Add API for SSH proxy (#3646)

Part-of: #3644

Author: un-def

pyproject.toml

@@ -126,6 +126,7 @@ markers = [
 ]
 env = [
     "DSTACK_CLI_RICH_FORCE_TERMINAL=0",
+    "DSTACK_SSHPROXY_API_TOKEN=test-token",
 ]
 filterwarnings = [
     # testcontainers modules use deprecated decorators – nothing we can do:
@@ -142,6 +143,7 @@ dev = [
     "pytest-httpbin>=2.1.0",
     "pytest-socket>=0.7.0",
     "pytest-env>=1.1.0",
+    "pytest-unordered>=0.7.0",
     "httpbin>=0.10.2", # indirect to make compatible with Werkzeug 3
     "requests-mock>=1.12.1",
     "openai>=1.68.2",

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -130,8 +130,8 @@ def run_job(
         commands = get_docker_commands(
             [run.run_spec.ssh_key_pub.strip(), project_ssh_public_key.strip()]
         )
-        # There is a one jump pod per Kubernetes backend that is used
-        # as an ssh proxy jump to connect to all other services in Kubernetes.
+        # 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()
         jump_pod_name = f"dstack-{run.project_name}-ssh-jump-pod"
         jump_pod_service_name = _get_pod_service_name(jump_pod_name)

src/dstack/_internal/server/app.py

@@ -45,6 +45,7 @@
     runs,
     secrets,
     server,
+    sshproxy,
     templates,
     users,
     volumes,
@@ -255,6 +256,7 @@ def register_routes(app: FastAPI, ui: bool = True):
     app.include_router(events.root_router)
     app.include_router(templates.router)
     app.include_router(exports.project_router)
+    app.include_router(sshproxy.router)
 
     @app.exception_handler(ForbiddenError)
     async def forbidden_error_handler(request: Request, exc: ForbiddenError):

src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py

@@ -59,6 +59,7 @@
     emit_job_status_change_event,
     get_job_provisioning_data,
     get_job_runtime_data,
+    get_job_spec,
 )
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
@@ -803,7 +804,7 @@ async def _detach_volumes_from_job_instance(
     jpd: JobProvisioningData,
     run_termination_reason: Optional[RunTerminationReason],
 ) -> _VolumeDetachResult:
-    job_spec = JobSpec.__response__.parse_raw(job_model.job_spec_data)
+    job_spec = get_job_spec(job_model)
     backend = await backends_services.get_project_backend_by_type(
         project=instance_model.project,
         backend_type=jpd.backend,

src/dstack/_internal/server/background/scheduled_tasks/probes.py

@@ -12,7 +12,7 @@
 from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.errors import SSHError
-from dstack._internal.core.models.runs import JobSpec, JobStatus, ProbeSpec
+from dstack._internal.core.models.runs import JobStatus, ProbeSpec
 from dstack._internal.core.services.ssh.tunnel import (
     SSH_DEFAULT_OPTIONS,
     IPSocket,
@@ -21,6 +21,7 @@
 )
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import InstanceModel, JobModel, ProbeModel
+from dstack._internal.server.services.jobs import get_job_spec
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.ssh import container_ssh_tunnel
@@ -71,7 +72,7 @@ async def process_probes():
                 if probe.job.status != JobStatus.RUNNING:
                     probe.active = False
                 else:
-                    job_spec: JobSpec = JobSpec.__response__.parse_raw(probe.job.job_spec_data)
+                    job_spec = get_job_spec(probe.job)
                     probe_spec = job_spec.probes[probe.probe_num]
                     if probe_spec.until_ready and probe.success_streak >= probe_spec.ready_after:
                         probe.active = False
@@ -148,7 +149,7 @@ async def _get_service_replica_client(job: JobModel) -> AsyncGenerator[AsyncClie
         **SSH_DEFAULT_OPTIONS,
         "ConnectTimeout": str(int(SSH_CONNECT_TIMEOUT.total_seconds())),
     }
-    job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
+    job_spec = get_job_spec(job)
     with TemporaryDirectory() as temp_dir:
         app_socket_path = (Path(temp_dir) / "replica.sock").absolute()
         async with container_ssh_tunnel(

src/dstack/_internal/server/background/scheduled_tasks/running_jobs.py

@@ -34,7 +34,6 @@
     JobTerminationReason,
     ProbeSpec,
     Run,
-    RunSpec,
     RunStatus,
 )
 from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, VolumeMountPoint
@@ -67,6 +66,7 @@
     find_job,
     get_job_attached_volumes,
     get_job_runtime_data,
+    get_job_spec,
     is_master_job,
     job_model_to_job_submission,
     switch_job_status,
@@ -82,6 +82,7 @@
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
 from dstack._internal.server.services.runs import (
+    get_run_spec,
     is_job_ready,
     run_model_to_run,
 )
@@ -734,7 +735,7 @@ def _process_provisioning_with_shim(
     Returns:
         is successful
     """
-    job_spec = JobSpec.__response__.parse_raw(job_model.job_spec_data)
+    job_spec = get_job_spec(job_model)
 
     shim_client = client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
 
@@ -982,7 +983,7 @@ def _terminate_if_inactivity_duration_exceeded(
     job_model: JobModel,
     no_connections_secs: Optional[int],
 ) -> None:
-    conf = RunSpec.__response__.parse_raw(run_model.run_spec).configuration
+    conf = get_run_spec(run_model).configuration
     if not isinstance(conf, DevEnvironmentConfiguration) or not isinstance(
         conf.inactivity_duration, int
     ):

src/dstack/_internal/server/background/scheduled_tasks/runs.py

@@ -13,7 +13,6 @@
 from dstack._internal.core.models.profiles import RetryEvent, StopCriteria
 from dstack._internal.core.models.runs import (
     Job,
-    JobSpec,
     JobStatus,
     JobTerminationReason,
     Run,
@@ -33,6 +32,7 @@
 from dstack._internal.server.services import events
 from dstack._internal.server.services.jobs import (
     find_job,
+    get_job_spec,
     get_job_specs_from_run_spec,
     group_jobs_by_replica_latest,
     is_master_job,
@@ -531,7 +531,7 @@ async def _handle_run_replicas(
             if job.status.is_finished():
                 continue
             try:
-                job_spec = JobSpec.__response__.parse_raw(job.job_spec_data)
+                job_spec = get_job_spec(job)
                 existing_group_names.add(job_spec.replica_group)
             except Exception:
                 continue
@@ -647,7 +647,7 @@ async def _update_jobs_to_new_deployment_in_place(
         replica_group_name = None
 
         if replicas:
-            job_spec = JobSpec.__response__.parse_raw(job_models[0].job_spec_data)
+            job_spec = get_job_spec(job_models[0])
             replica_group_name = job_spec.replica_group
 
         # FIXME: Handle getting image configuration errors or skip it.
@@ -662,7 +662,7 @@ async def _update_jobs_to_new_deployment_in_place(
         )
         can_update_all_jobs = True
         for old_job_model, new_job_spec in zip(job_models, new_job_specs):
-            old_job_spec = JobSpec.__response__.parse_raw(old_job_model.job_spec_data)
+            old_job_spec = get_job_spec(old_job_model)
             if new_job_spec != old_job_spec:
                 can_update_all_jobs = False
                 break

src/dstack/_internal/server/background/scheduled_tasks/terminating_jobs.py

@@ -39,6 +39,7 @@
 from dstack._internal.server.services.jobs import (
     get_job_provisioning_data,
     get_job_runtime_data,
+    get_job_spec,
     switch_job_status,
 )
 from dstack._internal.server.services.locking import get_locker
@@ -356,7 +357,7 @@ async def _detach_volumes_from_job_instance(
     instance_model: InstanceModel,
     volume_models: list[VolumeModel],
 ) -> bool:
-    job_spec = JobSpec.__response__.parse_raw(job_model.job_spec_data)
+    job_spec = get_job_spec(job_model)
     backend = await backends_services.get_project_backend_by_type(
         project=project,
         backend_type=jpd.backend,

src/dstack/_internal/server/routers/sshproxy.py

@@ -0,0 +1,39 @@
+import os
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.errors import ResourceNotExistsError
+from dstack._internal.server.db import get_session
+from dstack._internal.server.schemas.sshproxy import GetUpstreamRequest, GetUpstreamResponse
+from dstack._internal.server.security.permissions import AlwaysForbidden, ServiceAccount
+from dstack._internal.server.services.sshproxy import get_upstream_response
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
+
+if _token := os.getenv("DSTACK_SSHPROXY_API_TOKEN"):
+    _auth = ServiceAccount(_token)
+else:
+    _auth = AlwaysForbidden()
+
+
+router = APIRouter(
+    prefix="/api/sshproxy",
+    tags=["sshproxy"],
+    responses=get_base_api_additional_responses(),
+    dependencies=[Depends(_auth)],
+)
+
+
+@router.post("/get_upstream", response_model=GetUpstreamResponse)
+async def get_upstream(
+    body: GetUpstreamRequest,
+    session: Annotated[AsyncSession, Depends(get_session)],
+):
+    response = await get_upstream_response(session=session, upstream_id=body.id)
+    if response is None:
+        raise ResourceNotExistsError()
+    return CustomORJSONResponse(response)

src/dstack/_internal/server/schemas/sshproxy.py

@@ -0,0 +1,27 @@
+from typing import Annotated
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class GetUpstreamRequest(CoreModel):
+    # The format of id is intentionally not limited to UUID to allow further extensions
+    id: str
+
+
+class UpstreamHost(CoreModel):
+    host: Annotated[str, Field(description="The hostname or IP address")]
+    port: Annotated[int, Field(description="The SSH port")]
+    user: Annotated[str, Field(description="The user to log in")]
+    private_key: Annotated[str, Field(description="The private key in OpenSSH file format")]
+
+
+class GetUpstreamResponse(CoreModel):
+    hosts: Annotated[
+        list[UpstreamHost],
+        Field(description="The chain of SSH hosts, the jump host(s) first, the target host last"),
+    ]
+    authorized_keys: Annotated[
+        list[str], Field(description="The list of authorized public keys in OpenSSH file format")
+    ]