+
+```yaml
+type: dev-environment
+image: dstackai/dstack
+dstack: true
+env:
+ - DSTACK_TOKEN
+init:
+ - dstack ps
+```
+
+
+
+> 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 000ad7de8f..22116567cb 100644
--- a/mkdocs/docs/concepts/services.md
+++ b/mkdocs/docs/concepts/services.md
@@ -1048,6 +1048,29 @@ 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
+```
+
+
+
+> 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 d3a3515b82..6592e431b1 100644
--- a/mkdocs/docs/concepts/tasks.md
+++ b/mkdocs/docs/concepts/tasks.md
@@ -864,6 +864,27 @@ 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`.
+### `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.
+
+
+
+```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`.
+
--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
diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md
index 86a7dd051d..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.
@@ -168,6 +169,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/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py
index 4b57db1d47..693e3252e2 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, ServiceConfiguration),
+ )
+ and not run_spec.configuration.dstack
+ ):
+ configuration_excludes["dstack"] = True
+
if isinstance(run_spec.configuration, ServiceConfiguration):
if run_spec.configuration.probes:
probe_excludes: IncludeExcludeDictType = {}
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..126fba380a 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 ConfigurationWithDstackParams(CoreModel):
+ dstack: 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,
+ ConfigurationWithDstackParams,
DevEnvironmentConfigurationParams,
generate_dual_core_model(DevEnvironmentConfigurationConfig),
):
@@ -773,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
@@ -793,6 +813,7 @@ class TaskConfiguration(
BaseRunConfiguration,
ConfigurationWithCommandsParams,
ConfigurationWithPortsParams,
+ ConfigurationWithDstackParams,
TaskConfigurationParams,
generate_dual_core_model(TaskConfigurationConfig),
):
@@ -1338,6 +1359,7 @@ class ServiceConfiguration(
ProfileParams,
BaseRunConfiguration,
ConfigurationWithCommandsParams,
+ ConfigurationWithDstackParams,
ServiceConfigurationParams,
generate_dual_core_model(ServiceConfigurationConfig),
):
diff --git a/src/dstack/_internal/core/services/api_client.py b/src/dstack/_internal/core/services/api_client.py
index a3a2c097b2..72d2e3e999 100644
--- a/src/dstack/_internal/core/services/api_client.py
+++ b/src/dstack/_internal/core/services/api_client.py
@@ -1,15 +1,40 @@
+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 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/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..2a7cf456fb 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,10 @@ 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) 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
@@ -614,6 +626,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 +793,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 +906,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 +976,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, "dstack", 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..c4959a41c7 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._dstack():
+ env.setdefault(DSTACK_SERVER_URL_ENV, DSTACK_RUN_SERVER_URL)
+ return env
+
+ 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/jobs/server_connection.py b/src/dstack/_internal/server/services/jobs/server_connection.py
new file mode 100644
index 0000000000..ccfed04827
--- /dev/null
+++ b/src/dstack/_internal/server/services/jobs/server_connection.py
@@ -0,0 +1,228 @@
+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)
+
+
+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."""
+
+ 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))
+ # 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()
+ # 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():
+ # 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
+ 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()
+
+ async def ensure(
+ self,
+ job: JobModel,
+ job_runtime_data: Optional[JobRuntimeData],
+ ) -> bool:
+ 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():
+ 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 = 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:
+ job_ids = set(self._connections).union(self._failure_started_at)
+ await asyncio.gather(*(self.remove(job_id) for job_id in job_ids))
+
+ 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:
+ 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..93cf2f7238 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, "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:
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/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/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/_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/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..04095777ef 100644
--- a/src/tests/_internal/core/models/test_configurations.py
+++ b/src/tests/_internal/core/models/test_configurations.py
@@ -16,6 +16,45 @@
class TestParseConfiguration:
+ @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_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(
+ {
+ "type": "task",
+ "commands": ["true"],
+ "dstack": 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/models/test_runs.py b/src/tests/_internal/core/models/test_runs.py
index 7968845f4e..973aaa9140 100644
--- a/src/tests/_internal/core/models/test_runs.py
+++ b/src/tests/_internal/core/models/test_runs.py
@@ -1,6 +1,12 @@
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,
+ ServiceConfiguration,
+ TaskConfiguration,
+)
from dstack._internal.core.models.profiles import RetryEvent
from dstack._internal.core.models.runs import (
JobStatus,
@@ -23,6 +29,23 @@ 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", "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(
+ "configuration"
+ )
+
+ assert isinstance(configuration_excludes, dict)
+ assert ("dstack" in configuration_excludes) is not dstack
+
+
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()
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..94668481c4
--- /dev/null
+++ b/src/tests/_internal/core/services/test_api_client.py
@@ -0,0 +1,55 @@
+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"
+
+ 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"
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..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
@@ -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
@@ -977,6 +980,184 @@ 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"],
+ dstack=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."
+ "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,
+ ) 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()
+ # 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()
+ await session.refresh(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"],
+ dstack=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,
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..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
@@ -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"],
+ dstack=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/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py
index fdc5a8f355..4c7ddd857b 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": [],
+ "dstack": 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": [],
+ "dstack": False,
"python": "3.13" if not docker else None,
"nvcc": None,
"registry_auth": None,
@@ -3462,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(
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..6fab966f5e 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"], 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 "dstack" 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"],
+ dstack=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..fb855f40db
--- /dev/null
+++ b/src/tests/_internal/server/services/jobs/test_server_connection.py
@@ -0,0 +1,219 @@
+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.forwarded_sockets = []
+ tunnel.reverse_forwarded_sockets = []
+ 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
+
+
+@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_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)
+ # 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()
+ 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),
+ )
+ ]
+
+ @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..e7db24226b 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"], dstack=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"], dstack=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()