xpk 0.0.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (92) hide show
  1. xpk/__init__.py +15 -0
  2. xpk/api/__init__.py +15 -0
  3. xpk/api/storage_crd.yaml +52 -0
  4. xpk/commands/__init__.py +15 -0
  5. xpk/commands/batch.py +131 -0
  6. xpk/commands/cluster.py +808 -0
  7. xpk/commands/cluster_gcluster.py +269 -0
  8. xpk/commands/common.py +44 -0
  9. xpk/commands/config.py +29 -0
  10. xpk/commands/info.py +243 -0
  11. xpk/commands/inspector.py +357 -0
  12. xpk/commands/job.py +199 -0
  13. xpk/commands/kind.py +283 -0
  14. xpk/commands/kjob_common.py +44 -0
  15. xpk/commands/run.py +128 -0
  16. xpk/commands/shell.py +140 -0
  17. xpk/commands/storage.py +267 -0
  18. xpk/commands/version.py +27 -0
  19. xpk/commands/workload.py +889 -0
  20. xpk/core/__init__.py +15 -0
  21. xpk/core/blueprint/__init__.py +15 -0
  22. xpk/core/blueprint/blueprint_definitions.py +62 -0
  23. xpk/core/blueprint/blueprint_generator.py +708 -0
  24. xpk/core/capacity.py +185 -0
  25. xpk/core/cluster.py +564 -0
  26. xpk/core/cluster_private.py +200 -0
  27. xpk/core/commands.py +356 -0
  28. xpk/core/config.py +179 -0
  29. xpk/core/docker_container.py +225 -0
  30. xpk/core/docker_image.py +210 -0
  31. xpk/core/docker_manager.py +308 -0
  32. xpk/core/docker_resources.py +350 -0
  33. xpk/core/filestore.py +251 -0
  34. xpk/core/gcloud_context.py +196 -0
  35. xpk/core/gcluster_manager.py +176 -0
  36. xpk/core/gcsfuse.py +50 -0
  37. xpk/core/kjob.py +444 -0
  38. xpk/core/kueue.py +358 -0
  39. xpk/core/monitoring.py +134 -0
  40. xpk/core/nap.py +361 -0
  41. xpk/core/network.py +377 -0
  42. xpk/core/nodepool.py +581 -0
  43. xpk/core/pathways.py +377 -0
  44. xpk/core/ray.py +222 -0
  45. xpk/core/remote_state/__init__.py +15 -0
  46. xpk/core/remote_state/fuse_remote_state.py +99 -0
  47. xpk/core/remote_state/remote_state_client.py +38 -0
  48. xpk/core/resources.py +238 -0
  49. xpk/core/scheduling.py +253 -0
  50. xpk/core/storage.py +581 -0
  51. xpk/core/system_characteristics.py +1432 -0
  52. xpk/core/vertex.py +105 -0
  53. xpk/core/workload.py +341 -0
  54. xpk/core/workload_decorators/__init__.py +15 -0
  55. xpk/core/workload_decorators/rdma_decorator.py +129 -0
  56. xpk/core/workload_decorators/storage_decorator.py +52 -0
  57. xpk/core/workload_decorators/tcpxo_decorator.py +190 -0
  58. xpk/main.py +75 -0
  59. xpk/parser/__init__.py +15 -0
  60. xpk/parser/batch.py +43 -0
  61. xpk/parser/cluster.py +662 -0
  62. xpk/parser/common.py +259 -0
  63. xpk/parser/config.py +49 -0
  64. xpk/parser/core.py +135 -0
  65. xpk/parser/info.py +64 -0
  66. xpk/parser/inspector.py +65 -0
  67. xpk/parser/job.py +147 -0
  68. xpk/parser/kind.py +95 -0
  69. xpk/parser/run.py +47 -0
  70. xpk/parser/shell.py +59 -0
  71. xpk/parser/storage.py +316 -0
  72. xpk/parser/validators.py +39 -0
  73. xpk/parser/version.py +23 -0
  74. xpk/parser/workload.py +726 -0
  75. xpk/templates/__init__.py +15 -0
  76. xpk/templates/storage.yaml +13 -0
  77. xpk/utils/__init__.py +15 -0
  78. xpk/utils/console.py +55 -0
  79. xpk/utils/file.py +82 -0
  80. xpk/utils/gcs_utils.py +125 -0
  81. xpk/utils/kubectl.py +57 -0
  82. xpk/utils/network.py +168 -0
  83. xpk/utils/objects.py +88 -0
  84. xpk/utils/templates.py +28 -0
  85. xpk/utils/validation.py +80 -0
  86. xpk/utils/yaml.py +30 -0
  87. xpk-0.0.1.dist-info/LICENSE +202 -0
  88. xpk-0.0.1.dist-info/METADATA +1498 -0
  89. xpk-0.0.1.dist-info/RECORD +92 -0
  90. xpk-0.0.1.dist-info/WHEEL +5 -0
  91. xpk-0.0.1.dist-info/entry_points.txt +2 -0
  92. xpk-0.0.1.dist-info/top_level.txt +1 -0
