From 53f877694b475f1e1f9072d60f82ba5f97e4f0a9 Mon Sep 17 00:00:00 2001
From: Rohin Bhasin <bhasin.rohin@gmail.com>
Date: Wed, 4 Sep 2024 23:05:38 -0400
Subject: [PATCH] Docker exec directly into container.

---
 runhouse/resources/hardware/cluster.py            | 11 ++++++++---
 runhouse/resources/hardware/sky_command_runner.py |  6 ++++--
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/runhouse/resources/hardware/cluster.py b/runhouse/resources/hardware/cluster.py
index 77c20c8c7..b83ecc008 100644
--- a/runhouse/resources/hardware/cluster.py
+++ b/runhouse/resources/hardware/cluster.py
@@ -438,7 +438,9 @@ def is_shared(self) -> bool:
 
         return f"{self._creds.name}/" in ssh_creds.get("ssh_private_key", "")
 
-    def _command_runner(self, node: Optional[str] = None) -> "CommandRunner":
+    def _command_runner(
+        self, node: Optional[str] = None, use_docker_exec: Optional[bool] = False
+    ) -> "CommandRunner":
         from runhouse.resources.hardware.sky_command_runner import (
             SkyKubernetesRunner,
             SkySSHRunner,
@@ -473,7 +475,8 @@ def _command_runner(self, node: Optional[str] = None) -> "CommandRunner":
                 ssh_private_key=ssh_credentials.get("ssh_private_key"),
                 ssh_proxy_command=ssh_credentials.get("ssh_proxy_command"),
                 ssh_control_name=ssh_control_name,
-                docker_user=self.docker_user,
+                docker_user=self.docker_user if not use_docker_exec else None,
+                use_docker_exec=use_docker_exec,
             )
 
         return runner
@@ -1519,7 +1522,9 @@ def _run_commands_with_runner(
         ssh_credentials.pop("private_key", None)
         ssh_credentials.pop("public_key", None)
 
-        runner = self._command_runner(node=node)
+        runner = self._command_runner(
+            node=node, use_docker_exec=self.docker_user is not None
+        )
 
         env_var_prefix = (
             " ".join(f"{key}={val}" for key, val in env_vars.items())
diff --git a/runhouse/resources/hardware/sky_command_runner.py b/runhouse/resources/hardware/sky_command_runner.py
index bff7503c4..1ee4bca95 100644
--- a/runhouse/resources/hardware/sky_command_runner.py
+++ b/runhouse/resources/hardware/sky_command_runner.py
@@ -66,6 +66,7 @@ def __init__(
         docker_user: Optional[str] = None,
         disable_control_master: Optional[bool] = False,
         local_bind_port: Optional[int] = None,
+        use_docker_exec: Optional[bool] = False,
     ):
         super().__init__(
             node,
@@ -81,6 +82,7 @@ def __init__(
         self.docker_user = docker_user
         self.local_bind_port = local_bind_port
         self.remote_bind_port = None
+        self.use_docker_exec = use_docker_exec
 
     def _ssh_base_command(
         self,
@@ -176,9 +178,9 @@ def run(
         if quiet_ssh:  # RH MODIFIED
             base_ssh_command.append("-q")
 
-        if self.docker_user:  # RH MODIFIED
+        if self.use_docker_exec:  # RH MODIFIED
             cmd = " ".join(cmd) if isinstance(cmd, list) else cmd
-            cmd = f"conda deactivate && {cmd}"
+            cmd = f"sudo docker exec {DEFAULT_DOCKER_CONTAINER_NAME} bash -c {shlex.quote(cmd)}"
 
         command_str = self._get_command_to_run(
             cmd,