dstack 0.18.42__py3-none-any.whl → 0.18.44__py3-none-any.whl

dstack/_internal/server/services/{logs.py → logs/aws.py}

@@ -1,26 +1,26 @@
-import atexit
-import base64
 import itertools
 import operator
-from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
-from pathlib import Path
-from typing import Iterator, List, Optional, Set, Tuple, TypedDict, Union
+from typing import Iterator, List, Optional, Set, Tuple, TypedDict
 from uuid import UUID
 
-from dstack._internal.core.errors import DstackError
 from dstack._internal.core.models.logs import (
     JobSubmissionLogs,
     LogEvent,
     LogEventSource,
     LogProducer,
 )
-from dstack._internal.server import settings
 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.utils.common import run_async
+from dstack._internal.server.services.logs.base import (
+    LogStorage,
+    LogStorageError,
+    b64encode_raw_message,
+    datetime_to_unix_time_ms,
+    unix_time_ms_to_datetime,
+)
 from dstack._internal.utils.logging import get_logger
 
 BOTO_AVAILABLE = True
@@ -33,30 +33,6 @@ except ImportError:
 logger = get_logger(__name__)
 
 
-class LogStorageError(DstackError):
-    pass
-
-
-class LogStorage(ABC):
-    @abstractmethod
-    def poll_logs(self, project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
-        pass
-
-    @abstractmethod
-    def write_logs(
-        self,
-        project: ProjectModel,
-        run_name: str,
-        job_submission_id: UUID,
-        runner_logs: List[RunnerLogEvent],
-        job_logs: List[RunnerLogEvent],
-    ) -> None:
-        pass
-
-    def close(self) -> None:
-        pass
-
-
 class _CloudWatchLogEvent(TypedDict):
     timestamp: int  # unix time in milliseconds
     message: str
@@ -119,7 +95,7 @@ class CloudWatchLogStorage(LogStorage):
             cw_events_iter = iter(cw_events)
         logs = [
             LogEvent(
-                timestamp=_unix_time_ms_to_datetime(cw_event["timestamp"]),
+                timestamp=unix_time_ms_to_datetime(cw_event["timestamp"]),
                 log_source=LogEventSource.STDOUT,
                 message=cw_event["message"],
             )
@@ -138,11 +114,11 @@ class CloudWatchLogStorage(LogStorage):
         if request.start_time:
             # XXX: Since callers use start_time/end_time for pagination, one millisecond is added
             # to avoid an infinite loop because startTime boundary is inclusive.
-            parameters["startTime"] = _datetime_to_unix_time_ms(request.start_time) + 1
+            parameters["startTime"] = datetime_to_unix_time_ms(request.start_time) + 1
         if request.end_time:
             # No need to substract one millisecond in this case, though, seems that endTime is
             # exclusive, that is, time interval boundaries are [startTime, entTime)
-            parameters["endTime"] = _datetime_to_unix_time_ms(request.end_time)
+            parameters["endTime"] = datetime_to_unix_time_ms(request.end_time)
         response = self._client.get_log_events(**parameters)
         events: List[_CloudWatchLogEvent] = response["events"]
         if start_from_head or events:
@@ -294,7 +270,7 @@ class CloudWatchLogStorage(LogStorage):
     ) -> _CloudWatchLogEvent:
         return {
             "timestamp": runner_log_event.timestamp,
-            "message": _b64encode_raw_message(runner_log_event.message),
+            "message": b64encode_raw_message(runner_log_event.message),
         }
 
     @contextmanager
@@ -339,174 +315,3 @@ class CloudWatchLogStorage(LogStorage):
         producer: LogProducer,
     ) -> str:
         return f"{project_name}/{run_name}/{job_submission_id}/{producer.value}"
