[Docker] Add --gpus all for docker when GPU is available (#3833)

* Add --gpus all for when GPU is available * Get rid of --gpus all
skypilot-org · Aug 21, 2024 · fa60f7a · fa60f7a
1 parent 87e55ab
commit fa60f7a
Show file tree

Hide file tree

Showing 5 changed files with 1 addition and 13 deletions.
diff --git a/sky/provision/docker_utils.py b/sky/provision/docker_utils.py
@@ -381,7 +381,7 @@ def _configure_runtime(self, run_options: List[str]) -> List[str]:
         if 'nvidia-container-runtime' in runtime_output:
             try:
                 self._run('nvidia-smi', log_err_when_fail=False)
-                return run_options + ['--runtime=nvidia']
+                return run_options + ['--runtime=nvidia', '--gpus all']
             except Exception as e:  # pylint: disable=broad-except
                 logger.debug(
                     'Nvidia Container Runtime is present in the docker image'

diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
@@ -11,9 +11,6 @@ docker:
   container_name: {{docker_container_name}}
   run_options:
     - --ulimit nofile=1048576:1048576
-    {%- if custom_resources is not none %}
-      --gpus all
-    {%- endif %}
     {%- for run_option in docker_run_options %}
     - {{run_option}}
     {%- endfor %}

diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
@@ -11,9 +11,6 @@ docker:
   container_name: {{docker_container_name}}
   run_options:
     - --ulimit nofile=1048576:1048576
-    {%- if custom_resources is not none %}
-      --gpus all
-    {%- endif %}
     {%- for run_option in docker_run_options %}
     - {{run_option}}
     {%- endfor %}

diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
@@ -12,9 +12,6 @@ docker:
   container_name: {{docker_container_name}}
   run_options:
     - --ulimit nofile=1048576:1048576
-    {%- if gpu is not none %}
-      --gpus all
-    {%- endif %}
     {%- for run_option in docker_run_options %}
     - {{run_option}}
     {%- endfor %}

diff --git a/sky/templates/paperspace-ray.yml.j2 b/sky/templates/paperspace-ray.yml.j2
@@ -11,9 +11,6 @@ docker:
   container_name: {{docker_container_name}}
   run_options:
     - --ulimit nofile=1048576:1048576
-    {%- if custom_resources is not none %}
-      --gpus all
-    {%- endif %}
     {%- for run_option in docker_run_options %}
     - {{run_option}}
     {%- endfor %}