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