[Nebius] Add Docker support for Nebius cloud (#5334)

SalikovAlex · SeungjinYang · cblmemo · web-flow · commit 2b5774c849f3 · 2025-04-29T00:19:56.000Z
* Add Docker support for Nebius cloud

This commit introduces Docker compatibility for Nebius, enabling users to specify Docker images and runtime options. It includes updates to templates, resource management, and Docker daemon initialization to ensure smooth GPU integration and proper runtime setup. Relevant tests have been adjusted to reflect the new functionality.

* Enable Nebius support for H100 GPUs and CUDA &gt;= 12.0

Updated templates and tests to add Nebius compatibility for H100 GPUs and CUDA versions 12.0 and above. Adjusted smoke test parameters and logic to incorporate these changes while maintaining backward compatibility.

* fix typo

* fix typo

* Update formatting and clarify Docker installation logic

Adjusted spacing in templates for consistency. Added inline documentation to explain the Docker installation process for specific OS images, optimizing deployment speed by selectively installing Docker only when required.

* Fix GPU runtime option and refine Nebius template logic

Corrected the comment in `nebius.py` regarding GPU runtime options and adjusted YAML logic to handle docker installation conditionally. Improved formatting and removed redundant lines to enhance clarity and maintainability.

* Apply suggestions from code review

Co-authored-by: Tian Xia &lt;cblmemo@gmail.com&gt;

---------

Co-authored-by: Seung Jin &lt;seungjin219@gmail.com&gt;
Co-authored-by: Tian Xia &lt;cblmemo@gmail.com&gt;
diff --git a/sky/clouds/nebius.py b/sky/clouds/nebius.py
@@ -59,10 +59,6 @@ class Nebius(clouds.Cloud):
             ('Spot is not supported, as Nebius API does not implement spot.'),
         clouds.CloudImplementationFeatures.CLONE_DISK_FROM_CLUSTER:
             (f'Migrating disk is currently not supported on {_REPR}.'),
-        clouds.CloudImplementationFeatures.DOCKER_IMAGE:
-            (f'Docker image is currently not supported on {_REPR}. '
-             'You can try running docker command inside the '
-             '`run` section in task.yaml.'),
         clouds.CloudImplementationFeatures.CUSTOM_DISK_TIER:
             (f'Custom disk tier is currently not supported on {_REPR}.'),
         clouds.CloudImplementationFeatures.HIGH_AVAILABILITY_CONTROLLERS:
@@ -213,7 +209,8 @@ def make_deploy_resources_variables(
         else:
             raise RuntimeError('Unsupported instance type for Nebius cloud:'
                                f' {resources.instance_type}')
-        return {
+
+        resources_vars = {
             'instance_type': resources.instance_type,
             'custom_resources': custom_resources,
             'region': region.name,
@@ -222,6 +219,14 @@ def make_deploy_resources_variables(
             'zones': None,
         }
 
+        if acc_dict is not None:
+            # Nebius cloud's docker runtime information does not contain
+            # 'nvidia-container-runtime', causing no GPU option to be added to
+            # the docker run command. We patch this by adding it here.
+            resources_vars['docker_run_options'] = ['--gpus all']
+
+        return resources_vars
+
     def _get_feasible_launchable_resources(
         self, resources: 'resources_lib.Resources'
     ) -> 'resources_utils.FeasibleResources':
diff --git a/sky/provision/docker_utils.py b/sky/provision/docker_utils.py
@@ -28,6 +28,9 @@
 # the command.
 DOCKER_PERMISSION_DENIED_STR = ('permission denied while trying to connect to '
                                 'the Docker daemon socket')
+
+DOCKER_SOCKET_NOT_READY_STR = ('Is the docker daemon running?')
+
 _DOCKER_SOCKET_WAIT_TIMEOUT_SECONDS = 30
 
 
@@ -173,22 +176,25 @@ def _run(self,
                 stream_logs=False,
                 separate_stderr=separate_stderr,
                 log_path=self.log_path)
-            if (DOCKER_PERMISSION_DENIED_STR in stdout + stderr and
-                    wait_for_docker_daemon):
-                if time.time() - start > _DOCKER_SOCKET_WAIT_TIMEOUT_SECONDS:
-                    if rc == 0:
-                        # Set returncode to 1 if failed to connect to docker
-                        # daemon after timeout.
-                        rc = 1
-                    break
-                # Close the cached connection to make the permission update of
-                # ssh user take effect, e.g. usermod -aG docker $USER, called
-                # by cloud-init of Azure.
-                self.runner.close_cached_connection()
-                logger.info('Failed to connect to docker daemon. It might be '
-                            'initializing, retrying in 5 seconds...')
-                time.sleep(5)
-                continue
+            if (DOCKER_PERMISSION_DENIED_STR in stdout + stderr or
+                    DOCKER_SOCKET_NOT_READY_STR in stdout + stderr):
+                if wait_for_docker_daemon:
+                    if time.time(
+                    ) - start > _DOCKER_SOCKET_WAIT_TIMEOUT_SECONDS:
+                        if rc == 0:
+                            # Set returncode to 1 if failed to connect to docker
+                            # daemon after timeout.
+                            rc = 1
+                        break
+                    # Close the cached connection to make the permission update
+                    # of ssh user take effect, e.g. usermod -aG docker $USER,
+                    # called by cloud-init of Azure.
+                    self.runner.close_cached_connection()
+                    logger.info(
+                        'Failed to connect to docker daemon. It might be '
+                        'initializing, retrying in 5 seconds...')
+                    time.sleep(5)
+                    continue
             break
         subprocess_utils.handle_returncode(
             rc,
diff --git a/sky/templates/nebius-ray.yml.j2 b/sky/templates/nebius-ray.yml.j2
@@ -10,6 +10,27 @@ provider:
   module: sky.provision.nebius
   region: "{{region}}"
 
+{%- if docker_image is not none %}
+docker:
+  image: {{docker_image}}
+  container_name: {{docker_container_name}}
+  run_options:
+    - --ulimit nofile=1048576:1048576
+    {%- for run_option in docker_run_options %}
+    - {{run_option}}
+    {%- endfor %}
+  {%- if docker_login_config is not none %}
+  docker_login_config:
+    username: |-
+      {{docker_login_config.username}}
+    password: |-
+      {{docker_login_config.password}}
+    server: |-
+      {{docker_login_config.server}}
+  {%- endif %}
+{%- endif %}
+
+
 auth:
   ssh_user: ubuntu
   ssh_private_key: {{ssh_private_key}}
@@ -22,6 +43,48 @@ available_node_types:
       ImageId: {{image_id}}
       DiskSize: {{disk_size}}
       UserData: |
+        {%- if docker_image is not none %}
+        runcmd:
+          - sudo sed -i 's/^#\?AllowTcpForwarding.*/AllowTcpForwarding yes/' /etc/ssh/sshd_config
+          - systemctl restart sshd
+        {%- endif %}
+
+        {# Two available OS images:
+               1. ubuntu22.04-driverless - requires Docker installation
+               2. ubuntu22.04-cuda12 - comes with Docker pre-installed
+            To optimize deployment speed, Docker is only installed when using ubuntu22.04-driverless #}
+        {%- if docker_image is not none and image_id == 'ubuntu22.04-driverless' %}
+        apt:
+          sources:
+            docker.list:
+              source: deb [arch=amd64] https://download.docker.com/linux/ubuntu $RELEASE stable
+              keyid: 9DC858229FC7DD38854AE2D88D81803C0EBFCD88
+
+        packages:
+          - apt-transport-https
+          - ca-certificates
+          - curl
+          - gnupg-agent
+          - software-properties-common
+          - docker-ce
+          - docker-ce-cli
+          - containerd.io
+
+        # Enable ipv4 forwarding, required on CIS hardened machines
+        write_files:
+          - path: /etc/sysctl.d/enabled_ipv4_forwarding.conf
+            content: |
+              net.ipv4.conf.all.forwarding=1
+
+        # create the docker group
+        groups:
+          - docker
+
+        # Add default auto created user to docker group
+        system_info:
+          default_user:
+            groups: [docker]
+        {%- endif %}
         users:
           - name: skypilot:ssh_user
             shell: /bin/bash
diff --git a/tests/smoke_tests/test_cluster_job.py b/tests/smoke_tests/test_cluster_job.py
@@ -46,9 +46,8 @@
 @pytest.mark.no_scp  # SCP does not have T4 gpus. Run test_scp_job_queue instead
 @pytest.mark.no_paperspace  # Paperspace does not have T4 gpus.
 @pytest.mark.no_oci  # OCI does not have T4 gpus
-@pytest.mark.no_nebius  # Nebius does not support T4 GPUs
 @pytest.mark.resource_heavy
-@pytest.mark.parametrize('accelerator', [{'do': 'H100'}])
+@pytest.mark.parametrize('accelerator', [{'do': 'H100', 'nebius': 'H100'}])
 def test_job_queue(generic_cloud: str, accelerator: Dict[str, str]):
     accelerator = accelerator.get(generic_cloud, 'T4')
     name = smoke_tests_utils.get_cluster_name()
@@ -85,8 +84,7 @@ def test_job_queue(generic_cloud: str, accelerator: Dict[str, str]):
 @pytest.mark.no_scp  # Doesn't support SCP for now
 @pytest.mark.no_oci  # Doesn't support OCI for now
 @pytest.mark.no_kubernetes  # Doesn't support Kubernetes for now
-@pytest.mark.no_nebius  # Nebius does not support Docker
-@pytest.mark.parametrize('accelerator', [{'do': 'H100'}])
+@pytest.mark.parametrize('accelerator', [{'do': 'H100', 'nebius': 'H100'}])
 @pytest.mark.parametrize(
     'image_id',
     [
@@ -109,6 +107,11 @@ def test_job_queue_with_docker(generic_cloud: str, image_id: str,
     name = smoke_tests_utils.get_cluster_name() + image_id[len('docker:'):][:4]
     total_timeout_minutes = 40 if generic_cloud == 'azure' else 15
     time_to_sleep = 300 if generic_cloud == 'azure' else 200
+    # Nebius support Cuda >= 12.0
+    if (image_id == 'docker:nvidia/cuda:11.8.0-devel-ubuntu18.04' and
+            generic_cloud == 'nebius'):
+        image_id = 'docker:nvidia/cuda:12.1.0-devel-ubuntu18.04'
+
     test = smoke_tests_utils.Test(
         'job_queue_with_docker',
         [
@@ -139,7 +142,7 @@ def test_job_queue_with_docker(generic_cloud: str, image_id: str,
             f'sky logs {name} 5 --status',
             f'sky logs {name} 6 --status',
             # Make sure it is still visible after an stop & start cycle.
-            f'sky exec {name} --image-id {image_id} nvidia-smi | grep "Tesla T4"',
+            f'sky exec {name} --image-id {image_id} nvidia-smi | grep -i "{accelerator}"',
             f'sky logs {name} 7 --status'
         ],
         f'sky down -y {name}',
@@ -228,8 +231,7 @@ def test_scp_job_queue():
 @pytest.mark.no_oci  # OCI Cloud does not have T4 gpus.
 @pytest.mark.no_vast  # Vast does not support num_nodes > 1 yet
 @pytest.mark.no_kubernetes  # Kubernetes not support num_nodes > 1 yet
-@pytest.mark.no_nebius  # Nebius does not have T4 gpus.
-@pytest.mark.parametrize('accelerator', [{'do': 'H100'}])
+@pytest.mark.parametrize('accelerator', [{'do': 'H100', 'nebius': 'H100'}])
 def test_job_queue_multinode(generic_cloud: str, accelerator: Dict[str, str]):
     accelerator = accelerator.get(generic_cloud, 'T4')
     name = smoke_tests_utils.get_cluster_name()
@@ -383,7 +385,6 @@ def test_ibm_job_queue_multinode():
 @pytest.mark.no_scp  # Doesn't support SCP for now
 @pytest.mark.no_oci  # Doesn't support OCI for now
 @pytest.mark.no_kubernetes  # Doesn't support Kubernetes for now
-@pytest.mark.no_nebius  # Nebius does not support Docker
 # TODO(zhwu): we should fix this for kubernetes
 def test_docker_preinstalled_package(generic_cloud: str):
     name = smoke_tests_utils.get_cluster_name()
@@ -462,9 +463,8 @@ def test_multi_echo(generic_cloud: str):
 @pytest.mark.no_lambda_cloud  # Lambda Cloud does not have V100 gpus
 @pytest.mark.no_ibm  # IBM cloud currently doesn't provide public image with CUDA
 @pytest.mark.no_scp  # SCP does not have V100 (16GB) GPUs. Run test_scp_huggingface instead.
-@pytest.mark.no_nebius  # Nebius does not have T4 gpus for now
 @pytest.mark.resource_heavy
-@pytest.mark.parametrize('accelerator', [{'do': 'H100'}])
+@pytest.mark.parametrize('accelerator', [{'do': 'H100', 'nebius': 'H100'}])
 def test_huggingface(generic_cloud: str, accelerator: Dict[str, str]):
     accelerator = accelerator.get(generic_cloud, 'T4')
     name = smoke_tests_utils.get_cluster_name()
@@ -1399,9 +1399,8 @@ def test_cancel_azure():
 @pytest.mark.no_paperspace  # Paperspace has `gnome-shell` on nvidia-smi
 @pytest.mark.no_scp  # SCP does not support num_nodes > 1 yet
 @pytest.mark.no_vast  # Vast does not support num_nodes > 1 yet
-@pytest.mark.no_nebius  # Nebius Cloud does not work with Cuda 11
 @pytest.mark.resource_heavy
-@pytest.mark.parametrize('accelerator', [{'do': 'H100'}])
+@pytest.mark.parametrize('accelerator', [{'do': 'H100', 'nebius': 'H100'}])
 def test_cancel_pytorch(generic_cloud: str, accelerator: Dict[str, str]):
     accelerator = accelerator.get(generic_cloud, 'T4')
     name = smoke_tests_utils.get_cluster_name()
diff --git a/tests/smoke_tests/test_mount_and_storage.py b/tests/smoke_tests/test_mount_and_storage.py
@@ -387,7 +387,6 @@ def test_kubernetes_context_switch():
 # credentials for getting azure storage commands, even though the API server
 # is running remotely. We should fix this.
 @pytest.mark.no_vast  # Requires AWS
-@pytest.mark.no_nebius  # Docker image is currently not supported on Nebius.
 @pytest.mark.resource_heavy
 @pytest.mark.parametrize(
     'image_id',
@@ -557,7 +556,7 @@ def test_ibm_storage_mounts():
                          [constants.SKY_IGNORE_FILE, constants.GIT_IGNORE_FILE])
 def test_ignore_exclusions(generic_cloud: str, ignore_file: str):
     """Tests that .skyignore patterns correctly exclude files when using sky launch and sky jobs launch.
-    
+
     Creates a temporary directory with various files and folders, adds a .skyignore file
     that excludes specific files and folders, then verifies the exclusions work properly
     when using sky launch and sky jobs launch commands.
