From 8d6966e5f44d599ba2afdf8945ae032cc1b93de9 Mon Sep 17 00:00:00 2001
From: Victor Skvortsov <vds003@gmail.com>
Date: Mon, 20 Jan 2025 13:42:42 +0500
Subject: [PATCH] Fix local backend

---
 runner/cmd/shim/main.go                               |  2 --
 runner/internal/shim/runner.go                        |  3 +++
 src/dstack/_internal/core/services/ssh/attach.py      | 11 +++++++++++
 .../server/background/tasks/process_running_jobs.py   |  2 +-
 src/dstack/_internal/server/services/runner/ssh.py    | 10 ++++------
 5 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/runner/cmd/shim/main.go b/runner/cmd/shim/main.go
index 234623259..cf37776c3 100644
--- 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
index 6f5285d05..56dd2b1a7 100644
--- 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
index 103a8c90b..58f23a73f 100644
--- 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
index 0451d0546..40bb2fbef 100644
--- 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
index 925270cba..044c5d3bc 100644
--- 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