move helper functions to hardware utils

runhouse/resources/hardware/sky_ssh_runner.py

@@ -24,6 +24,7 @@
     SshMode,
 )
 
+from runhouse.resources.hardware.utils import _ssh_base_command
 
 # Get rid of the constant "Found credentials in shared credentials file: ~/.aws/credentials" message
 try:
@@ -102,39 +103,17 @@ def __init__(
     def _ssh_base_command(
         self, *, ssh_mode: SshMode, port_forward: Optional[List[int]]
     ) -> List[str]:
-        ssh = ["ssh"]
-        if ssh_mode == SshMode.NON_INTERACTIVE:
-            # Disable pseudo-terminal allocation. Otherwise, the output of
-            # ssh will be corrupted by the user's input.
-            ssh += ["-T"]
-        else:
-            # Force pseudo-terminal allocation for interactive/login mode.
-            ssh += ["-tt"]
-        if port_forward is not None:
-            # RH MODIFIED: Accept port int (to forward same port) or pair of ports
-            for fwd in port_forward:
-                if isinstance(fwd, int):
-                    local, remote = fwd, fwd
-                else:
-                    local, remote = fwd
-                logger.debug(f"Forwarding port {local} to port {remote} on localhost.")
-                ssh += ["-L", f"{local}:localhost:{remote}"]
-        if self._docker_ssh_proxy_command is not None:
-            docker_ssh_proxy_command = self._docker_ssh_proxy_command(ssh)
-        else:
-            docker_ssh_proxy_command = None
-        return (
-            ssh
-            + ssh_options_list(
-                self.ssh_private_key,
-                self.ssh_control_name,
-                ssh_proxy_command=self._ssh_proxy_command,
-                docker_ssh_proxy_command=docker_ssh_proxy_command,
-                # TODO change to None like before?
-                port=self.port,
-                disable_control_master=self.disable_control_master,
-            )
-            + [f"{self.ssh_user}@{self.ip}"]
+        return _ssh_base_command(
+            address=self.ip,
+            ssh_user=self.ssh_user,
+            ssh_private_key=self.ssh_private_key,
+            ssh_control_name=self.ssh_control_name,
+            ssh_proxy_command=self._ssh_proxy_command,
+            ssh_port=self.port,
+            docker_user=self.docker_user,
+            disable_control_master=self.disable_control_master,
+            ssh_mode=ssh_mode,
+            port_forward=port_forward,
         )
 
     def run(

runhouse/resources/hardware/utils.py

@@ -3,16 +3,17 @@
 
 from enum import Enum
 from pathlib import Path
-from typing import Dict
+from typing import Dict, List, Optional
 
 from runhouse.constants import (
     CLUSTER_CONFIG_PATH,
     EMPTY_DEFAULT_ENV_NAME,
     RESERVED_SYSTEM_NAMES,
 )
+
+from runhouse.logger import logger
 from runhouse.resources.envs.utils import _get_env_from, run_setup_command
-from runhouse.resources.hardware.sky.command_runner import SshMode
-from runhouse.resources.hardware.sky_ssh_runner import SkySSHRunner
+from runhouse.resources.hardware.sky.command_runner import ssh_options_list, SshMode
 
 
 class ServerConnectionType(str, Enum):
@@ -142,14 +143,80 @@ def _run_ssh_command(
     ssh_private_key: str,
     docker_user: str,
 ):
-    runner = SkySSHRunner(
-        ip=address,
+    ssh_command = _ssh_base_command(
+        address=address,
         ssh_user=ssh_user,
-        port=ssh_port,
         ssh_private_key=ssh_private_key,
+        ssh_mode=SshMode.INTERACTIVE,
+        ssh_port=ssh_port,
         docker_user=docker_user,
     )
-    ssh_command = runner._ssh_base_command(
-        ssh_mode=SshMode.INTERACTIVE, port_forward=None
-    )
     subprocess.run(ssh_command)
+
+
+def _docker_ssh_proxy_command(
+    address: str,
+    ssh_user: str,
+    ssh_private_key: str,
+):
+    return lambda ssh: " ".join(
+        ssh
+        + ssh_options_list(ssh_private_key, None)
+        + ["-W", "%h:%p", f"{ssh_user}@{address}"]
+    )
+
+
+# Adapted from SkyPilot Command Runner
+def _ssh_base_command(
+    address: str,
+    ssh_user: str,
+    ssh_private_key: str,
+    ssh_control_name: Optional[str] = "__default__",
+    ssh_proxy_command: Optional[str] = None,
+    ssh_port: int = 22,
+    docker_user: Optional[str] = None,
+    disable_control_master: Optional[bool] = False,
+    ssh_mode: SshMode = SshMode.INTERACTIVE,
+    port_forward: Optional[List[int]] = None,
+    timeout: Optional[int] = 30,
+):
+    ssh = ["ssh"]
+    if ssh_mode == SshMode.NON_INTERACTIVE:
+        # Disable pseudo-terminal allocation. Otherwise, the output of
+        # ssh will be corrupted by the user's input.
+        ssh += ["-T"]
+    else:
+        # Force pseudo-terminal allocation for interactive/login mode.
+        ssh += ["-tt"]
+
+    if port_forward is not None:
+        # RH MODIFIED: Accept port int (to forward same port) or pair of ports
+        for fwd in port_forward:
+            if isinstance(fwd, int):
+                local, remote = fwd, fwd
+            else:
+                local, remote = fwd
+            logger.debug(f"Forwarding port {local} to port {remote} on localhost.")
+            ssh += ["-L", f"{local}:localhost:{remote}"]
+
+    if docker_user:
+        docker_ssh_proxy_command = _docker_ssh_proxy_command(
+            address, ssh_user, ssh_private_key
+        )(["ssh"])
+    else:
+        docker_ssh_proxy_command = None
+
+    return (
+        ssh
+        + ssh_options_list(
+            ssh_private_key,
+            ssh_control_name,
+            ssh_proxy_command=ssh_proxy_command,
+            docker_ssh_proxy_command=docker_ssh_proxy_command,
+            # TODO change to None like before?
+            port=ssh_port,
+            timeout=timeout,
+            disable_control_master=disable_control_master,
+        )
+        + [f"{ssh_user}@{address}"]
+    )