dstack 0.19.17__py3-none-any.whl → 0.19.19__py3-none-any.whl

dstack/_internal/server/services/gateways/client.py

@@ -7,9 +7,9 @@ from pydantic import parse_obj_as
 
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.errors import GatewayError
-from dstack._internal.core.models.configurations import RateLimit
+from dstack._internal.core.models.configurations import RateLimit, ServiceConfiguration
 from dstack._internal.core.models.instances import SSHConnectionParams
-from dstack._internal.core.models.runs import JobSubmission, Run
+from dstack._internal.core.models.runs import JobSpec, JobSubmission, Run, get_service_port
 from dstack._internal.proxy.gateway.schemas.stats import ServiceStats
 from dstack._internal.server import settings
 
@@ -80,13 +80,15 @@ class GatewayClient:
     async def register_replica(
         self,
         run: Run,
+        job_spec: JobSpec,
         job_submission: JobSubmission,
         ssh_head_proxy: Optional[SSHConnectionParams],
         ssh_head_proxy_private_key: Optional[str],
     ):
+        assert isinstance(run.run_spec.configuration, ServiceConfiguration)
         payload = {
             "job_id": job_submission.id.hex,
-            "app_port": run.run_spec.configuration.port.container_port,
+            "app_port": get_service_port(job_spec, run.run_spec.configuration),
             "ssh_head_proxy": ssh_head_proxy.dict() if ssh_head_proxy is not None else None,
             "ssh_head_proxy_private_key": ssh_head_proxy_private_key,
         }

dstack/_internal/server/services/instances.py

@@ -106,6 +106,14 @@ def get_instance_requirements(instance_model: InstanceModel) -> Requirements:
     return Requirements.__response__.parse_raw(instance_model.requirements)
 
 
