@@ -94,7 +94,7 @@ locals {
94
94
for b in local . nightly_builds :
95
95
format (" %s_%s%s%s" ,
96
96
b. python_version ,
97
- " tpuvm" ,
97
+ b . accelerator == " tpu " ? " tpuvm" : format ( " cuda_%s " , b . cuda_version ) ,
98
98
try (b. cxx11_abi == " 0" , false ) ? " _precxx11" : " " ,
99
99
try (b. bundle_libtpu == " 1" , false ) ? " _libtpu" : " "
100
100
) => b
@@ -106,7 +106,7 @@ locals {
106
106
format (" r%s_%s_%s%s%s" ,
107
107
replace (b. package_version , " +" , " _" ),
108
108
b. python_version ,
109
- " tpuvm" ,
109
+ b . accelerator == " tpu " ? " tpuvm" : format ( " cuda_%s " , b . cuda_version ) ,
110
110
try (b. cxx11_abi == " 0" , false ) ? " _precxx11" : " " ,
111
111
try (b. bundle_libtpu == " 1" , false ) ? " _libtpu" : " "
112
112
) => b
@@ -136,13 +136,20 @@ module "nightly_builds" {
136
136
]
137
137
138
138
description = join (" " , [
139
- " Builds nightly xla:nightly_${ each . key } ' TPU docker image and " ,
140
- " corresponding wheels for PyTorch/XLA." ,
139
+ " Builds nightly xla:nightly_${ each . key } ' ${
140
+ each . value . accelerator == " tpu"
141
+ ? " TPU"
142
+ : format (" CUDA %s" , each. value . cuda_version )
143
+ } docker image and corresponding wheels for PyTorch/XLA." ,
141
144
" Trigger managed by Terraform setup in" ,
142
145
" infra/tpu-pytorch-releases/artifacts_builds.tf."
143
146
])
144
147
145
- wheels_dest = " ${ module . releases_storage_bucket . url } /wheels/tpuvm"
148
+ wheels_dest = " ${ module . releases_storage_bucket . url } /wheels/${
149
+ each . value . accelerator == " tpu"
150
+ ? " tpuvm"
151
+ : " cuda/${ each . value . cuda_version } "
152
+ } "
146
153
wheels_srcs = [" /dist/*.whl" ]
147
154
build_args = {
148
155
python_version = each.value.python_version
@@ -176,13 +183,20 @@ module "versioned_builds" {
176
183
image_tags = [each . key ]
177
184
178
185
description = join (" " , [
179
- " Builds official xla:${ each . key } ' TPU docker image and " ,
180
- " corresponding wheels for PyTorch/XLA." ,
186
+ " Builds official xla:${ each . key } ' ${
187
+ each . value . accelerator == " tpu"
188
+ ? " TPU"
189
+ : format (" CUDA %s" , each. value . cuda_version )
190
+ } docker image and corresponding wheels for PyTorch/XLA." ,
181
191
" Trigger managed by Terraform setup in" ,
182
192
" infra/tpu-pytorch-releases/artifacts_builds.tf."
183
193
])
184
194
185
- wheels_dest = " ${ module . releases_storage_bucket . url } /wheels/tpuvm"
195
+ wheels_dest = " ${ module . releases_storage_bucket . url } /wheels/${
196
+ each . value . accelerator == " tpu"
197
+ ? " tpuvm"
198
+ : " cuda/${ each . value . cuda_version } "
199
+ } "
186
200
wheels_srcs = [" /dist/*.whl" ]
187
201
# Pass docker build args to infra/ansible/Dockerfile, other than `ansible_vars`.
188
202
build_args = {
0 commit comments