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/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/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/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()
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/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,
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()