Added support for AWS Trainium accelerator (#2690)

* added support for Trainium accelerator * formatting issue * formatting issue --------- Co-authored-by: Zhanghao Wu <[email protected]>
skypilot-org · Nov 17, 2023 · 623c10d · 623c10d
1 parent 6b1bbc9
commit 623c10d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/sky/utils/accelerator_registry.py b/sky/utils/accelerator_registry.py
@@ -25,6 +25,7 @@
     'A10G',
     'Gaudi HL-205',
     'Inferentia',
+    'Trainium',
     'K520',
     'K80',
     'M60',

diff --git a/tests/test_list_accelerators.py b/tests/test_list_accelerators.py
@@ -8,6 +8,7 @@ def test_list_accelerators():
     assert 'V100' in result, result
     assert 'tpu-v3-8' in result, result
     assert 'Inferentia' not in result, result
+    assert 'Trainium' not in result, result
     assert 'A100-80GB' in result, result
 
 
@@ -16,6 +17,7 @@ def test_list_ccelerators_all():
     assert 'V100' in result, result
     assert 'tpu-v3-8' in result, result
     assert 'Inferentia' in result, result
+    assert 'Trainium' in result, result
     assert 'A100-80GB' in result, result
-Original file line number
+Diff line change
@@ Expand Up / @@ -25,6 +25,7 @@ @@
         'A10G',
         'Gaudi HL-205',
         'Inferentia',
+        'Trainium',
         'K520',
         'K80',
         'M60',
@@ Expand Down @@