Fix accidentally removed updates to cloud provider

sky/skylet/providers/aws/node_provider.py

@@ -24,6 +24,9 @@
     resource_cache,
     client_cache,
 )
+from sky.skylet.providers.command_runner import SkyDockerCommandRunner
+
+from ray.autoscaler._private.command_runner import SSHCommandRunner
 from ray.autoscaler._private.cli_logger import cli_logger, cf
 from ray.autoscaler._private.constants import BOTO_MAX_RETRIES, BOTO_CREATE_MAX_RETRIES
 from ray.autoscaler._private.log_timer import LogTimer
@@ -714,6 +717,30 @@ def fillout_available_node_types_resources(
                 )
         return cluster_config
 
+    def get_command_runner(
+        self,
+        log_prefix,
+        node_id,
+        auth_config,
+        cluster_name,
+        process_runner,
+        use_internal_ip,
+        docker_config=None,
+    ):
+        common_args = {
+            "log_prefix": log_prefix,
+            "node_id": node_id,
+            "provider": self,
+            "auth_config": auth_config,
+            "cluster_name": cluster_name,
+            "process_runner": process_runner,
+            "use_internal_ip": use_internal_ip,
+        }
+        if docker_config and docker_config["container_name"] != "":
+            return SkyDockerCommandRunner(docker_config, **common_args)
+        else:
+            return SSHCommandRunner(**common_args)
+
 
 class AWSNodeProviderV2(AWSNodeProvider):
     """Same as V1, except head and workers use a SkyPilot IAM role.

sky/skylet/providers/azure/azure-vm-template.json

@@ -117,6 +117,12 @@
             "metadata": {
                 "description": "OS disk tier."
             }
+        },
+        "cloudInitSetupCommands": {
+            "type": "string",
+            "metadata": {
+                "description": "Base64 encoded cloud-init setup commands."
+            }
         }
     },
     "variables": {
@@ -260,7 +266,8 @@
                                 }
                             ]
                         }
-                    }
+                    },
+                    "customData": "[parameters('cloudInitSetupCommands')]"
                 },
                 "priority": "[parameters('priority')]",
                 "billingProfile": "[parameters('billingProfile')]"

sky/skylet/providers/azure/config.py

@@ -81,6 +81,32 @@ def _configure_resource_group(config):
     with open(template_path, "r") as template_fp:
         template = json.load(template_fp)
 
+    # Setup firewall rules for ports
+    nsg_resource = None
+    for resource in template["resources"]:
+        if resource["type"] == "Microsoft.Network/networkSecurityGroups":
+            nsg_resource = resource
+            break
+    assert nsg_resource is not None, "Could not find NSG resource in template"
+    ports = config["provider"].get("ports", None)
+    if ports is not None:
+        ports = [str(port) for port in ports if port != 22]
+        nsg_resource["properties"]["securityRules"].append(
+            {
+                "name": "user-ports",
+                "properties": {
+                    "priority": 1001,
+                    "protocol": "TCP",
+                    "access": "Allow",
+                    "direction": "Inbound",
+                    "sourceAddressPrefix": "*",
+                    "sourcePortRange": "*",
+                    "destinationAddressPrefix": "*",
+                    "destinationPortRanges": ports,
+                },
+            }
+        )
+
     logger.info("Using cluster name: %s", config["cluster_name"])
 
     # set unique id for resources in this cluster

sky/skylet/providers/azure/node_provider.py

@@ -15,6 +15,9 @@
     bootstrap_azure,
     get_azure_sdk_function,
 )
+from sky.skylet.providers.command_runner import SkyDockerCommandRunner
+
+from ray.autoscaler._private.command_runner import SSHCommandRunner
 from ray.autoscaler.node_provider import NodeProvider
 from ray.autoscaler.tags import (
     TAG_RAY_CLUSTER_NAME,
@@ -419,3 +422,27 @@ def _get_cached_node(self, node_id):
     @staticmethod
     def bootstrap_config(cluster_config):
         return bootstrap_azure(cluster_config)
+
+    def get_command_runner(
+        self,
+        log_prefix,
+        node_id,
+        auth_config,
+        cluster_name,
+        process_runner,
+        use_internal_ip,
+        docker_config=None,
+    ):
+        common_args = {
+            "log_prefix": log_prefix,
+            "node_id": node_id,
+            "provider": self,
+            "auth_config": auth_config,
+            "cluster_name": cluster_name,
+            "process_runner": process_runner,
+            "use_internal_ip": use_internal_ip,
+        }
+        if docker_config and docker_config["container_name"] != "":
+            return SkyDockerCommandRunner(docker_config, **common_args)
+        else:
+            return SSHCommandRunner(**common_args)

sky/skylet/providers/gcp/node_provider.py

@@ -12,7 +12,9 @@
     construct_clients_from_provider_config,
     get_node_type,
 )
+from sky.skylet.providers.command_runner import SkyDockerCommandRunner
 
+from ray.autoscaler._private.command_runner import SSHCommandRunner
 from ray.autoscaler.tags import (
     TAG_RAY_LAUNCH_CONFIG,
     TAG_RAY_NODE_KIND,
@@ -364,3 +366,27 @@ def _get_cached_node(self, node_id: str) -> GCPNode:
     @staticmethod
     def bootstrap_config(cluster_config):
         return bootstrap_gcp(cluster_config)
+
+    def get_command_runner(
+        self,
+        log_prefix,
+        node_id,
+        auth_config,
+        cluster_name,
+        process_runner,
+        use_internal_ip,
+        docker_config=None,
+    ):
+        common_args = {
+            "log_prefix": log_prefix,
+            "node_id": node_id,
+            "provider": self,
+            "auth_config": auth_config,
+            "cluster_name": cluster_name,
+            "process_runner": process_runner,
+            "use_internal_ip": use_internal_ip,
+        }
+        if docker_config and docker_config["container_name"] != "":
+            return SkyDockerCommandRunner(docker_config, **common_args)
+        else:
+            return SSHCommandRunner(**common_args)