xpk/core/system_characteristics.py ADDED
@@ -0,0 +1,1432 @@
1
+ """
2
+ Copyright 2023 Google LLC
3
+
4
+ Licensed under the Apache License, Version 2.0 (the "License");
5
+ you may not use this file except in compliance with the License.
6
+ You may obtain a copy of the License at
7
+
8
+ https://www.apache.org/licenses/LICENSE-2.0
9
+
10
+ Unless required by applicable law or agreed to in writing, software
11
+ distributed under the License is distributed on an "AS IS" BASIS,
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ See the License for the specific language governing permissions and
14
+ limitations under the License.
15
+ """
16
+
17
+ from dataclasses import dataclass
18
+
19
+ AcceleratorType = {'TPU': 1, 'GPU': 2, 'CPU': 3}
20
+
21
+
22
+ @dataclass
23
+ class AcceleratorCharacteristics:
24
+ resource_type: str
25
+ accelerator_label: str
26
+ machine_label: str
27
+
28
+
29
+ AcceleratorTypeToAcceleratorCharacteristics = {
30
+ # TPU
31
+ AcceleratorType['TPU']: AcceleratorCharacteristics(
32
+ 'google.com/tpu',
33
+ 'cloud.google.com/gke-tpu-accelerator',
34
+ 'cloud.google.com/gke-tpu-topology',
35
+ ),
36
+ # GPU
37
+ AcceleratorType['GPU']: AcceleratorCharacteristics(
38
+ 'nvidia.com/gpu',
39
+ 'cloud.google.com/gke-accelerator',
40
+ 'cloud.google.com/gce-machine-type',
41
+ ),
42
+ # CPU
43
+ AcceleratorType['CPU']: AcceleratorCharacteristics(
44
+ 'cpu', '', 'cloud.google.com/gke-nodepool'
45
+ ),
46
+ }
47
+
48
+
49
+ @dataclass
50
+ class SystemCharacteristics:
51
+ topology: str
52
+ vms_per_slice: int
53
+ gke_accelerator: str
54
+ gce_machine_type: str
55
+ chips_per_vm: int
56
+ accelerator_type: AcceleratorType # type: ignore
57
+ device_type: str
58
+
59
+
60
+ def get_system_characteristics(
61
+ args,
62
+ ) -> tuple[SystemCharacteristics | None, int]:
63
+ """Get system characteristics based on user provided arguments.
64
+
65
+ Args:
66
+ args: user provided arguments for running the command.
67
+
68
+ Returns:
69
+ Tuple with string with the system characteristics and
70
+ int of 0 if successful and 1 otherwise.
71
+ """
72
+ device_type = args.tpu_type if args.tpu_type else args.device_type
73
+ return get_system_characteristics_by_device_type(device_type)
74
+
75
+
76
+ def get_system_characteristics_by_device_type(
77
+ device_type,
78
+ ) -> tuple[SystemCharacteristics | None, int]:
79
+ """Get system characteristics based on device_type.
80
+
81
+ Args:
82
+ device_type: device_type for running the command.
83
+
84
+ Returns:
85
+ Tuple with string with the system characteristics and
86
+ int of 0 if successful and 1 otherwise.
87
+ """
88
+ if device_type in UserFacingNameToSystemCharacteristics:
89
+ return UserFacingNameToSystemCharacteristics[device_type], 0
90
+ else:
91
+ return None, 1
92
+
93
+
94
+ ################### Subcommand Helper Functions #############################
95
+ """ !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
96
+ IF YOU MODIFY THE BELOW UserFacingNameToSystemCharacteristics MAP YOU SHOULD
97
+ ALSO ADD CORRESPONDING MODIFICATIONS TO UserFacingNameToSystemCharacteristics
98
+ IN MaxText/accelerator_to_spec_map.py !!!!! """
99
+ # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv
100
+ UserFacingNameToSystemCharacteristics = {
101
+ # GPU system characteristics
102
+ # l4-$CHIPSc
103
+ 'l4-1': SystemCharacteristics(
104
+ 'N/A',
105
+ 1,
106
+ 'nvidia-l4',
107
+ 'g2-standard-12',
108
+ 1,
109
+ AcceleratorType['GPU'],
110
+ 'l4-1',
111
+ ),
112
+ 'l4-2': SystemCharacteristics(
113
+ 'N/A',
114
+ 1,
115
+ 'nvidia-l4',
116
+ 'g2-standard-24',
117
+ 2,
118
+ AcceleratorType['GPU'],
119
+ 'l4-2',
120
+ ),
121
+ 'l4-4': SystemCharacteristics(
122
+ 'N/A',
123
+ 1,
124
+ 'nvidia-l4',
125
+ 'g2-standard-48',
126
+ 4,
127
+ AcceleratorType['GPU'],
128
+ 'l4-4',
129
+ ),
130
+ 'l4-8': SystemCharacteristics(
131
+ 'N/A',
132
+ 1,
133
+ 'nvidia-l4',
134
+ 'g2-standard-96',
135
+ 8,
136
+ AcceleratorType['GPU'],
137
+ 'l4-8',
138
+ ),
139
+ # A100-40gb-$CHIPSc
140
+ 'a100-40gb-1': SystemCharacteristics(
141
+ 'N/A',
142
+ 1,
143
+ 'nvidia-tesla-a100',
144
+ 'a2-highgpu-1g',
145
+ 1,
146
+ AcceleratorType['GPU'],
147
+ 'a100-40gb-1',
148
+ ),
149
+ 'a100-40gb-2': SystemCharacteristics(
150
+ 'N/A',
151
+ 1,
152
+ 'nvidia-tesla-a100',
153
+ 'a2-highgpu-2g',
154
+ 2,
155
+ AcceleratorType['GPU'],
156
+ 'a100-40gb-2',
157
+ ),
158
+ 'a100-40gb-4': SystemCharacteristics(
159
+ 'N/A',
160
+ 1,
161
+ 'nvidia-tesla-a100',
162
+ 'a2-highgpu-4g',
163
+ 4,
164
+ AcceleratorType['GPU'],
165
+ 'a100-40gb-4',
166
+ ),
167
+ 'a100-40gb-8': SystemCharacteristics(
168
+ 'N/A',
169
+ 1,
170
+ 'nvidia-tesla-a100',
171
+ 'a2-highgpu-8g',
172
+ 8,
173
+ AcceleratorType['GPU'],
174
+ 'a100-40gb-8',
175
+ ),
176
+ 'h200-141gb-8': SystemCharacteristics(
177
+ 'N/A',
178
+ 1,
179
+ 'nvidia-h200-141gb',
180
+ 'a3-ultragpu-8g',
181
+ 8,
182
+ AcceleratorType['GPU'],
183
+ 'h200-141gb-8',
184
+ ),
185
+ # H100-80gb-$CHIPS
186
+ 'h100-80gb-8': SystemCharacteristics(
187
+ 'N/A',
188
+ 1,
189
+ 'nvidia-h100-80gb',
190
+ 'a3-highgpu-8g',
191
+ 8,
192
+ AcceleratorType['GPU'],
193
+ 'h100-80gb-8',
194
+ ),
195
+ # H100-mega-80gb-$CHIPS
196
+ 'h100-mega-80gb-8': SystemCharacteristics(
197
+ 'N/A',
198
+ 1,
199
+ 'nvidia-h100-mega-80gb',
200
+ 'a3-megagpu-8g',
201
+ 8,
202
+ AcceleratorType['GPU'],
203
+ 'h100-mega-80gb-8',
204
+ ),
205
+ # TPU system characteristics
206
+ # v6e
207
+ 'v6e-1': SystemCharacteristics(
208
+ '1x1',
209
+ 1,
210
+ 'tpu-v6e-slice',
211
+ 'ct6e-standard-1t',
212
+ 1,
213
+ AcceleratorType['TPU'],
214
+ 'v6e-1',
215
+ ),
216
+ 'v6e-4': SystemCharacteristics(
217
+ '2x2',
218
+ 1,
219
+ 'tpu-v6e-slice',
220
+ 'ct6e-standard-4t',
221
+ 4,
222
+ AcceleratorType['TPU'],
223
+ 'v6e-4',
224
+ ),
225
+ 'v6e-8': SystemCharacteristics(
226
+ '2x4',
227
+ 2,
228
+ 'tpu-v6e-slice',
229
+ 'ct6e-standard-4t',
230
+ 4,
231
+ AcceleratorType['TPU'],
232
+ 'v6e-8',
233
+ ),
234
+ 'v6e-16': SystemCharacteristics(
235
+ '4x4',
236
+ 4,
237
+ 'tpu-v6e-slice',
238
+ 'ct6e-standard-4t',
239
+ 4,
240
+ AcceleratorType['TPU'],
241
+ 'v6e-16',
242
+ ),
243
+ 'v6e-32': SystemCharacteristics(
244
+ '4x8',
245
+ 8,
246
+ 'tpu-v6e-slice',
247
+ 'ct6e-standard-4t',
248
+ 4,
249
+ AcceleratorType['TPU'],
250
+ 'v6e-32',
251
+ ),
252
+ 'v6e-64': SystemCharacteristics(
253
+ '8x8',
254
+ 16,
255
+ 'tpu-v6e-slice',
256
+ 'ct6e-standard-4t',
257
+ 4,
258
+ AcceleratorType['TPU'],
259
+ 'v6e-64',
260
+ ),
261
+ 'v6e-128': SystemCharacteristics(
262
+ '8x16',
263
+ 32,
264
+ 'tpu-v6e-slice',
265
+ 'ct6e-standard-4t',
266
+ 4,
267
+ AcceleratorType['TPU'],
268
+ 'v6e-128',
269
+ ),
270
+ 'v6e-256': SystemCharacteristics(
271
+ '16x16',
272
+ 64,
273
+ 'tpu-v6e-slice',
274
+ 'ct6e-standard-4t',
275
+ 4,
276
+ AcceleratorType['TPU'],
277
+ 'v6e-256',
278
+ ),
279
+ # v5p
280
+ 'v5p-8': SystemCharacteristics(
281
+ '2x2x1',
282
+ 1,
283
+ 'tpu-v5p-slice',
284
+ 'ct5p-hightpu-4t',
285
+ 4,
286
+ AcceleratorType['TPU'],
287
+ 'v5p-8',
288
+ ),
289
+ 'v5p-16': SystemCharacteristics(
290
+ '2x2x2',
291
+ 2,
292
+ 'tpu-v5p-slice',
293
+ 'ct5p-hightpu-4t',
294
+ 4,
295
+ AcceleratorType['TPU'],
296
+ 'v5p-16',
297
+ ),
298
+ 'v5p-32': SystemCharacteristics(
299
+ '2x2x4',
300
+ 4,
301
+ 'tpu-v5p-slice',
302
+ 'ct5p-hightpu-4t',
303
+ 4,
304
+ AcceleratorType['TPU'],
305
+ 'v5p-32',
306
+ ),
307
+ 'v5p-64': SystemCharacteristics(
308
+ '2x4x4',
309
+ 8,
310
+ 'tpu-v5p-slice',
311
+ 'ct5p-hightpu-4t',
312
+ 4,
313
+ AcceleratorType['TPU'],
314
+ 'v5p-64',
315
+ ),
316
+ 'v5p-128': SystemCharacteristics(
317
+ '4x4x4',
318
+ 16,
319
+ 'tpu-v5p-slice',
320
+ 'ct5p-hightpu-4t',
321
+ 4,
322
+ AcceleratorType['TPU'],
323
+ 'v5p-128',
324
+ ),
325
+ 'v5p-256': SystemCharacteristics(
326
+ '4x4x8',
327
+ 32,
328
+ 'tpu-v5p-slice',
329
+ 'ct5p-hightpu-4t',
330
+ 4,
331
+ AcceleratorType['TPU'],
332
+ 'v5p-256',
333
+ ),
334
+ 'v5p-384': SystemCharacteristics(
335
+ '4x4x12',
336
+ 48,
337
+ 'tpu-v5p-slice',
338
+ 'ct5p-hightpu-4t',
339
+ 4,
340
+ AcceleratorType['TPU'],
341
+ 'v5p-384',
342
+ ),
343
+ 'v5p-512': SystemCharacteristics(
344
+ '4x8x8',
345
+ 64,
346
+ 'tpu-v5p-slice',
347
+ 'ct5p-hightpu-4t',
348
+ 4,
349
+ AcceleratorType['TPU'],
350
+ 'v5p-512',
351
+ ),
352
+ 'v5p-640': SystemCharacteristics(
353
+ '4x4x20',
354
+ 80,
355
+ 'tpu-v5p-slice',
356
+ 'ct5p-hightpu-4t',
357
+ 4,
358
+ AcceleratorType['TPU'],
359
+ 'v5p-640',
360
+ ),
361
+ 'v5p-768': SystemCharacteristics(
362
+ '4x8x12',
363
+ 96,
364
+ 'tpu-v5p-slice',
365
+ 'ct5p-hightpu-4t',
366
+ 4,
367
+ AcceleratorType['TPU'],
368
+ 'v5p-768',
369
+ ),
370
+ 'v5p-896': SystemCharacteristics(
371
+ '4x4x28',
372
+ 112,
373
+ 'tpu-v5p-slice',
374
+ 'ct5p-hightpu-4t',
375
+ 4,
376
+ AcceleratorType['TPU'],
377
+ 'v5p-896',
378
+ ),
379
+ 'v5p-1024': SystemCharacteristics(
380
+ '8x8x8',
381
+ 128,
382
+ 'tpu-v5p-slice',
383
+ 'ct5p-hightpu-4t',
384
+ 4,
385
+ AcceleratorType['TPU'],
386
+ 'v5p-1024',
387
+ ),
388
+ 'v5p-1152': SystemCharacteristics(
389
+ '4x12x12',
390
+ 144,
391
+ 'tpu-v5p-slice',
392
+ 'ct5p-hightpu-4t',
393
+ 4,
394
+ AcceleratorType['TPU'],
395
+ 'v5p-1152',
396
+ ),
397
+ 'v5p-1280': SystemCharacteristics(
398
+ '4x8x20',
399
+ 160,
400
+ 'tpu-v5p-slice',
401
+ 'ct5p-hightpu-4t',
402
+ 4,
403
+ AcceleratorType['TPU'],
404
+ 'v5p-1280',
405
+ ),
406
+ 'v5p-1408': SystemCharacteristics(
407
+ '4x4x44',
408
+ 176,
409
+ 'tpu-v5p-slice',
410
+ 'ct5p-hightpu-4t',
411
+ 4,
412
+ AcceleratorType['TPU'],
413
+ 'v5p-1408',
414
+ ),
415
+ 'v5p-1536': SystemCharacteristics(
416
+ '8x8x12',
417
+ 192,
418
+ 'tpu-v5p-slice',
419
+ 'ct5p-hightpu-4t',
420
+ 4,
421
+ AcceleratorType['TPU'],
422
+ 'v5p-1536',
423
+ ),
424
+ 'v5p-1664': SystemCharacteristics(
425
+ '4x4x52',
426
+ 208,
427
+ 'tpu-v5p-slice',
428
+ 'ct5p-hightpu-4t',
429
+ 4,
430
+ AcceleratorType['TPU'],
431
+ 'v5p-1664',
432
+ ),
433
+ 'v5p-1792': SystemCharacteristics(
434
+ '4x8x28',
435
+ 224,
436
+ 'tpu-v5p-slice',
437
+ 'ct5p-hightpu-4t',
438
+ 4,
439
+ AcceleratorType['TPU'],
440
+ 'v5p-1792',
441
+ ),
442
+ 'v5p-1920': SystemCharacteristics(
443
+ '4x12x20',
444
+ 240,
445
+ 'tpu-v5p-slice',
446
+ 'ct5p-hightpu-4t',
447
+ 4,
448
+ AcceleratorType['TPU'],
449
+ 'v5p-1920',
450
+ ),
451
+ 'v5p-2048': SystemCharacteristics(
452
+ '8x8x16',
453
+ 256,
454
+ 'tpu-v5p-slice',
455
+ 'ct5p-hightpu-4t',
456
+ 4,
457
+ AcceleratorType['TPU'],
458
+ 'v5p-2048',
459
+ ),
460
+ 'v5p-2176': SystemCharacteristics(
461
+ '4x4x68',
462
+ 272,
463
+ 'tpu-v5p-slice',
464
+ 'ct5p-hightpu-4t',
465
+ 4,
466
+ AcceleratorType['TPU'],
467
+ 'v5p-2176',
468
+ ),
469
+ 'v5p-2304': SystemCharacteristics(
470
+ '8x12x12',
471
+ 288,
472
+ 'tpu-v5p-slice',
473
+ 'ct5p-hightpu-4t',
474
+ 4,
475
+ AcceleratorType['TPU'],
476
+ 'v5p-2304',
477
+ ),
478
+ 'v5p-2432': SystemCharacteristics(
479
+ '4x4x76',
480
+ 304,
481
+ 'tpu-v5p-slice',
482
+ 'ct5p-hightpu-4t',
483
+ 4,
484
+ AcceleratorType['TPU'],
485
+ 'v5p-2432',
486
+ ),
487
+ 'v5p-2560': SystemCharacteristics(
488
+ '8x8x20',
489
+ 320,
490
+ 'tpu-v5p-slice',
491
+ 'ct5p-hightpu-4t',
492
+ 4,
493
+ AcceleratorType['TPU'],
494
+ 'v5p-2560',
495
+ ),
496
+ 'v5p-2688': SystemCharacteristics(
497
+ '4x12x28',
498
+ 336,
499
+ 'tpu-v5p-slice',
500
+ 'ct5p-hightpu-4t',
501
+ 4,
502
+ AcceleratorType['TPU'],
503
+ 'v5p-2688',
504
+ ),
505
+ 'v5p-2816': SystemCharacteristics(
506
+ '4x8x44',
507
+ 352,
508
+ 'tpu-v5p-slice',
509
+ 'ct5p-hightpu-4t',
510
+ 4,
511
+ AcceleratorType['TPU'],
512
+ 'v5p-2816',
513
+ ),
514
+ 'v5p-2944': SystemCharacteristics(
515
+ '4x4x92',
516
+ 368,
517
+ 'tpu-v5p-slice',
518
+ 'ct5p-hightpu-4t',
519
+ 4,
520
+ AcceleratorType['TPU'],
521
+ 'v5p-2944',
522
+ ),
523
+ 'v5p-3072': SystemCharacteristics(
524
+ '8x12x16',
525
+ 384,
526
+ 'tpu-v5p-slice',
527
+ 'ct5p-hightpu-4t',
528
+ 4,
529
+ AcceleratorType['TPU'],
530
+ 'v5p-3072',
531
+ ),
532
+ 'v5p-3200': SystemCharacteristics(
533
+ '4x20x20',
534
+ 400,
535
+ 'tpu-v5p-slice',
536
+ 'ct5p-hightpu-4t',
537
+ 4,
538
+ AcceleratorType['TPU'],
539
+ 'v5p-3200',
540
+ ),
541
+ 'v5p-3328': SystemCharacteristics(
542
+ '4x8x52',
543
+ 416,
544
+ 'tpu-v5p-slice',
545
+ 'ct5p-hightpu-4t',
546
+ 4,
547
+ AcceleratorType['TPU'],
548
+ 'v5p-3328',
549
+ ),
550
+ 'v5p-3456': SystemCharacteristics(
551
+ '12x12x12',
552
+ 432,
553
+ 'tpu-v5p-slice',
554
+ 'ct5p-hightpu-4t',
555
+ 4,
556
+ AcceleratorType['TPU'],
557
+ 'v5p-3456',
558
+ ),
559
+ 'v5p-3584': SystemCharacteristics(
560
+ '8x8x28',
561
+ 448,
562
+ 'tpu-v5p-slice',
563
+ 'ct5p-hightpu-4t',
564
+ 4,
565
+ AcceleratorType['TPU'],
566
+ 'v5p-3584',
567
+ ),
568
+ 'v5p-3712': SystemCharacteristics(
569
+ '4x4x116',
570
+ 464,
571
+ 'tpu-v5p-slice',
572
+ 'ct5p-hightpu-4t',
573
+ 4,
574
+ AcceleratorType['TPU'],
575
+ 'v5p-3712',
576
+ ),
577
+ 'v5p-3840': SystemCharacteristics(
578
+ '8x12x20',
579
+ 480,
580
+ 'tpu-v5p-slice',
581
+ 'ct5p-hightpu-4t',
582
+ 4,
583
+ AcceleratorType['TPU'],
584
+ 'v5p-3840',
585
+ ),
586
+ 'v5p-3968': SystemCharacteristics(
587
+ '4x4x124',
588
+ 496,
589
+ 'tpu-v5p-slice',
590
+ 'ct5p-hightpu-4t',
591
+ 4,
592
+ AcceleratorType['TPU'],
593
+ 'v5p-3968',
594
+ ),
595
+ 'v5p-4096': SystemCharacteristics(
596
+ '8x16x16',
597
+ 512,
598
+ 'tpu-v5p-slice',
599
+ 'ct5p-hightpu-4t',
600
+ 4,
601
+ AcceleratorType['TPU'],
602
+ 'v5p-4096',
603
+ ),
604
+ 'v5p-4224': SystemCharacteristics(
605
+ '4x12x44',
606
+ 528,
607
+ 'tpu-v5p-slice',
608
+ 'ct5p-hightpu-4t',
609
+ 4,
610
+ AcceleratorType['TPU'],
611
+ 'v5p-4224',
612
+ ),
613
+ 'v5p-4352': SystemCharacteristics(
614
+ '4x8x68',
615
+ 544,
616
+ 'tpu-v5p-slice',
617
+ 'ct5p-hightpu-4t',
618
+ 4,
619
+ AcceleratorType['TPU'],
620
+ 'v5p-4352',
621
+ ),
622
+ 'v5p-4480': SystemCharacteristics(
623
+ '4x20x28',
624
+ 560,
625
+ 'tpu-v5p-slice',
626
+ 'ct5p-hightpu-4t',
627
+ 4,
628
+ AcceleratorType['TPU'],
629
+ 'v5p-4480',
630
+ ),
631
+ 'v5p-4608': SystemCharacteristics(
632
+ '12x12x16',
633
+ 576,
634
+ 'tpu-v5p-slice',
635
+ 'ct5p-hightpu-4t',
636
+ 4,
637
+ AcceleratorType['TPU'],
638
+ 'v5p-4608',
639
+ ),
640
+ 'v5p-4736': SystemCharacteristics(
641
+ '4x4x148',
642
+ 592,
643
+ 'tpu-v5p-slice',
644
+ 'ct5p-hightpu-4t',
645
+ 4,
646
+ AcceleratorType['TPU'],
647
+ 'v5p-4736',
648
+ ),
649
+ 'v5p-4864': SystemCharacteristics(
650
+ '4x8x76',
651
+ 608,
652
+ 'tpu-v5p-slice',
653
+ 'ct5p-hightpu-4t',
654
+ 4,
655
+ AcceleratorType['TPU'],
656
+ 'v5p-4864',
657
+ ),
658
+ 'v5p-4992': SystemCharacteristics(
659
+ '4x12x52',
660
+ 624,
661
+ 'tpu-v5p-slice',
662
+ 'ct5p-hightpu-4t',
663
+ 4,
664
+ AcceleratorType['TPU'],
665
+ 'v5p-4992',
666
+ ),
667
+ 'v5p-5120': SystemCharacteristics(
668
+ '8x16x20',
669
+ 640,
670
+ 'tpu-v5p-slice',
671
+ 'ct5p-hightpu-4t',
672
+ 4,
673
+ AcceleratorType['TPU'],
674
+ 'v5p-5120',
675
+ ),
676
+ 'v5p-5248': SystemCharacteristics(
677
+ '4x4x164',
678
+ 656,
679
+ 'tpu-v5p-slice',
680
+ 'ct5p-hightpu-4t',
681
+ 4,
682
+ AcceleratorType['TPU'],
683
+ 'v5p-5248',
684
+ ),
685
+ 'v5p-5376': SystemCharacteristics(
686
+ '8x12x28',
687
+ 672,
688
+ 'tpu-v5p-slice',
689
+ 'ct5p-hightpu-4t',
690
+ 4,
691
+ AcceleratorType['TPU'],
692
+ 'v5p-5376',
693
+ ),
694
+ 'v5p-5504': SystemCharacteristics(
695
+ '4x4x172',
696
+ 688,
697
+ 'tpu-v5p-slice',
698
+ 'ct5p-hightpu-4t',
699
+ 4,
700
+ AcceleratorType['TPU'],
701
+ 'v5p-5504',
702
+ ),
703
+ 'v5p-5632': SystemCharacteristics(
704
+ '8x8x44',
705
+ 704,
706
+ 'tpu-v5p-slice',
707
+ 'ct5p-hightpu-4t',
708
+ 4,
709
+ AcceleratorType['TPU'],
710
+ 'v5p-5632',
711
+ ),
712
+ 'v5p-5760': SystemCharacteristics(
713
+ '12x12x20',
714
+ 720,
715
+ 'tpu-v5p-slice',
716
+ 'ct5p-hightpu-4t',
717
+ 4,
718
+ AcceleratorType['TPU'],
719
+ 'v5p-5760',
720
+ ),
721
+ 'v5p-5888': SystemCharacteristics(
722
+ '4x8x92',
723
+ 736,
724
+ 'tpu-v5p-slice',
725
+ 'ct5p-hightpu-4t',
726
+ 4,
727
+ AcceleratorType['TPU'],
728
+ 'v5p-5888',
729
+ ),
730
+ 'v5p-6016': SystemCharacteristics(
731
+ '4x4x188',
732
+ 752,
733
+ 'tpu-v5p-slice',
734
+ 'ct5p-hightpu-4t',
735
+ 4,
736
+ AcceleratorType['TPU'],
737
+ 'v5p-6016',
738
+ ),
739
+ 'v5p-6144': SystemCharacteristics(
740
+ '12x16x16',
741
+ 768,
742
+ 'tpu-v5p-slice',
743
+ 'ct5p-hightpu-4t',
744
+ 4,
745
+ AcceleratorType['TPU'],
746
+ 'v5p-6144',
747
+ ),
748
+ 'v5p-6272': SystemCharacteristics(
749
+ '4x28x28',
750
+ 784,
751
+ 'tpu-v5p-slice',
752
+ 'ct5p-hightpu-4t',
753
+ 4,
754
+ AcceleratorType['TPU'],
755
+ 'v5p-6272',
756
+ ),
757
+ 'v5p-6400': SystemCharacteristics(
758
+ '8x20x20',
759
+ 800,
760
+ 'tpu-v5p-slice',
761
+ 'ct5p-hightpu-4t',
762
+ 4,
763
+ AcceleratorType['TPU'],
764
+ 'v5p-6400',
765
+ ),
766
+ 'v5p-6528': SystemCharacteristics(
767
+ '4x12x68',
768
+ 816,
769
+ 'tpu-v5p-slice',
770
+ 'ct5p-hightpu-4t',
771
+ 4,
772
+ AcceleratorType['TPU'],
773
+ 'v5p-6528',
774
+ ),
775
+ 'v5p-6656': SystemCharacteristics(
776
+ '8x8x52',
777
+ 832,
778
+ 'tpu-v5p-slice',
779
+ 'ct5p-hightpu-4t',
780
+ 4,
781
+ AcceleratorType['TPU'],
782
+ 'v5p-6656',
783
+ ),
784
+ 'v5p-6784': SystemCharacteristics(
785
+ '4x4x212',
786
+ 848,
787
+ 'tpu-v5p-slice',
788
+ 'ct5p-hightpu-4t',
789
+ 4,
790
+ AcceleratorType['TPU'],
791
+ 'v5p-6784',
792
+ ),
793
+ 'v5p-6912': SystemCharacteristics(
794
+ '12x12x24',
795
+ 864,
796
+ 'tpu-v5p-slice',
797
+ 'ct5p-hightpu-4t',
798
+ 4,
799
+ AcceleratorType['TPU'],
800
+ 'v5p-6912',
801
+ ),
802
+ 'v5p-7040': SystemCharacteristics(
803
+ '4x20x44',
804
+ 880,
805
+ 'tpu-v5p-slice',
806
+ 'ct5p-hightpu-4t',
807
+ 4,
808
+ AcceleratorType['TPU'],
809
+ 'v5p-7040',
810
+ ),
811
+ 'v5p-7168': SystemCharacteristics(
812
+ '8x16x28',
813
+ 896,
814
+ 'tpu-v5p-slice',
815
+ 'ct5p-hightpu-4t',
816
+ 4,
817
+ AcceleratorType['TPU'],
818
+ 'v5p-7168',
819
+ ),
820
+ 'v5p-7296': SystemCharacteristics(
821
+ '4x12x76',
822
+ 912,
823
+ 'tpu-v5p-slice',
824
+ 'ct5p-hightpu-4t',
825
+ 4,
826
+ AcceleratorType['TPU'],
827
+ 'v5p-7296',
828
+ ),
829
+ 'v5p-7424': SystemCharacteristics(
830
+ '4x8x116',
831
+ 928,
832
+ 'tpu-v5p-slice',
833
+ 'ct5p-hightpu-4t',
834
+ 4,
835
+ AcceleratorType['TPU'],
836
+ 'v5p-7424',
837
+ ),
838
+ 'v5p-7552': SystemCharacteristics(
839
+ '4x4x236',
840
+ 944,
841
+ 'tpu-v5p-slice',
842
+ 'ct5p-hightpu-4t',
843
+ 4,
844
+ AcceleratorType['TPU'],
845
+ 'v5p-7552',
846
+ ),
847
+ 'v5p-7680': SystemCharacteristics(
848
+ '12x16x20',
849
+ 960,
850
+ 'tpu-v5p-slice',
851
+ 'ct5p-hightpu-4t',
852
+ 4,
853
+ AcceleratorType['TPU'],
854
+ 'v5p-7680',
855
+ ),
856
+ 'v5p-7808': SystemCharacteristics(
857
+ '4x4x244',
858
+ 976,
859
+ 'tpu-v5p-slice',
860
+ 'ct5p-hightpu-4t',
861
+ 4,
862
+ AcceleratorType['TPU'],
863
+ 'v5p-7808',
864
+ ),
865
+ 'v5p-7936': SystemCharacteristics(
866
+ '4x8x124',
867
+ 992,
868
+ 'tpu-v5p-slice',
869
+ 'ct5p-hightpu-4t',
870
+ 4,
871
+ AcceleratorType['TPU'],
872
+ 'v5p-7936',
873
+ ),
874
+ 'v5p-8064': SystemCharacteristics(
875
+ '12x12x28',
876
+ 1008,
877
+ 'tpu-v5p-slice',
878
+ 'ct5p-hightpu-4t',
879
+ 4,
880
+ AcceleratorType['TPU'],
881
+ 'v5p-8064',
882
+ ),
883
+ 'v5p-8192': SystemCharacteristics(
884
+ '16x16x16',
885
+ 1024,
886
+ 'tpu-v5p-slice',
887
+ 'ct5p-hightpu-4t',
888
+ 4,
889
+ AcceleratorType['TPU'],
890
+ 'v5p-8192',
891
+ ),
892
+ 'v5p-8320': SystemCharacteristics(
893
+ '4x20x52',
894
+ 1040,
895
+ 'tpu-v5p-slice',
896
+ 'ct5p-hightpu-4t',
897
+ 4,
898
+ AcceleratorType['TPU'],
899
+ 'v5p-8320',
900
+ ),
901
+ 'v5p-8448': SystemCharacteristics(
902
+ '8x12x44',
903
+ 1056,
904
+ 'tpu-v5p-slice',
905
+ 'ct5p-hightpu-4t',
906
+ 4,
907
+ AcceleratorType['TPU'],
908
+ 'v5p-8448',
909
+ ),
910
+ 'v5p-8704': SystemCharacteristics(
911
+ '8x8x68',
912
+ 1088,
913
+ 'tpu-v5p-slice',
914
+ 'ct5p-hightpu-4t',
915
+ 4,
916
+ AcceleratorType['TPU'],
917
+ 'v5p-8704',
918
+ ),
919
+ 'v5p-8832': SystemCharacteristics(
920
+ '4x12x92',
921
+ 1104,
922
+ 'tpu-v5p-slice',
923
+ 'ct5p-hightpu-4t',
924
+ 4,
925
+ AcceleratorType['TPU'],
926
+ 'v5p-8832',
927
+ ),
928
+ 'v5p-8960': SystemCharacteristics(
929
+ '8x20x28',
930
+ 1120,
931
+ 'tpu-v5p-slice',
932
+ 'ct5p-hightpu-4t',
933
+ 4,
934
+ AcceleratorType['TPU'],
935
+ 'v5p-8960',
936
+ ),
937
+ 'v5p-9216': SystemCharacteristics(
938
+ '12x16x24',
939
+ 1152,
940
+ 'tpu-v5p-slice',
941
+ 'ct5p-hightpu-4t',
942
+ 4,
943
+ AcceleratorType['TPU'],
944
+ 'v5p-9216',
945
+ ),
946
+ 'v5p-9472': SystemCharacteristics(
947
+ '4x8x148',
948
+ 1184,
949
+ 'tpu-v5p-slice',
950
+ 'ct5p-hightpu-4t',
951
+ 4,
952
+ AcceleratorType['TPU'],
953
+ 'v5p-9472',
954
+ ),
955
+ 'v5p-9600': SystemCharacteristics(
956
+ '12x20x20',
957
+ 1200,
958
+ 'tpu-v5p-slice',
959
+ 'ct5p-hightpu-4t',
960
+ 4,
961
+ AcceleratorType['TPU'],
962
+ 'v5p-9600',
963
+ ),
964
+ 'v5p-9728': SystemCharacteristics(
965
+ '8x8x76',
966
+ 1216,
967
+ 'tpu-v5p-slice',
968
+ 'ct5p-hightpu-4t',
969
+ 4,
970
+ AcceleratorType['TPU'],
971
+ 'v5p-9728',
972
+ ),
973
+ 'v5p-9856': SystemCharacteristics(
974
+ '4x28x44',
975
+ 1232,
976
+ 'tpu-v5p-slice',
977
+ 'ct5p-hightpu-4t',
978
+ 4,
979
+ AcceleratorType['TPU'],
980
+ 'v5p-9856',
981
+ ),
982
+ 'v5p-9984': SystemCharacteristics(
983
+ '8x12x52',
984
+ 1248,
985
+ 'tpu-v5p-slice',
986
+ 'ct5p-hightpu-4t',
987
+ 4,
988
+ AcceleratorType['TPU'],
989
+ 'v5p-9984',
990
+ ),
991
+ 'v5p-10240': SystemCharacteristics(
992
+ '16x16x20',
993
+ 1280,
994
+ 'tpu-v5p-slice',
995
+ 'ct5p-hightpu-4t',
996
+ 4,
997
+ AcceleratorType['TPU'],
998
+ 'v5p-10240',
999
+ ),
1000
+ 'v5p-10368': SystemCharacteristics(
1001
+ '12x12x36',
1002
+ 1296,
1003
+ 'tpu-v5p-slice',
1004
+ 'ct5p-hightpu-4t',
1005
+ 4,
1006
+ AcceleratorType['TPU'],
1007
+ 'v5p-10368',
1008
+ ),
1009
+ 'v5p-10496': SystemCharacteristics(
1010
+ '4x8x164',
1011
+ 1312,
1012
+ 'tpu-v5p-slice',
1013
+ 'ct5p-hightpu-4t',
1014
+ 4,
1015
+ AcceleratorType['TPU'],
1016
+ 'v5p-10496',
1017
+ ),
1018
+ 'v5p-10752': SystemCharacteristics(
1019
+ '12x16x28',
1020
+ 1344,
1021
+ 'tpu-v5p-slice',
1022
+ 'ct5p-hightpu-4t',
1023
+ 4,
1024
+ AcceleratorType['TPU'],
1025
+ 'v5p-10752',
1026
+ ),
1027
+ 'v5p-10880': SystemCharacteristics(
1028
+ '4x20x68',
1029
+ 1360,
1030
+ 'tpu-v5p-slice',
1031
+ 'ct5p-hightpu-4t',
1032
+ 4,
1033
+ AcceleratorType['TPU'],
1034
+ 'v5p-10880',
1035
+ ),
1036
+ 'v5p-11008': SystemCharacteristics(
1037
+ '4x8x172',
1038
+ 1376,
1039
+ 'tpu-v5p-slice',
1040
+ 'ct5p-hightpu-4t',
1041
+ 4,
1042
+ AcceleratorType['TPU'],
1043
+ 'v5p-11008',
1044
+ ),
1045
+ 'v5p-11136': SystemCharacteristics(
1046
+ '4x12x116',
1047
+ 1392,
1048
+ 'tpu-v5p-slice',
1049
+ 'ct5p-hightpu-4t',
1050
+ 4,
1051
+ AcceleratorType['TPU'],
1052
+ 'v5p-11136',
1053
+ ),
1054
+ 'v5p-11264': SystemCharacteristics(
1055
+ '8x16x44',
1056
+ 1408,
1057
+ 'tpu-v5p-slice',
1058
+ 'ct5p-hightpu-4t',
1059
+ 4,
1060
+ AcceleratorType['TPU'],
1061
+ 'v5p-11264',
1062
+ ),
1063
+ 'v5p-11520': SystemCharacteristics(
1064
+ '12x20x24',
1065
+ 1440,
1066
+ 'tpu-v5p-slice',
1067
+ 'ct5p-hightpu-4t',
1068
+ 4,
1069
+ AcceleratorType['TPU'],
1070
+ 'v5p-11520',
1071
+ ),
1072
+ 'v5p-11648': SystemCharacteristics(
1073
+ '4x28x52',
1074
+ 1456,
1075
+ 'tpu-v5p-slice',
1076
+ 'ct5p-hightpu-4t',
1077
+ 4,
1078
+ AcceleratorType['TPU'],
1079
+ 'v5p-11648',
1080
+ ),
1081
+ 'v5p-11776': SystemCharacteristics(
1082
+ '8x8x92',
1083
+ 1472,
1084
+ 'tpu-v5p-slice',
1085
+ 'ct5p-hightpu-4t',
1086
+ 4,
1087
+ AcceleratorType['TPU'],
1088
+ 'v5p-11776',
1089
+ ),
1090
+ 'v5p-11904': SystemCharacteristics(
1091
+ '4x12x124',
1092
+ 1488,
1093
+ 'tpu-v5p-slice',
1094
+ 'ct5p-hightpu-4t',
1095
+ 4,
1096
+ AcceleratorType['TPU'],
1097
+ 'v5p-11904',
1098
+ ),
1099
+ 'v5p-12032': SystemCharacteristics(
1100
+ '4x8x188',
1101
+ 1504,
1102
+ 'tpu-v5p-slice',
1103
+ 'ct5p-hightpu-4t',
1104
+ 4,
1105
+ AcceleratorType['TPU'],
1106
+ 'v5p-12032',
1107
+ ),
1108
+ 'v5p-12160': SystemCharacteristics(
1109
+ '4x20x76',
1110
+ 1520,
1111
+ 'tpu-v5p-slice',
1112
+ 'ct5p-hightpu-4t',
1113
+ 4,
1114
+ AcceleratorType['TPU'],
1115
+ 'v5p-12160',
1116
+ ),
1117
+ 'v5p-12288': SystemCharacteristics(
1118
+ '16x16x24',
1119
+ 1536,
1120
+ 'tpu-v5p-slice',
1121
+ 'ct5p-hightpu-4t',
1122
+ 4,
1123
+ AcceleratorType['TPU'],
1124
+ 'v5p-12288',
1125
+ ),
1126
+ 'v5p-13824': SystemCharacteristics(
1127
+ '12x24x24',
1128
+ 1728,
1129
+ 'tpu-v5p-slice',
1130
+ 'ct5p-hightpu-4t',
1131
+ 4,
1132
+ AcceleratorType['TPU'],
1133
+ 'v5p-13824',
1134
+ ),
1135
+ 'v5p-17920': SystemCharacteristics(
1136
+ '16x20x28',
1137
+ 2240,
1138
+ 'tpu-v5p-slice',
1139
+ 'ct5p-hightpu-4t',
1140
+ 4,
1141
+ AcceleratorType['TPU'],
1142
+ 'v5p-17920',
1143
+ ),
1144
+ # v5litepod
1145
+ 'v5litepod-8': SystemCharacteristics(
1146
+ '2x4',
1147
+ 2,
1148
+ 'tpu-v5-lite-podslice',
1149
+ 'ct5lp-hightpu-4t',
1150
+ 8,
1151
+ AcceleratorType['TPU'],
1152
+ 'v5litepod-8',
1153
+ ),
1154
+ 'v5litepod-16': SystemCharacteristics(
1155
+ '4x4',
1156
+ 4,
1157
+ 'tpu-v5-lite-podslice',
1158
+ 'ct5lp-hightpu-4t',
1159
+ 4,
1160
+ AcceleratorType['TPU'],
1161
+ 'v5litepod-16',
1162
+ ),
1163
+ 'v5litepod-32': SystemCharacteristics(
1164
+ '4x8',
1165
+ 8,
1166
+ 'tpu-v5-lite-podslice',
1167
+ 'ct5lp-hightpu-4t',
1168
+ 4,
1169
+ AcceleratorType['TPU'],
1170
+ 'v5litepod-32',
1171
+ ),
1172
+ 'v5litepod-64': SystemCharacteristics(
1173
+ '8x8',
1174
+ 16,
1175
+ 'tpu-v5-lite-podslice',
1176
+ 'ct5lp-hightpu-4t',
1177
+ 4,
1178
+ AcceleratorType['TPU'],
1179
+ 'v5litepod-64',
1180
+ ),
1181
+ 'v5litepod-128': SystemCharacteristics(
1182
+ '8x16',
1183
+ 32,
1184
+ 'tpu-v5-lite-podslice',
1185
+ 'ct5lp-hightpu-4t',
1186
+ 4,
1187
+ AcceleratorType['TPU'],
1188
+ 'v5litepod-128',
1189
+ ),
1190
+ 'v5litepod-256': SystemCharacteristics(
1191
+ '16x16',
1192
+ 64,
1193
+ 'tpu-v5-lite-podslice',
1194
+ 'ct5lp-hightpu-4t',
1195
+ 4,
1196
+ AcceleratorType['TPU'],
1197
+ 'v5litepod-256',
1198
+ ),
1199
+ # v4
1200
+ 'v4-8': SystemCharacteristics(
1201
+ '2x2x1',
1202
+ 1,
1203
+ 'tpu-v4-podslice',
1204
+ 'ct4p-hightpu-4t',
1205
+ 4,
1206
+ AcceleratorType['TPU'],
1207
+ 'v4-8',
1208
+ ),
1209
+ 'v4-16': SystemCharacteristics(
1210
+ '2x2x2',
1211
+ 2,
1212
+ 'tpu-v4-podslice',
1213
+ 'ct4p-hightpu-4t',
1214
+ 4,
1215
+ AcceleratorType['TPU'],
1216
+ 'v4-16',
1217
+ ),
1218
+ 'v4-32': SystemCharacteristics(
1219
+ '2x2x4',
1220
+ 4,
1221
+ 'tpu-v4-podslice',
1222
+ 'ct4p-hightpu-4t',
1223
+ 4,
1224
+ AcceleratorType['TPU'],
1225
+ 'v4-32',
1226
+ ),
1227
+ 'v4-64': SystemCharacteristics(
1228
+ '2x4x4',
1229
+ 8,
1230
+ 'tpu-v4-podslice',
1231
+ 'ct4p-hightpu-4t',
1232
+ 4,
1233
+ AcceleratorType['TPU'],
1234
+ 'v4-64',
1235
+ ),
1236
+ 'v4-128': SystemCharacteristics(
1237
+ '4x4x4',
1238
+ 16,
1239
+ 'tpu-v4-podslice',
1240
+ 'ct4p-hightpu-4t',
1241
+ 4,
1242
+ AcceleratorType['TPU'],
1243
+ 'v4-128',
1244
+ ),
1245
+ 'v4-256': SystemCharacteristics(
1246
+ '4x4x8',
1247
+ 32,
1248
+ 'tpu-v4-podslice',
1249
+ 'ct4p-hightpu-4t',
1250
+ 4,
1251
+ AcceleratorType['TPU'],
1252
+ 'v4-256',
1253
+ ),
1254
+ 'v4-512': SystemCharacteristics(
1255
+ '4x8x8',
1256
+ 64,
1257
+ 'tpu-v4-podslice',
1258
+ 'ct4p-hightpu-4t',
1259
+ 4,
1260
+ AcceleratorType['TPU'],
1261
+ 'v4-512',
1262
+ ),
1263
+ 'v4-1024': SystemCharacteristics(
1264
+ '8x8x8',
1265
+ 128,
1266
+ 'tpu-v4-podslice',
1267
+ 'ct4p-hightpu-4t',
1268
+ 4,
1269
+ AcceleratorType['TPU'],
1270
+ 'v4-1024',
1271
+ ),
1272
+ 'v4-1536': SystemCharacteristics(
1273
+ '8x8x12',
1274
+ 192,
1275
+ 'tpu-v4-podslice',
1276
+ 'ct4p-hightpu-4t',
1277
+ 4,
1278
+ AcceleratorType['TPU'],
1279
+ 'v4-1536',
1280
+ ),
1281
+ 'v4-2048': SystemCharacteristics(
1282
+ '8x8x16',
1283
+ 256,
1284
+ 'tpu-v4-podslice',
1285
+ 'ct4p-hightpu-4t',
1286
+ 4,
1287
+ AcceleratorType['TPU'],
1288
+ 'v4-2048',
1289
+ ),
1290
+ 'v4-4096': SystemCharacteristics(
1291
+ '8x16x16',
1292
+ 512,
1293
+ 'tpu-v4-podslice',
1294
+ 'ct4p-hightpu-4t',
1295
+ 4,
1296
+ AcceleratorType['TPU'],
1297
+ 'v4-4096',
1298
+ ),
1299
+ # CPU system characteristics.
1300
+ # Note that chips_per_vm is actually the number of vCPUs in that CPU.
1301
+ # There are no chips in CPUs.
1302
+ # m1-megamem-#vCPUs-#VMs
1303
+ 'm1-megamem-96-1': SystemCharacteristics(
1304
+ 'N/A',
1305
+ 1,
1306
+ 'N/A',
1307
+ 'm1-megamem-96',
1308
+ 96,
1309
+ AcceleratorType['CPU'],
1310
+ 'm1-megamem-96-1',
1311
+ ),
1312
+ # n2-standard-#vCPUs-#VMs
1313
+ 'n2-standard-64-1': SystemCharacteristics(
1314
+ 'N/A',
1315
+ 1,
1316
+ 'N/A',
1317
+ 'n2-standard-64',
1318
+ 64,
1319
+ AcceleratorType['CPU'],
1320
+ 'n2-standard-64-1',
1321
+ ),
1322
+ 'n2-standard-32-1': SystemCharacteristics(
1323
+ 'N/A',
1324
+ 1,
1325
+ 'N/A',
1326
+ 'n2-standard-32',
1327
+ 32,
1328
+ AcceleratorType['CPU'],
1329
+ 'n2-standard-32-1',
1330
+ ),
1331
+ 'n2-standard-32-2': SystemCharacteristics(
1332
+ 'N/A',
1333
+ 2,
1334
+ 'N/A',
1335
+ 'n2-standard-32',
1336
+ 32,
1337
+ AcceleratorType['CPU'],
1338
+ 'n2-standard-32-2',
1339
+ ),
1340
+ 'n2-standard-32-4': SystemCharacteristics(
1341
+ 'N/A',
1342
+ 4,
1343
+ 'N/A',
1344
+ 'n2-standard-32',
1345
+ 32,
1346
+ AcceleratorType['CPU'],
1347
+ 'n2-standard-32-4',
1348
+ ),
1349
+ 'n2-standard-32-8': SystemCharacteristics(
1350
+ 'N/A',
1351
+ 8,
1352
+ 'N/A',
1353
+ 'n2-standard-32',
1354
+ 32,
1355
+ AcceleratorType['CPU'],
1356
+ 'n2-standard-32-8',
1357
+ ),
1358
+ 'n2-standard-32-16': SystemCharacteristics(
1359
+ 'N/A',
1360
+ 16,
1361
+ 'N/A',
1362
+ 'n2-standard-32',
1363
+ 32,
1364
+ AcceleratorType['CPU'],
1365
+ 'n2-standard-32-16',
1366
+ ),
1367
+ 'n2-standard-32-32': SystemCharacteristics(
1368
+ 'N/A',
1369
+ 32,
1370
+ 'N/A',
1371
+ 'n2-standard-32',
1372
+ 32,
1373
+ AcceleratorType['CPU'],
1374
+ 'n2-standard-32-32',
1375
+ ),
1376
+ 'n2-standard-32-64': SystemCharacteristics(
1377
+ 'N/A',
1378
+ 64,
1379
+ 'N/A',
1380
+ 'n2-standard-32',
1381
+ 32,
1382
+ AcceleratorType['CPU'],
1383
+ 'n2-standard-32-64',
1384
+ ),
1385
+ 'n2-standard-32-128': SystemCharacteristics(
1386
+ 'N/A',
1387
+ 128,
1388
+ 'N/A',
1389
+ 'n2-standard-32',
1390
+ 32,
1391
+ AcceleratorType['CPU'],
1392
+ 'n2-standard-32-128',
1393
+ ),
1394
+ 'n2-standard-32-256': SystemCharacteristics(
1395
+ 'N/A',
1396
+ 256,
1397
+ 'N/A',
1398
+ 'n2-standard-32',
1399
+ 32,
1400
+ AcceleratorType['CPU'],
1401
+ 'n2-standard-32-256',
1402
+ ),
1403
+ 'n2-standard-32-512': SystemCharacteristics(
1404
+ 'N/A',
1405
+ 512,
1406
+ 'N/A',
1407
+ 'n2-standard-32',
1408
+ 32,
1409
+ AcceleratorType['CPU'],
1410
+ 'n2-standard-32-512',
1411
+ ),
1412
+ 'n2-standard-32-1024': SystemCharacteristics(
1413
+ 'N/A',
1414
+ 1024,
1415
+ 'N/A',
1416
+ 'n2-standard-32',
1417
+ 32,
1418
+ AcceleratorType['CPU'],
1419
+ 'n2-standard-32-1024',
1420
+ ),
1421
+ 'n2-standard-32-2048': SystemCharacteristics(
1422
+ 'N/A',
1423
+ 2048,
1424
+ 'N/A',
1425
+ 'n2-standard-32',
1426
+ 32,
1427
+ AcceleratorType['CPU'],
1428
+ 'n2-standard-32-2048',
1429
+ ),
1430
+ }
1431
+ """ If you modify UserFacingNameToSystemCharacteristics you should also modify
1432
+ the corresponding Map in MaxText/accelerator_to_spec_map.py """