diff --git a/src/dstack/_internal/cli/services/configurators/fleet.py b/src/dstack/_internal/cli/services/configurators/fleet.py index 0e3dbdd80..97ff87277 100644 --- a/src/dstack/_internal/cli/services/configurators/fleet.py +++ b/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]: diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 75f4c6773..d3273ba9c 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/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,16 +96,19 @@ 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( @@ -112,6 +116,7 @@ def get_plan( 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 diff --git a/src/dstack/_internal/utils/ssh.py b/src/dstack/_internal/utils/ssh.py index 8b180127e..d8a701233 100644 --- a/src/dstack/_internal/utils/ssh.py +++ b/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) diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py index 18448aab7..018500cd3 100644 --- a/src/dstack/api/_public/runs.py +++ b/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, diff --git a/src/tests/_internal/cli/services/configurators/test_fleet.py b/src/tests/_internal/cli/services/configurators/test_fleet.py index 80bc215e9..35c6a313f 100644 --- a/src/tests/_internal/cli/services/configurators/test_fleet.py +++ b/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)) diff --git a/src/tests/_internal/cli/services/configurators/test_run.py b/src/tests/_internal/cli/services/configurators/test_run.py index aa77dec68..a382ab299 100644 --- a/src/tests/_internal/cli/services/configurators/test_run.py +++ b/src/tests/_internal/cli/services/configurators/test_run.py @@ -1,6 +1,7 @@ import argparse +from pathlib import Path from textwrap import dedent -from typing import List, Tuple +from typing import List, Optional, Tuple from unittest.mock import Mock import pytest @@ -11,7 +12,7 @@ ServiceConfigurator, render_run_spec_diff, ) -from dstack._internal.core.errors import ConfigurationError +from dstack._internal.core.errors import CLIError, ConfigurationError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import RegistryAuth from dstack._internal.core.models.configurations import ( @@ -24,6 +25,7 @@ from dstack._internal.core.models.envs import Env from dstack._internal.core.models.profiles import Profile from dstack._internal.server.testing.common import get_run_spec +from tests._internal.utils.test_ssh import PRIVATE_KEY, PUBLIC_KEY, PUBLIC_KEY_NO_COMMENT _TENSTORRENT_ACCELERATOR_NAMES = tuple( sorted({gpu.name for gpu in KNOWN_TENSTORRENT_ACCELERATORS}) @@ -168,7 +170,8 @@ def test_composes_get_plan_and_apply_plan(self, monkeypatch): apply_plan = Mock() monkeypatch.setattr(ServiceConfigurator, "get_plan", get_plan) monkeypatch.setattr(ServiceConfigurator, "apply_plan", apply_plan) - conf, command_args, configurator_args = Mock(), Mock(), Mock() + conf, command_args = Mock(), Mock() + configurator_args = argparse.Namespace(ssh_identity_file=None) ServiceConfigurator(api_client=Mock()).apply_configuration( conf, "svc.dstack.yml", command_args, configurator_args @@ -178,12 +181,14 @@ def test_composes_get_plan_and_apply_plan(self, monkeypatch): conf=conf, configuration_path="svc.dstack.yml", configurator_args=configurator_args, + ssh_key_pub=None, ) apply_plan.assert_called_once_with( run_plan=run_plan, repo=repo, command_args=command_args, configurator_args=configurator_args, + ssh_identity_file=None, ) @@ -269,3 +274,45 @@ def test_no_diff(self): old = get_run_spec(run_name="test", repo_id="test") new = get_run_spec(run_name="test", repo_id="test") assert render_run_spec_diff(old, new) is None + + +class TestGetSSHKey: + def get_ssh_key( + self, ssh_identity_file: Optional[Path] + ) -> tuple[Optional[str], Optional[Path]]: + configurator_args = argparse.Namespace(ssh_identity_file=ssh_identity_file) + return ServiceConfigurator(api_client=Mock()).get_ssh_key(configurator_args) + + def test_returns_none_if_not_specified(self): + assert self.get_ssh_key(None) == (None, 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 self.get_ssh_key(private_key_path) == (PUBLIC_KEY, private_key_path) + + def test_generates_public_key(self, tmp_path: Path): + private_key_path = tmp_path / "id_ed25519" + private_key_path.write_text(PRIVATE_KEY) + + assert self.get_ssh_key(private_key_path) == (PUBLIC_KEY_NO_COMMENT, private_key_path) + + 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"): + self.get_ssh_key(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"): + self.get_ssh_key(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"): + self.get_ssh_key(key_path) diff --git a/src/tests/_internal/utils/test_ssh.py b/src/tests/_internal/utils/test_ssh.py index 6dbca3786..0512d9bf0 100644 --- a/src/tests/_internal/utils/test_ssh.py +++ b/src/tests/_internal/utils/test_ssh.py @@ -6,16 +6,39 @@ import pytest from dstack._internal.compat import IS_WINDOWS +from dstack._internal.utils import crypto from dstack._internal.utils.path import FilePath from dstack._internal.utils.ssh import ( check_required_ssh_version, + find_ssh_util, include_ssh_config, normalize_path, + pkey_from_str, + resolve_ssh_key, update_ssh_config, ) pytestmark = pytest.mark.windows +PRIVATE_KEY = """\ +-----BEGIN OPENSSH PRIVATE KEY----- +b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW +QyNTUxOQAAACDTpsdE/oQUWieowottWWHtjVtxGUvHtDHcJfsmSbfpAwAAAJDfsiip37Io +qQAAAAtzc2gtZWQyNTUxOQAAACDTpsdE/oQUWieowottWWHtjVtxGUvHtDHcJfsmSbfpAw +AAAEDD+JQrRu/CGiOsZTV8yXAukWWMwQeJSsRZvS36UpQRvdOmx0T+hBRaJ6jCi21ZYe2N +W3EZS8e0Mdwl+yZJt+kDAAAAC3Rlc3RAZHN0YWNrAQI= +-----END OPENSSH PRIVATE KEY----- +""" +PUBLIC_KEY_NO_COMMENT = ( + "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD" +) +PUBLIC_KEY = f"{PUBLIC_KEY_NO_COMMENT} test@dstack\n" +# A valid public key of a type paramiko cannot construct a PKey for +UNSUPPORTED_TYPE_PUBLIC_KEY = ( + "sk-ssh-ed25519@openssh.com AAAAGnNrLXNzaC1lZDI1NTE5QG9wZW5zc2guY29tAAAAIAABAgMEBQYHCAkKCwwN" + "Dg8QERITFBUWFxgZGhscHR4fAAAABHNzaDo= test@dstack\n" +) + class TestNormalizePath: @pytest.mark.skipif(IS_WINDOWS, reason="POSIX OpenSSH home semantics") @@ -105,3 +128,113 @@ def test_ssh_version_on_windows_below_8_4(self, mock_run): ) self.assertFalse(check_required_ssh_version()) + + +class TestResolveSSHKey: + def test_private_key_with_public_key_file(self, tmp_path: Path): + private_key_path = tmp_path / "id_ed25519" + private_key_path.write_text(PRIVATE_KEY) + public_key_path = tmp_path / "id_ed25519.pub" + public_key_path.write_text(PUBLIC_KEY) + + assert resolve_ssh_key(private_key_path) == ( + PUBLIC_KEY, + public_key_path, + PRIVATE_KEY, + private_key_path, + ) + + def test_returns_public_key_file_contents_as_is(self, tmp_path: Path): + private_key_path = tmp_path / "id_ed25519" + private_key_path.write_text(PRIVATE_KEY) + # A ".pub" file is not validated, its contents is passed through verbatim + (tmp_path / "id_ed25519.pub").write_text(" not a key at all ") + + public_key, _, _, _ = resolve_ssh_key(private_key_path) + + assert public_key == " not a key at all " + + def test_private_key_without_public_key_file(self, tmp_path: Path): + private_key_path = tmp_path / "id_ed25519" + private_key_path.write_text(PRIVATE_KEY) + + assert resolve_ssh_key(private_key_path) == ( + PUBLIC_KEY_NO_COMMENT, + None, + PRIVATE_KEY, + private_key_path, + ) + + def test_public_key(self, tmp_path: Path): + public_key_path = tmp_path / "id_ed25519.pub" + public_key_path.write_text(PUBLIC_KEY) + + assert resolve_ssh_key(public_key_path) == (PUBLIC_KEY, public_key_path, None, None) + + def test_public_key_of_unsupported_type(self, tmp_path: Path): + public_key_path = tmp_path / "id_sk_ed25519.pub" + public_key_path.write_text(UNSUPPORTED_TYPE_PUBLIC_KEY) + + assert resolve_ssh_key(public_key_path) == ( + UNSUPPORTED_TYPE_PUBLIC_KEY, + public_key_path, + None, + None, + ) + + def test_accepts_str_path(self, tmp_path: Path): + private_key_path = tmp_path / "id_ed25519" + private_key_path.write_text(PRIVATE_KEY) + + _, _, _, returned_path = resolve_ssh_key(str(private_key_path)) + + assert returned_path == private_key_path + + @pytest.mark.skipif(IS_WINDOWS, reason="POSIX OpenSSH home semantics") + def test_expands_user(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HOME", str(tmp_path)) + private_key_path = tmp_path / "id_ed25519" + private_key_path.write_text(PRIVATE_KEY) + + _, _, _, returned_path = resolve_ssh_key("~/id_ed25519") + + assert returned_path == private_key_path + + def test_finds_public_key_file_for_dotted_key_name(self, tmp_path: Path): + private_key_path = tmp_path / "my.key" + private_key_path.write_text(PRIVATE_KEY) + public_key_path = tmp_path / "my.key.pub" + public_key_path.write_text(PUBLIC_KEY) + + assert resolve_ssh_key(private_key_path) == ( + PUBLIC_KEY, + public_key_path, + PRIVATE_KEY, + private_key_path, + ) + + @pytest.mark.skipif(find_ssh_util("ssh-keygen") is None, reason="requires ssh-keygen") + def test_converts_pkcs8_private_key_to_pem(self, tmp_path: Path): + # dstack generates PKCS#8 keys, paramiko only reads PEM and OpenSSH ones + private_key_bytes, public_key_bytes = crypto.generate_rsa_key_pair_bytes() + private_key_path = tmp_path / "id_rsa" + private_key_path.write_bytes(private_key_bytes) + + public_key, public_key_path, private_key, _ = resolve_ssh_key(private_key_path) + + assert public_key_path is None + assert public_key == public_key_bytes.decode().rsplit(" ", 1)[0] + assert private_key.startswith("-----BEGIN RSA PRIVATE KEY-----") + assert pkey_from_str(private_key) + + @pytest.mark.parametrize("contents", ["", "garbage", "ssh-ed25519 not-base64"]) + def test_raises_on_invalid_key(self, tmp_path: Path, contents: str): + key_path = tmp_path / "id_ed25519" + key_path.write_text(contents) + + with pytest.raises(ValueError, match="Unsupported key type or invalid key"): + resolve_ssh_key(key_path) + + def test_raises_if_key_does_not_exist(self, tmp_path: Path): + with pytest.raises(OSError): + resolve_ssh_key(tmp_path / "id_ed25519") diff --git a/src/tests/api/test_runs.py b/src/tests/api/test_runs.py index 337750051..34017acf3 100644 --- a/src/tests/api/test_runs.py +++ b/src/tests/api/test_runs.py @@ -1,7 +1,12 @@ import base64 import uuid from datetime import datetime, timezone +from pathlib import Path +from typing import Optional +import pytest + +from dstack._internal.core.errors import ConfigurationError from dstack._internal.core.models.configurations import TaskConfiguration from dstack._internal.core.models.logs import JobSubmissionLogs, LogEvent, LogEventSource from dstack._internal.core.models.resources import ResourcesSpec @@ -17,6 +22,7 @@ from dstack._internal.core.models.runs import Run as RunModel from dstack._internal.server.schemas.logs import PollLogsRequest from dstack.api._public.runs import Run, RunCollection +from tests._internal.utils.test_ssh import PRIVATE_KEY, PUBLIC_KEY, PUBLIC_KEY_NO_COMMENT class _RunsAPI: @@ -35,6 +41,20 @@ def __init__(self): self.runs = _RunsAPI() +class _PlansAPI: + def __init__(self): + self.run_specs: list[RunSpec] = [] + + def get_plan(self, project_name: str, run_spec: RunSpec, **kwargs) -> str: + self.run_specs.append(run_spec) + return "run-plan" + + +class _PlansAPIClient: + def __init__(self): + self.runs = _PlansAPI() + + class TestRunCollectionList: def test_default_list_fallback_limits_job_submissions(self): api_client = _APIClient() @@ -166,3 +186,65 @@ def test_returns_logs_of_requested_replica(self): run = _get_run(run_model, logs_api) assert b"".join(run.logs(replica_num=1)) == b"replica 1\n" + + +class TestRunCollectionGetRunPlan: + def _get_run_plan( + self, ssh_identity_file: Optional[Path] = None, ssh_key_pub: Optional[str] = None + ) -> RunSpec: + api_client = _PlansAPIClient() + runs = RunCollection(api_client=api_client, project="main", client=None) + + assert ( + runs.get_run_plan( + configuration=TaskConfiguration(commands=["echo hello"], image="ubuntu:latest"), + ssh_identity_file=ssh_identity_file, + ssh_key_pub=ssh_key_pub, + ) + == "run-plan" + ) + + return api_client.runs.run_specs[0] + + 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 self._get_run_plan(private_key_path).ssh_key_pub == PUBLIC_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 self._get_run_plan(private_key_path).ssh_key_pub == PUBLIC_KEY_NO_COMMENT + + def test_accepts_public_key(self, tmp_path: Path): + public_key_path = tmp_path / "id_ed25519.pub" + public_key_path.write_text(PUBLIC_KEY) + + assert self._get_run_plan(public_key_path).ssh_key_pub == PUBLIC_KEY + + def test_uses_user_key_if_no_key_given(self): + assert self._get_run_plan().ssh_key_pub is None + + def test_uses_public_key_as_is(self): + assert self._get_run_plan(ssh_key_pub=PUBLIC_KEY).ssh_key_pub == PUBLIC_KEY + + def test_raises_if_both_public_key_and_identity_file_given(self, tmp_path: Path): + private_key_path = tmp_path / "id_ed25519" + private_key_path.write_text(PRIVATE_KEY) + + with pytest.raises(ConfigurationError, match="mutually exclusive"): + self._get_run_plan(ssh_identity_file=private_key_path, ssh_key_pub=PUBLIC_KEY) + + def test_raises_if_key_does_not_exist(self, tmp_path: Path): + with pytest.raises(ConfigurationError, match="Unable to read the SSH key"): + self._get_run_plan(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(ConfigurationError, match="Unsupported or invalid SSH key"): + self._get_run_plan(key_path)