From 497cac2b198da76b42a1a0d5ed60b715c4177696 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Tue, 1 Sep 2026 14:57:04 +0000 Subject: [PATCH] Honor sshproxy enforcement in container-based backends VM-based backends receive the container's authorized keys uniformly, via dstack-shim task submission, which already omits the user key when DSTACK_SERVER_SSHPROXY_ENFORCED confines users to the SSH proxy. Container-based backends built the list ad-hoc inside run_job() and run_jobs() from run_spec.ssh_key_pub, so the user key reached the container regardless of the setting. Add an extra_authorized_keys argument to Compute.run_job() and run_jobs(). The caller builds it once with get_extra_authorized_keys(), which applies the setting; Kubernetes, Runpod, Slurm and Vast.ai combine it with the project key through the new build_authorized_keys(), which requires the project key to be valid and skips extra keys that cannot be parsed. The argument is deliberately unvalidated at the call site, as most backends are shim-based and ignore it. The shim path now derives its key list from the same helper so that the two cannot drift. --- .../core/backends/base/authorized_keys.py | 52 +++++++++++++++ .../_internal/core/backends/base/compute.py | 20 ++++++ .../core/backends/kubernetes/compute.py | 66 +++++++++++++------ .../_internal/core/backends/runpod/compute.py | 19 ++---- .../_internal/core/backends/slurm/compute.py | 10 ++- .../_internal/core/backends/vastai/compute.py | 5 +- .../background/pipeline_tasks/jobs_running.py | 10 +-- .../pipeline_tasks/jobs_submitted.py | 4 ++ .../server/services/jobs/__init__.py | 16 +++++ .../core/backends/vastai/test_compute.py | 10 ++- 10 files changed, 171 insertions(+), 41 deletions(-) diff --git a/src/dstack/_internal/core/backends/base/authorized_keys.py b/src/dstack/_internal/core/backends/base/authorized_keys.py index 15fcad132..3c72604ee 100644 --- a/src/dstack/_internal/core/backends/base/authorized_keys.py +++ b/src/dstack/_internal/core/backends/base/authorized_keys.py @@ -2,6 +2,7 @@ from typing import Optional from uuid import uuid4 +from dstack._internal.core.errors import ComputeError from dstack._internal.utils.logging import get_logger from dstack._internal.utils.ssh import parse_public_key @@ -20,6 +21,57 @@ DSTACK_PUBLIC_KEY_MARKER = "# added by dstack" +def build_authorized_keys( + project_ssh_public_key: str, + extra_authorized_keys: list[str], +) -> list[str]: + """ + Builds the list of public keys to authorize on a job container. + + Args: + project_ssh_public_key: The project public key, always authorized. Must be valid -- + the server cannot reach the container without it. + extra_authorized_keys: The public keys to authorize in addition to the project key, + as passed to `Compute.run_job()`. Untrusted -- keys that cannot be parsed are + skipped with a warning, one bad key does not keep the rest out. + + Returns: + The normalized keys in OpenSSH disk format, the project key first. + + Raises: + ComputeError: The project key is invalid. + """ + project_authorized_keys = normalize_authorized_keys([project_ssh_public_key]) + if not project_authorized_keys: + raise ComputeError("Invalid project SSH key") + return project_authorized_keys + normalize_authorized_keys(extra_authorized_keys) + + +def normalize_authorized_keys(authorized_keys: list[str]) -> list[str]: + """ + Rebuilds the given public keys from their parsed fields. + + Every returned entry is a single `type blob [comment]` line with the comment whitespace + normalized, so that nothing unvalidated reaches a command or a file. + + Args: + authorized_keys: The public keys in OpenSSH disk format. + + Returns: + The normalized keys, in the original order. Keys that cannot be parsed are skipped with + a warning, therefore the result may be shorter than the input, or empty. + """ + normalized: list[str] = [] + for authorized_key in authorized_keys: + try: + key = parse_public_key(authorized_key) + except ValueError as e: + logger.warning("Failed to parse authorized key: %r: %s", authorized_key, e) + continue + normalized.append(str(key)) + return normalized + + def get_add_authorized_keys_script( authorized_keys: list[str], *, diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 234b3a309..6062acc3c 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -148,11 +148,21 @@ def run_job( volumes: List[Volume], placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> JobProvisioningData: """ Launches a new instance for the job. It should return `JobProvisioningData` ASAP. If required to wait to get the IP address or SSH port, return partially filled `JobProvisioningData` and implement `update_provisioning_data()`. + + `extra_authorized_keys` are the public keys to authorize on the job container in addition + to `project_ssh_public_key`, as decided by the caller -- typically the user key, or + nothing if the server does not let the user connect to the container directly. The project + key is never among them, and the keys are not validated; pass them to + `base.authorized_keys.build_authorized_keys()` to get the complete, validated list to + authorize. Only Computes that add the keys themselves need this argument; VM-based + (shim-based) Computes ignore it, as the server submits the keys to the shim once the + instance is up. """ pass @@ -390,10 +400,14 @@ def run_job( volumes: List[Volume], placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> JobProvisioningData: """ The default `run_job()` implementation for all backends that support `create_instance()`. Override only if custom `run_job()` behavior is required. + + `extra_authorized_keys` is ignored -- all such backends are VM-based, and the server + submits the keys to the shim later, see `Compute.run_job()`. """ instance_config = InstanceConfiguration( project_name=run.project_name, @@ -442,7 +456,13 @@ def run_jobs( project_ssh_private_key: str, placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> ComputeGroupProvisioningData: + """ + Launches a compute group -- instances created all at once via the provider API -- running + one job per instance. See `Compute.run_job()` for the arguments shared with it, including + `extra_authorized_keys`. + """ pass @abstractmethod diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index d21784f3e..cc159b8bd 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -15,7 +15,9 @@ from typing_extensions import Self from dstack._internal.core.backends.base.authorized_keys import ( + build_authorized_keys, get_add_authorized_keys_script, + normalize_authorized_keys, ) from dstack._internal.core.backends.base.compute import ( Compute, @@ -126,7 +128,15 @@ class Operator(str, Enum): class KubernetesBackendData(CoreModel): jump_pod_name: str jump_pod_service_name: str - user_ssh_public_key: str + # TODO: remove `user_ssh_public_key` once servers before 0.21.4 are no longer supported. + user_ssh_public_key: str = "" + """Superseded by `extra_authorized_keys`, but still written -- as the first extra key, or an + empty string if there are none -- because servers before 0.21.4 require this field and may + read this record during a rolling deployment. Read only if `extra_authorized_keys` is empty, + that is, if such a server wrote the record. + """ + extra_authorized_keys: list[str] = [] + """Empty in records written before the field was introduced, see `user_ssh_public_key`.""" @classmethod def load(cls, raw: str) -> Self: @@ -193,6 +203,7 @@ def run_job( volumes: list[Volume], placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> JobProvisioningData: cluster = self.region_cluster_map.get(instance_offer.region) if cluster is None: @@ -202,6 +213,13 @@ def run_job( api = client.CoreV1Api(cluster.api_client) namespace = cluster.namespace + # The jump pod, created below only if it does not exist yet, gets the project key alone. + # This job's extra keys are added to it later by update_provisioning_data(), which runs + # for every job, whether or not the job created the pod. The job pod gets both at once. + project_authorized_keys = build_authorized_keys(project_ssh_public_key, []) + extra_authorized_keys = normalize_authorized_keys(extra_authorized_keys) + authorized_keys = project_authorized_keys + extra_authorized_keys + # There is one jump pod per project that is used as an ssh proxy jump to connect # to all job pods of the same project. # The service is created here and configured later in update_provisioning_data() @@ -214,7 +232,7 @@ def run_job( jump_pod_name=jump_pod_name, jump_pod_service_name=jump_pod_service_name, jump_pod_port=cluster.proxy_jump.port, - project_ssh_public_key=project_ssh_public_key.strip(), + authorized_keys=project_authorized_keys, ) pod_name = generate_unique_instance_name_for_job( @@ -256,8 +274,6 @@ def run_job( should_delete_manually_if_failed=True, ) - assert run.run_spec.ssh_key_pub is not None - authorized_keys = [run.run_spec.ssh_key_pub.strip(), project_ssh_public_key.strip()] _create_job_pod( api=api, namespace=namespace, @@ -329,7 +345,9 @@ def run_job( backend_data = KubernetesBackendData( jump_pod_name=jump_pod_name, jump_pod_service_name=jump_pod_service_name, - user_ssh_public_key=run.run_spec.ssh_key_pub.strip(), + extra_authorized_keys=extra_authorized_keys, + # For compatibility with server replicas < 0.21.4 + user_ssh_public_key=next(iter(extra_authorized_keys), ""), ) return JobProvisioningData( @@ -369,6 +387,11 @@ def update_provisioning_data( if provisioning_data.backend_data is not None: # Before running a job, ensure the jump pod is running and has user's public SSH key. backend_data = KubernetesBackendData.load(provisioning_data.backend_data) + extra_authorized_keys = backend_data.extra_authorized_keys + # TODO: remove the fallback once servers before 0.21.4 are no longer supported. + if not extra_authorized_keys and backend_data.user_ssh_public_key: + # Written by a server that predates the `extra_authorized_keys` field. + extra_authorized_keys = [backend_data.user_ssh_public_key] ssh_proxy = _check_and_configure_jump_pod_service( api=api, namespace=namespace, @@ -376,7 +399,7 @@ def update_provisioning_data( jump_pod_service_name=backend_data.jump_pod_service_name, jump_pod_hostname=cluster.proxy_jump.hostname, project_ssh_private_key=project_ssh_private_key, - user_ssh_public_key=backend_data.user_ssh_public_key, + authorized_keys=extra_authorized_keys, ) if ssh_proxy is None: # Jump pod is not ready yet @@ -837,7 +860,7 @@ def _create_jump_pod_service_if_not_exists( jump_pod_name: str, jump_pod_service_name: str, jump_pod_port: Optional[int], - project_ssh_public_key: str, + authorized_keys: list[str], ) -> None: base_labels = build_base_labels( component="ssh-proxy", @@ -914,7 +937,7 @@ def _create_jump_pod_service_if_not_exists( ) if not tolerations: logger.warning("No appropriate node found, the jump pod may never be scheduled") - commands = _get_jump_pod_commands(authorized_keys=[project_ssh_public_key]) + commands = _get_jump_pod_commands(authorized_keys) pod = client.V1Pod( metadata=client.V1ObjectMeta( name=jump_pod_name, @@ -977,7 +1000,7 @@ def _check_and_configure_jump_pod_service( jump_pod_service_name: str, jump_pod_hostname: Optional[str], project_ssh_private_key: str, - user_ssh_public_key: str, + authorized_keys: list[str], ) -> Optional[SSHConnectionParams]: jump_pod = api.read_namespaced_pod( namespace=namespace, @@ -1031,13 +1054,9 @@ def _check_and_configure_jump_pod_service( if (jump_pod_port := jump_pod_service_ports[0].node_port) is None: raise ProvisioningError("Jump pod service %s port is not set", jump_pod_service_name) - ssh_exit_status, ssh_output = _run_ssh_command( - hostname=jump_pod_hostname, - port=jump_pod_port, - username=JUMP_POD_USER, - ssh_private_key=project_ssh_private_key, - command=get_add_authorized_keys_script( - [user_ssh_public_key], + if authorized_keys: + command = get_add_authorized_keys_script( + authorized_keys, # The project key, written by _get_jump_pod_commands(), carries no marker, and # marking only some of our entries defeats the purpose of the marker. It is # redundant on the jump pod anyway -- every entry there is ours @@ -1045,7 +1064,16 @@ def _check_and_configure_jump_pod_service( # command= in authorized_keys is equivalent to ForceCommand in sshd_config # By forcing the /bin/false command we only allow proxy jumping, no shell access options='command="/bin/false"', - ), + ) + else: + # No keys to add, but the connection itself is still the check that sshd is up + command = "true" + ssh_exit_status, ssh_output = _run_ssh_command( + hostname=jump_pod_hostname, + port=jump_pod_port, + username=JUMP_POD_USER, + ssh_private_key=project_ssh_private_key, + command=command, ) if ssh_exit_status != 0: logger.debug( @@ -1072,11 +1100,11 @@ def _check_and_configure_jump_pod_service( def _get_jump_pod_commands(authorized_keys: list[str]) -> list[str]: - authorized_keys_content = "\n".join(authorized_keys).strip() + authorized_keys_content = "\n".join(authorized_keys) commands = [ "mkdir -p ~/.ssh", "chmod 700 ~/.ssh", - f"echo '{authorized_keys_content}' > ~/.ssh/authorized_keys", + f"echo {shlex.quote(authorized_keys_content)} > ~/.ssh/authorized_keys", "chmod 600 ~/.ssh/authorized_keys", # regenerate host keys "rm -rf /etc/ssh/ssh_host_*", diff --git a/src/dstack/_internal/core/backends/runpod/compute.py b/src/dstack/_internal/core/backends/runpod/compute.py index 390b42600..63d677f99 100644 --- a/src/dstack/_internal/core/backends/runpod/compute.py +++ b/src/dstack/_internal/core/backends/runpod/compute.py @@ -4,6 +4,7 @@ from datetime import timedelta from typing import Callable, List, Optional +from dstack._internal.core.backends.base.authorized_keys import build_authorized_keys from dstack._internal.core.backends.base.backend import Compute from dstack._internal.core.backends.base.compute import ( ComputeWithAllOffersCached, @@ -34,7 +35,6 @@ InstanceAvailability, InstanceConfiguration, InstanceOfferWithAvailability, - SSHKey, ) from dstack._internal.core.models.placement import PlacementGroup from dstack._internal.core.models.resources import Memory, Range @@ -135,20 +135,17 @@ def run_job( volumes: List[Volume], placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> JobProvisioningData: - assert run.run_spec.ssh_key_pub is not None instance_config = InstanceConfiguration( project_name=run.project_name, instance_name=get_job_instance_name(run, job), - ssh_keys=[ - SSHKey(public=run.run_spec.ssh_key_pub.strip()), - SSHKey(public=project_ssh_public_key.strip()), - ], + ssh_keys=[], user=run.user, ) pod_name = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN) - authorized_keys = instance_config.get_public_keys() + authorized_keys = build_authorized_keys(project_ssh_public_key, extra_authorized_keys) memory_size = round(instance_offer.instance.resources.memory_mib / 1024) disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024) @@ -251,6 +248,7 @@ def run_jobs( project_ssh_private_key: str, placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> ComputeGroupProvisioningData: master_job_configuration = job_configurations[0] master_job = master_job_configuration.job @@ -259,15 +257,12 @@ def run_jobs( instance_config = InstanceConfiguration( project_name=run.project_name, instance_name=get_job_instance_name(run, master_job), - ssh_keys=[ - SSHKey(public=get_or_error(run.run_spec.ssh_key_pub).strip()), - SSHKey(public=project_ssh_public_key.strip()), - ], + ssh_keys=[], user=run.user, ) pod_name = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN) - authorized_keys = instance_config.get_public_keys() + authorized_keys = build_authorized_keys(project_ssh_public_key, extra_authorized_keys) disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024) network_volume_id = None diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index 2fa9ab86b..44dfa8cf1 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -6,6 +6,7 @@ from typing import Optional from dstack._internal.core.backends.base.authorized_keys import ( + build_authorized_keys, get_add_authorized_keys_script, ) from dstack._internal.core.backends.base.compute import ( @@ -140,6 +141,7 @@ def run_job( volumes: list[Volume], placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> JobProvisioningData: # run_job provisions a single dstack job → one Slurm node. Do not fall # back to jobs_per_replica (total across hetero groups). @@ -148,6 +150,7 @@ def run_job( job=job, instance_offer=instance_offer, project_ssh_public_key=project_ssh_public_key, + extra_authorized_keys=extra_authorized_keys, requirements=requirements, node_count=1, ) @@ -162,6 +165,7 @@ def run_jobs( project_ssh_private_key: str, placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> ComputeGroupProvisioningData: master_job = job_configurations[0].job return self._run_slurm_job( @@ -169,6 +173,7 @@ def run_jobs( job=master_job, instance_offer=instance_offer, project_ssh_public_key=project_ssh_public_key, + extra_authorized_keys=extra_authorized_keys, requirements=requirements, node_count=len(job_configurations), ) @@ -192,6 +197,7 @@ def _run_slurm_job( job: Job, instance_offer: InstanceOfferWithAvailability, project_ssh_public_key: str, + extra_authorized_keys: list[str], requirements: Requirements, node_count: int, ) -> ComputeGroupProvisioningData: @@ -214,8 +220,8 @@ def _run_slurm_job( max_length=SLURM_JOB_NAME_MAX_LENGTH, ) - assert run.run_spec.ssh_key_pub is not None - authorized_keys = [project_ssh_public_key.strip(), run.run_spec.ssh_key_pub.strip()] + # The same keys go to the login node, for proxy jumping, and to the container + authorized_keys = build_authorized_keys(project_ssh_public_key, extra_authorized_keys) # Slurm --nodes for this call (1 from run_job, len(batch) from run_jobs). resources_spec = requirements.resources diff --git a/src/dstack/_internal/core/backends/vastai/compute.py b/src/dstack/_internal/core/backends/vastai/compute.py index a9e080d41..aff025e82 100644 --- a/src/dstack/_internal/core/backends/vastai/compute.py +++ b/src/dstack/_internal/core/backends/vastai/compute.py @@ -5,6 +5,7 @@ from gpuhunt.providers.vastai import VastAIProvider from typing_extensions import assert_never +from dstack._internal.core.backends.base.authorized_keys import build_authorized_keys from dstack._internal.core.backends.base.backend import Compute from dstack._internal.core.backends.base.compute import ( ComputeWithFilteredOffersCached, @@ -123,13 +124,13 @@ def run_job( volumes: List[Volume], placement_group: Optional[PlacementGroup], requirements: Requirements, + extra_authorized_keys: list[str], ) -> JobProvisioningData: instance_name = generate_unique_instance_name_for_job( run, job, max_length=MAX_INSTANCE_NAME_LEN ) - assert run.run_spec.ssh_key_pub is not None commands = get_docker_commands( - [run.run_spec.ssh_key_pub.strip(), project_ssh_public_key.strip()] + build_authorized_keys(project_ssh_public_key, extra_authorized_keys) ) offer_backend_data = validate_extra_ignore( VastAIOfferBackendData, instance_offer.backend_data diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py index 0a4b90bac..2cfc889ff 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -91,6 +91,7 @@ from dstack._internal.server.services.jobs import ( emit_job_status_change_event, find_job, + get_extra_authorized_keys, get_job_attached_volumes, get_job_runtime_data, get_job_spec, @@ -816,14 +817,15 @@ async def _process_provisioning_status( fmt(context.job_model), context.job_submission.age, ) - public_keys = [context.project.ssh_public_key.strip()] + extra_authorized_keys = get_extra_authorized_keys(context.run.run_spec) + public_keys = [context.project.ssh_public_key.strip(), *extra_authorized_keys] + # Host access, unlike container access, is all or nothing -- the user key is added to + # the host only if they are allowed to bypass the SSH proxy ssh_user: Optional[str] = None user_ssh_key: Optional[str] = None if not server_settings.SSHPROXY_ENFORCED: ssh_user = job_provisioning_data.username - assert context.run.run_spec.ssh_key_pub is not None - user_ssh_key = context.run.run_spec.ssh_key_pub.strip() - public_keys.append(user_ssh_key) + user_ssh_key = get_or_error(context.run.run_spec.ssh_key_pub).strip() success = await run_async( _process_provisioning_with_shim, server_ssh_private_keys, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index 8923e5f06..4a1796847 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -104,6 +104,7 @@ check_can_attach_job_volumes, find_job, find_jobs, + get_extra_authorized_keys, get_job_configured_volume_models, get_job_configured_volumes, get_job_runtime_data, @@ -2418,6 +2419,7 @@ async def _provision_new_capacity( instance_mounts=check_run_spec_requires_instance_mounts(run.run_spec), placement_group=placement_group_model_to_placement_group_optional(placement_group_model), ) + extra_authorized_keys = get_extra_authorized_keys(run.run_spec) offers_iter = iter(offers) offers_tried = 0 offers_taken = 0 @@ -2492,6 +2494,7 @@ async def _provision_new_capacity( project_ssh_private_key, placement_group_model_to_placement_group_optional(placement_group_model), requirements, + extra_authorized_keys, ) return _ProvisionNewCapacityResult( provisioning_data=compute_group_provisioning_data, @@ -2516,6 +2519,7 @@ async def _provision_new_capacity( offer_volumes, placement_group_model_to_placement_group_optional(placement_group_model), requirements, + extra_authorized_keys, ) return _ProvisionNewCapacityResult( provisioning_data=job_provisioning_data, diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index e9f8eeca2..dafacc2ee 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -288,6 +288,22 @@ def get_job_runtime_data(job_model: JobModel) -> Optional[JobRuntimeData]: return validate_json_extra_ignore(JobRuntimeData, job_model.job_runtime_data) +def get_extra_authorized_keys(run_spec: RunSpec) -> list[str]: + """ + Returns the public keys, besides the project key, to authorize on a job's container. + + The user key is included unless `DSTACK_SERVER_SSHPROXY_ENFORCED` is set, in which case the + user is only let in through the SSH proxy, which authenticates them itself and connects with + the project key. + + The keys are unvalidated -- whoever writes them to a container must normalize them first, + see `backends.base.authorized_keys`. + """ + if settings.SSHPROXY_ENFORCED: + return [] + return [common.get_or_error(run_spec.ssh_key_pub).strip()] + + def _get_image_pull_progress(job_model: JobModel) -> Optional[ImagePullProgress]: if job_model.image_pull_progress is None: return None diff --git a/src/tests/_internal/core/backends/vastai/test_compute.py b/src/tests/_internal/core/backends/vastai/test_compute.py index ee5b823d8..3c6b1c71f 100644 --- a/src/tests/_internal/core/backends/vastai/test_compute.py +++ b/src/tests/_internal/core/backends/vastai/test_compute.py @@ -47,9 +47,14 @@ def _offer( ) +# build_authorized_keys() rejects the project key unless it actually parses +PROJECT_SSH_PUBLIC_KEY = ( + "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD project" +) + + def _run_job(compute: VastAICompute, offer: InstanceOfferWithAvailability): run = MagicMock() - run.run_spec.ssh_key_pub = "ssh-rsa AAAA test" job = MagicMock() job.job_spec.image_name = "dstackai/base:latest" job.job_spec.registry_auth = None @@ -67,11 +72,12 @@ def _run_job(compute: VastAICompute, offer: InstanceOfferWithAvailability): run=run, job=job, instance_offer=offer, - project_ssh_public_key="ssh-rsa BBBB project", + project_ssh_public_key=PROJECT_SSH_PUBLIC_KEY, project_ssh_private_key="private-key", volumes=[], placement_group=None, requirements=_requirements(), + extra_authorized_keys=[], )