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