+def get_instance_remote_connection_info(
+    instance_model: InstanceModel,
+) -> Optional[RemoteConnectionInfo]:
+    if instance_model.remote_connection_info is None:
+        return None
+    return RemoteConnectionInfo.__response__.parse_raw(instance_model.remote_connection_info)
+
+
 def get_instance_ssh_private_keys(instance_model: InstanceModel) -> tuple[str, Optional[str]]:
     """
     Returns a pair of SSH private keys: host key and optional proxy jump key.

dstack/_internal/server/services/jobs/__init__.py

@@ -134,6 +134,8 @@ def job_model_to_job_submission(job_model: JobModel) -> JobSubmission:
     finished_at = None
     if job_model.status.is_finished():
         finished_at = last_processed_at
+    status_message = _get_job_status_message(job_model)
+    error = _get_job_error(job_model)
     return JobSubmission(
         id=job_model.id,
         submission_num=job_model.submission_num,
@@ -143,11 +145,13 @@ def job_model_to_job_submission(job_model: JobModel) -> JobSubmission:
         finished_at=finished_at,
         inactivity_secs=job_model.inactivity_secs,
         status=job_model.status,
+        status_message=status_message,
         termination_reason=job_model.termination_reason,
         termination_reason_message=job_model.termination_reason_message,
         exit_status=job_model.exit_status,
         job_provisioning_data=job_provisioning_data,
         job_runtime_data=get_job_runtime_data(job_model),
+        error=error,
     )
 
 
@@ -289,6 +293,19 @@ async def process_terminating_job(
     # so that stuck volumes don't prevent the instance from terminating.
     job_model.instance_id = None
     instance_model.last_job_processed_at = common.get_current_datetime()
+
+    volume_names = (
+        jrd.volume_names
+        if jrd and jrd.volume_names
+        else [va.volume.name for va in instance_model.volume_attachments]
+    )
+    if volume_names:
+        volumes = await list_project_volume_models(
+            session=session, project=instance_model.project, names=volume_names
+        )
+        for volume in volumes:
+            volume.last_job_processed_at = common.get_current_datetime()
+
     logger.info(
         "%s: instance '%s' has been released, new status is %s",
         fmt(job_model),
@@ -693,3 +710,31 @@ def _get_job_mount_point_attached_volume(
             continue
         return volume
     raise ServerClientError("Failed to find an eligible volume for the mount point")
+
+
+def _get_job_status_message(job_model: JobModel) -> str:
+    if job_model.status == JobStatus.DONE:
+        return "exited (0)"
+    elif job_model.status == JobStatus.FAILED:
+        if job_model.termination_reason == JobTerminationReason.CONTAINER_EXITED_WITH_ERROR:
+            return f"exited ({job_model.exit_status})"
+        elif (
+            job_model.termination_reason == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
+        ):
+            return "no offers"
+        elif job_model.termination_reason == JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY:
+            return "interrupted"
+        else:
+            return "error"
+    elif job_model.status == JobStatus.TERMINATED:
+        if job_model.termination_reason == JobTerminationReason.TERMINATED_BY_USER:
+            return "stopped"
+        elif job_model.termination_reason == JobTerminationReason.ABORTED_BY_USER:
+            return "aborted"
+    return job_model.status.value
+
+
+def _get_job_error(job_model: JobModel) -> Optional[str]:
+    if job_model.termination_reason is None:
+        return None
+    return job_model.termination_reason.to_error()

dstack/_internal/server/services/jobs/configurators/base.py

@@ -1,5 +1,6 @@
 import shlex
 import sys
+import threading
 from abc import ABC, abstractmethod
 from pathlib import PurePosixPath
 from typing import Dict, List, Optional, Union
@@ -14,6 +15,7 @@ from dstack._internal.core.models.configurations import (
     PortMapping,
     PythonVersion,
     RunConfigurationType,
+    ServiceConfiguration,
 )
 from dstack._internal.core.models.profiles import (
     DEFAULT_STOP_DURATION,
@@ -152,6 +154,7 @@ class JobConfigurator(ABC):
             repo_data=self.run_spec.repo_data,
             repo_code_hash=self.run_spec.repo_code_hash,
             file_archives=self.run_spec.file_archives,
+            service_port=self._service_port(),
         )
         return job_spec
 
@@ -305,6 +308,11 @@ class JobConfigurator(ABC):
             )
         return self._job_ssh_key
 
+    def _service_port(self) -> Optional[int]:
+        if isinstance(self.run_spec.configuration, ServiceConfiguration):
+            return self.run_spec.configuration.port.container_port
+        return None
+
 
 def interpolate_job_volumes(
     run_volumes: List[Union[MountPoint, str]],
@@ -354,7 +362,10 @@ def _join_shell_commands(commands: List[str]) -> str:
     return " && ".join(commands)
 
 
-@cached(TTLCache(maxsize=2048, ttl=80))
+@cached(
+    cache=TTLCache(maxsize=2048, ttl=80),
+    lock=threading.Lock(),
+)
 def _get_image_config(image: str, registry_auth: Optional[RegistryAuth]) -> ImageConfig:
     try:
         return get_image_config(image, registry_auth).config

dstack/_internal/server/services/locking.py

@@ -1,8 +1,10 @@
 import asyncio
+import collections.abc
 import hashlib
+from abc import abstractmethod
 from asyncio import Lock
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, Dict, List, Set, Tuple, TypeVar, Union
+from typing import AsyncGenerator, Iterable, Iterator, Protocol, TypeVar, Union
 
 from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
@@ -10,23 +12,54 @@ from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 KeyT = TypeVar("KeyT")
 
 
-class ResourceLocker:
-    def __init__(self):
-        self.namespace_to_locks_map: Dict[str, Tuple[Lock, set]] = {}
+class LocksetLock(Protocol):
+    async def acquire(self) -> bool: ...
+    def release(self) -> None: ...
+    async def __aenter__(self): ...
+    async def __aexit__(self, exc_type, exc, tb): ...
+
+
+T = TypeVar("T")
+
 
-    def get_lockset(self, namespace: str) -> Tuple[Lock, set]:
+class Lockset(Protocol[T]):
+    def __contains__(self, item: T) -> bool: ...
+    def __iter__(self) -> Iterator[T]: ...
+    def __len__(self) -> int: ...
+    def add(self, item: T) -> None: ...
+    def discard(self, item: T) -> None: ...
+    def update(self, other: Iterable[T]) -> None: ...
+    def difference_update(self, other: Iterable[T]) -> None: ...
+
+
+class ResourceLocker:
+    @abstractmethod
+    def get_lockset(self, namespace: str) -> tuple[LocksetLock, Lockset]:
         """
         Returns a lockset containing locked resources for in-memory locking.
         Also returns a lock that guards the lockset.
         """
-        return self.namespace_to_locks_map.setdefault(namespace, (Lock(), set()))
+        pass
 
+    @abstractmethod
     @asynccontextmanager
-    async def lock_ctx(self, namespace: str, keys: List[KeyT]):
+    async def lock_ctx(self, namespace: str, keys: list[KeyT]):
         """
         Acquires locks for all keys in namespace.
         The keys must be sorted to prevent deadlock.
         """
+        yield
+
+
+class InMemoryResourceLocker(ResourceLocker):
+    def __init__(self):
+        self.namespace_to_locks_map: dict[str, tuple[Lock, set]] = {}
+
+    def get_lockset(self, namespace: str) -> tuple[Lock, set]:
+        return self.namespace_to_locks_map.setdefault(namespace, (Lock(), set()))
+
+    @asynccontextmanager
+    async def lock_ctx(self, namespace: str, keys: list[KeyT]):
         lock, lockset = self.get_lockset(namespace)
         try:
             await _wait_to_lock_many(lock, lockset, keys)
@@ -35,6 +68,56 @@ class ResourceLocker:
             lockset.difference_update(keys)
 
 
+class DummyAsyncLock:
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
+    async def acquire(self):
+        return True
+
+    def release(self):
+        pass
+
+
+class DummySet(collections.abc.MutableSet):
+    def __contains__(self, item):
+        return False
+
+    def __iter__(self):
+        return iter(())
+
+    def __len__(self):
+        return 0
+
+    def add(self, value):
+        pass
+
+    def discard(self, value):
+        pass
+
+    def update(self, other):
+        pass
+
+    def difference_update(self, other):
+        pass
+
+
+class DummyResourceLocker(ResourceLocker):
+    def __init__(self):
+        self.lock = DummyAsyncLock()
+        self.lockset = DummySet()
+
+    def get_lockset(self, namespace: str) -> tuple[DummyAsyncLock, DummySet]:
+        return self.lock, self.lockset
+
+    @asynccontextmanager
+    async def lock_ctx(self, namespace: str, keys: list[KeyT]):
+        yield
+
+
 def string_to_lock_id(s: str) -> int:
     return int(hashlib.sha256(s.encode()).hexdigest(), 16) % (2**63)
 
@@ -67,15 +150,21 @@ async def try_advisory_lock_ctx(
             await bind.execute(select(func.pg_advisory_unlock(string_to_lock_id(resource))))
 
 
-_locker = ResourceLocker()
+_in_memory_locker = InMemoryResourceLocker()
+_dummy_locker = DummyResourceLocker()
 
 
-def get_locker() -> ResourceLocker:
-    return _locker
+def get_locker(dialect_name: str) -> ResourceLocker:
+    if dialect_name == "sqlite":
+        return _in_memory_locker
+    # We could use an in-memory locker on Postgres
+    # but it can lead to unnecessary lock contention,
+    # so we use a dummy locker that does not take any locks.
+    return _dummy_locker
 
 
 async def _wait_to_lock_many(
-    lock: asyncio.Lock, locked: Set[KeyT], keys: List[KeyT], *, delay: float = 0.1
+    lock: asyncio.Lock, locked: set[KeyT], keys: list[KeyT], *, delay: float = 0.1
 ):
     """
     Retry locking until all the keys are locked.
