diff --git a/infra/tpu-pytorch-releases/artifacts.auto.tfvars b/infra/tpu-pytorch-releases/artifacts.auto.tfvars index 36d3b98c920..8a9da5595ab 100644 --- a/infra/tpu-pytorch-releases/artifacts.auto.tfvars +++ b/infra/tpu-pytorch-releases/artifacts.auto.tfvars @@ -27,14 +27,19 @@ nightly_builds = [ ] # TODO: Remove this after the 2.1 release -xrt_nightly_builds = [ +xrt_versioned_builds = [ { accelerator = "tpu" python_version = "3.10" + pytorch_git_rev = "v2.1.0-rc5" + package_version = "2.1.0rc5+xrt" }, { accelerator = "cuda" + python_version = "3.10" cuda_version = "12.0" + pytorch_git_rev = "v2.1.0-rc5" + package_version = "2.1.0rc5+xrt" }, ] diff --git a/infra/tpu-pytorch-releases/artifacts_builds.tf b/infra/tpu-pytorch-releases/artifacts_builds.tf index cba72bfe8b7..1fedd3ecefa 100644 --- a/infra/tpu-pytorch-releases/artifacts_builds.tf +++ b/infra/tpu-pytorch-releases/artifacts_builds.tf @@ -17,10 +17,12 @@ variable "nightly_builds" { } // TODO: Remove this after the 2.1 release -variable "xrt_nightly_builds" { +variable "xrt_versioned_builds" { type = list( object({ + package_version = string accelerator = string + pytorch_git_rev = optional(string, "") cuda_version = optional(string, "11.8") python_version = optional(string, "3.8") arch = optional(string, "amd64") @@ -58,9 +60,10 @@ locals { } // TODO: Remove this after the 2.1 release - xrt_nightly_builds_dict = { - for b in var.xrt_nightly_builds : - format("%s_%s", + xrt_versioned_builds_dict = { + for b in var.xrt_versioned_builds : + format("r%s_%s_%s", + replace(b.package_version, "+", "_"), b.python_version, b.accelerator == "tpu" ? "tpuvm" : format("cuda_%s", b.cuda_version) ) => b @@ -123,9 +126,9 @@ module "nightly_builds" { } // TODO: Remove this after the 2.1 release -module "xrt_nightly_builds" { +module "xrt_versioned_builds" { source = "../terraform_modules/xla_docker_build" - for_each = local.xrt_nightly_builds_dict + for_each = local.xrt_versioned_builds_dict ansible_vars = merge(each.value, { package_version = var.nightly_package_version @@ -136,16 +139,16 @@ module "xrt_nightly_builds" { trigger_on_schedule = { schedule = "0 0 * * *", branch = "xrt" } - trigger_name = "nightly-xrt-${replace(each.key, "/[_.]/", "-")}" + trigger_name = replace(each.key, "/[_.]/", "-") image_name = "xla" image_tags = [ - "nightly_xrt_${each.key}", + each.key, # Append _YYYYMMDD suffix to nightly image name. - "nightly_xrt_${each.key}_$(date +%Y%m%d)", + "${each.key}_$(date +%Y%m%d)", ] description = join(" ", [ - "Builds nightly xla:nightly_${each.key}' ${ + "Builds nightly xla:${each.key}' ${ each.value.accelerator == "tpu" ? "TPU" : format("CUDA %s", each.value.cuda_version)