Fix dstack apply --ssh-identity

src/dstack/_internal/cli/services/configurators/fleet.py

@@ -38,7 +38,7 @@
 from dstack._internal.utils.common import local_time
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.nested_list import NestedList, NestedListItem
-from dstack._internal.utils.ssh import convert_ssh_key_to_pem, generate_public_key, pkey_from_str
+from dstack._internal.utils.ssh import resolve_ssh_key
 from dstack.api.utils import load_profile
 
 logger = get_logger(__name__)
@@ -354,22 +354,15 @@ def _preprocess_spec(spec: FleetSpec):
 def _resolve_ssh_key(ssh_key_path: Optional[str]) -> Optional[SSHKey]:
     if ssh_key_path is None:
         return None
-    ssh_key_path_obj = Path(ssh_key_path).expanduser()
     try:
-        private_key = convert_ssh_key_to_pem(ssh_key_path_obj.read_text())
-        try:
-            pub_key = ssh_key_path_obj.with_suffix(".pub").read_text()
-        except FileNotFoundError:
-            pub_key = generate_public_key(pkey_from_str(private_key))
-        return SSHKey(public=pub_key, private=private_key)
+        public_key, _, private_key, _ = resolve_ssh_key(ssh_key_path)
     except OSError as e:
-        logger.debug("Got OSError: %s", repr(e))
-        console.print(f"[error]Unable to read the SSH key at {ssh_key_path}[/]")
-        exit()
+        raise CLIError(f"Unable to read the SSH key at {ssh_key_path}") from e
     except ValueError as e:
-        logger.debug("Key type is not supported", repr(e))
-        console.print("[error]Key type is not supported[/]")
-        exit()
+        raise CLIError(f"Unsupported or invalid SSH key at {ssh_key_path}") from e
+    if private_key is None:
+        raise CLIError(f"Expected a private key at {ssh_key_path}, got a public key")
+    return SSHKey(public=public_key, private=private_key)
 
 
 def _render_fleet_spec_diff(old_spec: FleetSpec, new_spec: FleetSpec) -> Optional[str]:

src/dstack/_internal/cli/services/configurators/run.py

@@ -75,6 +75,7 @@
 from dstack._internal.utils.nested_list import NestedList, NestedListItem
 from dstack._internal.utils.nodes_interpolator import is_valid_groups_ip_ref
 from dstack._internal.utils.path import is_absolute_posix_path
+from dstack._internal.utils.ssh import resolve_ssh_key
 from dstack.api._public.runs import Run
 
 _BIND_ADDRESS_ARG = "bind_address"
@@ -95,23 +96,27 @@ def apply_configuration(
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
     ):
+        ssh_key_pub, ssh_identity_file = self.get_ssh_key(configurator_args)
         run_plan, repo = self.get_plan(
             conf=conf,
             configuration_path=configuration_path,
             configurator_args=configurator_args,
+            ssh_key_pub=ssh_key_pub,
         )
         return self.apply_plan(
             run_plan=run_plan,
             repo=repo,
             command_args=command_args,
             configurator_args=configurator_args,
+            ssh_identity_file=ssh_identity_file,
         )
 
     def get_plan(
         self,
         conf: RunConfigurationT,
         configuration_path: str,
         configurator_args: argparse.Namespace,
+        ssh_key_pub: Optional[str],
     ) -> tuple[RunPlan, Repo]:
         """Apply CLI arguments and validation, then return the run plan and its repo."""
         if configurator_args.repo and configurator_args.no_repo:
