From a85f194db657084a0d913b3cf53f91bc93d43668 Mon Sep 17 00:00:00 2001
From: Zhanghao Wu <zhanghao.wu@outlook.com>
Date: Wed, 3 Jul 2024 15:44:52 +0000
Subject: [PATCH] Allow opening ports for cudo

---
 sky/provision/cudo/__init__.py |  3 ++-
 sky/provision/cudo/instance.py | 10 ++++++++++
 sky/utils/controller_utils.py  | 11 +++++------
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/sky/provision/cudo/__init__.py b/sky/provision/cudo/__init__.py
index bbdc96413a8..c4587bfdfa7 100644
--- a/sky/provision/cudo/__init__.py
+++ b/sky/provision/cudo/__init__.py
@@ -3,6 +3,7 @@
 from sky.provision.cudo.config import bootstrap_instances
 from sky.provision.cudo.instance import cleanup_ports
 from sky.provision.cudo.instance import get_cluster_info
+from sky.provision.cudo.instance import open_ports
 from sky.provision.cudo.instance import query_instances
 from sky.provision.cudo.instance import run_instances
 from sky.provision.cudo.instance import stop_instances
@@ -11,4 +12,4 @@
 
 __all__ = ('bootstrap_instances', 'run_instances', 'stop_instances',
            'terminate_instances', 'wait_instances', 'get_cluster_info',
-           'cleanup_ports', 'query_instances')
+           'cleanup_ports', 'query_instances', 'open_ports')
diff --git a/sky/provision/cudo/instance.py b/sky/provision/cudo/instance.py
index 71ada577e53..fcc7b8167cb 100644
--- a/sky/provision/cudo/instance.py
+++ b/sky/provision/cudo/instance.py
@@ -220,6 +220,16 @@ def query_instances(
     return statuses
 
 
+def open_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    del cluster_name_on_cloud, ports, provider_config
+    # Cudo has all ports open by default. Nothing to do here.
+    return
+
+
 def cleanup_ports(
     cluster_name_on_cloud: str,
     ports: List[str],
diff --git a/sky/utils/controller_utils.py b/sky/utils/controller_utils.py
index ba65d4b664a..d605d7babad 100644
--- a/sky/utils/controller_utils.py
+++ b/sky/utils/controller_utils.py
@@ -247,6 +247,11 @@ def _get_cloud_dependencies_installation_commands(
                 '/bin/linux/amd64/kubectl" && '
                 'sudo install -o root -g root -m 0755 '
                 'kubectl /usr/local/bin/kubectl))')
+        elif isinstance(cloud, clouds.Cudo):
+            commands.append(
+                f'echo -en "\\r{prefix_str}Cudo{empty_str}" && '
+                'pip list | grep cudo-compute > /dev/null 2>&1 || '
+                'pip install "cudo-compute>=0.1.10" > /dev/null 2>&1')
         if controller == Controllers.JOBS_CONTROLLER:
             if isinstance(cloud, clouds.IBM):
                 commands.append(
@@ -263,12 +268,6 @@ def _get_cloud_dependencies_installation_commands(
                     f'echo -en "\\r{prefix_str}RunPod{empty_str}" && '
                     'pip list | grep runpod > /dev/null 2>&1 || '
                     'pip install "runpod>=1.5.1" > /dev/null 2>&1')
-            elif isinstance(cloud, clouds.Cudo):
-                # cudo doesn't support open port
-                commands.append(
-                    f'echo -en "\\r{prefix_str}Cudo{empty_str}" && '
-                    'pip list | grep cudo-compute > /dev/null 2>&1 || '
-                    'pip install "cudo-compute>=0.1.8" > /dev/null 2>&1')
     if (cloudflare.NAME
             in storage_lib.get_cached_enabled_storage_clouds_or_refresh()):
         commands.append(f'echo -en "\\r{prefix_str}Cloudflare{empty_str}" && ' +