From 979de8ef1ac3dc65bfbc741a42f14114faadef0b Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Sun, 12 Jul 2026 00:40:29 +0200 Subject: [PATCH 01/14] Add server access to tasks and dev environments --- mkdocs/docs/concepts/dev-environments.md | 20 ++ mkdocs/docs/concepts/tasks.md | 19 ++ mkdocs/docs/reference/env.md | 11 + pyproject.toml | 2 +- src/dstack/_internal/core/consts.py | 8 + .../_internal/core/models/configurations.py | 14 ++ .../_internal/core/services/api_client.py | 12 + src/dstack/_internal/server/app.py | 2 + .../background/pipeline_tasks/jobs_running.py | 45 ++++ .../pipeline_tasks/jobs_terminating.py | 4 + .../services/jobs/configurators/base.py | 12 +- .../server/services/jobs/server_connection.py | 207 ++++++++++++++++++ .../server/services/runner/client.py | 6 +- src/dstack/_internal/server/services/ssh.py | 10 +- src/dstack/api/server/__init__.py | 3 + .../core/models/test_configurations.py | 33 +++ .../core/services/test_api_client.py | 36 +++ .../pipeline_tasks/test_running_jobs.py | 76 +++++++ .../pipeline_tasks/test_terminating_jobs.py | 40 ++++ .../services/jobs/configurators/test_task.py | 41 ++++ .../services/jobs/test_server_connection.py | 163 ++++++++++++++ .../server/services/runner/test_client.py | 74 ++++++- src/tests/api/test_client.py | 30 +++ 23 files changed, 861 insertions(+), 7 deletions(-) create mode 100644 src/dstack/_internal/server/services/jobs/server_connection.py create mode 100644 src/tests/_internal/core/services/test_api_client.py create mode 100644 src/tests/_internal/server/services/jobs/test_server_connection.py create mode 100644 src/tests/api/test_client.py diff --git a/mkdocs/docs/concepts/dev-environments.md b/mkdocs/docs/concepts/dev-environments.md index 2e4bb73f11..113909568b 100644 --- a/mkdocs/docs/concepts/dev-environments.md +++ b/mkdocs/docs/concepts/dev-environments.md @@ -642,6 +642,26 @@ The `schedule` property can be combined with `max_duration` or `utilization_poli By default, `dstack` uses on-demand instances. However, you can change that via the [`spot_policy`](../reference/dstack.yml/dev-environment.md#spot_policy) property. It accepts `spot`, `on-demand`, and `auto`. +### Server access + +Set `server` to `true` when a dev environment needs to use the dstack CLI. dstack configures the +server and current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` +explicitly. + +
+ +```yaml +type: dev-environment +image: dstackai/dstack +server: true +env: + - DSTACK_TOKEN +init: + - dstack ps +``` + +
+ --8<-- "docs/concepts/snippets/manage-fleets.ext" !!! info "Reference" diff --git a/mkdocs/docs/concepts/tasks.md b/mkdocs/docs/concepts/tasks.md index d3a3515b82..2d61b5838a 100644 --- a/mkdocs/docs/concepts/tasks.md +++ b/mkdocs/docs/concepts/tasks.md @@ -864,6 +864,25 @@ schedule: By default, `dstack` uses on-demand instances. However, you can change that via the [`spot_policy`](../reference/dstack.yml/task.md#spot_policy) property. It accepts `spot`, `on-demand`, and `auto`. +### Server access + +Set `server` to `true` when a task needs to use the dstack CLI. dstack configures the server and +current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` explicitly. + +
+ +```yaml +type: task +image: dstackai/dstack +server: true +env: + - DSTACK_TOKEN +commands: + - dstack ps +``` + +
+ --8<-- "docs/concepts/snippets/manage-fleets.ext" !!! info "Reference" diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 86a7dd051d..36a9f65d1c 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -168,6 +168,17 @@ slows down processing and may cause CPU spikes due to frequent SSH-connection es The following environment variables are supported by the CLI. +- `DSTACK_TOKEN`{ #DSTACK_TOKEN } – The user token used by the CLI. Set `DSTACK_TOKEN`, + `DSTACK_SERVER_URL`, and `DSTACK_PROJECT` together to use the CLI without a project in + `~/.dstack/config.yml`, or to override the configured server, project, and user. + + ```shell + DSTACK_SERVER_URL=https://server.example.com \ + DSTACK_PROJECT=main \ + DSTACK_TOKEN=your-token \ + dstack ps + ``` + - `DSTACK_CLI_LOG_LEVEL`{ #DSTACK_CLI_LOG_LEVEL } – Sets the logging level for CLI output to stdout. Defaults to `INFO`. Example: diff --git a/pyproject.toml b/pyproject.toml index b61f974ed3..0ed8a8c21b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ classifiers = [ dependencies = [ "pyyaml", "requests", + "requests-unixsocket>=0.4.1", "typing-extensions>=4.0.0", "cryptography", "packaging", @@ -187,7 +188,6 @@ server = [ "aiorwlock", "aiocache", "httpx>=0.28.0", - "requests-unixsocket>=0.4.1", "jinja2", "watchfiles", "sqlalchemy[asyncio]>=2.0.0", diff --git a/src/dstack/_internal/core/consts.py b/src/dstack/_internal/core/consts.py index cb68f9be82..9b60381bdc 100644 --- a/src/dstack/_internal/core/consts.py +++ b/src/dstack/_internal/core/consts.py @@ -1,8 +1,16 @@ +from urllib.parse import quote + # shim (runs on the host) HTTP API port DSTACK_SHIM_HTTP_PORT = 10998 # runner (runs inside a container) HTTP API port DSTACK_RUNNER_HTTP_PORT = 10999 # ssh server (runs alongside the runner inside a container) listen port DSTACK_RUNNER_SSH_PORT = 10022 +# Private socket created inside jobs that request access to the dstack server. +DSTACK_RUN_SERVER_SOCKET_PATH = "/run/dstack/server.sock" +DSTACK_RUN_SERVER_URL = f"http+unix://{quote(DSTACK_RUN_SERVER_SOCKET_PATH, safe='')}" +DSTACK_PROJECT_ENV = "DSTACK_PROJECT" +DSTACK_SERVER_URL_ENV = "DSTACK_SERVER_URL" +DSTACK_TOKEN_ENV = "DSTACK_TOKEN" # legacy AWS, Azure, GCP, and OCI image for older GPUs DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES = "0.10" diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 2aab13233e..3d43ef6fe0 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -693,6 +693,18 @@ def check_image_or_commands_present(cls, values): return values +class ConfigurationWithServerParams(CoreModel): + server: Annotated[ + bool, + Field( + description=( + "Make the dstack server accessible inside the run. " + "No authentication credentials are provided" + ) + ), + ] = False + + class DevEnvironmentConfigurationParams(CoreModel): ide: Annotated[ Optional[Union[Literal["vscode"], Literal["cursor"], Literal["windsurf"], Literal["zed"]]], @@ -762,6 +774,7 @@ class DevEnvironmentConfiguration( ProfileParams, BaseRunConfiguration, ConfigurationWithPortsParams, + ConfigurationWithServerParams, DevEnvironmentConfigurationParams, generate_dual_core_model(DevEnvironmentConfigurationConfig), ): @@ -793,6 +806,7 @@ class TaskConfiguration( BaseRunConfiguration, ConfigurationWithCommandsParams, ConfigurationWithPortsParams, + ConfigurationWithServerParams, TaskConfigurationParams, generate_dual_core_model(TaskConfigurationConfig), ): diff --git a/src/dstack/_internal/core/services/api_client.py b/src/dstack/_internal/core/services/api_client.py index a3a2c097b2..f1718be97b 100644 --- a/src/dstack/_internal/core/services/api_client.py +++ b/src/dstack/_internal/core/services/api_client.py @@ -1,11 +1,23 @@ +import os from typing import Optional, Tuple import dstack._internal.core.services.configs as configs +from dstack._internal.core.consts import ( + DSTACK_PROJECT_ENV, + DSTACK_SERVER_URL_ENV, + DSTACK_TOKEN_ENV, +) from dstack._internal.core.errors import ConfigurationError from dstack.api.server import APIClient def get_api_client(project_name: Optional[str] = None) -> Tuple[APIClient, str]: + env_project_name = project_name or os.getenv(DSTACK_PROJECT_ENV) + server_url = os.getenv(DSTACK_SERVER_URL_ENV) + token = os.getenv(DSTACK_TOKEN_ENV) + if env_project_name is not None and server_url is not None and token is not None: + return APIClient(server_url, token), env_project_name + config = configs.ConfigManager() project = config.get_project_config(project_name) if project is None: diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py index 0f02806aa4..eb701673e5 100644 --- a/src/dstack/_internal/server/app.py +++ b/src/dstack/_internal/server/app.py @@ -53,6 +53,7 @@ ) from dstack._internal.server.services.config import ServerConfigManager from dstack._internal.server.services.gateways import gateway_connections_pool +from dstack._internal.server.services.jobs.server_connection import job_server_connections_pool from dstack._internal.server.services.locking import advisory_lock_ctx from dstack._internal.server.services.projects import get_or_create_default_project from dstack._internal.server.services.proxy.deps import ServerProxyDependencyInjector @@ -213,6 +214,7 @@ async def lifespan(app: FastAPI): if pipeline_manager is not None: await pipeline_manager.drain() await gateway_connections_pool.remove_all() + await job_server_connections_pool.remove_all() service_conn_pool = await get_injector_from_app(app).get_service_connection_pool() await service_conn_pool.remove_all() if settings.SERVER_SSH_POOL_ENABLED: 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 61599172b5..a12c33e827 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,9 @@ is_master_job, job_model_to_job_submission, ) +from dstack._internal.server.services.jobs.server_connection import ( + job_server_connections_pool, +) from dstack._internal.server.services.locking import get_locker from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.metrics import get_job_metrics @@ -474,6 +477,11 @@ async def _process_running_job(context: _ProcessContext) -> _ProcessResult: context=context, startup_context=startup_context, result=result ) elif context.job_model.status == JobStatus.RUNNING: + if _server_access_enabled(context): + await job_server_connections_pool.ensure( + context.job_model, + context.job_submission.job_runtime_data, + ) await _process_running_status(context=context, result=result) if _get_result_status(context.job_model, result) == JobStatus.RUNNING: @@ -485,6 +493,8 @@ async def _process_running_job(context: _ProcessContext) -> _ProcessResult: ) await _maybe_register_replica(context=context, result=result) await _check_gpu_utilization(context=context, result=result) + elif _server_access_enabled(context): + await job_server_connections_pool.remove(context.job_model.id) return result @@ -614,6 +624,7 @@ async def _refetch_locked_job_model( JobModel.lock_token == item.lock_token, ) .options(joinedload(JobModel.instance).joinedload(InstanceModel.project)) + .options(joinedload(JobModel.project)) .options(joinedload(JobModel.probes).load_only(ProbeModel.success_streak)) .options( joinedload(JobModel.run).load_only(RunModel.id, RunModel.run_spec, RunModel.status) @@ -780,6 +791,8 @@ async def _process_provisioning_status( None, ) if runner_availability == _RunnerAvailability.AVAILABLE: + if not await _ensure_job_server_connection(context, result): + return file_archives = await _get_job_file_archives( archive_mappings=context.job.job_spec.file_archives, user=context.run_model.user, @@ -891,6 +904,8 @@ async def _process_pulling_status( return if runner_availability == _RunnerAvailability.AVAILABLE: + if not await _ensure_job_server_connection(context, result): + return file_archives = await _get_job_file_archives( archive_mappings=context.job.job_spec.file_archives, user=context.run_model.user, @@ -959,6 +974,36 @@ async def _process_running_status( _handle_instance_unreachable(context, result, job_provisioning_data) +async def _ensure_job_server_connection( + context: _ProcessContext, + result: _ProcessResult, +) -> bool: + if not _server_access_enabled(context): + return True + connected = await job_server_connections_pool.ensure( + context.job_model, + _get_result_job_runtime_data(context.job_model, result), + ) + if connected: + return True + + if job_server_connections_pool.retry_timed_out( + context.job_model.id, + JOB_DISCONNECTED_RETRY_TIMEOUT.total_seconds(), + ): + _terminate_job( + job_model=context.job_model, + job_update_map=result.job_update_map, + termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, + termination_reason_message="Could not establish dstack server access", + ) + return False + + +def _server_access_enabled(context: _ProcessContext) -> bool: + return bool(getattr(context.run.run_spec.configuration, "server", False)) + + async def _apply_process_result( item: JobRunningPipelineItem, job_model: JobModel, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py index be8d80948d..02f78a9136 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py @@ -64,6 +64,9 @@ get_job_spec, stop_runner, ) +from dstack._internal.server.services.jobs.server_connection import ( + job_server_connections_pool, +) from dstack._internal.server.services.locking import get_locker from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.pipelines import PipelineHinterProtocol @@ -268,6 +271,7 @@ async def process(self, item: JobTerminatingPipelineItem): return if job_model.volumes_detached_at is None: + await job_server_connections_pool.remove(job_model.id) result = await _process_terminating_job( job_model=job_model, instance_model=instance_model, diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 761745247d..d3e652d2c5 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/src/dstack/_internal/server/services/jobs/configurators/base.py @@ -9,6 +9,10 @@ from cachetools import TTLCache, cached from dstack._internal import settings +from dstack._internal.core.consts import ( + DSTACK_RUN_SERVER_URL, + DSTACK_SERVER_URL_ENV, +) from dstack._internal.core.errors import DockerRegistryError, ServerClientError from dstack._internal.core.models.common import RegistryAuth from dstack._internal.core.models.configurations import ( @@ -274,7 +278,13 @@ def _app_specs(self) -> List[AppSpec]: return specs def _env(self) -> Dict[str, str]: - return self.run_spec.configuration.env.as_dict() + env = self.run_spec.configuration.env.as_dict() + if self._server(): + env.setdefault(DSTACK_SERVER_URL_ENV, DSTACK_RUN_SERVER_URL) + return env + + def _server(self) -> bool: + return bool(getattr(self.run_spec.configuration, "server", False)) def _home_dir(self) -> Optional[str]: return self.run_spec.configuration.home_dir diff --git a/src/dstack/_internal/server/services/jobs/server_connection.py b/src/dstack/_internal/server/services/jobs/server_connection.py new file mode 100644 index 0000000000..56f45bd88d --- /dev/null +++ b/src/dstack/_internal/server/services/jobs/server_connection.py @@ -0,0 +1,207 @@ +import asyncio +import shlex +import shutil +import time +import uuid +from pathlib import Path +from typing import Optional +from weakref import WeakValueDictionary + +import httpx + +from dstack._internal.core.consts import DSTACK_RUN_SERVER_SOCKET_PATH +from dstack._internal.core.errors import SSHError +from dstack._internal.core.models.runs import JobRuntimeData +from dstack._internal.core.services.ssh.tunnel import ( + SSH_DEFAULT_OPTIONS, + IPSocket, + SocketPair, + SSHTunnel, + UnixSocket, +) +from dstack._internal.server import settings +from dstack._internal.server.models import JobModel +from dstack._internal.server.services.ssh import get_container_ssh_credentials +from dstack._internal.utils.logging import get_logger +from dstack._internal.utils.path import make_tmp_symlink_to_dir + +logger = get_logger(__name__) + +CONNECTIONS_DIR = settings.SERVER_DIR_PATH / "job-server-connections" +_MIN_ALIVE_CHECK_INTERVAL = 30 +_PROBE_TIMEOUT = 3 +_REMOTE_SOCKET_PATH = Path(DSTACK_RUN_SERVER_SOCKET_PATH) + + +class JobServerConnection: + """A private reverse SSH tunnel from one job to the dstack server.""" + + def __init__(self, job: JobModel, job_runtime_data: Optional[JobRuntimeData]) -> None: + self.job_id = job.id + self._last_verified_at = 0.0 + # Keep the control socket discoverable across server process restarts. The temporary + # symlink keeps its effective path below OpenSSH's Unix-socket length limit. + self._connection_dir = CONNECTIONS_DIR / str(job.id) + self._connection_dir.mkdir(parents=True, exist_ok=True) + self._temp_dir, effective_dir = make_tmp_symlink_to_dir( + self._connection_dir, + "connection", + ) + self._control_socket_path = effective_dir / "control.sock" + self._probe_socket_path = effective_dir / "probe.sock" + self._real_control_socket_path = self._connection_dir / "control.sock" + + hosts = get_container_ssh_credentials(job, job_runtime_data=job_runtime_data) + target, identity = hosts[-1] + self._tunnel = SSHTunnel( + destination=f"{target.username}@{target.hostname}", + port=target.port, + identity=identity, + control_sock_path=self._control_socket_path, + ssh_proxies=hosts[:-1], + options={ + **SSH_DEFAULT_OPTIONS, + "ConnectTimeout": str(settings.SERVER_SSH_CONNECT_TIMEOUT), + "ServerAliveInterval": "10", + "ServerAliveCountMax": "3", + "ControlPersist": "2m", + }, + batch_mode=True, + ) + + async def open(self) -> None: + if self._real_control_socket_path.exists(): + if await self._tunnel.acheck() and await self._server_is_reachable(): + self._last_verified_at = time.monotonic() + return + await self._tunnel.aclose() + self._real_control_socket_path.unlink(missing_ok=True) + + await self._tunnel.aopen() + try: + remote_dir = shlex.quote(str(_REMOTE_SOCKET_PATH.parent)) + remote_socket = shlex.quote(str(_REMOTE_SOCKET_PATH)) + # A new server owner replaces the stable path, making an orphaned forward unreachable. + await self._tunnel.aexec( + f"mkdir -p {remote_dir} && chmod 755 {remote_dir} && rm -f {remote_socket}" + ) + self._tunnel.reverse_forwarded_sockets = [ + SocketPair( + local=IPSocket(host="127.0.0.1", port=settings.SERVER_PORT), + remote=UnixSocket(path=_REMOTE_SOCKET_PATH), + ) + ] + # Probe through the job socket itself: a socket path can remain after its listener + # becomes unreachable. + self._tunnel.forwarded_sockets = [ + SocketPair( + local=UnixSocket(path=self._probe_socket_path), + remote=UnixSocket(path=_REMOTE_SOCKET_PATH), + ) + ] + await self._tunnel.aopen() + # The socket carries no credentials. World access inside the isolated job container + # lets configurations using a non-root `user` reach it as well. + await self._tunnel.aexec(f"chmod 666 {remote_socket}") + if not await self._server_is_reachable(): + raise SSHError("dstack server is not reachable from the job") + except Exception: + await self._tunnel.aclose() + raise + self._last_verified_at = time.monotonic() + + async def is_alive(self) -> bool: + if not self._control_socket_path.exists(): + return False + now = time.monotonic() + if now - self._last_verified_at < _MIN_ALIVE_CHECK_INTERVAL: + return True + if not await self._tunnel.acheck() or not await self._server_is_reachable(): + return False + self._last_verified_at = now + return True + + async def close(self) -> None: + await self._tunnel.aclose() + shutil.rmtree(self._connection_dir, ignore_errors=True) + + async def _server_is_reachable(self) -> bool: + transport = httpx.AsyncHTTPTransport(uds=str(self._probe_socket_path)) + try: + async with httpx.AsyncClient( + transport=transport, + timeout=_PROBE_TIMEOUT, + ) as client: + response = await client.get("http://localhost/healthcheck") + except httpx.HTTPError: + return False + return response.status_code == 200 + + +class JobServerConnectionsPool: + def __init__(self) -> None: + self._connections: dict[uuid.UUID, JobServerConnection] = {} + self._failure_started_at: dict[uuid.UUID, float] = {} + self._locks: WeakValueDictionary[uuid.UUID, asyncio.Lock] = WeakValueDictionary() + self._pool_lock = asyncio.Lock() + + async def ensure( + self, + job: JobModel, + job_runtime_data: Optional[JobRuntimeData], + ) -> bool: + lock = await self._get_lock(job.id) + async with lock: + connection = self._connections.get(job.id) + if connection is not None and await connection.is_alive(): + self._failure_started_at.pop(job.id, None) + return True + if connection is not None: + await self._close(connection) + self._connections.pop(job.id, None) + + connection = JobServerConnection(job, job_runtime_data) + try: + await connection.open() + except SSHError as e: + logger.warning("Failed to establish server access for job %s: %s", job.id, e) + await self._close(connection) + self._failure_started_at.setdefault(job.id, time.monotonic()) + return False + self._connections[job.id] = connection + self._failure_started_at.pop(job.id, None) + return True + + def retry_timed_out(self, job_id: uuid.UUID, timeout: float) -> bool: + failure_started_at = self._failure_started_at.get(job_id) + if failure_started_at is None: + return False + return time.monotonic() - failure_started_at > timeout + + async def remove(self, job_id: uuid.UUID) -> None: + lock = await self._get_lock(job_id) + async with lock: + connection = self._connections.pop(job_id, None) + if connection is not None: + await self._close(connection) + self._failure_started_at.pop(job_id, None) + shutil.rmtree(CONNECTIONS_DIR / str(job_id), ignore_errors=True) + + async def remove_all(self) -> None: + async with self._pool_lock: + job_ids = set(self._connections).union(self._failure_started_at) + await asyncio.gather(*(self.remove(job_id) for job_id in job_ids)) + + async def _get_lock(self, job_id: uuid.UUID) -> asyncio.Lock: + async with self._pool_lock: + return self._locks.setdefault(job_id, asyncio.Lock()) + + @staticmethod + async def _close(connection: JobServerConnection) -> None: + try: + await connection.close() + except Exception: + logger.exception("Failed to close server access for job %s", connection.job_id) + + +job_server_connections_pool = JobServerConnectionsPool() diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index 7ccc2b1af7..a0f34b158c 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -11,6 +11,7 @@ import requests_unixsocket from typing_extensions import Self +from dstack._internal.core.consts import DSTACK_PROJECT_ENV from dstack._internal.core.errors import DstackError from dstack._internal.core.models.common import CoreModel, NetworkMode from dstack._internal.core.models.envs import Env @@ -126,7 +127,8 @@ def submit_job( # API modification. Both layers are merged into a deep-copied job_spec # so the shared spec object held by the caller is not mutated. job_spec = job.job_spec - if instance_env is not None or router_env is not None: + server_access = bool(getattr(run.run_spec.configuration, "server", False)) + if instance_env is not None or router_env is not None or server_access: merged_env: Dict[str, str] = {} if instance_env is not None: if isinstance(instance_env, Env): @@ -136,6 +138,8 @@ def submit_job( merged_env.update(job_spec.env) if router_env is not None: merged_env.update(router_env) + if server_access: + merged_env.setdefault(DSTACK_PROJECT_ENV, run.project_name) job_spec = job_spec.copy(deep=True) job_spec.env = merged_env quota = server_settings.SERVER_LOG_QUOTA_PER_JOB_HOUR diff --git a/src/dstack/_internal/server/services/ssh.py b/src/dstack/_internal/server/services/ssh.py index 9d07263885..3a14b97d72 100644 --- a/src/dstack/_internal/server/services/ssh.py +++ b/src/dstack/_internal/server/services/ssh.py @@ -1,7 +1,9 @@ from collections.abc import Iterable +from typing import Optional from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT from dstack._internal.core.models.instances import SSHConnectionParams +from dstack._internal.core.models.runs import JobRuntimeData from dstack._internal.core.services.ssh.tunnel import SSH_DEFAULT_OPTIONS, SocketPair, SSHTunnel from dstack._internal.server.models import JobModel from dstack._internal.server.services.instances import get_instance_remote_connection_info @@ -10,7 +12,10 @@ from dstack._internal.utils.path import FileContent -def get_container_ssh_credentials(job: JobModel) -> list[tuple[SSHConnectionParams, FileContent]]: +def get_container_ssh_credentials( + job: JobModel, + job_runtime_data: Optional[JobRuntimeData] = None, +) -> list[tuple[SSHConnectionParams, FileContent]]: """ Returns the information needed to connect to the SSH server inside the job container. @@ -21,6 +26,7 @@ def get_container_ssh_credentials(job: JobModel) -> list[tuple[SSHConnectionPara Args: job: `JobModel` with `project`, `instance` and `instance.project` fields loaded. + job_runtime_data: Runtime data to use before it has been persisted to `job`. Returns: A list of hosts credentials as (host's `SSHConnectionParams`, private key's `FileContent`) @@ -51,7 +57,7 @@ def get_container_ssh_credentials(job: JobModel) -> list[tuple[SSHConnectionPara instance_project_key = FileContent(instance.project.ssh_private_key) hosts.append((instance_proxy, instance_project_key)) ssh_port = DSTACK_RUNNER_SSH_PORT - jrd = get_job_runtime_data(job) + jrd = job_runtime_data or get_job_runtime_data(job) if jrd is not None and jrd.ports is not None: ssh_port = jrd.ports.get(ssh_port, ssh_port) target_host = SSHConnectionParams( diff --git a/src/dstack/api/server/__init__.py b/src/dstack/api/server/__init__.py index 82b009863a..d1130713d6 100644 --- a/src/dstack/api/server/__init__.py +++ b/src/dstack/api/server/__init__.py @@ -5,6 +5,7 @@ from typing import Dict, List, Optional, Type import requests +import requests_unixsocket from dstack import version from dstack._internal.core.errors import ( @@ -64,6 +65,8 @@ def __init__(self, base_url: str, token: Optional[str] = None): """ self._base_url = base_url.rstrip("/") self._s = requests.session() + if self._base_url.startswith("http+unix://"): + self._s.mount("http+unix://", requests_unixsocket.UnixAdapter()) self._token = None if token is not None: self._token = token diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index 364624e597..a10f18e4c1 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -16,6 +16,39 @@ class TestParseConfiguration: + @pytest.mark.parametrize("configuration_type", ["task", "dev-environment"]) + def test_server_access_supported(self, configuration_type: str): + conf = {"type": configuration_type, "server": True} + if configuration_type == "task": + conf["commands"] = ["true"] + + parsed = parse_run_configuration(conf) + + assert parsed.server is True + + def test_server_access_not_supported_for_service(self): + with pytest.raises(ConfigurationError, match="extra fields not permitted"): + parse_run_configuration( + { + "type": "service", + "commands": ["python3 -m http.server"], + "port": 8000, + "server": True, + } + ) + + def test_server_access_allows_unresolved_passthrough_env(self): + parsed = parse_run_configuration( + { + "type": "task", + "commands": ["true"], + "server": True, + "env": ["DSTACK_TOKEN"], + } + ) + + assert "DSTACK_TOKEN" in parsed.env + def test_service_model_probes_none_when_omitted(self): """When model is set but probes omitted, probes should remain None. The default probe is generated server-side in the job configurator.""" diff --git a/src/tests/_internal/core/services/test_api_client.py b/src/tests/_internal/core/services/test_api_client.py new file mode 100644 index 0000000000..efbb747e49 --- /dev/null +++ b/src/tests/_internal/core/services/test_api_client.py @@ -0,0 +1,36 @@ +from unittest.mock import patch + +from dstack._internal.core.models.config import ProjectConfig +from dstack._internal.core.services.api_client import get_api_client + + +class TestGetAPIClient: + def test_uses_complete_environment_config(self, monkeypatch): + monkeypatch.setenv("DSTACK_SERVER_URL", "http+unix://%2Frun%2Fdstack%2Fserver.sock") + monkeypatch.setenv("DSTACK_PROJECT", "main") + monkeypatch.setenv("DSTACK_TOKEN", "token") + + with patch("dstack._internal.core.services.api_client.configs.ConfigManager") as manager: + client, project_name = get_api_client() + + manager.assert_not_called() + assert client.base_url == "http+unix://%2Frun%2Fdstack%2Fserver.sock" + assert project_name == "main" + + def test_incomplete_environment_config_uses_config_file(self, monkeypatch): + monkeypatch.setenv("DSTACK_SERVER_URL", "http+unix://%2Frun%2Fdstack%2Fserver.sock") + monkeypatch.setenv("DSTACK_PROJECT", "environment-project") + monkeypatch.delenv("DSTACK_TOKEN", raising=False) + project = ProjectConfig( + name="configured-project", + url="https://server.example.com", + token="configured-token", + default=True, + ) + + with patch("dstack._internal.core.services.api_client.configs.ConfigManager") as manager: + manager.return_value.get_project_config.return_value = project + client, project_name = get_api_client() + + assert client.base_url == "https://server.example.com" + assert project_name == "configured-project" diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index c3070ae75a..d2c09b1bd0 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -977,6 +977,82 @@ async def test_pulling_shim_runner_not_ready( await session.refresh(job) assert job.status == JobStatus.PULLING + async def test_pulling_waits_for_requested_server_access_before_starting_job( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ssh_tunnel_mock: Mock, + shim_client_mock: Mock, + runner_client_mock: Mock, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + commands=["true"], + server=True, + ), + ), + ) + instance = await create_instance( + session=session, project=project, status=InstanceStatus.BUSY + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.PULLING, + submitted_at=get_current_datetime(), + job_provisioning_data=get_job_provisioning_data(dockerized=True), + job_runtime_data=get_job_runtime_data(network_mode="bridge", ports=None), + instance=instance, + instance_assigned=True, + ) + shim_client_mock.get_task.return_value.status = TaskStatus.RUNNING + shim_client_mock.get_task.return_value.ports = [ + PortMapping(container=10022, host=32771), + PortMapping(container=10999, host=32772), + ] + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running." + "job_server_connections_pool.ensure", + new_callable=AsyncMock, + return_value=False, + ) as ensure_server_connection_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running." + "job_server_connections_pool.retry_timed_out", + return_value=False, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running._get_job_file_archives", + new_callable=AsyncMock, + ) as get_job_file_archives_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running._get_job_code", + new_callable=AsyncMock, + ) as get_job_code_mock, + ): + await _process_job(session, worker, job) + + ensure_server_connection_mock.assert_awaited_once() + runner_client_mock.submit_job.assert_not_called() + get_job_file_archives_mock.assert_not_awaited() + get_job_code_mock.assert_not_awaited() + await session.refresh(job) + assert job.status == JobStatus.PULLING + assert job.disconnected_at is None + async def test_pulling_shim_uses_runtime_port_mapping_for_runner_calls( self, test_db, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py index af34b1913e..9c79cb5e7b 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py @@ -9,6 +9,7 @@ from sqlalchemy.orm import joinedload from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.configurations import TaskConfiguration from dstack._internal.core.models.instances import InstanceStatus from dstack._internal.core.models.runs import JobStatus, JobTerminationReason from dstack._internal.core.models.volumes import VolumeStatus @@ -32,6 +33,7 @@ get_instance_offer_with_availability, get_job_provisioning_data, get_job_runtime_data, + get_run_spec, get_volume_configuration, get_volume_provisioning_data, list_events, @@ -248,6 +250,44 @@ async def test_fetch_returns_oldest_jobs_first_up_to_limit( @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @pytest.mark.usefixtures("image_config_mock") class TestJobTerminatingWorker: + async def test_closes_requested_server_access( + self, test_db, session: AsyncSession, worker: JobTerminatingWorker + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + commands=["true"], + server=True, + ), + ), + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.TERMINATING, + termination_reason=JobTerminationReason.TERMINATED_BY_USER, + ) + _lock_job(job) + await session.commit() + + with patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating." + "job_server_connections_pool.remove", + new_callable=AsyncMock, + ) as remove_server_connection_mock: + await worker.process(_job_to_pipeline_item(job)) + + remove_server_connection_mock.assert_awaited_once_with(job.id) + async def test_stops_job_gracefully_before_terminating_container( self, test_db, session: AsyncSession, worker: JobTerminatingWorker ): diff --git a/src/tests/_internal/server/services/jobs/configurators/test_task.py b/src/tests/_internal/server/services/jobs/configurators/test_task.py index 3c80bf226f..2c89c5c65e 100644 --- a/src/tests/_internal/server/services/jobs/configurators/test_task.py +++ b/src/tests/_internal/server/services/jobs/configurators/test_task.py @@ -37,6 +37,47 @@ async def test_multi_node(self): assert job_specs[1].ssh_key == JobSSHKey(private="private1", public="public1") +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_config_mock") +class TestServerAccess: + async def test_adds_transport_without_credentials(self): + configuration = TaskConfiguration(image="debian", commands=["true"], server=True) + run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration) + configurator = TaskJobConfigurator(run_spec) + + job_spec = (await configurator.get_job_specs(replica_num=0))[0] + + assert "server" not in job_spec.dict() + assert job_spec.env == { + "DSTACK_SERVER_URL": "http+unix://%2Frun%2Fdstack%2Fserver.sock", + } + + async def test_disabled_by_default(self): + configuration = TaskConfiguration(image="debian", commands=["true"]) + run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration) + configurator = TaskJobConfigurator(run_spec) + + job_spec = (await configurator.get_job_specs(replica_num=0))[0] + + assert "DSTACK_SERVER_URL" not in job_spec.env + + async def test_preserves_explicit_transport_env(self): + configuration = TaskConfiguration( + image="debian", + commands=["true"], + server=True, + env={ + "DSTACK_SERVER_URL": "https://server.example.com", + }, + ) + run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration) + configurator = TaskJobConfigurator(run_spec) + + job_spec = (await configurator.get_job_specs(replica_num=0))[0] + + assert job_spec.env["DSTACK_SERVER_URL"] == "https://server.example.com" + + @pytest.mark.asyncio @pytest.mark.usefixtures("image_config_mock") class TestCommands: diff --git a/src/tests/_internal/server/services/jobs/test_server_connection.py b/src/tests/_internal/server/services/jobs/test_server_connection.py new file mode 100644 index 0000000000..c8e5a8b9a1 --- /dev/null +++ b/src/tests/_internal/server/services/jobs/test_server_connection.py @@ -0,0 +1,163 @@ +import asyncio +import socket +import uuid +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import pytest + +from dstack._internal.core.errors import SSHError +from dstack._internal.core.models.instances import SSHConnectionParams +from dstack._internal.core.services.ssh.tunnel import IPSocket, SocketPair, UnixSocket +from dstack._internal.server.services.jobs import server_connection +from dstack._internal.server.services.jobs.server_connection import ( + JobServerConnection, + JobServerConnectionsPool, +) +from dstack._internal.utils.path import FileContent + + +@pytest.fixture +def tunnel_mock(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(server_connection, "CONNECTIONS_DIR", tmp_path) + tunnel = MagicMock() + tunnel.acheck = AsyncMock(return_value=False) + tunnel.aopen = AsyncMock() + tunnel.aexec = AsyncMock(return_value="") + tunnel.aclose = AsyncMock() + tunnel_class = Mock(return_value=tunnel) + monkeypatch.setattr(server_connection, "SSHTunnel", tunnel_class) + monkeypatch.setattr( + server_connection, + "get_container_ssh_credentials", + Mock( + return_value=[ + ( + SSHConnectionParams(hostname="job.example.com", username="root", port=10022), + FileContent("private-key"), + ) + ] + ), + ) + return tunnel, tunnel_class + + +class TestJobServerConnection: + @pytest.mark.asyncio + async def test_opens_private_reverse_socket(self, tunnel_mock): + tunnel, tunnel_class = tunnel_mock + job = Mock(id=uuid.uuid4()) + connection = JobServerConnection(job, job_runtime_data=None) + connection._server_is_reachable = AsyncMock(return_value=True) + + await connection.open() + + tunnel_class.assert_called_once() + assert tunnel.aopen.await_count == 2 + assert tunnel.reverse_forwarded_sockets == [ + SocketPair( + local=IPSocket(host="127.0.0.1", port=server_connection.settings.SERVER_PORT), + remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH), + ) + ] + assert tunnel.forwarded_sockets == [ + SocketPair( + local=UnixSocket(path=connection._probe_socket_path), + remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH), + ) + ] + commands = [call.args[0] for call in tunnel.aexec.await_args_list] + assert commands == [ + "mkdir -p /run/dstack && chmod 755 /run/dstack && rm -f /run/dstack/server.sock", + "chmod 666 /run/dstack/server.sock", + ] + + @pytest.mark.asyncio + async def test_reuses_live_tunnel_with_existing_socket(self, tunnel_mock): + tunnel, _ = tunnel_mock + tunnel.acheck.return_value = True + job = Mock(id=uuid.uuid4()) + connection = JobServerConnection(job, job_runtime_data=None) + connection._server_is_reachable = AsyncMock(return_value=True) + connection._real_control_socket_path.touch() + + await connection.open() + + tunnel.aopen.assert_not_awaited() + connection._server_is_reachable.assert_awaited_once() + + @pytest.mark.asyncio + async def test_replaces_tunnel_with_stale_socket(self, tunnel_mock): + tunnel, _ = tunnel_mock + tunnel.acheck.return_value = True + job = Mock(id=uuid.uuid4()) + connection = JobServerConnection(job, job_runtime_data=None) + connection._server_is_reachable = AsyncMock(side_effect=[False, True]) + connection._real_control_socket_path.touch() + + await connection.open() + + tunnel.aclose.assert_awaited_once() + assert tunnel.aopen.await_count == 2 + + @pytest.mark.asyncio + async def test_stale_probe_socket_is_not_reachable(self, tunnel_mock): + job = Mock(id=uuid.uuid4()) + connection = JobServerConnection(job, job_runtime_data=None) + stale_socket = socket.socket(socket.AF_UNIX) + stale_socket.bind(str(connection._probe_socket_path)) + stale_socket.close() + + assert not await connection._server_is_reachable() + + @pytest.mark.asyncio + async def test_healthcheck_through_probe_socket(self, tunnel_mock): + job = Mock(id=uuid.uuid4()) + connection = JobServerConnection(job, job_runtime_data=None) + + async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + await reader.readuntil(b"\r\n\r\n") + writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") + await writer.drain() + writer.close() + + server = await asyncio.start_unix_server(handle, path=connection._probe_socket_path) + try: + assert await connection._server_is_reachable() + finally: + server.close() + await server.wait_closed() + + +class TestJobServerConnectionsPool: + @pytest.mark.asyncio + async def test_reuses_healthy_connection(self): + job = Mock(id=uuid.uuid4()) + connection = MagicMock() + connection.job_id = job.id + connection.open = AsyncMock() + connection.is_alive = AsyncMock(return_value=True) + connection.close = AsyncMock() + pool = JobServerConnectionsPool() + + with patch.object(server_connection, "JobServerConnection", return_value=connection): + assert await pool.ensure(job, job_runtime_data=None) + assert await pool.ensure(job, job_runtime_data=None) + + connection.open.assert_awaited_once() + connection.is_alive.assert_awaited_once() + + @pytest.mark.asyncio + async def test_returns_false_when_tunnel_cannot_open(self): + job = Mock(id=uuid.uuid4()) + connection = MagicMock() + connection.job_id = job.id + connection.open = AsyncMock(side_effect=SSHError("connection failed")) + connection.close = AsyncMock() + pool = JobServerConnectionsPool() + + with patch.object(server_connection, "JobServerConnection", return_value=connection): + with patch.object(server_connection.time, "monotonic", side_effect=[10, 131]): + assert not await pool.ensure(job, job_runtime_data=None) + assert pool.retry_timed_out(job.id, timeout=120) + + connection.close.assert_awaited_once() diff --git a/src/tests/_internal/server/services/runner/test_client.py b/src/tests/_internal/server/services/runner/test_client.py index 588c231a19..dd7805b5b2 100644 --- a/src/tests/_internal/server/services/runner/test_client.py +++ b/src/tests/_internal/server/services/runner/test_client.py @@ -1,14 +1,17 @@ import uuid from collections.abc import Generator +from datetime import datetime, timezone from typing import Optional import pytest import requests_mock -from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT +from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT, DSTACK_SHIM_HTTP_PORT from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import NetworkMode +from dstack._internal.core.models.configurations import TaskConfiguration from dstack._internal.core.models.resources import Memory +from dstack._internal.core.models.runs import ClusterInfo, Job, JobSpec, JobSubmission, Run from dstack._internal.core.models.volumes import ( InstanceMountPoint, VolumeAttachment, @@ -25,11 +28,16 @@ TaskStatus, ) from dstack._internal.server.services.runner.client import ( + RunnerClient, ShimClient, ShimHTTPError, _parse_version, ) -from dstack._internal.server.testing.common import get_volume, get_volume_configuration +from dstack._internal.server.testing.common import ( + get_run_spec, + get_volume, + get_volume_configuration, +) class BaseShimClientTest: @@ -65,6 +73,68 @@ def assert_request( assert req.json() == json +class TestRunnerClientSubmitJob(BaseShimClientTest): + def test_adds_default_project_for_server_access(self, adapter: requests_mock.Adapter): + adapter.register_uri("POST", "/api/submit", json={}) + run_spec = get_run_spec( + repo_id="repo", configuration=TaskConfiguration(commands=["true"], server=True) + ) + run = Run.construct(id=uuid.uuid4(), project_name="main", run_spec=run_spec) + job = Job.construct( + job_spec=JobSpec.construct(env={"DSTACK_TOKEN": "token"}), + job_submissions=[ + JobSubmission.construct( + id=uuid.uuid4(), + submitted_at=datetime.now(timezone.utc), + ) + ], + ) + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + client.submit_job( + run=run, + job=job, + cluster_info=ClusterInfo(job_ips=[], master_job_ip="", gpus_per_job=0), + secrets={}, + repo_credentials=None, + ) + + assert adapter.last_request is not None + assert adapter.last_request.json()["job_spec"]["env"] == { + "DSTACK_PROJECT": "main", + "DSTACK_TOKEN": "token", + } + assert job.job_spec.env == {"DSTACK_TOKEN": "token"} + + def test_preserves_explicit_project_for_server_access(self, adapter: requests_mock.Adapter): + adapter.register_uri("POST", "/api/submit", json={}) + run_spec = get_run_spec( + repo_id="repo", configuration=TaskConfiguration(commands=["true"], server=True) + ) + run = Run.construct(id=uuid.uuid4(), project_name="main", run_spec=run_spec) + job = Job.construct( + job_spec=JobSpec.construct(env={"DSTACK_PROJECT": "other"}), + job_submissions=[ + JobSubmission.construct( + id=uuid.uuid4(), + submitted_at=datetime.now(timezone.utc), + ) + ], + ) + client = RunnerClient(port=DSTACK_RUNNER_HTTP_PORT) + + client.submit_job( + run=run, + job=job, + cluster_info=ClusterInfo(job_ips=[], master_job_ip="", gpus_per_job=0), + secrets={}, + repo_credentials=None, + ) + + assert adapter.last_request is not None + assert adapter.last_request.json()["job_spec"]["env"]["DSTACK_PROJECT"] == "other" + + class TestShimClientNegotiate(BaseShimClientTest): @pytest.mark.parametrize( ["expected_shim_version", "expected_api_version"], diff --git a/src/tests/api/test_client.py b/src/tests/api/test_client.py new file mode 100644 index 0000000000..44b5e3f769 --- /dev/null +++ b/src/tests/api/test_client.py @@ -0,0 +1,30 @@ +from unittest.mock import Mock, patch + +from dstack.api.server import APIClient + + +class TestAPIClientTransport: + def test_mounts_unix_socket_adapter(self): + session = Mock() + adapter = Mock() + with ( + patch("dstack.api.server.requests.session", return_value=session), + patch("dstack.api.server.requests_unixsocket.UnixAdapter", return_value=adapter), + ): + client = APIClient("http+unix://%2Frun%2Fdstack%2Fserver.sock", token="token") + + assert client.base_url == "http+unix://%2Frun%2Fdstack%2Fserver.sock" + session.mount.assert_called_once_with("http+unix://", adapter) + session.headers.update.assert_any_call({"Authorization": "Bearer token"}) + + def test_http_transport_does_not_mount_unix_socket_adapter(self): + session = Mock() + with ( + patch("dstack.api.server.requests.session", return_value=session), + patch("dstack.api.server.requests_unixsocket.UnixAdapter") as adapter_class, + ): + client = APIClient("https://server.example.com/") + + assert client.base_url == "https://server.example.com" + adapter_class.assert_not_called() + session.mount.assert_not_called() From ebe8aa358273d8d304f791bf01943bc98725035d Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Sun, 12 Jul 2026 00:50:56 +0200 Subject: [PATCH 02/14] Update run response expectations --- src/tests/_internal/server/routers/test_runs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index fdc5a8f355..0fdeea68af 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -173,6 +173,7 @@ def get_dev_env_run_plan_dict( "privileged": privileged, "init": [], "ports": [], + "server": False, "python": "3.13" if not docker else None, "nvcc": None, "registry_auth": None, @@ -424,6 +425,7 @@ def get_dev_env_run_dict( "privileged": privileged, "init": [], "ports": [], + "server": False, "python": "3.13" if not docker else None, "nvcc": None, "registry_auth": None, From eec89915840e49c3b7ea49e070d3801a6829e956 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Sun, 12 Jul 2026 00:58:54 +0200 Subject: [PATCH 03/14] Preserve compatibility with older servers --- .../_internal/core/compatibility/runs.py | 15 +++++++++++++- src/tests/_internal/core/models/test_runs.py | 20 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py index 4b57db1d47..eefc326f92 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -6,7 +6,11 @@ IncludeExcludeDictType, IncludeExcludeSetType, ) -from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.configurations import ( + DevEnvironmentConfiguration, + ServiceConfiguration, + TaskConfiguration, +) from dstack._internal.core.models.routers import SGLangServiceRouterConfig from dstack._internal.core.models.runs import ( DEFAULT_PROBE_UNTIL_READY, @@ -89,6 +93,15 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType: if run_spec.configuration.backend_options is None: configuration_excludes["backend_options"] = True + if ( + isinstance( + run_spec.configuration, + (DevEnvironmentConfiguration, TaskConfiguration), + ) + and not run_spec.configuration.server + ): + configuration_excludes["server"] = True + if isinstance(run_spec.configuration, ServiceConfiguration): if run_spec.configuration.probes: probe_excludes: IncludeExcludeDictType = {} diff --git a/src/tests/_internal/core/models/test_runs.py b/src/tests/_internal/core/models/test_runs.py index 7968845f4e..0fb2300d5c 100644 --- a/src/tests/_internal/core/models/test_runs.py +++ b/src/tests/_internal/core/models/test_runs.py @@ -1,6 +1,11 @@ import pytest from pydantic import ValidationError +from dstack._internal.core.compatibility.runs import get_run_spec_excludes +from dstack._internal.core.models.configurations import ( + DevEnvironmentConfiguration, + TaskConfiguration, +) from dstack._internal.core.models.profiles import RetryEvent from dstack._internal.core.models.runs import ( JobStatus, @@ -23,6 +28,21 @@ def test_run_termination_reason_to_status_works_with_all_enum_variants(): assert isinstance(run_status, RunStatus) +@pytest.mark.parametrize("configuration_type", ["task", "dev-environment"]) +@pytest.mark.parametrize("server", [False, True]) +def test_server_access_run_spec_compatibility(configuration_type: str, server: bool): + if configuration_type == "task": + configuration = TaskConfiguration(commands=["true"], server=server) + else: + configuration = DevEnvironmentConfiguration(server=server) + configuration_excludes = get_run_spec_excludes(RunSpec(configuration=configuration)).get( + "configuration" + ) + + assert isinstance(configuration_excludes, dict) + assert ("server" in configuration_excludes) is not server + + def test_job_termination_reason_to_status_works_with_all_enum_variants(): for job_termination_reason in JobTerminationReason: job_status = job_termination_reason.to_status() From 8d0dfd4c8ad28e10fdab93c12a8f34c2e88636d3 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 12:36:31 +0200 Subject: [PATCH 04/14] Fix server access retry timeout for starting jobs Removing the job server connection on PROVISIONING/PULLING iterations reset the failure time tracked by the pool, so retry_timed_out() could never fire and jobs with a failing tunnel were never terminated by the server. Remove the connection only when the job leaves RUNNING; jobs_terminating covers the final cleanup. Co-Authored-By: Claude Fable 5 --- .../background/pipeline_tasks/jobs_running.py | 4 +- .../pipeline_tasks/test_running_jobs.py | 105 ++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) 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 a12c33e827..e126cdf3c6 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -493,7 +493,9 @@ async def _process_running_job(context: _ProcessContext) -> _ProcessResult: ) await _maybe_register_replica(context=context, result=result) await _check_gpu_utilization(context=context, result=result) - elif _server_access_enabled(context): + elif _server_access_enabled(context) and context.job_model.status == JobStatus.RUNNING: + # Removing on PROVISIONING/PULLING iterations would reset the failure time + # tracked by the pool, breaking retry_timed_out() await job_server_connections_pool.remove(context.job_model.id) return result diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index d2c09b1bd0..c0d4f8537a 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -39,6 +39,7 @@ from dstack._internal.core.services.ssh.tunnel import SSHTunnel from dstack._internal.server import settings as server_settings from dstack._internal.server.background.pipeline_tasks.jobs_running import ( + JOB_DISCONNECTED_RETRY_TIMEOUT, ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS, JobRunningFetcher, JobRunningPipeline, @@ -61,6 +62,8 @@ PullResponse, TaskStatus, ) +from dstack._internal.server.services.jobs import server_connection +from dstack._internal.server.services.jobs.server_connection import job_server_connections_pool from dstack._internal.server.services.runner.client import RunnerClient, ShimClient from dstack._internal.server.services.runs.replicas import RouterEnvStatus from dstack._internal.server.services.volumes import volume_model_to_volume @@ -1034,6 +1037,11 @@ async def test_pulling_waits_for_requested_server_access_before_starting_job( "job_server_connections_pool.retry_timed_out", return_value=False, ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running." + "job_server_connections_pool.remove", + new_callable=AsyncMock, + ) as remove_server_connection_mock, patch( "dstack._internal.server.background.pipeline_tasks.jobs_running._get_job_file_archives", new_callable=AsyncMock, @@ -1046,6 +1054,8 @@ async def test_pulling_waits_for_requested_server_access_before_starting_job( await _process_job(session, worker, job) ensure_server_connection_mock.assert_awaited_once() + # Removing the connection here would reset the failure time tracked by the pool + remove_server_connection_mock.assert_not_awaited() runner_client_mock.submit_job.assert_not_called() get_job_file_archives_mock.assert_not_awaited() get_job_code_mock.assert_not_awaited() @@ -1053,6 +1063,101 @@ async def test_pulling_waits_for_requested_server_access_before_starting_job( assert job.status == JobStatus.PULLING assert job.disconnected_at is None + async def test_provisioning_server_access_failure_terminates_job_after_retry_timeout( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + commands=["true"], + server=True, + ), + ), + ) + instance = await create_instance( + session=session, project=project, status=InstanceStatus.BUSY + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.PROVISIONING, + submitted_at=get_current_datetime(), + job_provisioning_data=get_job_provisioning_data(dockerized=False), + job_runtime_data=get_job_runtime_data(), + instance=instance, + instance_assigned=True, + ) + last_processed_at = job.last_processed_at + monkeypatch.setattr(server_connection, "CONNECTIONS_DIR", tmp_path) + failing_connection = MagicMock() + failing_connection.job_id = job.id + failing_connection.open = AsyncMock(side_effect=SSHError("cannot open tunnel")) + failing_connection.close = AsyncMock() + monkeypatch.setattr( + server_connection, "JobServerConnection", Mock(return_value=failing_connection) + ) + + try: + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel"), + patch.object(RunnerClient, "_healthcheck") as healthcheck_mock, + patch.object(RunnerClient, "submit_job") as submit_job_mock, + ): + healthcheck_mock.return_value = HealthcheckResponse( + service="dstack-runner", version="0.0.1.dev2" + ) + await _process_job(session, worker, job) + + submit_job_mock.assert_not_called() + await session.refresh(job) + assert job.status == JobStatus.PROVISIONING + failure_started_at = job_server_connections_pool._failure_started_at.get(job.id) + assert failure_started_at is not None + + # The connection failure time must survive further processing iterations, + # otherwise the retry timeout can never elapse + job.last_processed_at = last_processed_at + await session.commit() + await _process_job(session, worker, job) + + submit_job_mock.assert_not_called() + await session.refresh(job) + assert job.status == JobStatus.PROVISIONING + assert ( + job_server_connections_pool._failure_started_at.get(job.id) + == failure_started_at + ) + + job_server_connections_pool._failure_started_at[job.id] = ( + failure_started_at - JOB_DISCONNECTED_RETRY_TIMEOUT.total_seconds() - 1 + ) + job.last_processed_at = last_processed_at + await session.commit() + await _process_job(session, worker, job) + + submit_job_mock.assert_not_called() + await session.refresh(job) + assert job.status == JobStatus.TERMINATING + assert job.termination_reason == JobTerminationReason.TERMINATED_BY_SERVER + assert job.termination_reason_message == "Could not establish dstack server access" + finally: + job_server_connections_pool._connections.pop(job.id, None) + job_server_connections_pool._failure_started_at.pop(job.id, None) + async def test_pulling_shim_uses_runtime_port_mapping_for_runner_calls( self, test_db, From 065c6abfe3160734b2abdf3af518c7c0bafdadf7 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 12:36:32 +0200 Subject: [PATCH 05/14] Allow DSTACK_TOKEN to override the configured token DSTACK_TOKEN alone now overrides the token of the project configured in config.yml. DSTACK_SERVER_URL still takes effect only together with DSTACK_TOKEN, and a missing token or project now produces an explicit error instead of a silent fallback when no config is available. Co-Authored-By: Claude Fable 5 --- .../_internal/core/services/api_client.py | 15 ++++++++++++++- .../core/services/test_api_client.py | 19 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/src/dstack/_internal/core/services/api_client.py b/src/dstack/_internal/core/services/api_client.py index f1718be97b..72d2e3e999 100644 --- a/src/dstack/_internal/core/services/api_client.py +++ b/src/dstack/_internal/core/services/api_client.py @@ -15,13 +15,26 @@ def get_api_client(project_name: Optional[str] = None) -> Tuple[APIClient, str]: env_project_name = project_name or os.getenv(DSTACK_PROJECT_ENV) server_url = os.getenv(DSTACK_SERVER_URL_ENV) token = os.getenv(DSTACK_TOKEN_ENV) - if env_project_name is not None and server_url is not None and token is not None: + if server_url is not None and token is not None: + if env_project_name is None: + raise ConfigurationError( + f"{DSTACK_SERVER_URL_ENV} and {DSTACK_TOKEN_ENV} are set," + f" but the project is not specified." + f" Set {DSTACK_PROJECT_ENV} or use --project" + ) return APIClient(server_url, token), env_project_name config = configs.ConfigManager() project = config.get_project_config(project_name) if project is None: + if server_url is not None: + raise ConfigurationError( + f"{DSTACK_SERVER_URL_ENV} is set, but {DSTACK_TOKEN_ENV} is not set" + ) if project_name is not None: raise ConfigurationError(f"Project {project_name} is not configured") raise ConfigurationError("No default project, specify project name") + if token is not None: + # DSTACK_TOKEN overrides the configured token + return APIClient(project.url, token), project.name return APIClient(project.url, project.token), project.name diff --git a/src/tests/_internal/core/services/test_api_client.py b/src/tests/_internal/core/services/test_api_client.py index efbb747e49..94668481c4 100644 --- a/src/tests/_internal/core/services/test_api_client.py +++ b/src/tests/_internal/core/services/test_api_client.py @@ -34,3 +34,22 @@ def test_incomplete_environment_config_uses_config_file(self, monkeypatch): assert client.base_url == "https://server.example.com" assert project_name == "configured-project" + + def test_token_alone_overrides_configured_token(self, monkeypatch): + monkeypatch.delenv("DSTACK_SERVER_URL", raising=False) + monkeypatch.delenv("DSTACK_PROJECT", raising=False) + monkeypatch.setenv("DSTACK_TOKEN", "environment-token") + project = ProjectConfig( + name="configured-project", + url="https://server.example.com", + token="configured-token", + default=True, + ) + + with patch("dstack._internal.core.services.api_client.configs.ConfigManager") as manager: + manager.return_value.get_project_config.return_value = project + client, project_name = get_api_client() + + assert client.base_url == "https://server.example.com" + assert client._s.headers["Authorization"] == "Bearer environment-token" + assert project_name == "configured-project" From 2c29b3ab0518c63ef56a4bc9e8ad9519706d9ae6 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 13:23:26 +0200 Subject: [PATCH 06/14] Rename `server` to `dstack` in run configurations Enabling dstack access inside a run now reads like Docker-in-Docker (`docker: true`): `dstack: true`. No migration needed: the field is stored in run spec JSON only; pre-release DBs with the old field can be reset. Co-Authored-By: Claude Fable 5 --- mkdocs/docs/concepts/dev-environments.md | 4 ++-- mkdocs/docs/concepts/tasks.md | 4 ++-- src/dstack/_internal/core/compatibility/runs.py | 4 ++-- src/dstack/_internal/core/models/configurations.py | 8 ++++---- .../server/background/pipeline_tasks/jobs_running.py | 2 +- .../server/services/jobs/configurators/base.py | 6 +++--- src/dstack/_internal/server/services/runner/client.py | 2 +- src/tests/_internal/core/models/test_configurations.py | 8 ++++---- src/tests/_internal/core/models/test_runs.py | 10 +++++----- .../background/pipeline_tasks/test_running_jobs.py | 4 ++-- .../background/pipeline_tasks/test_terminating_jobs.py | 2 +- src/tests/_internal/server/routers/test_runs.py | 4 ++-- .../server/services/jobs/configurators/test_task.py | 6 +++--- .../_internal/server/services/runner/test_client.py | 4 ++-- 14 files changed, 34 insertions(+), 34 deletions(-) diff --git a/mkdocs/docs/concepts/dev-environments.md b/mkdocs/docs/concepts/dev-environments.md index 113909568b..f960a67c6e 100644 --- a/mkdocs/docs/concepts/dev-environments.md +++ b/mkdocs/docs/concepts/dev-environments.md @@ -644,7 +644,7 @@ via the [`spot_policy`](../reference/dstack.yml/dev-environment.md#spot_policy) ### Server access -Set `server` to `true` when a dev environment needs to use the dstack CLI. dstack configures the +Set `dstack` to `true` when a dev environment needs to use the dstack CLI. dstack configures the server and current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` explicitly. @@ -653,7 +653,7 @@ explicitly. ```yaml type: dev-environment image: dstackai/dstack -server: true +dstack: true env: - DSTACK_TOKEN init: diff --git a/mkdocs/docs/concepts/tasks.md b/mkdocs/docs/concepts/tasks.md index 2d61b5838a..78ed975281 100644 --- a/mkdocs/docs/concepts/tasks.md +++ b/mkdocs/docs/concepts/tasks.md @@ -866,7 +866,7 @@ via the [`spot_policy`](../reference/dstack.yml/task.md#spot_policy) property. I ### Server access -Set `server` to `true` when a task needs to use the dstack CLI. dstack configures the server and +Set `dstack` to `true` when a task needs to use the dstack CLI. dstack configures the server and current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` explicitly.
@@ -874,7 +874,7 @@ current project automatically. To run authenticated commands, pass `DSTACK_TOKEN ```yaml type: task image: dstackai/dstack -server: true +dstack: true env: - DSTACK_TOKEN commands: diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py index eefc326f92..1b76ede880 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -98,9 +98,9 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType: run_spec.configuration, (DevEnvironmentConfiguration, TaskConfiguration), ) - and not run_spec.configuration.server + and not run_spec.configuration.dstack ): - configuration_excludes["server"] = True + configuration_excludes["dstack"] = True if isinstance(run_spec.configuration, ServiceConfiguration): if run_spec.configuration.probes: diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 3d43ef6fe0..c0830b82fd 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -693,8 +693,8 @@ def check_image_or_commands_present(cls, values): return values -class ConfigurationWithServerParams(CoreModel): - server: Annotated[ +class ConfigurationWithDstackParams(CoreModel): + dstack: Annotated[ bool, Field( description=( @@ -774,7 +774,7 @@ class DevEnvironmentConfiguration( ProfileParams, BaseRunConfiguration, ConfigurationWithPortsParams, - ConfigurationWithServerParams, + ConfigurationWithDstackParams, DevEnvironmentConfigurationParams, generate_dual_core_model(DevEnvironmentConfigurationConfig), ): @@ -806,7 +806,7 @@ class TaskConfiguration( BaseRunConfiguration, ConfigurationWithCommandsParams, ConfigurationWithPortsParams, - ConfigurationWithServerParams, + ConfigurationWithDstackParams, TaskConfigurationParams, generate_dual_core_model(TaskConfigurationConfig), ): 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 e126cdf3c6..2a7cf456fb 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -1003,7 +1003,7 @@ async def _ensure_job_server_connection( def _server_access_enabled(context: _ProcessContext) -> bool: - return bool(getattr(context.run.run_spec.configuration, "server", False)) + return bool(getattr(context.run.run_spec.configuration, "dstack", False)) async def _apply_process_result( diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index d3e652d2c5..c4959a41c7 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/src/dstack/_internal/server/services/jobs/configurators/base.py @@ -279,12 +279,12 @@ def _app_specs(self) -> List[AppSpec]: def _env(self) -> Dict[str, str]: env = self.run_spec.configuration.env.as_dict() - if self._server(): + if self._dstack(): env.setdefault(DSTACK_SERVER_URL_ENV, DSTACK_RUN_SERVER_URL) return env - def _server(self) -> bool: - return bool(getattr(self.run_spec.configuration, "server", False)) + def _dstack(self) -> bool: + return bool(getattr(self.run_spec.configuration, "dstack", False)) def _home_dir(self) -> Optional[str]: return self.run_spec.configuration.home_dir diff --git a/src/dstack/_internal/server/services/runner/client.py b/src/dstack/_internal/server/services/runner/client.py index a0f34b158c..93cf2f7238 100644 --- a/src/dstack/_internal/server/services/runner/client.py +++ b/src/dstack/_internal/server/services/runner/client.py @@ -127,7 +127,7 @@ def submit_job( # API modification. Both layers are merged into a deep-copied job_spec # so the shared spec object held by the caller is not mutated. job_spec = job.job_spec - server_access = bool(getattr(run.run_spec.configuration, "server", False)) + server_access = bool(getattr(run.run_spec.configuration, "dstack", False)) if instance_env is not None or router_env is not None or server_access: merged_env: Dict[str, str] = {} if instance_env is not None: diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index a10f18e4c1..cbe7a31fd0 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -18,13 +18,13 @@ class TestParseConfiguration: @pytest.mark.parametrize("configuration_type", ["task", "dev-environment"]) def test_server_access_supported(self, configuration_type: str): - conf = {"type": configuration_type, "server": True} + conf = {"type": configuration_type, "dstack": True} if configuration_type == "task": conf["commands"] = ["true"] parsed = parse_run_configuration(conf) - assert parsed.server is True + assert parsed.dstack is True def test_server_access_not_supported_for_service(self): with pytest.raises(ConfigurationError, match="extra fields not permitted"): @@ -33,7 +33,7 @@ def test_server_access_not_supported_for_service(self): "type": "service", "commands": ["python3 -m http.server"], "port": 8000, - "server": True, + "dstack": True, } ) @@ -42,7 +42,7 @@ def test_server_access_allows_unresolved_passthrough_env(self): { "type": "task", "commands": ["true"], - "server": True, + "dstack": True, "env": ["DSTACK_TOKEN"], } ) diff --git a/src/tests/_internal/core/models/test_runs.py b/src/tests/_internal/core/models/test_runs.py index 0fb2300d5c..463e30cb1a 100644 --- a/src/tests/_internal/core/models/test_runs.py +++ b/src/tests/_internal/core/models/test_runs.py @@ -29,18 +29,18 @@ def test_run_termination_reason_to_status_works_with_all_enum_variants(): @pytest.mark.parametrize("configuration_type", ["task", "dev-environment"]) -@pytest.mark.parametrize("server", [False, True]) -def test_server_access_run_spec_compatibility(configuration_type: str, server: bool): +@pytest.mark.parametrize("dstack", [False, True]) +def test_server_access_run_spec_compatibility(configuration_type: str, dstack: bool): if configuration_type == "task": - configuration = TaskConfiguration(commands=["true"], server=server) + configuration = TaskConfiguration(commands=["true"], dstack=dstack) else: - configuration = DevEnvironmentConfiguration(server=server) + configuration = DevEnvironmentConfiguration(dstack=dstack) configuration_excludes = get_run_spec_excludes(RunSpec(configuration=configuration)).get( "configuration" ) assert isinstance(configuration_excludes, dict) - assert ("server" in configuration_excludes) is not server + assert ("dstack" in configuration_excludes) is not dstack def test_job_termination_reason_to_status_works_with_all_enum_variants(): diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index c0d4f8537a..95a93d6e37 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -1002,7 +1002,7 @@ async def test_pulling_waits_for_requested_server_access_before_starting_job( configuration=TaskConfiguration( image="debian", commands=["true"], - server=True, + dstack=True, ), ), ) @@ -1084,7 +1084,7 @@ async def test_provisioning_server_access_failure_terminates_job_after_retry_tim configuration=TaskConfiguration( image="debian", commands=["true"], - server=True, + dstack=True, ), ), ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py index 9c79cb5e7b..0ecb64cd24 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py @@ -266,7 +266,7 @@ async def test_closes_requested_server_access( configuration=TaskConfiguration( image="debian", commands=["true"], - server=True, + dstack=True, ), ), ) diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 0fdeea68af..16d507d1a6 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -173,7 +173,7 @@ def get_dev_env_run_plan_dict( "privileged": privileged, "init": [], "ports": [], - "server": False, + "dstack": False, "python": "3.13" if not docker else None, "nvcc": None, "registry_auth": None, @@ -425,7 +425,7 @@ def get_dev_env_run_dict( "privileged": privileged, "init": [], "ports": [], - "server": False, + "dstack": False, "python": "3.13" if not docker else None, "nvcc": None, "registry_auth": None, diff --git a/src/tests/_internal/server/services/jobs/configurators/test_task.py b/src/tests/_internal/server/services/jobs/configurators/test_task.py index 2c89c5c65e..6fab966f5e 100644 --- a/src/tests/_internal/server/services/jobs/configurators/test_task.py +++ b/src/tests/_internal/server/services/jobs/configurators/test_task.py @@ -41,13 +41,13 @@ async def test_multi_node(self): @pytest.mark.usefixtures("image_config_mock") class TestServerAccess: async def test_adds_transport_without_credentials(self): - configuration = TaskConfiguration(image="debian", commands=["true"], server=True) + configuration = TaskConfiguration(image="debian", commands=["true"], dstack=True) run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration) configurator = TaskJobConfigurator(run_spec) job_spec = (await configurator.get_job_specs(replica_num=0))[0] - assert "server" not in job_spec.dict() + assert "dstack" not in job_spec.dict() assert job_spec.env == { "DSTACK_SERVER_URL": "http+unix://%2Frun%2Fdstack%2Fserver.sock", } @@ -65,7 +65,7 @@ async def test_preserves_explicit_transport_env(self): configuration = TaskConfiguration( image="debian", commands=["true"], - server=True, + dstack=True, env={ "DSTACK_SERVER_URL": "https://server.example.com", }, diff --git a/src/tests/_internal/server/services/runner/test_client.py b/src/tests/_internal/server/services/runner/test_client.py index dd7805b5b2..e7db24226b 100644 --- a/src/tests/_internal/server/services/runner/test_client.py +++ b/src/tests/_internal/server/services/runner/test_client.py @@ -77,7 +77,7 @@ class TestRunnerClientSubmitJob(BaseShimClientTest): def test_adds_default_project_for_server_access(self, adapter: requests_mock.Adapter): adapter.register_uri("POST", "/api/submit", json={}) run_spec = get_run_spec( - repo_id="repo", configuration=TaskConfiguration(commands=["true"], server=True) + repo_id="repo", configuration=TaskConfiguration(commands=["true"], dstack=True) ) run = Run.construct(id=uuid.uuid4(), project_name="main", run_spec=run_spec) job = Job.construct( @@ -109,7 +109,7 @@ def test_adds_default_project_for_server_access(self, adapter: requests_mock.Ada def test_preserves_explicit_project_for_server_access(self, adapter: requests_mock.Adapter): adapter.register_uri("POST", "/api/submit", json={}) run_spec = get_run_spec( - repo_id="repo", configuration=TaskConfiguration(commands=["true"], server=True) + repo_id="repo", configuration=TaskConfiguration(commands=["true"], dstack=True) ) run = Run.construct(id=uuid.uuid4(), project_name="main", run_spec=run_spec) job = Job.construct( From c95885587484b9c9565d4ee7bbf1a444df40ed48 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 13:30:14 +0200 Subject: [PATCH 07/14] Rename docs section to `dstack` inside `dstack` Co-Authored-By: Claude Fable 5 --- mkdocs/docs/concepts/dev-environments.md | 2 +- mkdocs/docs/concepts/tasks.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mkdocs/docs/concepts/dev-environments.md b/mkdocs/docs/concepts/dev-environments.md index f960a67c6e..e03ac7dd17 100644 --- a/mkdocs/docs/concepts/dev-environments.md +++ b/mkdocs/docs/concepts/dev-environments.md @@ -642,7 +642,7 @@ The `schedule` property can be combined with `max_duration` or `utilization_poli By default, `dstack` uses on-demand instances. However, you can change that via the [`spot_policy`](../reference/dstack.yml/dev-environment.md#spot_policy) property. It accepts `spot`, `on-demand`, and `auto`. -### Server access +### `dstack` inside `dstack` Set `dstack` to `true` when a dev environment needs to use the dstack CLI. dstack configures the server and current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` diff --git a/mkdocs/docs/concepts/tasks.md b/mkdocs/docs/concepts/tasks.md index 78ed975281..e9ab6e791a 100644 --- a/mkdocs/docs/concepts/tasks.md +++ b/mkdocs/docs/concepts/tasks.md @@ -864,7 +864,7 @@ schedule: By default, `dstack` uses on-demand instances. However, you can change that via the [`spot_policy`](../reference/dstack.yml/task.md#spot_policy) property. It accepts `spot`, `on-demand`, and `auto`. -### Server access +### `dstack` inside `dstack` Set `dstack` to `true` when a task needs to use the dstack CLI. dstack configures the server and current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` explicitly. From 93d455d813b3b4d999c95b13c8586f58d884792e Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 14:32:50 +0200 Subject: [PATCH 08/14] Follow the server bind address for job server access The reverse forward targeted 127.0.0.1 while the server may be bound to a specific address via DSTACK_SERVER_HOST, making the forward target unreachable. Derive the target from the bind address and include it in the error when the server is not reachable through the tunnel. Co-Authored-By: Claude Fable 5 --- .../server/services/jobs/server_connection.py | 18 ++++++++++++++-- .../services/jobs/test_server_connection.py | 21 +++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/src/dstack/_internal/server/services/jobs/server_connection.py b/src/dstack/_internal/server/services/jobs/server_connection.py index 56f45bd88d..5b29879d42 100644 --- a/src/dstack/_internal/server/services/jobs/server_connection.py +++ b/src/dstack/_internal/server/services/jobs/server_connection.py @@ -33,6 +33,16 @@ _REMOTE_SOCKET_PATH = Path(DSTACK_RUN_SERVER_SOCKET_PATH) +def _get_server_socket() -> IPSocket: + # The server may be bound to a specific address, making loopback unreachable + host = settings.SERVER_HOST + if not host or host in ("0.0.0.0", "localhost"): + host = "127.0.0.1" + elif host == "::": + host = "::1" + return IPSocket(host=host, port=settings.SERVER_PORT) + + class JobServerConnection: """A private reverse SSH tunnel from one job to the dstack server.""" @@ -85,9 +95,10 @@ async def open(self) -> None: await self._tunnel.aexec( f"mkdir -p {remote_dir} && chmod 755 {remote_dir} && rm -f {remote_socket}" ) + server_socket = _get_server_socket() self._tunnel.reverse_forwarded_sockets = [ SocketPair( - local=IPSocket(host="127.0.0.1", port=settings.SERVER_PORT), + local=server_socket, remote=UnixSocket(path=_REMOTE_SOCKET_PATH), ) ] @@ -104,7 +115,10 @@ async def open(self) -> None: # lets configurations using a non-root `user` reach it as well. await self._tunnel.aexec(f"chmod 666 {remote_socket}") if not await self._server_is_reachable(): - raise SSHError("dstack server is not reachable from the job") + raise SSHError( + "dstack server is not reachable from the job" + f" (forward target {server_socket.render()})" + ) except Exception: await self._tunnel.aclose() raise diff --git a/src/tests/_internal/server/services/jobs/test_server_connection.py b/src/tests/_internal/server/services/jobs/test_server_connection.py index c8e5a8b9a1..da8f65cc8e 100644 --- a/src/tests/_internal/server/services/jobs/test_server_connection.py +++ b/src/tests/_internal/server/services/jobs/test_server_connection.py @@ -41,6 +41,27 @@ def tunnel_mock(tmp_path, monkeypatch: pytest.MonkeyPatch): return tunnel, tunnel_class +@pytest.mark.parametrize( + ("server_host", "expected_host"), + [ + ("localhost", "127.0.0.1"), + ("127.0.0.1", "127.0.0.1"), + ("0.0.0.0", "127.0.0.1"), + ("", "127.0.0.1"), + ("::", "::1"), + ("10.0.0.5", "10.0.0.5"), + ], +) +def test_get_server_socket_follows_server_bind_host( + monkeypatch: pytest.MonkeyPatch, server_host: str, expected_host: str +): + monkeypatch.setattr(server_connection.settings, "SERVER_HOST", server_host) + + assert server_connection._get_server_socket() == IPSocket( + host=expected_host, port=server_connection.settings.SERVER_PORT + ) + + class TestJobServerConnection: @pytest.mark.asyncio async def test_opens_private_reverse_socket(self, tunnel_mock): From e45abbf41ee19ec83845352c651504469ec096e7 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 20:42:28 +0200 Subject: [PATCH 09/14] Drop redundant pool lock in JobServerConnectionsPool _get_lock and remove_all only run await-free expressions, so the extra lock adds nothing under the single-threaded event loop. Make _get_lock synchronous and drop the lock. Co-Authored-By: Claude Fable 5 --- .../server/services/jobs/server_connection.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/dstack/_internal/server/services/jobs/server_connection.py b/src/dstack/_internal/server/services/jobs/server_connection.py index 5b29879d42..ce56b1c2fd 100644 --- a/src/dstack/_internal/server/services/jobs/server_connection.py +++ b/src/dstack/_internal/server/services/jobs/server_connection.py @@ -157,14 +157,13 @@ def __init__(self) -> None: self._connections: dict[uuid.UUID, JobServerConnection] = {} self._failure_started_at: dict[uuid.UUID, float] = {} self._locks: WeakValueDictionary[uuid.UUID, asyncio.Lock] = WeakValueDictionary() - self._pool_lock = asyncio.Lock() async def ensure( self, job: JobModel, job_runtime_data: Optional[JobRuntimeData], ) -> bool: - lock = await self._get_lock(job.id) + lock = self._get_lock(job.id) async with lock: connection = self._connections.get(job.id) if connection is not None and await connection.is_alive(): @@ -193,7 +192,7 @@ def retry_timed_out(self, job_id: uuid.UUID, timeout: float) -> bool: return time.monotonic() - failure_started_at > timeout async def remove(self, job_id: uuid.UUID) -> None: - lock = await self._get_lock(job_id) + lock = self._get_lock(job_id) async with lock: connection = self._connections.pop(job_id, None) if connection is not None: @@ -202,13 +201,12 @@ async def remove(self, job_id: uuid.UUID) -> None: shutil.rmtree(CONNECTIONS_DIR / str(job_id), ignore_errors=True) async def remove_all(self) -> None: - async with self._pool_lock: - job_ids = set(self._connections).union(self._failure_started_at) + job_ids = set(self._connections).union(self._failure_started_at) await asyncio.gather(*(self.remove(job_id) for job_id in job_ids)) - async def _get_lock(self, job_id: uuid.UUID) -> asyncio.Lock: - async with self._pool_lock: - return self._locks.setdefault(job_id, asyncio.Lock()) + def _get_lock(self, job_id: uuid.UUID) -> asyncio.Lock: + # setdefault is atomic under the single-threaded event loop, so no extra lock is needed + return self._locks.setdefault(job_id, asyncio.Lock()) @staticmethod async def _close(connection: JobServerConnection) -> None: From e428da5e6e8d13d49a9db0ef640db65e4b628b72 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 21:06:32 +0200 Subject: [PATCH 10/14] Reject dstack access combined with inactivity_duration The persistent server connection counts as SSH activity, so a dev environment with both dstack access and inactivity_duration would never be detected as inactive. Reject the combination at configuration time instead of silently breaking inactivity_duration. Co-Authored-By: Claude Fable 5 --- src/dstack/_internal/core/models/configurations.py | 7 +++++++ .../_internal/core/models/test_configurations.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index c0830b82fd..5491292f38 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -786,6 +786,13 @@ def validate_entrypoint(cls, v: Optional[str]) -> Optional[str]: raise ValueError("entrypoint is not supported for dev-environment") return v + @root_validator + def validate_dstack_and_inactivity_duration(cls, values): + if values.get("dstack") and values.get("inactivity_duration") is not None: + # The persistent server connection counts as activity, so inactivity is never detected + raise ValueError("`dstack` is not supported together with `inactivity_duration`") + return values + class TaskConfigurationParams(CoreModel): nodes: Annotated[int, Field(description="Number of nodes", ge=1)] = 1 diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index cbe7a31fd0..b0a27fcff8 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -37,6 +37,20 @@ def test_server_access_not_supported_for_service(self): } ) + def test_server_access_not_supported_with_inactivity_duration(self): + with pytest.raises(ConfigurationError, match="inactivity_duration"): + parse_run_configuration( + { + "type": "dev-environment", + "dstack": True, + "inactivity_duration": "1h", + } + ) + + def test_server_access_and_inactivity_duration_allowed_separately(self): + parse_run_configuration({"type": "dev-environment", "dstack": True}) + parse_run_configuration({"type": "dev-environment", "inactivity_duration": "1h"}) + def test_server_access_allows_unresolved_passthrough_env(self): parsed = parse_run_configuration( { From 42c8f6e84b0ffe2a2d71a8d605bcdf5bd1b17d2f Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 13 Jul 2026 23:18:30 +0200 Subject: [PATCH 11/14] Adopt an existing job server socket instead of overwriting it With multiple server replicas, each replica used to overwrite the job's /run/dstack/server.sock on every open, so ownership churned between replicas (they repeatedly stole the socket from each other) and every overwrite briefly broke access. Probe the socket first and (re)create the reverse forward only when it is missing or unreachable; otherwise adopt the existing one. This keeps a single stable owner. If the owner replica dies, another takes over on its next health check. Co-Authored-By: Claude Fable 5 --- .../server/services/jobs/server_connection.py | 43 +++++++------ .../services/jobs/test_server_connection.py | 61 +++++++++++++++---- 2 files changed, 74 insertions(+), 30 deletions(-) diff --git a/src/dstack/_internal/server/services/jobs/server_connection.py b/src/dstack/_internal/server/services/jobs/server_connection.py index ce56b1c2fd..ccfed04827 100644 --- a/src/dstack/_internal/server/services/jobs/server_connection.py +++ b/src/dstack/_internal/server/services/jobs/server_connection.py @@ -91,17 +91,6 @@ async def open(self) -> None: try: remote_dir = shlex.quote(str(_REMOTE_SOCKET_PATH.parent)) remote_socket = shlex.quote(str(_REMOTE_SOCKET_PATH)) - # A new server owner replaces the stable path, making an orphaned forward unreachable. - await self._tunnel.aexec( - f"mkdir -p {remote_dir} && chmod 755 {remote_dir} && rm -f {remote_socket}" - ) - server_socket = _get_server_socket() - self._tunnel.reverse_forwarded_sockets = [ - SocketPair( - local=server_socket, - remote=UnixSocket(path=_REMOTE_SOCKET_PATH), - ) - ] # Probe through the job socket itself: a socket path can remain after its listener # becomes unreachable. self._tunnel.forwarded_sockets = [ @@ -111,14 +100,34 @@ async def open(self) -> None: ) ] await self._tunnel.aopen() - # The socket carries no credentials. World access inside the isolated job container - # lets configurations using a non-root `user` reach it as well. - await self._tunnel.aexec(f"chmod 666 {remote_socket}") + # With multiple server replicas, do not overwrite a socket that another replica already + # serves: create the reverse forward only when the current one is missing or + # unreachable. This keeps a single stable owner and avoids ownership churn (replicas + # repeatedly stealing the socket) and the brief unavailability of overwriting a live one. if not await self._server_is_reachable(): - raise SSHError( - "dstack server is not reachable from the job" - f" (forward target {server_socket.render()})" + # A new server owner replaces the stable path, making an orphaned forward + # unreachable. + await self._tunnel.aexec( + f"mkdir -p {remote_dir} && chmod 755 {remote_dir} && rm -f {remote_socket}" ) + server_socket = _get_server_socket() + self._tunnel.reverse_forwarded_sockets = [ + SocketPair( + local=server_socket, + remote=UnixSocket(path=_REMOTE_SOCKET_PATH), + ) + ] + # The probe forward is already established; do not request it again. + self._tunnel.forwarded_sockets = [] + await self._tunnel.aopen() + # The socket carries no credentials. World access inside the isolated job container + # lets configurations using a non-root `user` reach it as well. + await self._tunnel.aexec(f"chmod 666 {remote_socket}") + if not await self._server_is_reachable(): + raise SSHError( + "dstack server is not reachable from the job" + f" (forward target {server_socket.render()})" + ) except Exception: await self._tunnel.aclose() raise diff --git a/src/tests/_internal/server/services/jobs/test_server_connection.py b/src/tests/_internal/server/services/jobs/test_server_connection.py index da8f65cc8e..fb855f40db 100644 --- a/src/tests/_internal/server/services/jobs/test_server_connection.py +++ b/src/tests/_internal/server/services/jobs/test_server_connection.py @@ -20,6 +20,8 @@ def tunnel_mock(tmp_path, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(server_connection, "CONNECTIONS_DIR", tmp_path) tunnel = MagicMock() + tunnel.forwarded_sockets = [] + tunnel.reverse_forwarded_sockets = [] tunnel.acheck = AsyncMock(return_value=False) tunnel.aopen = AsyncMock() tunnel.aexec = AsyncMock(return_value="") @@ -64,33 +66,66 @@ def test_get_server_socket_follows_server_bind_host( class TestJobServerConnection: @pytest.mark.asyncio - async def test_opens_private_reverse_socket(self, tunnel_mock): + async def test_becomes_owner_when_socket_unreachable(self, tunnel_mock): tunnel, tunnel_class = tunnel_mock job = Mock(id=uuid.uuid4()) connection = JobServerConnection(job, job_runtime_data=None) - connection._server_is_reachable = AsyncMock(return_value=True) + # The probe after the local forward is unreachable (no healthy owner) -> become owner; + # the probe after the reverse forward confirms reachability. + connection._server_is_reachable = AsyncMock(side_effect=[False, True]) + + snapshots = [] + + async def record_aopen(): + snapshots.append( + (list(tunnel.forwarded_sockets), list(tunnel.reverse_forwarded_sockets)) + ) + + tunnel.aopen.side_effect = record_aopen await connection.open() tunnel_class.assert_called_once() - assert tunnel.aopen.await_count == 2 - assert tunnel.reverse_forwarded_sockets == [ - SocketPair( - local=IPSocket(host="127.0.0.1", port=server_connection.settings.SERVER_PORT), - remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH), - ) + probe_pair = SocketPair( + local=UnixSocket(path=connection._probe_socket_path), + remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH), + ) + reverse_pair = SocketPair( + local=IPSocket(host="127.0.0.1", port=server_connection.settings.SERVER_PORT), + remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH), + ) + # master (no forwards) -> add probe forward -> add reverse forward + assert snapshots == [ + ([], []), + ([probe_pair], []), + ([], [reverse_pair]), + ] + commands = [call.args[0] for call in tunnel.aexec.await_args_list] + assert commands == [ + "mkdir -p /run/dstack && chmod 755 /run/dstack && rm -f /run/dstack/server.sock", + "chmod 666 /run/dstack/server.sock", ] + + @pytest.mark.asyncio + async def test_adopts_existing_healthy_socket(self, tunnel_mock): + tunnel, _ = tunnel_mock + job = Mock(id=uuid.uuid4()) + connection = JobServerConnection(job, job_runtime_data=None) + # A healthy owner (another replica) already serves the socket. + connection._server_is_reachable = AsyncMock(return_value=True) + + await connection.open() + + # master + probe forward only; the reverse forward is not (re)created + assert tunnel.aopen.await_count == 2 + assert tunnel.reverse_forwarded_sockets == [] + tunnel.aexec.assert_not_awaited() assert tunnel.forwarded_sockets == [ SocketPair( local=UnixSocket(path=connection._probe_socket_path), remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH), ) ] - commands = [call.args[0] for call in tunnel.aexec.await_args_list] - assert commands == [ - "mkdir -p /run/dstack && chmod 755 /run/dstack && rm -f /run/dstack/server.sock", - "chmod 666 /run/dstack/server.sock", - ] @pytest.mark.asyncio async def test_reuses_live_tunnel_with_existing_socket(self, tunnel_mock): From e69acf8b327ca526834748372efdbdb9aa32b379 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 14 Jul 2026 00:35:55 +0200 Subject: [PATCH 12/14] Support dstack access for services Allow `dstack: true` on service configurations, mirroring tasks and dev environments. Each replica is a separate job that gets its own `/run/dstack/server.sock`, so multiple replicas behave like multi-node job replicas. Exclude the field from run specs sent to older servers, same as for tasks and dev environments. Co-Authored-By: Claude Fable 5 --- mkdocs/docs/concepts/services.md | 21 +++++++++++++++++++ .../_internal/core/compatibility/runs.py | 2 +- .../_internal/core/models/configurations.py | 1 + .../core/models/test_configurations.py | 16 ++++---------- src/tests/_internal/core/models/test_runs.py | 5 ++++- 5 files changed, 31 insertions(+), 14 deletions(-) diff --git a/mkdocs/docs/concepts/services.md b/mkdocs/docs/concepts/services.md index 000ad7de8f..6d965debbc 100644 --- a/mkdocs/docs/concepts/services.md +++ b/mkdocs/docs/concepts/services.md @@ -1048,6 +1048,27 @@ resources: gpu: H100:1 ``` +### `dstack` inside `dstack` + +Set `dstack` to `true` when a service needs to use the dstack CLI. dstack configures the server and +current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` explicitly. + +
+ +```yaml +type: service +image: dstackai/dstack +dstack: true +port: 8000 +env: + - DSTACK_TOKEN +commands: + - dstack ps + - python -m http.server 8000 +``` + +
+ ### Environment variables
diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py index 1b76ede880..693e3252e2 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -96,7 +96,7 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType: if ( isinstance( run_spec.configuration, - (DevEnvironmentConfiguration, TaskConfiguration), + (DevEnvironmentConfiguration, TaskConfiguration, ServiceConfiguration), ) and not run_spec.configuration.dstack ): diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 5491292f38..126fba380a 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -1359,6 +1359,7 @@ class ServiceConfiguration( ProfileParams, BaseRunConfiguration, ConfigurationWithCommandsParams, + ConfigurationWithDstackParams, ServiceConfigurationParams, generate_dual_core_model(ServiceConfigurationConfig), ): diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index b0a27fcff8..04095777ef 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -16,27 +16,19 @@ class TestParseConfiguration: - @pytest.mark.parametrize("configuration_type", ["task", "dev-environment"]) + @pytest.mark.parametrize("configuration_type", ["task", "dev-environment", "service"]) def test_server_access_supported(self, configuration_type: str): conf = {"type": configuration_type, "dstack": True} if configuration_type == "task": conf["commands"] = ["true"] + elif configuration_type == "service": + conf["commands"] = ["true"] + conf["port"] = 8000 parsed = parse_run_configuration(conf) assert parsed.dstack is True - def test_server_access_not_supported_for_service(self): - with pytest.raises(ConfigurationError, match="extra fields not permitted"): - parse_run_configuration( - { - "type": "service", - "commands": ["python3 -m http.server"], - "port": 8000, - "dstack": True, - } - ) - def test_server_access_not_supported_with_inactivity_duration(self): with pytest.raises(ConfigurationError, match="inactivity_duration"): parse_run_configuration( diff --git a/src/tests/_internal/core/models/test_runs.py b/src/tests/_internal/core/models/test_runs.py index 463e30cb1a..973aaa9140 100644 --- a/src/tests/_internal/core/models/test_runs.py +++ b/src/tests/_internal/core/models/test_runs.py @@ -4,6 +4,7 @@ from dstack._internal.core.compatibility.runs import get_run_spec_excludes from dstack._internal.core.models.configurations import ( DevEnvironmentConfiguration, + ServiceConfiguration, TaskConfiguration, ) from dstack._internal.core.models.profiles import RetryEvent @@ -28,11 +29,13 @@ def test_run_termination_reason_to_status_works_with_all_enum_variants(): assert isinstance(run_status, RunStatus) -@pytest.mark.parametrize("configuration_type", ["task", "dev-environment"]) +@pytest.mark.parametrize("configuration_type", ["task", "dev-environment", "service"]) @pytest.mark.parametrize("dstack", [False, True]) def test_server_access_run_spec_compatibility(configuration_type: str, dstack: bool): if configuration_type == "task": configuration = TaskConfiguration(commands=["true"], dstack=dstack) + elif configuration_type == "service": + configuration = ServiceConfiguration(commands=["true"], port=8000, dstack=dstack) else: configuration = DevEnvironmentConfiguration(dstack=dstack) configuration_excludes = get_run_spec_excludes(RunSpec(configuration=configuration)).get( From 1c541bbf651e28ea0389af63c82e0e228fa513e6 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 14 Jul 2026 10:55:58 +0200 Subject: [PATCH 13/14] Add DSTACK_FORBID_DSTACK_IN_RUNS server setting Let operators forbid `dstack: true` in run configurations, mirroring DSTACK_FORBID_SERVICES_WITHOUT_GATEWAY. When set, submitting a run with dstack access enabled is rejected server-side. Co-Authored-By: Claude Fable 5 --- mkdocs/docs/reference/env.md | 1 + .../server/services/runs/__init__.py | 5 ++++ src/dstack/_internal/server/settings.py | 1 + .../_internal/server/routers/test_runs.py | 30 +++++++++++++++++++ 4 files changed, 37 insertions(+) diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 36a9f65d1c..ec344fbe87 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -124,6 +124,7 @@ For more details on the options below, refer to the [server deployment](../guide - `DSTACK_DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE`{ #DSTACK_DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE } – Request body size limit for services running with a gateway, in bytes. Defaults to 64 MiB. - `DSTACK_SERVICE_CLIENT_TIMEOUT`{ #DSTACK_SERVICE_CLIENT_TIMEOUT } – Timeout in seconds for HTTP requests sent from the in-server proxy and gateways to service replicas. Defaults to 60. - `DSTACK_FORBID_SERVICES_WITHOUT_GATEWAY`{ #DSTACK_FORBID_SERVICES_WITHOUT_GATEWAY } – Forbids registering new services without a gateway if set to any value. +- `DSTACK_FORBID_DSTACK_IN_RUNS`{ #DSTACK_FORBID_DSTACK_IN_RUNS } – Forbids submitting runs with `dstack: true` (dstack server access inside runs) if set to any value. - `DSTACK_SERVER_CODE_UPLOAD_LIMIT`{ #DSTACK_SERVER_CODE_UPLOAD_LIMIT } - The repo size limit when uploading diffs or local repos, in bytes. Set to `0` to disable size limits. Defaults to `2MiB`. - `DSTACK_SERVER_S3_BUCKET`{ #DSTACK_SERVER_S3_BUCKET } - The bucket that repo diffs will be uploaded to if set. If unset, diffs are uploaded to the database. - `DSTACK_SERVER_S3_BUCKET_REGION`{ #DSTACK_SERVER_S3_BUCKET_REGION } - The region of the S3 Bucket. diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index e7877605d5..e2cb927446 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -40,6 +40,7 @@ ) from dstack._internal.core.models.users import GlobalRole from dstack._internal.core.services.diff import format_diff_fields_for_event +from dstack._internal.server import settings as server_settings from dstack._internal.server.db import get_db, is_db_postgres, is_db_sqlite from dstack._internal.server.models import ( FleetModel, @@ -1157,6 +1158,10 @@ async def _validate_run( project: ProjectModel, run_spec: RunSpec, ): + if server_settings.FORBID_DSTACK_IN_RUNS and getattr(run_spec.configuration, "dstack", False): + raise ServerClientError( + "This dstack-server installation forbids `dstack: true` in run configurations." + ) await _validate_run_volumes( session=session, project=project, diff --git a/src/dstack/_internal/server/settings.py b/src/dstack/_internal/server/settings.py index 27a97a6db5..671eacef07 100644 --- a/src/dstack/_internal/server/settings.py +++ b/src/dstack/_internal/server/settings.py @@ -139,6 +139,7 @@ USER_PROJECT_DEFAULT_QUOTA = int(os.getenv("DSTACK_USER_PROJECT_DEFAULT_QUOTA", 10)) FORBID_SERVICES_WITHOUT_GATEWAY = os.getenv("DSTACK_FORBID_SERVICES_WITHOUT_GATEWAY") is not None +FORBID_DSTACK_IN_RUNS = os.getenv("DSTACK_FORBID_DSTACK_IN_RUNS") is not None SERVER_CODE_UPLOAD_LIMIT = int(os.getenv("DSTACK_SERVER_CODE_UPLOAD_LIMIT", 2 * 2**20)) diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 16d507d1a6..4c7ddd857b 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -3464,6 +3464,36 @@ async def test_returns_400_if_bad_run_name( ) assert response.status_code == 400 + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_400_if_dstack_in_runs_forbidden( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session=session, global_role=GlobalRole.USER) + project = await create_project(session=session, owner=user) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + repo = await create_repo(session=session, project_id=project.id) + run_dict = get_dev_env_run_dict( + project_name=project.name, + username=user.name, + run_name="test-run", + repo_id=repo.name, + ) + run_dict["run_spec"]["configuration"]["dstack"] = True + body = {"run_spec": run_dict["run_spec"]} + with patch( + "dstack._internal.server.services.runs.server_settings.FORBID_DSTACK_IN_RUNS", True + ): + response = await client.post( + f"/api/project/{project.name}/runs/submit", + headers=get_auth_headers(user.token), + json=body, + ) + assert response.status_code == 400 + assert "forbids" in response.json()["detail"][0]["msg"] + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) async def test_returns_400_if_repo_does_not_exist( From f8f969d56c74ceb219665971870895a914af0711 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 14 Jul 2026 12:35:30 +0200 Subject: [PATCH 14/14] Document dstack in dstack in Protips and note apply/attach Add a 'dstack in dstack' Protips section, and note in each dstack-access section that runs can also submit runs with dstack apply and attach to them with dstack attach, not just inspect them. Co-Authored-By: Claude Fable 5 --- mkdocs/docs/concepts/dev-environments.md | 2 ++ mkdocs/docs/concepts/services.md | 2 ++ mkdocs/docs/concepts/tasks.md | 2 ++ mkdocs/docs/guides/protips.md | 20 ++++++++++++++++++++ 4 files changed, 26 insertions(+) diff --git a/mkdocs/docs/concepts/dev-environments.md b/mkdocs/docs/concepts/dev-environments.md index e03ac7dd17..e8ddb70456 100644 --- a/mkdocs/docs/concepts/dev-environments.md +++ b/mkdocs/docs/concepts/dev-environments.md @@ -662,6 +662,8 @@ init:
+> Besides inspecting runs, you can submit new runs with `dstack apply` and attach to them with `dstack attach`. + --8<-- "docs/concepts/snippets/manage-fleets.ext" !!! info "Reference" diff --git a/mkdocs/docs/concepts/services.md b/mkdocs/docs/concepts/services.md index 6d965debbc..22116567cb 100644 --- a/mkdocs/docs/concepts/services.md +++ b/mkdocs/docs/concepts/services.md @@ -1069,6 +1069,8 @@ commands:
+> Besides inspecting runs, you can submit new runs with `dstack apply` and attach to them with `dstack attach`. + ### Environment variables
diff --git a/mkdocs/docs/concepts/tasks.md b/mkdocs/docs/concepts/tasks.md index e9ab6e791a..6592e431b1 100644 --- a/mkdocs/docs/concepts/tasks.md +++ b/mkdocs/docs/concepts/tasks.md @@ -883,6 +883,8 @@ commands:
+> Besides inspecting runs, you can submit new runs with `dstack apply` and attach to them with `dstack attach`. + --8<-- "docs/concepts/snippets/manage-fleets.ext" !!! info "Reference" diff --git a/mkdocs/docs/guides/protips.md b/mkdocs/docs/guides/protips.md index 2b5c7c5b4c..02995f0c74 100644 --- a/mkdocs/docs/guides/protips.md +++ b/mkdocs/docs/guides/protips.md @@ -188,6 +188,26 @@ Set `docker` to `true` to enable the `docker` CLI in your dev environment, e.g., optional: true ``` +## dstack in dstack + +Set `dstack` to `true` when a run needs to use the dstack CLI. dstack configures the server and current project automatically. To run authenticated commands, pass `DSTACK_TOKEN` explicitly. + +
+ +```yaml +type: task +image: dstackai/dstack +dstack: true +env: + - DSTACK_TOKEN +commands: + - dstack ps +``` + +
+ +> Besides inspecting runs, you can submit new runs with `dstack apply` and attach to them with `dstack attach`. + ## Fleets ### Creation policy