Repository: skypilot-org/skypilot
Pull request: [Setup] unset CUDA_VISIBLE_DEVICES for detach setup
Closes #1403.

Tested:
- [x] `tests/run_smoke_tests.sh test_env_check`

Problem statement:
[Setup] `--detach-setup` will have CUDA_VISIBLE_DEVICES set causing inconsistency
From the user:
> Is there a recent change that sets CUDA_VISIBLE_DEVICES to empty string during cluster setup? I have a setup logic that configures the cuda compute capability compile flags based on this, it used to work, but suddenly stopped working after updating to the latest Sky on master today.

Hints from issue discussion:

Commit message: Add setting CUDA_VISIBLE_DEVICES test

Patch to review:
diff --git a/examples/env_check.yaml b/examples/env_check.yaml
index 30adaefb3b0..7d2ed57ab8f 100644
--- a/examples/env_check.yaml
+++ b/examples/env_check.yaml
@@ -8,6 +8,9 @@ workdir: .
 setup: |
   echo "here"
   echo "export TEST_VAR=test" >> ~/.bashrc
+  [[ -v CUDA_VISIBLE_DEVICES ]] && exit 1 || exit 0
+  export CUDA_VISIBLE_DEVICES=1
+  [[ -v CUDA_VISIBLE_DEVICES ]] && exit 0 || exit 1
 
 run: |
   if [[ -z "${TEST_VAR}" ]]; then
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 8082dbd7cad..7e55b5275ac 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -208,6 +208,9 @@ def add_prologue(self,
             self._code += [
                 textwrap.dedent(f"""\
                 _SETUP_CPUS = 0.0001
+                # Unset CUDA_VISIBLE_DEVICES (set by ray), since the setup command
+                # may need to detect the number of GPUs on the machine.
+                setup_cmd = 'unset CUDA_VISIBLE_DEVICES; ' + setup_cmd
                 job_lib.set_status({job_id!r}, job_lib.JobStatus.SETTING_UP)
                 print({_CTRL_C_TIP_MESSAGE!r}, file=sys.stderr, flush=True)
                 total_num_nodes = len(ray.nodes())
@@ -2041,6 +2044,7 @@ def _setup(self, handle: ResourceHandle, task: task_lib.Task,
             f.write(setup_script)
             f.flush()
             setup_sh_path = f.name
+            print(setup_sh_path)
             setup_file = os.path.basename(setup_sh_path)
             # Sync the setup script up and run it.
             ip_list = backend_utils.get_node_ips(
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 9f4a736bfb2..2c7ec4576f5 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -229,7 +229,7 @@ def test_env_check():
     test = Test(
         'env_check',
         [
-            f'sky launch -y -c {name} examples/env_check.yaml',
+            f'sky launch -y -c {name} --detach-setup examples/env_check.yaml',
             f'sky logs {name} 1 --status',  # Ensure the job succeeded.
         ],
         f'sky down -y {name}',


Code review comments:
diff_hunk: @@ -2041,6 +2044,7 @@ def _setup(self, handle: ResourceHandle, task: task_lib.Task,
             f.write(setup_script)
             f.flush()
             setup_sh_path = f.name
+            print(setup_sh_path)
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
remnant?

@author:
Good catch! Removed.

diff_hunk: @@ -208,6 +208,9 @@ def add_prologue(self,
             self._code += [
                 textwrap.dedent(f"""\
                 _SETUP_CPUS = 0.0001
+                # Unset CUDA_VISIBLE_DEVICES (set by ray), since the setup command
+                # may need to detect the number of GPUs on the machine.
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
```suggestion
                # The setup command will be run as a ray task with num_cpus=_SETUP_CPUS as the requirement; this means Ray will set CUDA_VISIBLE_DEVICES to an empty string. We unset it so that user setup command may properly use this env var.
```