-
-
-class FileLogStorage(LogStorage):
-    root: Path
-
-    def __init__(self, root: Union[Path, str, None] = None) -> None:
-        if root is None:
-            self.root = settings.SERVER_DIR_PATH
-        else:
-            self.root = Path(root)
-
-    def poll_logs(self, project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
-        # TODO Respect request.limit to support pagination
-        log_producer = LogProducer.RUNNER if request.diagnose else LogProducer.JOB
-        log_file_path = self._get_log_file_path(
-            project_name=project.name,
-            run_name=request.run_name,
-            job_submission_id=request.job_submission_id,
-            producer=log_producer,
-        )
-        logs = []
-        try:
-            with open(log_file_path) as f:
-                for line in f:
-                    log_event = LogEvent.__response__.parse_raw(line)
-                    if request.start_time and log_event.timestamp <= request.start_time:
-                        continue
-                    if request.end_time is None or log_event.timestamp < request.end_time:
-                        logs.append(log_event)
-                    else:
-                        break
-        except IOError:
-            pass
-        if request.descending:
-            logs = list(reversed(logs))
-        return JobSubmissionLogs(logs=logs)
-
-    def write_logs(
-        self,
-        project: ProjectModel,
-        run_name: str,
-        job_submission_id: UUID,
-        runner_logs: List[RunnerLogEvent],
-        job_logs: List[RunnerLogEvent],
-    ):
-        if len(runner_logs) > 0:
-            runner_log_file_path = self._get_log_file_path(
-                project.name, run_name, job_submission_id, LogProducer.RUNNER
-            )
-            self._write_logs(
-                log_file_path=runner_log_file_path,
-                log_events=runner_logs,
-            )
-        if len(job_logs) > 0:
-            job_log_file_path = self._get_log_file_path(
-                project.name, run_name, job_submission_id, LogProducer.JOB
-            )
-            self._write_logs(
-                log_file_path=job_log_file_path,
-                log_events=job_logs,
-            )
-
-    def _write_logs(self, log_file_path: Path, log_events: List[RunnerLogEvent]) -> None:
-        log_events_parsed = [self._runner_log_event_to_log_event(event) for event in log_events]
-        log_file_path.parent.mkdir(exist_ok=True, parents=True)
-        with open(log_file_path, "a") as f:
-            f.writelines(log.json() + "\n" for log in log_events_parsed)
-
-    def _get_log_file_path(
-        self,
-        project_name: str,
-        run_name: str,
-        job_submission_id: UUID,
-        producer: LogProducer,
-    ) -> Path:
-        return (
-            self.root
-            / "projects"
-            / project_name
-            / "logs"
-            / run_name
-            / str(job_submission_id)
-            / f"{producer.value}.log"
-        )
-
-    def _runner_log_event_to_log_event(self, runner_log_event: RunnerLogEvent) -> LogEvent:
-        return LogEvent(
-            timestamp=_unix_time_ms_to_datetime(runner_log_event.timestamp),
-            log_source=LogEventSource.STDOUT,
-            message=_b64encode_raw_message(runner_log_event.message),
-        )
-
-
-def _unix_time_ms_to_datetime(unix_time_ms: int) -> datetime:
-    return datetime.fromtimestamp(unix_time_ms / 1000, tz=timezone.utc)
-
-
-def _datetime_to_unix_time_ms(dt: datetime) -> int:
-    return int(dt.timestamp() * 1000)
-
-
-def _b64encode_raw_message(message: bytes) -> str:
-    return base64.b64encode(message).decode()
-
-
-_default_log_storage: Optional[LogStorage] = None
-
-
-def get_default_log_storage() -> LogStorage:
-    global _default_log_storage
-    if _default_log_storage is not None:
-        return _default_log_storage
-    if settings.SERVER_CLOUDWATCH_LOG_GROUP:
-        if BOTO_AVAILABLE:
-            try:
-                _default_log_storage = CloudWatchLogStorage(
-                    group=settings.SERVER_CLOUDWATCH_LOG_GROUP,
-                    region=settings.SERVER_CLOUDWATCH_LOG_REGION,
-                )
-            except LogStorageError as e:
-                logger.error("Failed to initialize CloudWatch Logs storage: %s", e)
-            else:
-                logger.debug("Using CloudWatch Logs storage")
-        else:
-            logger.error("Cannot use CloudWatch Logs storage, boto3 is not installed")
-    if _default_log_storage is None:
-        logger.debug("Using file-based storage")
-        _default_log_storage = FileLogStorage()
-    atexit.register(_default_log_storage.close)
-    return _default_log_storage
-
-
-def poll_logs(project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
-    return get_default_log_storage().poll_logs(project=project, request=request)
-
-
-def write_logs(
-    project: ProjectModel,
-    run_name: str,
-    job_submission_id: UUID,
-    runner_logs: List[RunnerLogEvent],
-    job_logs: List[RunnerLogEvent],
-) -> None:
-    return get_default_log_storage().write_logs(
-        project=project,
-        run_name=run_name,
-        job_submission_id=job_submission_id,
-        runner_logs=runner_logs,
-        job_logs=job_logs,
-    )
-
-
-async def poll_logs_async(project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
-    return await run_async(get_default_log_storage().poll_logs, project=project, request=request)
-
-
-async def write_logs_async(
-    project: ProjectModel,
-    run_name: str,
-    job_submission_id: UUID,
-    runner_logs: List[RunnerLogEvent],
-    job_logs: List[RunnerLogEvent],
-) -> None:
-    return await run_async(
-        get_default_log_storage().write_logs,
-        project=project,
-        run_name=run_name,
-        job_submission_id=job_submission_id,
-        runner_logs=runner_logs,
-        job_logs=job_logs,
-    )

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

@@ -0,0 +1,47 @@
+import base64
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+from typing import List
+from uuid import UUID
+
+from dstack._internal.core.errors import DstackError
+from dstack._internal.core.models.logs import JobSubmissionLogs
+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
+
+
+class LogStorageError(DstackError):
+    pass
+
+
+class LogStorage(ABC):
+    @abstractmethod
+    def poll_logs(self, project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
+        pass
+
+    @abstractmethod
+    def write_logs(
+        self,
+        project: ProjectModel,
+        run_name: str,
+        job_submission_id: UUID,
+        runner_logs: List[RunnerLogEvent],
+        job_logs: List[RunnerLogEvent],
+    ) -> None:
+        pass
+
+    def close(self) -> None:
+        pass
+
+
+def unix_time_ms_to_datetime(unix_time_ms: int) -> datetime:
+    return datetime.fromtimestamp(unix_time_ms / 1000, tz=timezone.utc)
+
+
+def datetime_to_unix_time_ms(dt: datetime) -> int:
+    return int(dt.timestamp() * 1000)
+
+
+def b64encode_raw_message(message: bytes) -> str:
+    return base64.b64encode(message).decode()

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

@@ -0,0 +1,110 @@
+from pathlib import Path
+from typing import List, Union
+from uuid import UUID
+
+from dstack._internal.core.models.logs import (
+    JobSubmissionLogs,
+    LogEvent,
+    LogEventSource,
+    LogProducer,
+)
+from dstack._internal.server import settings
+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.base import (
+    LogStorage,
+    b64encode_raw_message,
+    unix_time_ms_to_datetime,
+)
+
+
+class FileLogStorage(LogStorage):
+    root: Path
+
+    def __init__(self, root: Union[Path, str, None] = None) -> None:
+        if root is None:
+            self.root = settings.SERVER_DIR_PATH
+        else:
+            self.root = Path(root)
+
+    def poll_logs(self, project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
+        # TODO Respect request.limit to support pagination
+        log_producer = LogProducer.RUNNER if request.diagnose else LogProducer.JOB
+        log_file_path = self._get_log_file_path(
+            project_name=project.name,
+            run_name=request.run_name,
+            job_submission_id=request.job_submission_id,
+            producer=log_producer,
+        )
+        logs = []
+        try:
+            with open(log_file_path) as f:
+                for line in f:
+                    log_event = LogEvent.__response__.parse_raw(line)
+                    if request.start_time and log_event.timestamp <= request.start_time:
+                        continue
+                    if request.end_time is None or log_event.timestamp < request.end_time:
+                        logs.append(log_event)
+                    else:
+                        break
+        except IOError:
+            pass
+        if request.descending:
+            logs = list(reversed(logs))
+        return JobSubmissionLogs(logs=logs)
+
+    def write_logs(
+        self,
+        project: ProjectModel,
+        run_name: str,
+        job_submission_id: UUID,
+        runner_logs: List[RunnerLogEvent],
+        job_logs: List[RunnerLogEvent],
+    ):
+        if len(runner_logs) > 0:
+            runner_log_file_path = self._get_log_file_path(
+                project.name, run_name, job_submission_id, LogProducer.RUNNER
+            )
+            self._write_logs(
+                log_file_path=runner_log_file_path,
+                log_events=runner_logs,
+            )
+        if len(job_logs) > 0:
+            job_log_file_path = self._get_log_file_path(
+                project.name, run_name, job_submission_id, LogProducer.JOB
+            )
+            self._write_logs(
+                log_file_path=job_log_file_path,
+                log_events=job_logs,
+            )
+
+    def _write_logs(self, log_file_path: Path, log_events: List[RunnerLogEvent]) -> None:
+        log_events_parsed = [self._runner_log_event_to_log_event(event) for event in log_events]
+        log_file_path.parent.mkdir(exist_ok=True, parents=True)
+        with open(log_file_path, "a") as f:
+            f.writelines(log.json() + "\n" for log in log_events_parsed)
+
+    def _get_log_file_path(
+        self,
+        project_name: str,
+        run_name: str,
+        job_submission_id: UUID,
+        producer: LogProducer,
+    ) -> Path:
+        return (
+            self.root
+            / "projects"
+            / project_name
+            / "logs"
+            / run_name
+            / str(job_submission_id)
+            / f"{producer.value}.log"
+        )
+
+    def _runner_log_event_to_log_event(self, runner_log_event: RunnerLogEvent) -> LogEvent:
+        return LogEvent(
+            timestamp=unix_time_ms_to_datetime(runner_log_event.timestamp),
+            log_source=LogEventSource.STDOUT,
+            message=b64encode_raw_message(runner_log_event.message),
+        )

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

@@ -0,0 +1,165 @@
+import time
+from typing import Iterable, List
+from uuid import UUID
+
+from dstack._internal.core.errors import ServerClientError
+from dstack._internal.core.models.logs import (
+    JobSubmissionLogs,
+    LogEvent,
+    LogEventSource,
+    LogProducer,
+)
+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.base import (
+    LogStorage,
+    LogStorageError,
+    b64encode_raw_message,
+    unix_time_ms_to_datetime,
+)
+from dstack._internal.utils.common import batched
+from dstack._internal.utils.logging import get_logger
+
+GCP_LOGGING_AVAILABLE = True
+try:
+    import google.api_core.exceptions
+    import google.auth.exceptions
+    from google.cloud import logging
+except ImportError:
+    GCP_LOGGING_AVAILABLE = False
+
+
+logger = get_logger(__name__)
+
+
+class GCPLogStorage(LogStorage):
+    # Max expected message size from runner is 32KB.
+    # Max expected LogEntry size is 32KB + metadata < 50KB < 256KB limit.
+    # With MAX_BATCH_SIZE = 100, max write request size < 5MB < 10 MB limit.
+    # See: https://cloud.google.com/logging/quotas.
+    MAX_RUNNER_MESSAGE_SIZE = 32 * 1024
+    MAX_BATCH_SIZE = 100
+
+    # Use the same log name for all run logs so that it's easy to manage all dstack-related logs.
+    LOG_NAME = "dstack-run-logs"
+    # Logs from different jobs belong to different "streams".
+    # GCP Logging has no built-in concepts of streams, so we implement them with labels.
+    # It should be fast to filter by labels since labels are indexed by default
+    # (https://cloud.google.com/logging/docs/analyze/custom-index).
+
+    def __init__(self, project_id: str):
+        try:
+            self.client = logging.Client(project=project_id)
+            self.logger = self.client.logger(name=self.LOG_NAME)
+            self.logger.list_entries(max_results=1)
+            # Python client doesn't seem to support dry_run,
+            # so emit an empty log to check permissions.
+            self.logger.log_empty()
+        except google.auth.exceptions.DefaultCredentialsError:
+            raise LogStorageError("Default credentials not found")
+        except google.api_core.exceptions.NotFound:
+            raise LogStorageError(f"Project {project_id} not found")
+        except google.api_core.exceptions.PermissionDenied:
+            raise LogStorageError("Insufficient permissions")
+
+    def poll_logs(self, project: ProjectModel, request: PollLogsRequest) -> JobSubmissionLogs:
+        producer = LogProducer.RUNNER if request.diagnose else LogProducer.JOB
+        stream_name = self._get_stream_name(
+            project_name=project.name,
+            run_name=request.run_name,
+            job_submission_id=request.job_submission_id,
+            producer=producer,
+        )
+        log_filters = [f'labels.stream = "{stream_name}"']
+        if request.start_time:
+            log_filters.append(f'timestamp > "{request.start_time.isoformat()}"')
+        if request.end_time:
+            log_filters.append(f'timestamp < "{request.end_time.isoformat()}"')
+        log_filter = " AND ".join(log_filters)
+
+        order_by = logging.DESCENDING if request.descending else logging.ASCENDING
+        try:
+            entries: Iterable[logging.LogEntry] = self.logger.list_entries(
+                filter_=log_filter,
+                order_by=order_by,
+                max_results=request.limit,
+                # Specify max possible page_size (<=1000) to reduce number of API calls.
+                page_size=request.limit,
+            )
+            logs = [
+                LogEvent(
+                    timestamp=entry.timestamp,
+                    message=entry.payload["message"],
+                    log_source=LogEventSource.STDOUT,
+                )
+                for entry in entries
+            ]
+        except google.api_core.exceptions.ResourceExhausted as e:
+            logger.warning("GCP Logging exception: %s", repr(e))
+            # GCP Logging has severely low quota of 60 reads/min for entries.list
+            raise ServerClientError(
+                "GCP Logging read request limit exceeded."
+                " It's recommended to increase default entries.list request quota from 60 per minute."
+            )
+        # We intentionally make reading logs slow to prevent hitting GCP quota.
+        # This doesn't help with many concurrent clients but
+        # should help with one client reading all logs sequentially.
+        time.sleep(1)
+        return JobSubmissionLogs(logs=logs)
+
+    def write_logs(
+        self,
+        project: ProjectModel,
+        run_name: str,
+        job_submission_id: UUID,
+        runner_logs: List[RunnerLogEvent],
+        job_logs: List[RunnerLogEvent],
+    ):
+        producers_with_logs = [(LogProducer.RUNNER, runner_logs), (LogProducer.JOB, job_logs)]
+        for producer, producer_logs in producers_with_logs:
+            stream_name = self._get_stream_name(
+                project_name=project.name,
+                run_name=run_name,
+                job_submission_id=job_submission_id,
+                producer=producer,
+            )
+            self._write_logs_to_stream(
+                stream_name=stream_name,
+                logs=producer_logs,
+            )
+
+    def close(self):
+        self.client.close()
+
+    def _write_logs_to_stream(self, stream_name: str, logs: List[RunnerLogEvent]):
+        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)
+                    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:
+                        logger.error(
+                            "Stream %s: skipping event at %s, message exceeds max size: %d > %d",
+                            stream_name,
+                            timestamp.isoformat(),
+                            len(message),
+                            self.MAX_RUNNER_MESSAGE_SIZE,
+                        )
+                        continue
+                    batcher.log_struct(
+                        {
+                            "message": message,
+                        },
+                        labels={
+                            "stream": stream_name,
+                        },
+                        timestamp=timestamp,
+                    )
+                batcher.commit()
+
+    def _get_stream_name(
+        self, project_name: str, run_name: str, job_submission_id: UUID, producer: LogProducer
+    ) -> str:
+        return f"{project_name}-{run_name}-{job_submission_id}-{producer.value}"