@@ -83,14 +172,16 @@ async def _wait_to_lock_many(
     The keys must be sorted to prevent deadlock.
     """
     left_to_lock = keys.copy()
-    while len(left_to_lock) > 0:
+    while True:
         async with lock:
             locked_now_num = 0
             for key in left_to_lock:
                 if key in locked:
-                    # Someone already aquired the lock, wait
+                    # Someone already acquired the lock, wait
                     break
                 locked.add(key)
                 locked_now_num += 1
             left_to_lock = left_to_lock[locked_now_num:]
+        if not left_to_lock:
+            return
         await asyncio.sleep(delay)

dstack/_internal/server/services/logging.py

@@ -1,12 +1,14 @@
 from typing import Union
 
-from dstack._internal.server.models import JobModel, RunModel
+from dstack._internal.server.models import GatewayModel, JobModel, RunModel
 
 
-def fmt(model: Union[RunModel, JobModel]) -> str:
+def fmt(model: Union[RunModel, JobModel, GatewayModel]) -> str:
     """Consistent string representation of a model for logging."""
     if isinstance(model, RunModel):
         return f"run({model.id.hex[:6]}){model.run_name}"
     if isinstance(model, JobModel):
         return f"job({model.id.hex[:6]}){model.job_name}"
+    if isinstance(model, GatewayModel):
+        return f"gateway({model.id.hex[:6]}){model.name}"
     return str(model)

dstack/_internal/server/services/logs/__init__.py

@@ -8,7 +8,11 @@ from dstack._internal.server.models import ProjectModel
 from dstack._internal.server.schemas.logs import PollLogsRequest
 from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent
 from dstack._internal.server.services.logs.aws import BOTO_AVAILABLE, CloudWatchLogStorage
-from dstack._internal.server.services.logs.base import LogStorage, LogStorageError
+from dstack._internal.server.services.logs.base import (
+    LogStorage,
+    LogStorageError,
+    b64encode_raw_message,
+)
 from dstack._internal.server.services.logs.filelog import FileLogStorage
 from dstack._internal.server.services.logs.gcp import GCP_LOGGING_AVAILABLE, GCPLogStorage
 from dstack._internal.utils.common import run_async
@@ -75,4 +79,13 @@ def write_logs(
 
 
 async def poll_logs_async(project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
-    return await run_async(get_log_storage().poll_logs, project=project, request=request)
+    job_submission_logs = await run_async(
+        get_log_storage().poll_logs, project=project, request=request
+    )
+    # Logs are stored in plaintext but transmitted in base64 for API/CLI backward compatibility.
+    # Old logs stored in base64 are encoded twice for transmission and shown as base64 in CLI/UI.
+    # We live with that.
+    # TODO: Drop base64 encoding in 0.20.
+    for log_event in job_submission_logs.logs:
+        log_event.message = b64encode_raw_message(log_event.message.encode())
+    return job_submission_logs

dstack/_internal/server/services/logs/aws.py

@@ -17,7 +17,6 @@ from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent
 from dstack._internal.server.services.logs.base import (
     LogStorage,
     LogStorageError,
-    b64encode_raw_message,
     datetime_to_unix_time_ms,
     unix_time_ms_to_datetime,
 )
@@ -238,8 +237,7 @@ class CloudWatchLogStorage(LogStorage):
                 skipped_future_events += 1
                 continue
             cw_event = self._runner_log_event_to_cloudwatch_event(event)
-            # as message is base64-encoded, length in bytes = length in code points.
-            message_size = len(cw_event["message"]) + self.MESSAGE_OVERHEAD_SIZE
+            message_size = len(event.message) + self.MESSAGE_OVERHEAD_SIZE
             if message_size > self.MESSAGE_MAX_SIZE:
                 # we should never hit this limit, as we use `io.Copy` to copy from pty to logs,
                 # which under the hood uses 32KiB buffer, see runner/internal/executor/executor.go,
@@ -271,7 +269,7 @@ class CloudWatchLogStorage(LogStorage):
     ) -> _CloudWatchLogEvent:
         return {
             "timestamp": runner_log_event.timestamp,
-            "message": b64encode_raw_message(runner_log_event.message),
+            "message": runner_log_event.message.decode(errors="replace"),
         }
 
     @contextmanager

dstack/_internal/server/services/logs/filelog.py

@@ -2,6 +2,7 @@ from pathlib import Path
 from typing import List, Union
 from uuid import UUID
 
+from dstack._internal.core.errors import ServerClientError
 from dstack._internal.core.models.logs import (
     JobSubmissionLogs,
     LogEvent,
@@ -14,8 +15,6 @@ from dstack._internal.server.schemas.logs import PollLogsRequest
 from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent
 from dstack._internal.server.services.logs.base import (
     LogStorage,
-    LogStorageError,
-    b64encode_raw_message,
     unix_time_ms_to_datetime,
 )
 
@@ -30,9 +29,6 @@ class FileLogStorage(LogStorage):
             self.root = Path(root)
 
     def poll_logs(self, project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
-        if request.descending:
-            raise LogStorageError("descending: true is not supported")
-
         log_producer = LogProducer.RUNNER if request.diagnose else LogProducer.JOB
         log_file_path = self._get_log_file_path(
             project_name=project.name,
@@ -46,11 +42,11 @@ class FileLogStorage(LogStorage):
             try:
                 start_line = int(request.next_token)
                 if start_line < 0:
-                    raise LogStorageError(
+                    raise ServerClientError(
                         f"Invalid next_token: {request.next_token}. Must be a non-negative integer."
                     )
             except ValueError:
-                raise LogStorageError(
+                raise ServerClientError(
                     f"Invalid next_token: {request.next_token}. Must be a valid integer."
                 )
 
@@ -60,31 +56,41 @@ class FileLogStorage(LogStorage):
 
         try:
             with open(log_file_path) as f:
-                lines = f.readlines()
-
-            for i, line in enumerate(lines):
-                if current_line < start_line:
+                # Skip to start_line if needed
+                for _ in range(start_line):
+                    if f.readline() == "":
+                        # File is shorter than start_line
+                        return JobSubmissionLogs(logs=logs, next_token=next_token)
                     current_line += 1
-                    continue
 
-                log_event = LogEvent.__response__.parse_raw(line)
-                current_line += 1
+                # Read lines one by one
+                while True:
+                    line = f.readline()
+                    if line == "":  # EOF
+                        break
+
+                    current_line += 1
 
-                if request.start_time and log_event.timestamp <= request.start_time:
-                    continue
-                if request.end_time is not None and log_event.timestamp >= request.end_time:
-                    break
+                    try:
+                        log_event = LogEvent.__response__.parse_raw(line)
+                    except Exception:
+                        # Skip malformed lines
+                        continue
 
-                logs.append(log_event)
+                    if request.start_time and log_event.timestamp <= request.start_time:
+                        continue
+                    if request.end_time is not None and log_event.timestamp >= request.end_time:
+                        break
 
-                if len(logs) >= request.limit:
-                    # Only set next_token if there are more lines to read
-                    if current_line < len(lines):
-                        next_token = str(current_line)
-                    break
+                    logs.append(log_event)
 
-        except IOError as e:
-            raise LogStorageError(f"Failed to read log file {log_file_path}: {e}")
+                    if len(logs) >= request.limit:
+                        # Check if there are more lines to read
+                        if f.readline() != "":
+                            next_token = str(current_line)
+                        break
+        except FileNotFoundError:
+            pass
 
         return JobSubmissionLogs(logs=logs, next_token=next_token)
 
@@ -140,5 +146,5 @@ class FileLogStorage(LogStorage):
         return LogEvent(
             timestamp=unix_time_ms_to_datetime(runner_log_event.timestamp),
             log_source=LogEventSource.STDOUT,
-            message=b64encode_raw_message(runner_log_event.message),
+            message=runner_log_event.message.decode(errors="replace"),
         )

dstack/_internal/server/services/logs/gcp.py

@@ -14,7 +14,6 @@ from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent
 from dstack._internal.server.services.logs.base import (
     LogStorage,
     LogStorageError,
-    b64encode_raw_message,
     unix_time_ms_to_datetime,
 )
 from dstack._internal.utils.common import batched
@@ -137,15 +136,14 @@ class GCPLogStorage(LogStorage):
         with self.logger.batch() as batcher:
             for batch in batched(logs, self.MAX_BATCH_SIZE):
                 for log in batch:
-                    message = b64encode_raw_message(log.message)
+                    message = log.message.decode(errors="replace")
                     timestamp = unix_time_ms_to_datetime(log.timestamp)
-                    # as message is base64-encoded, length in bytes = length in code points
-                    if len(message) > self.MAX_RUNNER_MESSAGE_SIZE:
+                    if len(log.message) > self.MAX_RUNNER_MESSAGE_SIZE:
                         logger.error(
                             "Stream %s: skipping event at %s, message exceeds max size: %d > %d",
                             stream_name,
                             timestamp.isoformat(),
-                            len(message),
+                            len(log.message),
                             self.MAX_RUNNER_MESSAGE_SIZE,
                         )
                         continue

dstack/_internal/server/services/proxy/repo.py

@@ -12,10 +12,12 @@ from dstack._internal.core.models.configurations import ServiceConfiguration
 from dstack._internal.core.models.instances import RemoteConnectionInfo, SSHConnectionParams
 from dstack._internal.core.models.runs import (
     JobProvisioningData,
+    JobSpec,
     JobStatus,
     RunSpec,
     RunStatus,
     ServiceSpec,
+    get_service_port,
 )
 from dstack._internal.core.models.services import AnyModel
 from dstack._internal.proxy.lib.models import (
@@ -97,9 +99,10 @@ class ServerProxyRepo(BaseProxyRepo):
                 if rci.ssh_proxy is not None:
                     ssh_head_proxy = rci.ssh_proxy
                     ssh_head_proxy_private_key = get_or_error(rci.ssh_proxy_keys)[0].private
+            job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
             replica = Replica(
                 id=job.id.hex,
-                app_port=run_spec.configuration.port.container_port,
+                app_port=get_service_port(job_spec, run_spec.configuration),
                 ssh_destination=ssh_destination,
                 ssh_port=ssh_port,
                 ssh_proxy=ssh_proxy,