Repository: skypilot-org/skypilot
Pull request: [Core] Fix Inferentia job scheduling
<!-- Describe the changes in this PR -->
Fixes #2968. Bug fixed is when we are submitting a job request, do not request `GPU: <count>` resources. 


<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below): `pytest tests/test_smoke.py::test_inferentia --aws`
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
Job submission on Inferentia instances blocks
```
sky launch -t inf2.8xlarge -i60 --down echo hi                   
```
blocks:
```
..
I 01-09 15:41:36 provisioner.py:464] Successfully provisioned cluster: inf
I 01-09 15:41:42 cloud_vm_ray_backend.py:3320] Job submitted with Job ID: 1
I 01-09 23:41:44 log_lib.py:431] Start streaming logs for job 1.
INFO: Tip: use Ctrl-C to exit log streaming (task will not be killed).
INFO: Waiting for task resources on 1 node. This will block if the cluster is full.
```

`ray status` on the cluster shows the reason
```
Usage:
 0.0/32.0 CPU
 0.0/1.0 Inferentia
 0B/76.49GiB memory
 0B/36.77GiB object_store_memory

Demands:
 {'Inferentia': 1.0, 'GPU': 1.0, 'CPU': 0.5} * 1 (STRICT_SPREAD): 1+ pending placement groups
```
Namely, we demand `'GPU': 1.0` but `GPU` is not available as a resource.

Hints from issue discussion:

Commit message: Add smoke

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 1e6085ec222..e40ad1248ac 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -325,8 +325,8 @@ def add_gang_scheduling_placement_group_and_setup(
             gpu_dict = {'GPU': acc_count}
             # gpu_dict should be empty when the accelerator is not GPU.
             # FIXME: This is a hack to make sure that we do not reserve
-            # GPU when requesting TPU.
-            if 'tpu' in acc_name.lower():
+            # GPU when requesting non-GPU accelerators (e.g., TPU/Inferentia).
+            if 'tpu' in acc_name.lower() or 'inferentia' in acc_name.lower():
                 gpu_dict = {}
             for bundle in bundles:
                 bundle.update({
@@ -488,7 +488,8 @@ def add_ray_task(self,
             options.append(f'resources={json.dumps(ray_resources_dict)}')
 
             resources_key = list(ray_resources_dict.keys())[0]
-            if 'tpu' not in resources_key.lower():
+            if 'tpu' not in resources_key.lower(
+            ) and 'inferentia' not in resources_key.lower():
                 # `num_gpus` should be empty when the accelerator is not GPU.
                 # FIXME: use a set of GPU types, instead of 'tpu' in the key.
 
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index fb64cf5ba68..bfced150f85 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -1401,6 +1401,23 @@ def test_scp_huggingface(generic_cloud: str):
     run_one_test(test)
 
 
+# ---------- Inferentia. ----------
+@pytest.mark.aws
+def test_inferentia():
+    name = _get_cluster_name()
+    test = Test(
+        'test_inferentia',
+        [
+            f'sky launch -y -c {name} -t inf2.xlarge -- echo hi',
+            f'sky exec {name} --gpus Inferentia:1 echo hi',
+            f'sky logs {name} 1 --status',  # Ensure the job succeeded.
+            f'sky logs {name} 2 --status',  # Ensure the job succeeded.
+        ],
+        f'sky down -y {name}',
+    )
+    run_one_test(test)
+
+
 # ---------- TPU. ----------
 @pytest.mark.gcp
 @pytest.mark.tpu


Code review comments:
diff_hunk: @@ -325,8 +325,8 @@ def add_gang_scheduling_placement_group_and_setup(
             gpu_dict = {'GPU': acc_count}
             # gpu_dict should be empty when the accelerator is not GPU.
             # FIXME: This is a hack to make sure that we do not reserve
-            # GPU when requesting TPU.
-            if 'tpu' in acc_name.lower():
+            # GPU when requesting non-GPU accelerators (e.g., TPU/Inferentia).
+            if 'tpu' in acc_name.lower() or 'inferentia' in acc_name.lower():
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
This seems not very generalizable for the cases where more "non-GPU" accelerators are supported, such as: `Trainium`. Is it possible that we can directly add a `--num-gpus` in the `ray start` command where we start the ray cluster at the following place to force the ray cluster to contain the missing `GPU`: https://github.com/skypilot-org/skypilot/blob/6fd5bf298f9d830bb3fbafee0e131ef688142f86/sky/provision/instance_setup.py#L211

@author:
That seems reasonable. One concern is calling any accelerator "GPU" may be confusing to users (and it's only convenient for us).

Ray in recent releases took a different approach by not calling all accelerators "GPU". Instead there will be several resource names:
- https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/neuron.py#L36 
- https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/nvidia_gpu.py#L23
- https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/tpu.py#L80
- etc.

Wdyt?

@user1:
I think this will only change the representation in `ray status`, which should ideally not be exposed to users, so not affecting the user side. That said, we can merge the current PR changes first to unblock our users first, and do a further refactoring in the future for avoiding the hardcode here (either by using `GPU` or further distinguish them).

For the hardcode, we probably want to add `trainium` as well. : )

@author:
Sg, PTAL.