Add server access to tasks and dev environments

Author: Andrey Cheptsov

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.
+
+<div editor-title=".dstack.yml">
+
+```yaml
+type: dev-environment
+image: dstackai/dstack
+server: true
+env:
+  - DSTACK_TOKEN
+init:
+  - dstack ps
+```
+
+</div>
+
 --8<-- "docs/concepts/snippets/manage-fleets.ext"
 
 !!! info "Reference"

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.
+
+<div editor-title=".dstack.yml">
+
+```yaml
+type: task
+image: dstackai/dstack
+server: true
+env:
+  - DSTACK_TOKEN
+commands:
+  - dstack ps
+```
+
+</div>
+
 --8<-- "docs/concepts/snippets/manage-fleets.ext"
 
 !!! info "Reference"

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:

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",

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"

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),
 ):

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:

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:

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,

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,