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