Fix local backend (#2203)

* Fix local backend * Update runner readme
dstackai · Jan 20, 2025 · 1fdb91a · 1fdb91a
1 parent 3bd02ab
commit 1fdb91a
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 11 deletions.
diff --git a/runner/README.md b/runner/README.md
@@ -25,13 +25,12 @@ Here's the steps to build `dstack-shim` and `dstack-runner` and run `dstack` wit
 3. Start the shim:
 
     ```shell
-    ./shim --shim-home $RUNNER_DIR --runner-binary-path $COMPILED_RUNNER_PATH --ssh-key $DSTACK_PUBLIC_KEY
+    ./shim --shim-home $RUNNER_DIR --runner-binary-path $COMPILED_RUNNER_PATH
     ```
 
     Notes:
 
     * `$RUNNER_DIR` is any directory for storing runner files.
-    * `$DSTACK_PUBLIC_KEY` is `~/.dstack/ssh/id_rsa.pub` that allows the dstack CLI to connect to the ssh server inside the container.
 
 Now you can call shim API:
 

diff --git a/runner/cmd/shim/main.go b/runner/cmd/shim/main.go
@@ -65,7 +65,6 @@ func main() {
 				Name:        "runner-download-url",
 				Usage:       "Set runner's download URL",
 				Destination: &args.Runner.DownloadURL,
-				Required:    true,
 				EnvVars:     []string{"DSTACK_RUNNER_DOWNLOAD_URL"},
 			},
 			&cli.PathFlag{
@@ -106,7 +105,6 @@ func main() {
 			&cli.StringFlag{
 				Name:        "ssh-key",
 				Usage:       "Public SSH key",
-				Required:    true,
 				Destination: &args.Docker.ConcatinatedPublicSSHKeys,
 				EnvVars:     []string{"DSTACK_PUBLIC_SSH_KEY"},
 			},

diff --git a/runner/internal/shim/runner.go b/runner/internal/shim/runner.go
@@ -16,6 +16,9 @@ import (
 )
 
 func (c *CLIArgs) DownloadRunner(ctx context.Context) error {
+	if c.Runner.DownloadURL == "" {
+		return nil
+	}
 	err := downloadRunner(ctx, c.Runner.DownloadURL, c.Runner.BinaryPath, false)
 	if err != nil {
 		return gerrors.Wrap(err)

diff --git a/src/dstack/_internal/core/services/ssh/attach.py b/src/dstack/_internal/core/services/ssh/attach.py
@@ -131,6 +131,17 @@ def __init__(
             }
         else:
             self.container_config = None
+        if local_backend:
+            self.container_config = None
+            self.host_config = {
+                "HostName": hostname,
+                "Port": container_ssh_port,
+                "User": "root",  # TODO(#1535): support non-root images properly
+                "IdentityFile": self.identity_file,
+                "IdentitiesOnly": "yes",
+                "StrictHostKeyChecking": "no",
+                "UserKnownHostsFile": "/dev/null",
+            }
         if self.container_config is not None and get_ssh_client_info().supports_multiplexing:
             self.container_config.update(
                 {

diff --git a/src/dstack/_internal/server/background/tasks/process_running_jobs.py b/src/dstack/_internal/server/background/tasks/process_running_jobs.py
@@ -441,7 +441,7 @@ def _process_provisioning_with_shim(
             volume_mounts=volume_mounts,
             instance_mounts=instance_mounts,
             host_ssh_user=ssh_user,
-            host_ssh_keys=[ssh_key],
+            host_ssh_keys=[ssh_key] if ssh_key else [],
             container_ssh_keys=public_keys,
         )
     else:

diff --git a/src/dstack/_internal/server/services/runner/ssh.py b/src/dstack/_internal/server/services/runner/ssh.py
@@ -42,17 +42,15 @@ def wrapper(
             Returns:
                 is successful
             """
-
-            if job_provisioning_data.backend == BackendType.LOCAL:
-                # without SSH
-                port_map = {p: p for p in ports}
-                return func(port_map, *args, **kwargs)
-
             # container:host mapping
             container_ports_map = {port: port for port in ports}
             if job_runtime_data is not None and job_runtime_data.ports is not None:
                 container_ports_map.update(job_runtime_data.ports)
 
+            if job_provisioning_data.backend == BackendType.LOCAL:
+                # without SSH
+                return func(container_ports_map, *args, **kwargs)
+
             for attempt in range(retries):
                 last = attempt == retries - 1
                 # remote_host:local mapping