Add new logic for logs tracking via a thread with jsonl format

Author: polyaxon-ci

traceml/traceml/events/__init__.py

@@ -3,6 +3,7 @@
     get_event_assets_path,
     get_event_path,
     get_resource_path,
+    get_logs_path,
 )
 from traceml.events.schemas import (
     LoggedEventListSpec,

traceml/traceml/events/paths.py

@@ -3,6 +3,18 @@
 from clipped.utils.enums import get_enum_value
 
 
+def get_logs_path(
+    run_path: str,
+    filename: Optional[str] = None,
+    full_path: bool = True,
+) -> str:
+    _path = "{}/plxlogs".format(run_path)
+    if full_path:
+        filename = filename or "logs"
+        _path = "{}/{}.jsonl".format(_path, filename)
+    return _path
+
+
 def get_resource_path(
     run_path: str, kind: Optional[str] = None, name: Optional[str] = None
 ) -> str:

traceml/traceml/serialization/base.py

@@ -1,28 +1,51 @@
 import os
 
-from typing import List
+from typing import List, Union
 
 from clipped.utils.enums import get_enum_value
 from clipped.utils.paths import check_or_create_path, set_permissions
 
 from traceml.artifacts import V1ArtifactKind
+from traceml.logging import V1Log, V1Logs
 from traceml.events import (
     LoggedEventListSpec,
     LoggedEventSpec,
     get_event_path,
+    get_logs_path,
     get_resource_path,
 )
 
 
-class EventWriter:
+class BaseWriter:
+    """Base class for writing events and resources to files."""
+
+    def __init__(self, run_path: str):
+        self._run_path = run_path
+        self._closed = False
+
+    def flush(self):
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    def close(self):
+        self.flush()
+        self._closed = True
+
+    @property
+    def closed(self):
+        return self._closed
+
+    def write(self, events: List[LoggedEventSpec]):
+        raise NotImplementedError("Subclasses must implement this method.")
+
+
+class EventWriter(BaseWriter):
     EVENTS_BACKEND = "events"
     RESOURCES_BACKEND = "resources"
 
     def __init__(self, run_path: str, backend: str):
+        super().__init__(run_path)
         self._events_backend = backend
-        self._run_path = run_path
         self._files = {}  # type: dict[str, LoggedEventListSpec]
-        self._closed = False
 
     def _get_event_path(self, kind: str, name: str) -> str:
         if self._events_backend == self.EVENTS_BACKEND:
@@ -86,13 +109,50 @@ def flush(self):
                 self._append_events(events_spec)
             self._files[file_name].empty_events()
 
-    def close(self):
-        self.flush()
-        self._closed = True
 
-    @property
-    def closed(self):
-        return self._closed
+class LogWriter(BaseWriter):
+    def __init__(self, run_path: str):
+        super().__init__(run_path)
+        self._file = None  # type: V1Logs | None
+
+    def _get_event_path(self) -> str:
+        return get_logs_path(
+            run_path=self._run_path,
+        )
+
+    def _init_events(self, events_spec: V1Logs):
+        event_path = self._get_event_path()
+        # Check if the file exists otherwise initialize
+        if not os.path.exists(event_path):
+            check_or_create_path(event_path, is_dir=False)
+            with open(event_path, "w") as event_file:
+                event_file.write("")
+            set_permissions(event_path)
+
+    def _append_events(self, events_spec: V1Logs):
+        event_path = self._get_event_path()
+        with open(event_path, "a") as event_file:
+            event_file.write(events_spec.get_jsonl_events())
+
+    def _events_to_files(self, logs: List[V1Log]):
+        if self._file:
+            self._file.logs += logs
+        else:
+            self._file = V1Logs(logs=logs)
+            self._init_events(self._file)
+
+    def write(self, events: List[V1Log]):
+        if not events:
+            return
+        if isinstance(events, V1Log):
+            events = [events]
+        self._events_to_files(events)
+
+    def flush(self):
+        events_spec = self._file
+        if events_spec.logs:
+            self._append_events(events_spec)
+        self._file = V1Logs(logs=[])
 
 
 class BaseFileWriter:
@@ -110,14 +170,14 @@ def __init__(self, run_path: str):
     def run_path(self):
         return self._run_path
 
-    def add_event(self, event: LoggedEventSpec):
-        if not isinstance(event, LoggedEventSpec):
+    def add_event(self, event: Union[LoggedEventSpec, V1Log]):
+        if not isinstance(event, LoggedEventSpec) and not isinstance(event, V1Log):
             raise TypeError("Expected an LoggedEventSpec, " " but got %s" % type(event))
         self._async_writer.write(event)
 
-    def add_events(self, events: List[LoggedEventSpec]):
+    def add_events(self, events: List[Union[LoggedEventSpec, V1Log]]):
         for e in events:
-            if not isinstance(e, LoggedEventSpec):
+            if not isinstance(e, LoggedEventSpec) and not isinstance(e, V1Log):
                 raise TypeError("Expected an LoggedEventSpec, " " but got %s" % type(e))
         self._async_writer.write(events)
 

traceml/traceml/serialization/writer.py

@@ -7,13 +7,13 @@
 from clipped.utils.paths import check_or_create_path
 
 from traceml.events import LoggedEventSpec, get_asset_path, get_event_path
-from traceml.events.paths import get_resource_path
+from traceml.events.paths import get_resource_path, get_logs_path
 from traceml.processors.gpu_processor import can_log_gpu_resources, get_gpu_metrics
 from traceml.processors.psutil_processor import (
     can_log_psutil_resources,
     get_psutils_metrics,
 )
-from traceml.serialization.base import BaseFileWriter, EventWriter
+from traceml.serialization.base import BaseFileWriter, EventWriter, LogWriter
 
 
 class EventFileWriter(BaseFileWriter):
@@ -59,10 +59,33 @@ def __init__(self, run_path: str, max_queue_size: int = 20, flush_secs: int = 10
         )
 
 
+class LogsFileWriter(BaseFileWriter):
+    def __init__(self, run_path: str, max_queue_size: int = 20, flush_secs: int = 10):
+        """Creates a `LogsFileWriter`.
+
+        Args:
+          run_path: A string. Directory where events files will be written.
+          max_queue_size: Integer. Size of the queue for pending logs.
+          flush_secs: Number. How often, in seconds, to flush the
+            pending logs to disk.
+        """
+        super().__init__(run_path=run_path)
+
+        check_or_create_path(get_logs_path(run_path, False), is_dir=True)
+
+        self._async_writer = LogAsyncManager(
+            LogWriter(self._run_path),
+            max_queue_size,
+            flush_secs,
+        )
+
+
 class BaseAsyncManager:
     """Base manager for writing events to files by name by event kind."""
 
-    def __init__(self, event_writer: EventWriter, max_queue_size: int = 20):
+    def __init__(
+        self, event_writer: Union[EventWriter, LogWriter], max_queue_size: int = 20
+    ):
         """Writes events json spec to files asynchronously. An instance of this class
         holds a queue to keep the incoming data temporarily. Data passed to the
         `write` function will be put to the queue and the function returns
@@ -248,3 +271,18 @@ def run(self):
                     self._event_writer.write(data)
                     self._event_writer.flush()
                 self._next_flush_time = now + self._flush_secs
+
+
+class LogAsyncManager(BaseAsyncManager):
+    """Writes logs to local file."""
+
+    def __init__(
+        self, event_writer: LogWriter, max_queue_size: int = 20, flush_secs: int = 10
+    ):
+        super().__init__(event_writer=event_writer, max_queue_size=max_queue_size)
+        self._worker = EventWriterThread(
+            self._event_queue,
+            self._event_writer,
+            flush_secs,
+        )
+        self._worker.start()

traceml/traceml/tracking/run.py

@@ -37,10 +37,13 @@
 from traceml.artifacts import V1ArtifactKind
 from traceml.events import LoggedEventSpec, V1Event, V1EventSpan, get_asset_path
 from traceml.logger import logger
-from traceml.logging import V1Log, V1Logs
 from traceml.processors import events_processors
 from traceml.processors.logs_processor import end_log_processor, start_log_processor
-from traceml.serialization.writer import EventFileWriter, ResourceFileWriter
+from traceml.serialization.writer import (
+    EventFileWriter,
+    ResourceFileWriter,
+    LogsFileWriter,
+)
 
 
 class Run(RunClient):
@@ -139,11 +142,11 @@ def __init__(
             no_op=no_op,
         )
         track_logs = track_logs if track_logs is not None else self._is_offline
-        self._logs_history = V1Logs.construct(logs=[])
         self._artifacts_path = None
         self._outputs_path = None
         self._event_logger = None
         self._resource_logger = None
+        self._logs_logger = None
         self._sidecar = None
         self._exit_handler = None
         self._store_path = None
@@ -180,7 +183,8 @@ def __init__(
             self.log_code_ref()
 
         if (is_new or has_process_sidecar) and self._artifacts_path and track_logs:
-            start_log_processor(add_logs=self._add_logs)
+            self.set_run_logs_logger()
+            start_log_processor(add_logs=self._logs_logger.add_event)
 
         self._set_exit_handler(force=is_new or has_process_sidecar)
 
@@ -235,27 +239,6 @@ def _add_events(self, events: List[LoggedEventSpec]):
                 "the event logger was not configured properly".format(len(events))
             )
 
-    def _persist_logs_history(self):
-        if self._logs_history.logs and len(self._logs_history.logs) > 0:
-            logs_path = os.path.join(
-                self._artifacts_path,
-                "plxlogs",
-                "{}".format(datetime.timestamp(self._logs_history.logs[-1].timestamp)),
-            )
-            check_or_create_path(logs_path, is_dir=False)
-            with open(logs_path, "w") as logs_file:
-                logs_file.write(self._logs_history.to_json())
-            set_permissions(logs_path)
-
-    def _add_logs(self, log: V1Log):
-        if not log:
-            return
-        self._logs_history.logs.append(log)
-        if V1Logs.should_chunk(self._logs_history.logs):
-            self._persist_logs_history()
-            # Reset
-            self._logs_history = V1Logs.construct(logs=[])
-
     def create(self, **kwargs):
         raise NotImplementedError(
             "The tracking `Run` subclass does not allow to call "
@@ -459,6 +442,15 @@ def set_run_resource_logger(self):
         """
         self._resource_logger = ResourceFileWriter(run_path=self._artifacts_path)
 
+    @client_handler(check_no_op=True)
+    def set_run_logs_logger(self):
+        """Sets a logs logger.
+
+        > **Note**: This is only used during manual tracking, and it's not used by the `
+        > in-cluster` runs.
+        """
+        self._logs_logger = LogsFileWriter(run_path=self._artifacts_path)
+
     @client_handler(check_no_op=True)
     def set_run_process_sidecar(self):
         """Sets a sidecar process to sync artifacts.
@@ -1730,7 +1722,6 @@ def excepthook(exception, value, tb):
     def _end(self):
         self.log_succeeded()
         end_log_processor()
-        self._persist_logs_history()
         self._wait(sync_artifacts=True)
         if self._is_offline:
             self.persist_run(path=self._artifacts_path)
@@ -1754,6 +1745,8 @@ def _wait(self, sync_artifacts: bool = False):
             self._event_logger.close()
         if self._resource_logger:
             self._resource_logger.close()
+        if self._logs_logger:
+            self._logs_logger.close()
         if self._sidecar:
             self._sidecar.close()
         if self._results: