dstack 0.19.25__py3-none-any.whl → 0.19.26__py3-none-any.whl

dstack/_internal/core/models/resources.py

@@ -130,6 +130,12 @@ DEFAULT_GPU_COUNT = Range[int](min=1)
 
 
 class CPUSpec(CoreModel):
+    arch: Annotated[
+        Optional[gpuhunt.CPUArchitecture],
+        Field(description="The CPU architecture, one of: `x86`, `arm`"),
+    ] = None
+    count: Annotated[Range[int], Field(description="The number of CPU cores")] = DEFAULT_CPU_COUNT
+
     class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
@@ -138,12 +144,6 @@ class CPUSpec(CoreModel):
                 extra_types=[{"type": "integer"}, {"type": "string"}],
             )
 
-    arch: Annotated[
-        Optional[gpuhunt.CPUArchitecture],
-        Field(description="The CPU architecture, one of: `x86`, `arm`"),
-    ] = None
-    count: Annotated[Range[int], Field(description="The number of CPU cores")] = DEFAULT_CPU_COUNT
-
     @classmethod
     def __get_validators__(cls):
         yield cls.parse
@@ -191,22 +191,6 @@ class CPUSpec(CoreModel):
 
 
 class GPUSpec(CoreModel):
-    class Config(CoreModel.Config):
-        @staticmethod
-        def schema_extra(schema: Dict[str, Any]):
-            add_extra_schema_types(
-                schema["properties"]["count"],
-                extra_types=[{"type": "integer"}, {"type": "string"}],
-            )
-            add_extra_schema_types(
-                schema["properties"]["memory"],
-                extra_types=[{"type": "integer"}, {"type": "string"}],
-            )
-            add_extra_schema_types(
-                schema["properties"]["total_memory"],
-                extra_types=[{"type": "integer"}, {"type": "string"}],
-            )
-
     vendor: Annotated[
         Optional[gpuhunt.AcceleratorVendor],
         Field(
@@ -234,6 +218,26 @@ class GPUSpec(CoreModel):
         Field(description="The minimum compute capability of the GPU (e.g., `7.5`)"),
     ] = None
 
+    class Config(CoreModel.Config):
+        @staticmethod
+        def schema_extra(schema: Dict[str, Any]):
+            add_extra_schema_types(
+                schema["properties"]["count"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["name"],
+                extra_types=[{"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["memory"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+            add_extra_schema_types(
+                schema["properties"]["total_memory"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+
     @classmethod
     def __get_validators__(cls):
         yield cls.parse
@@ -314,6 +318,8 @@ class GPUSpec(CoreModel):
 
 
 class DiskSpec(CoreModel):
+    size: Annotated[Range[Memory], Field(description="Disk size")]
+
     class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
@@ -322,8 +328,6 @@ class DiskSpec(CoreModel):
                 extra_types=[{"type": "integer"}, {"type": "string"}],
             )
 
-    size: Annotated[Range[Memory], Field(description="Disk size")]
-
     @classmethod
     def __get_validators__(cls):
         yield cls._parse
@@ -340,6 +344,24 @@ DEFAULT_DISK = DiskSpec(size=Range[Memory](min=Memory.parse("100GB"), max=None))
 
 
 class ResourcesSpec(CoreModel):
+    # TODO: Remove Range[int] in 0.20. Range[int] for backward compatibility only.
+    cpu: Annotated[Union[CPUSpec, Range[int]], Field(description="The CPU requirements")] = (
+        CPUSpec()
+    )
+    memory: Annotated[Range[Memory], Field(description="The RAM size (e.g., `8GB`)")] = (
+        DEFAULT_MEMORY_SIZE
+    )
+    shm_size: Annotated[
+        Optional[Memory],
+        Field(
+            description="The size of shared memory (e.g., `8GB`). "
+            "If you are using parallel communicating processes (e.g., dataloaders in PyTorch), "
+            "you may need to configure this"
+        ),
+    ] = None
+    gpu: Annotated[Optional[GPUSpec], Field(description="The GPU requirements")] = None
+    disk: Annotated[Optional[DiskSpec], Field(description="The disk resources")] = DEFAULT_DISK
+
     class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
@@ -364,24 +386,6 @@ class ResourcesSpec(CoreModel):
                 extra_types=[{"type": "integer"}, {"type": "string"}],
             )
 
-    # TODO: Remove Range[int] in 0.20. Range[int] for backward compatibility only.
-    cpu: Annotated[Union[CPUSpec, Range[int]], Field(description="The CPU requirements")] = (
-        CPUSpec()
-    )
-    memory: Annotated[Range[Memory], Field(description="The RAM size (e.g., `8GB`)")] = (
-        DEFAULT_MEMORY_SIZE
-    )
-    shm_size: Annotated[
-        Optional[Memory],
-        Field(
-            description="The size of shared memory (e.g., `8GB`). "
-            "If you are using parallel communicating processes (e.g., dataloaders in PyTorch), "
-            "you may need to configure this"
-        ),
-    ] = None
-    gpu: Annotated[Optional[GPUSpec], Field(description="The GPU requirements")] = None
-    disk: Annotated[Optional[DiskSpec], Field(description="The disk resources")] = DEFAULT_DISK
-
     def pretty_format(self) -> str:
         # TODO: Remove in 0.20. Use self.cpu directly
         cpu = parse_obj_as(CPUSpec, self.cpu)

dstack/_internal/core/models/runs.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Type
+from urllib.parse import urlparse
 
 from pydantic import UUID4, Field, root_validator
 from typing_extensions import Annotated
@@ -483,6 +484,9 @@ class ServiceSpec(CoreModel):
     model: Optional[ServiceModelSpec] = None
     options: Dict[str, Any] = {}
 
+    def get_domain(self) -> Optional[str]:
+        return urlparse(self.url).hostname
+
 
 class RunStatus(str, Enum):
     PENDING = "pending"

dstack/_internal/core/services/configs/__init__.py

@@ -68,8 +68,8 @@ class ConfigManager:
         if len(self.config.projects) == 1:
             self.config.projects[0].default = True
 
-    def list_projects(self):
-        return [project.name for project in self.config.projects]
+    def list_project_configs(self) -> list[ProjectConfig]:
+        return self.config.projects
 
     def delete_project(self, name: str):
         self.config.projects = [p for p in self.config.projects if p.name != name]

dstack/_internal/core/services/profiles.py

@@ -37,10 +37,10 @@ def get_termination(
 ) -> Tuple[TerminationPolicy, int]:
     termination_policy = TerminationPolicy.DESTROY_AFTER_IDLE
     termination_idle_time = default_termination_idle_time
-    if profile.idle_duration is not None and int(profile.idle_duration) < 0:
+    if profile.idle_duration is not None and profile.idle_duration < 0:
         termination_policy = TerminationPolicy.DONT_DESTROY
     elif profile.idle_duration is not None:
         termination_idle_time = profile.idle_duration
     if termination_policy == TerminationPolicy.DONT_DESTROY:
         termination_idle_time = -1
-    return termination_policy, int(termination_idle_time)
+    return termination_policy, termination_idle_time

dstack/_internal/core/services/repos.py

@@ -2,7 +2,7 @@ import os
 from pathlib import Path
 from typing import Optional, Union
 
-import git
+import git.cmd
 import requests
 import yaml
 from git.exc import GitCommandError
@@ -24,6 +24,8 @@ logger = get_logger(__name__)
 gh_config_path = os.path.expanduser("~/.config/gh/hosts.yml")
 default_ssh_key = os.path.expanduser("~/.ssh/id_rsa")
 
+no_prompt_env = dict(GIT_TERMINAL_PROMPT="0")
+
 
 class InvalidRepoCredentialsError(DstackError):
     pass
@@ -84,7 +86,7 @@ def get_local_repo_credentials(
 
 def check_remote_repo_credentials_https(url: GitRepoURL, oauth_token: str) -> RemoteRepoCreds:
     try:
-        git.cmd.Git().ls_remote(url.as_https(oauth_token), env=dict(GIT_TERMINAL_PROMPT="0"))
+        git.cmd.Git().ls_remote(url.as_https(oauth_token), env=no_prompt_env)
     except GitCommandError:
         masked = len(oauth_token[:-4]) * "*" + oauth_token[-4:]
         raise InvalidRepoCredentialsError(
@@ -131,7 +133,7 @@ def get_default_branch(remote_url: str) -> Optional[str]:
     Get the default branch of a remote Git repository.
     """
     try:
-        output = git.cmd.Git().ls_remote("--symref", remote_url, "HEAD")
+        output = git.cmd.Git().ls_remote("--symref", remote_url, "HEAD", env=no_prompt_env)
         for line in output.splitlines():
             if line.startswith("ref:"):
                 return line.split()[1].split("/")[-1]

dstack/_internal/core/services/ssh/ports.py

@@ -74,7 +74,7 @@ class PortsLock:
         try:
             sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
             if IS_WINDOWS:
-                sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)  # type: ignore[attr-defined]
             sock.bind(("", port))
             return sock
         except socket.error as e:

dstack/_internal/proxy/lib/deps.py

@@ -21,12 +21,16 @@ class ProxyDependencyInjector(ABC):
     def __init__(self) -> None:
         self._service_conn_pool = ServiceConnectionPool()
 
+    # Abstract AsyncGenerator does not need async def since
+    # type checkers infer a different type without yield in body.
+    # https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+
     @abstractmethod
-    async def get_repo(self) -> AsyncGenerator[BaseProxyRepo, None]:
+    def get_repo(self) -> AsyncGenerator[BaseProxyRepo, None]:
         pass
 
     @abstractmethod
-    async def get_auth_provider(self) -> AsyncGenerator[BaseProxyAuthProvider, None]:
+    def get_auth_provider(self) -> AsyncGenerator[BaseProxyAuthProvider, None]:
         pass
 
     async def get_service_connection_pool(self) -> ServiceConnectionPool:

dstack/_internal/server/app.py

@@ -110,9 +110,11 @@ async def lifespan(app: FastAPI):
     _print_dstack_logo()
     if not check_required_ssh_version():
         logger.warning("OpenSSH 8.4+ is required. The dstack server may not work properly")
+    server_config_manager = None
+    server_config_loaded = False
     if settings.SERVER_CONFIG_ENABLED:
         server_config_manager = ServerConfigManager()
-        config_loaded = server_config_manager.load_config()
+        server_config_loaded = server_config_manager.load_config()
         # Encryption has to be configured before working with users and projects
         await server_config_manager.apply_encryption()
     async with get_session_ctx() as session:
@@ -126,11 +128,9 @@ async def lifespan(app: FastAPI):
                 session=session,
                 user=admin,
             )
-            if settings.SERVER_CONFIG_ENABLED:
-                server_config_dir = str(SERVER_CONFIG_FILE_PATH).replace(
-                    os.path.expanduser("~"), "~", 1
-                )
-                if not config_loaded:
+            if server_config_manager is not None:
+                server_config_dir = _get_server_config_dir()
+                if not server_config_loaded:
                     logger.info("Initializing the default configuration...", {"show_path": False})
                     await server_config_manager.init_config(session=session)
                     logger.info(
@@ -153,6 +153,7 @@ async def lifespan(app: FastAPI):
     )
     if settings.SERVER_S3_BUCKET is not None or settings.SERVER_GCS_BUCKET is not None:
         init_default_storage()
+    scheduler = None
     if settings.SERVER_BACKGROUND_PROCESSING_ENABLED:
         scheduler = start_background_tasks()
     else:
@@ -167,7 +168,7 @@ async def lifespan(app: FastAPI):
     for func in _ON_STARTUP_HOOKS:
         await func(app)
     yield
-    if settings.SERVER_BACKGROUND_PROCESSING_ENABLED:
+    if scheduler is not None:
         scheduler.shutdown()
     PROBES_SCHEDULER.shutdown(wait=False)
     await gateway_connections_pool.remove_all()
@@ -371,6 +372,18 @@ def _is_prometheus_request(request: Request) -> bool:
     return request.url.path.startswith("/metrics")
 
 
+def _sentry_traces_sampler(sampling_context: SamplingContext) -> float:
+    parent_sampling_decision = sampling_context["parent_sampled"]
+    if parent_sampling_decision is not None:
+        return float(parent_sampling_decision)
+    transaction_context = sampling_context["transaction_context"]
+    name = transaction_context.get("name")
+    if name is not None:
+        if name.startswith("background."):
+            return settings.SENTRY_TRACES_BACKGROUND_SAMPLE_RATE
+    return settings.SENTRY_TRACES_SAMPLE_RATE
+
+
 def _print_dstack_logo():
     console.print(
         """[purple]╱╱╭╮╱╱╭╮╱╱╱╱╱╱╭╮
@@ -387,13 +400,5 @@ def _print_dstack_logo():
     )
 
 
-def _sentry_traces_sampler(sampling_context: SamplingContext) -> float:
-    parent_sampling_decision = sampling_context["parent_sampled"]
-    if parent_sampling_decision is not None:
-        return float(parent_sampling_decision)
-    transaction_context = sampling_context["transaction_context"]
-    name = transaction_context.get("name")
-    if name is not None:
-        if name.startswith("background."):
-            return settings.SENTRY_TRACES_BACKGROUND_SAMPLE_RATE
-    return settings.SENTRY_TRACES_SAMPLE_RATE
+def _get_server_config_dir() -> str:
+    return str(SERVER_CONFIG_FILE_PATH).replace(os.path.expanduser("~"), "~", 1)

dstack/_internal/server/background/tasks/process_gateways.py

@@ -49,8 +49,8 @@ async def process_gateways():
             if gateway_model is None:
                 return
             lockset.add(gateway_model.id)
+        gateway_model_id = gateway_model.id
         try:
-            gateway_model_id = gateway_model.id
             initial_status = gateway_model.status
             if initial_status == GatewayStatus.SUBMITTED:
                 await _process_submitted_gateway(session=session, gateway_model=gateway_model)
@@ -165,6 +165,9 @@ async def _process_provisioning_gateway(
     )
     gateway_model = res.unique().scalar_one()
 
+    # Provisioning gateways must have compute.
+    assert gateway_model.gateway_compute is not None
+
     # FIXME: problems caused by blocking on connect_to_gateway_with_retry and configure_gateway:
     # - cannot delete the gateway before it is provisioned because the DB model is locked
     # - connection retry counter is reset on server restart

dstack/_internal/server/background/tasks/process_instances.py

@@ -85,8 +85,10 @@ from dstack._internal.server.services.instances import (
     get_instance_provisioning_data,
     get_instance_requirements,
     get_instance_ssh_private_keys,
+    remove_dangling_tasks_from_instance,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import is_divisible_into_blocks
 from dstack._internal.server.services.placement import (
     get_fleet_placement_group_models,
@@ -181,8 +183,8 @@ async def _process_next_instance():
             if instance is None:
                 return
             lockset.add(instance.id)
+        instance_model_id = instance.id
         try:
-            instance_model_id = instance.id
             await _process_instance(session=session, instance=instance)
         finally:
             lockset.difference_update([instance_model_id])
@@ -393,6 +395,7 @@ async def _add_remote(instance: InstanceModel) -> None:
         return
 
     region = instance.region
+    assert region is not None  # always set for ssh instances
     jpd = JobProvisioningData(
         backend=BackendType.REMOTE,
         instance_type=instance_type,
@@ -788,6 +791,7 @@ async def _check_instance(session: AsyncSession, instance: InstanceModel) -> Non
         ssh_private_keys,
         job_provisioning_data,
         None,
+        instance=instance,
         check_instance_health=check_instance_health,
     )
     if instance_check is False:
@@ -934,7 +938,7 @@ async def _wait_for_instance_provisioning_data(
 
 @runner_ssh_tunnel(ports=[DSTACK_SHIM_HTTP_PORT], retries=1)
 def _check_instance_inner(
-    ports: Dict[int, int], *, check_instance_health: bool = False
+    ports: Dict[int, int], *, instance: InstanceModel, check_instance_health: bool = False
 ) -> InstanceCheck:
     instance_health_response: Optional[InstanceHealthResponse] = None
     shim_client = runner_client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
@@ -954,6 +958,10 @@ def _check_instance_inner(
         args = (method.__func__.__name__, e.__class__.__name__, e)
         logger.exception(template, *args)
         return InstanceCheck(reachable=False, message=template % args)
+    try:
+        remove_dangling_tasks_from_instance(shim_client, instance)
+    except Exception as e:
+        logger.exception("%s: error removing dangling tasks: %s", fmt(instance), e)
     return runner_client.healthcheck_response_to_instance_check(
         healthcheck_response, instance_health_response
     )

dstack/_internal/server/background/tasks/process_probes.py

@@ -120,7 +120,7 @@ async def _execute_probe(probe: ProbeModel, probe_spec: ProbeSpec) -> bool:
                 method=probe_spec.method,
                 url="http://dstack" + probe_spec.url,
                 headers=[(h.name, h.value) for h in probe_spec.headers],
-                data=probe_spec.body,
+                content=probe_spec.body,
                 timeout=probe_spec.timeout,
                 follow_redirects=False,
             )

dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -128,9 +128,8 @@ async def _process_next_running_job():
             if job_model is None:
                 return
             lockset.add(job_model.id)
-
+        job_model_id = job_model.id
         try:
-            job_model_id = job_model.id
             await _process_running_job(session=session, job_model=job_model)
         finally:
             lockset.difference_update([job_model_id])
@@ -170,6 +169,11 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
 
     job = find_job(run.jobs, job_model.replica_num, job_model.job_num)
 
+    volumes = []
+    secrets = {}
+    cluster_info = None
+    repo_creds = None
+
     initial_status = job_model.status
     if initial_status in [JobStatus.PROVISIONING, JobStatus.PULLING]:
         # Wait until all other jobs in the replica are provisioned
@@ -257,6 +261,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 user_ssh_key,
             )
         else:
+            assert cluster_info is not None
             logger.debug(
                 "%s: process provisioning job without shim, age=%s",
                 fmt(job_model),
@@ -275,7 +280,6 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 repo=repo_model,
                 code_hash=_get_repo_code_hash(run, job),
             )
-
             success = await common_utils.run_async(
                 _submit_job_to_runner,
                 server_ssh_private_keys,
@@ -309,6 +313,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
 
     else:  # fails are not acceptable
         if initial_status == JobStatus.PULLING:
+            assert cluster_info is not None
             logger.debug(
                 "%s: process pulling job with shim, age=%s", fmt(job_model), job_submission.age
             )
@@ -341,7 +346,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 server_ssh_private_keys,
                 job_provisioning_data,
             )
-        elif initial_status == JobStatus.RUNNING:
+        else:
             logger.debug("%s: process running job, age=%s", fmt(job_model), job_submission.age)
             success = await common_utils.run_async(
                 _process_running,
@@ -632,6 +637,7 @@ def _process_pulling_with_shim(
         is successful
     """
     shim_client = client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
+    job_runtime_data = None
     if shim_client.is_api_v2_supported():  # raises error if shim is down, causes retry
         task = shim_client.get_task(job_model.id)
 

dstack/_internal/server/background/tasks/process_runs.py

@@ -129,8 +129,8 @@ async def _process_next_run():
             job_ids = [j.id for j in run_model.jobs]
             run_lockset.add(run_model.id)
             job_lockset.update(job_ids)
+        run_model_id = run_model.id
         try:
-            run_model_id = run_model.id
             await _process_run(session=session, run_model=run_model)
         finally:
             run_lockset.difference_update([run_model_id])