@@ -132,7 +137,7 @@ def get_plan(
                 repo=repo,
                 configuration_path=configuration_path,
                 profile=profile,
-                ssh_identity_file=configurator_args.ssh_identity_file,
+                ssh_key_pub=ssh_key_pub,
                 max_offers=configurator_args.max_offers,
                 full_offers=configurator_args.full_offers,
                 unallocated_resources=configurator_args.unallocated,
@@ -145,6 +150,7 @@ def apply_plan(
         repo: Repo,
         command_args: argparse.Namespace,
         configurator_args: argparse.Namespace,
+        ssh_identity_file: Optional[Path],
         plan_properties: Optional[Dict[str, str]] = None,
     ):
         """Apply a run plan using the standard CLI behavior."""
@@ -270,7 +276,10 @@ def apply_plan(
                     )
                     try:
                         try:
-                            attached = run.attach(bind_address=bind_address)
+                            attached = run.attach(
+                                ssh_identity_file=ssh_identity_file,
+                                bind_address=bind_address,
+                            )
                         except PortUsedError as e:
                             console.print(
                                 f"[error]Failed to attach: port [code]{e.port}[/code] is already in use."
@@ -543,6 +552,23 @@ def get_repo(
 
         return repo
 
+    def get_ssh_key(
+        self, configurator_args: argparse.Namespace
+    ) -> tuple[Optional[str], Optional[Path]]:
+        """Resolve the `--ssh-identity` argument to a (public key, private key path) pair."""
+        ssh_identity_file: Optional[Path] = configurator_args.ssh_identity_file
+        if ssh_identity_file is None:
+            return None, None
+        try:
+            public_key, _, _, private_key_path = resolve_ssh_key(ssh_identity_file)
+        except OSError as e:
+            raise CLIError(f"Unable to read the SSH key at {ssh_identity_file}") from e
+        except ValueError as e:
+            raise CLIError(f"Unsupported or invalid SSH key at {ssh_identity_file}") from e
+        if private_key_path is None:
+            raise CLIError(f"Expected a private key at {ssh_identity_file}, got a public key")
+        return public_key, private_key_path
+
 
 class RunWithPortsConfiguratorMixin:
     @classmethod

src/dstack/_internal/utils/ssh.py

@@ -5,6 +5,7 @@
 import subprocess
 import sys
 import tempfile
+from contextlib import suppress
 from pathlib import Path
 from typing import Dict, Optional, Union
 
@@ -268,6 +269,58 @@ def generate_public_key(private_key: PKey) -> str:
     return public_key
 
 
+def resolve_ssh_key(
+    path: PathLike,
+) -> Union[
+    tuple[str, Path, str, Path],
+    tuple[str, None, str, Path],
+    tuple[str, Path, None, None],
+]:
+    """
+    Resolves a private or public key path to key contents and paths.
+
+    If a private key is given, only supported private key types are allowed. PKCS#8 keys are
+    converted to PEM, so the returned private key may differ from the file contents. If a
+    corresponding ".pub" file exists, its contents is used as a public key without any validation
+    and its path is returned as the public key path, otherwise a public key is generated from the
+    private key and the public key path is None.
+
+    If a public key is given, any valid public key is allowed regardless of its type, and both
+    private key values are None. No corresponding private key (a file without ".pub" suffix) is
+    checked.
+
+    Args:
+        path: The private or public key path.
+
+    Returns:
+        A (public key, public key path, private key, private key path) tuple.
+
+    Raises:
+        OSError: Error reading key file(s).
+        ValueError: Unsupported or invalid private key or invalid public key.
+    """
+    path = Path(path).expanduser()
+    content = path.read_text()
+    private_key = convert_ssh_key_to_pem(content)
+    pkey: Optional[PKey] = None
+    with suppress(ValueError):
+        pkey = pkey_from_str(private_key)
+    if pkey is None:
+        # unsupported private key or public key or garbage
+        try:
+            PublicBlob.from_string(content)
+        except ValueError:
+            # unsupported private key or garbage
+            raise ValueError("Unsupported key type or invalid key")
+        # any valid public key, including unsupported (without matching SUPPORTED_KEY_TYPES PKey)
+        return content, path, None, None
+    # supported private key
+    public_key_path = path.with_name(path.name + ".pub")
+    if public_key_path.is_file():
+        return public_key_path.read_text(), public_key_path, private_key, path
+    return generate_public_key(pkey), None, private_key, path
+
+
 def check_required_ssh_version() -> bool:
     try:
         result = subprocess.run(["ssh", "-V"], capture_output=True, text=True)

src/dstack/api/_public/runs.py

@@ -50,6 +50,7 @@
 from dstack._internal.utils.files import create_file_archive
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import PathLike
+from dstack._internal.utils.ssh import resolve_ssh_key
 from dstack.api.server import APIClient
 
 logger = get_logger(__name__)
@@ -477,6 +478,7 @@ def get_run_plan(
         configuration_path: Optional[str] = None,
         repo_dir: Union[Deprecated, str, None] = Deprecated.PLACEHOLDER,
         ssh_identity_file: Optional[PathLike] = None,
+        ssh_key_pub: Optional[str] = None,
         max_offers: Optional[int] = None,
         full_offers: bool = False,
         unallocated_resources: bool = False,
@@ -493,10 +495,15 @@ def get_run_plan(
             profile: The profile to use for the run.
             configuration_path: The path to the configuration file. Omit if the configuration
                 is not loaded from a file.
-            ssh_identity_file: Path to the private SSH key file. The corresponding public key
-                (`.pub` file) is read and included in the run plan, allowing SSH access to the instances.
-                If the `.pub` file does not exist, it is generated automatically.
-                If ssh_identity_file is not specified, the user key is used.
+            ssh_identity_file: Path to a private or public SSH key file. The public key is
+                included in the run plan, allowing SSH access to the instances. If a private key
+                is given, its public key is read from the corresponding `.pub` file or, if there
+                is no such file, generated from the private key.
+                Mutually exclusive with ssh_key_pub.
+            ssh_key_pub: The public SSH key to include in the run plan, allowing SSH access to
+                the instances. Use it instead of ssh_identity_file if the key is not stored
+                on disk. Mutually exclusive with ssh_identity_file.
+                If neither ssh_key_pub nor ssh_identity_file is specified, the user key is used.
             max_offers: Maximum number of offers returned in the run plan.
             full_offers: Return full offers not adjusted by requirements.
             unallocated_resources: Subtract allocated resources to return only unallocated
@@ -536,10 +543,20 @@ def get_run_plan(
                 archive = self._api_client.files.upload_archive(hash=archive_hash, fp=fp)
             file_archives.append(FileArchiveMapping(id=archive.id, path=file_mapping.path))
 
+        if ssh_key_pub and ssh_identity_file:
+            raise ConfigurationError("ssh_key_pub and ssh_identity_file are mutually exclusive")
         if ssh_identity_file:
-            ssh_key_pub = Path(ssh_identity_file).with_suffix(".pub").read_text()
-        else:
-            ssh_key_pub = None  # using the server-managed user key
+            try:
+                ssh_key_pub, _, _, _ = resolve_ssh_key(ssh_identity_file)
+            except OSError as e:
+                raise ConfigurationError(
+                    f"Unable to read the SSH key at {ssh_identity_file}"
+                ) from e
+            except ValueError as e:
+                raise ConfigurationError(
+                    f"Unsupported or invalid SSH key at {ssh_identity_file}"
+                ) from e
+        # `ssh_key_pub` is None if neither is given: using the server-managed user key
         run_spec = RunSpec(
             run_name=configuration.name,
             repo_id=repo.repo_id,
@@ -609,6 +626,7 @@ def apply_configuration(
         configuration_path: Optional[str] = None,
         reserve_ports: bool = True,
         ssh_identity_file: Optional[PathLike] = None,
+        ssh_key_pub: Optional[str] = None,
     ) -> Run:
         """
         Apply the run configuration.
@@ -621,10 +639,15 @@ def apply_configuration(
             profile: The profile to use for the run.
             configuration_path: The path to the configuration file. Omit if the configuration is not loaded from a file.
             reserve_ports: Reserve local ports before applying. Use if you'll attach to the run.
-            ssh_identity_file: Path to the private SSH key file. The corresponding public key
-                (`.pub` file) is read and included in the run plan, allowing SSH access to the instances.
-                If the `.pub` file does not exist, it is generated automatically.
-                If ssh_identity_file is not specified, the user key is used.
+            ssh_identity_file: Path to a private or public SSH key file. The public key is
+                included in the run plan, allowing SSH access to the instances. If a private key
+                is given, its public key is read from the corresponding `.pub` file or, if there
+                is no such file, generated from the private key.
+                Mutually exclusive with ssh_key_pub.
+            ssh_key_pub: The public SSH key to include in the run plan, allowing SSH access to
+                the instances. Use it instead of ssh_identity_file if the key is not stored
+                on disk. Mutually exclusive with ssh_identity_file.
+                If neither ssh_key_pub nor ssh_identity_file is specified, the user key is used.
 
         Returns:
             Submitted run.
@@ -635,6 +658,7 @@ def apply_configuration(
             profile=profile,
             configuration_path=configuration_path,
             ssh_identity_file=ssh_identity_file,
+            ssh_key_pub=ssh_key_pub,
         )
         run = self.apply_plan(
             run_plan=run_plan,

src/tests/_internal/cli/services/configurators/test_fleet.py

@@ -1,5 +1,6 @@
 import argparse
 from datetime import datetime, timezone
+from pathlib import Path
 from textwrap import dedent
 from typing import List, Optional, Tuple
 from unittest.mock import Mock
@@ -12,8 +13,9 @@
 from dstack._internal.cli.services.configurators.fleet import (
     FleetConfigurator,
     _render_fleet_spec_diff,
+    _resolve_ssh_key,
 )
-from dstack._internal.core.errors import ConfigurationError
+from dstack._internal.core.errors import CLIError, ConfigurationError
 from dstack._internal.core.models.common import ApplyAction
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.fleets import (
@@ -25,7 +27,9 @@
     FleetStatus,
     InstanceGroupPlacement,
 )
+from dstack._internal.core.models.instances import SSHKey
 from dstack._internal.core.models.profiles import Profile
+from tests._internal.utils.test_ssh import PRIVATE_KEY, PUBLIC_KEY, PUBLIC_KEY_NO_COMMENT
 
 
 def create_conf() -> FleetConfiguration:
@@ -245,3 +249,43 @@ def test_no_diff(self):
         spec = get_cloud_fleet_spec()
 
         assert _render_fleet_spec_diff(spec, spec.model_copy(deep=True)) is None
+
+
+class TestResolveSSHKey:
+    def test_returns_none_if_no_path_given(self):
+        assert _resolve_ssh_key(None) is None
+
+    def test_uses_public_key_file(self, tmp_path: Path):
+        private_key_path = tmp_path / "id_ed25519"
+        private_key_path.write_text(PRIVATE_KEY)
+        (tmp_path / "id_ed25519.pub").write_text(PUBLIC_KEY)
+
+        assert _resolve_ssh_key(str(private_key_path)) == SSHKey(
+            public=PUBLIC_KEY, private=PRIVATE_KEY
+        )
+
+    def test_generates_public_key(self, tmp_path: Path):
+        private_key_path = tmp_path / "id_ed25519"
+        private_key_path.write_text(PRIVATE_KEY)
+
+        assert _resolve_ssh_key(str(private_key_path)) == SSHKey(
+            public=PUBLIC_KEY_NO_COMMENT, private=PRIVATE_KEY
+        )
+
+    def test_raises_if_public_key_given(self, tmp_path: Path):
+        public_key_path = tmp_path / "id_ed25519.pub"
+        public_key_path.write_text(PUBLIC_KEY)
+
+        with pytest.raises(CLIError, match="Expected a private key"):
+            _resolve_ssh_key(str(public_key_path))
+
+    def test_raises_if_key_does_not_exist(self, tmp_path: Path):
+        with pytest.raises(CLIError, match="Unable to read the SSH key"):
+            _resolve_ssh_key(str(tmp_path / "id_ed25519"))
+
+    def test_raises_if_key_type_is_not_supported(self, tmp_path: Path):
+        key_path = tmp_path / "id_ed25519"
+        key_path.write_text("garbage")
+
+        with pytest.raises(CLIError, match="Unsupported or invalid SSH key"):
+            _resolve_ssh_key(str(key_path))