Use user property as the default user for SSH (#2263)

Closes: #2262
dstackai · Feb 4, 2025 · 8aada88 · 8aada88
1 parent 4cc82c3
commit 8aada88
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/src/dstack/_internal/core/services/ssh/attach.py b/src/dstack/_internal/core/services/ssh/attach.py
@@ -56,6 +56,7 @@ def __init__(
         ssh_port: int,
         container_ssh_port: int,
         user: str,
+        container_user: str,
         id_rsa_path: PathLike,
         ports_lock: PortsLock,
         run_name: str,
@@ -74,7 +75,7 @@ def __init__(
         self.control_sock_path = FilePath(control_sock_path)
         self.identity_file = FilePath(id_rsa_path)
         self.tunnel = SSHTunnel(
-            destination=run_name,
+            destination=f"root@{run_name}",
             identity=self.identity_file,
             forwarded_sockets=ports_to_forwarded_sockets(
                 ports=self.ports,
@@ -91,7 +92,7 @@ def __init__(
             self.host_config = {
                 "HostName": hostname,
                 "Port": ssh_port,
-                "User": user,
+                "User": user if dockerized else container_user,
                 "IdentityFile": self.identity_file,
                 "IdentitiesOnly": "yes",
                 "StrictHostKeyChecking": "no",
@@ -111,7 +112,7 @@ def __init__(
             self.container_config = {
                 "HostName": "localhost",
                 "Port": container_ssh_port,
-                "User": "root",  # TODO(#1535): support non-root images properly
+                "User": container_user,
                 "IdentityFile": self.identity_file,
                 "IdentitiesOnly": "yes",
                 "StrictHostKeyChecking": "no",
@@ -122,7 +123,7 @@ def __init__(
             self.container_config = {
                 "HostName": hostname,
                 "Port": ssh_port,
-                "User": user,
+                "User": container_user,
                 "IdentityFile": self.identity_file,
                 "IdentitiesOnly": "yes",
                 "StrictHostKeyChecking": "no",
@@ -136,7 +137,7 @@ def __init__(
             self.host_config = {
                 "HostName": hostname,
                 "Port": container_ssh_port,
-                "User": "root",  # TODO(#1535): support non-root images properly
+                "User": container_user,
                 "IdentityFile": self.identity_file,
                 "IdentitiesOnly": "yes",
                 "StrictHostKeyChecking": "no",

diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py
@@ -324,11 +324,19 @@ def attach(
             if runtime_data is not None and runtime_data.ports is not None:
                 container_ssh_port = runtime_data.ports.get(container_ssh_port, container_ssh_port)
 
+            # TODO: get login name from runner in case it's not specified in the run configuration
+            # (i.e. the default image user is used, and it is not root)
+            if job.job_spec.user is not None and job.job_spec.user.username is not None:
+                container_user = job.job_spec.user.username
+            else:
+                container_user = "root"
+
             self._ssh_attach = SSHAttach(
                 hostname=provisioning_data.hostname,
                 ssh_port=provisioning_data.ssh_port,
                 container_ssh_port=container_ssh_port,
                 user=provisioning_data.username,
+                container_user=container_user,
                 id_rsa_path=ssh_identity_file,
                 ports_lock=self._ports_lock,
                 run_name=name,