Rework emmitter registration and other improvements

Implemented datadog specific log exporter
Implement a non-blocking log store flushing system
Add OTEL log store flavor

Author: stefannica

src/zenml/log_stores/__init__.py

@@ -23,6 +23,7 @@
 # OpenTelemetry log store
 from zenml.log_stores.otel.otel_flavor import (
     OtelLogStoreConfig,
+    OtelLogStoreFlavor,
 )
 from zenml.log_stores.otel.otel_log_store import OtelLogStore
 
@@ -50,4 +51,5 @@
     "DatadogLogStoreFlavor",
     "OtelLogStore",
     "OtelLogStoreConfig",
+    "OtelLogStoreFlavor",
 ]

src/zenml/log_stores/artifact/artifact_log_exporter.py

@@ -73,7 +73,7 @@ def export(self, batch: Sequence["LogData"]) -> LogExportResult:
                 if not attrs:
                     continue
 
-                log_uri = attrs.get("zenml.log_model.uri")
+                log_uri = attrs.get("zenml.log.uri")
                 if not log_uri or not isinstance(log_uri, str):
                     continue
 

src/zenml/log_stores/artifact/artifact_log_store.py

@@ -18,9 +18,11 @@
 from datetime import datetime
 from typing import (
     Any,
+    Dict,
     Iterator,
     List,
     Optional,
+    Type,
     cast,
 )
 from uuid import UUID
@@ -30,9 +32,16 @@
 from zenml.artifact_stores import BaseArtifactStore
 from zenml.enums import LoggingLevels, StackComponentType
 from zenml.exceptions import DoesNotExistException
-from zenml.log_stores.base_log_store import MAX_ENTRIES_PER_REQUEST
+from zenml.log_stores import BaseLogStore
+from zenml.log_stores.base_log_store import (
+    MAX_ENTRIES_PER_REQUEST,
+    BaseLogStoreEmitter,
+)
 from zenml.log_stores.otel.otel_flavor import OtelLogStoreConfig
-from zenml.log_stores.otel.otel_log_store import OtelLogStore
+from zenml.log_stores.otel.otel_log_store import (
+    OtelLogStore,
+    OtelLogStoreEmitter,
+)
 from zenml.logger import get_logger
 from zenml.models import LogsResponse
 from zenml.utils.io_utils import sanitize_remote_path
@@ -128,6 +137,9 @@ def _stream_logs_line_by_line(
     Raises:
         DoesNotExistException: If the artifact does not exist in the artifact store.
     """
+    if not artifact_store.exists(logs_uri):
+        return []
+
     if not artifact_store.isdir(logs_uri):
         # Single file case
         with artifact_store.open(logs_uri, "r") as file:
@@ -199,6 +211,31 @@ class ArtifactLogStoreConfig(OtelLogStoreConfig):
     """Configuration for the artifact log store."""
 
 
+class ArtifactLogStoreEmitter(OtelLogStoreEmitter):
+    """Artifact log store emitter."""
+
+    def __init__(
+        self,
+        name: str,
+        log_store: "BaseLogStore",
+        log_model: LogsResponse,
+        metadata: Dict[str, Any],
+    ) -> None:
+        """Initialize a log store emitter.
+
+        Args:
+            name: The name of the emitter.
+            log_store: The log store to emit logs to.
+            log_model: The log model associated with the emitter.
+            metadata: Additional metadata to attach to all log entries that will
+                be emitted by this emitter.
+        """
+        super().__init__(name, log_store, log_model, metadata)
+
+        if log_model.uri:
+            self._metadata["zenml.log.uri"] = log_model.uri
+
+
 class ArtifactLogStore(OtelLogStore):
     """Log store that saves logs to the artifact store.
 
@@ -220,6 +257,15 @@ def __init__(
         super().__init__(*args, **kwargs)
         self._artifact_store = artifact_store
 
+    @property
+    def emitter_class(self) -> Type[ArtifactLogStoreEmitter]:
+        """Class of the emitter.
+
+        Returns:
+            The class of the emitter.
+        """
+        return ArtifactLogStoreEmitter
+
     @classmethod
     def from_artifact_store(
         cls, artifact_store: "BaseArtifactStore"
@@ -236,7 +282,7 @@ def from_artifact_store(
             artifact_store=artifact_store,
             id=artifact_store.id,
             name="default",
-            config=ArtifactLogStoreConfig(),
+            config=ArtifactLogStoreConfig(endpoint=artifact_store.path),
             flavor="artifact",
             type=StackComponentType.LOG_STORE,
             user=artifact_store.user,
@@ -265,26 +311,20 @@ def get_exporter(self) -> "LogExporter":
 
         return ArtifactLogExporter(artifact_store=self._artifact_store)
 
-    def finalize(
+    def _finalize(
         self,
-        log_model: LogsResponse,
+        emitter: BaseLogStoreEmitter,
     ) -> None:
-        """Finalize the stream of log records associated with a log model.
+        """Finalize the stream of log records associated with an emitter.
 
         Args:
-            log_model: The log model to finalize.
+            emitter: The emitter to finalize.
         """
+        assert isinstance(emitter, ArtifactLogStoreEmitter)
         with self._lock:
-            if not self._provider or self._logger is None:
-                return
-
-        self._logger.emit(
-            body=END_OF_STREAM_MESSAGE,
-            attributes={
-                "zenml.log_model.id": str(log_model.id),
-                "zenml.log_model.uri": str(log_model.uri),
-            },
-        )
+            emitter.logger.emit(
+                body=END_OF_STREAM_MESSAGE,
+            )
 
     def fetch(
         self,

src/zenml/log_stores/base_log_store.py

@@ -15,7 +15,7 @@
 
 import logging
 import threading
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Type, cast
 
@@ -31,7 +31,73 @@ class BaseLogStoreConfig(StackComponentConfig):
     """Base configuration for all log stores."""
 
 
-class BaseLogStore(StackComponent):
+class BaseLogStoreEmitter:
+    """Base class for all ZenML log store emitters.
+
+    The emitter is the entry point for all log records to be emitted to the log
+    store. The process of emitting a log record is as follows:
+
+    1. instantiate a BaseLogStoreEmitter
+    2. create an emitter by calling log_store.register_emitter() and passing the
+    log model and optional metadata to be attached to each log record
+    3. emit the log record by calling emitter.emit() and passing the log record
+    4. deregister the emitter when all logs have been emitted by calling
+    emitter.deregister()
+    """
+
+    def __init__(
+        self,
+        name: str,
+        log_store: "BaseLogStore",
+        log_model: LogsResponse,
+        metadata: Dict[str, Any],
+    ) -> None:
+        """Initialize a log store emitter.
+
+        Args:
+            name: The name of the emitter.
+            log_store: The log store to emit logs to.
+            log_model: The log model associated with the emitter.
+            metadata: Additional metadata to attach to all log entries that will
+                be emitted by this emitter.
+        """
+        self._name = name
+        self._log_store = log_store
+        self._log_model = log_model
+        self._metadata = metadata
+
+    @property
+    def name(self) -> str:
+        """The name of the emitter.
+
+        Returns:
+            The name of the emitter.
+        """
+        return self._name
+
+    @property
+    def log_model(self) -> LogsResponse:
+        """The log model associated with the emitter.
+
+        Returns:
+            The log model.
+        """
+        return self._log_model
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """Emit a log record to the log store.
+
+        Args:
+            record: The log record to emit.
+        """
+        self._log_store._emit(self, record, metadata=self._metadata)
+
+    def deregister(self) -> None:
+        """Deregister the emitter from the log store."""
+        self._log_store.deregister_emitter(self)
+
+
+class BaseLogStore(StackComponent, ABC):
     """Base class for all ZenML log stores.
 
     A log store is responsible for collecting, storing, and retrieving logs
@@ -47,7 +113,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             **kwargs: Keyword arguments for the base class.
         """
         super().__init__(*args, **kwargs)
-        self._emitter_counter = 0
+        self._emitters: Dict[str, BaseLogStoreEmitter] = {}
         self._lock = threading.RLock()
 
     @property
@@ -59,56 +125,88 @@ def config(self) -> BaseLogStoreConfig:
         """
         return cast(BaseLogStoreConfig, self._config)
 
+    @property
+    def emitter_class(self) -> Type[BaseLogStoreEmitter]:
+        """Class of the emitter.
+
+        Returns:
+            The class of the emitter.
+        """
+        return BaseLogStoreEmitter
+
+    def register_emitter(
+        self, name: str, log_model: LogsResponse, metadata: Dict[str, Any]
+    ) -> BaseLogStoreEmitter:
+        """Register an emitter for the log store.
+
+        Args:
+            name: The name of the emitter.
+            log_model: The log model associated with the emitter.
+            metadata: Additional metadata to attach to the log entry.
+
+        Returns:
+            The emitter.
+        """
+        with self._lock:
+            emitter = self.emitter_class(name, self, log_model, metadata)
+            self._emitters[name] = emitter
+            return emitter
+
+    def deregister_emitter(self, emitter: BaseLogStoreEmitter) -> None:
+        """Deregister an emitter registered with the log store.
+
+        Args:
+            emitter: The emitter to deregister.
+        """
+        with self._lock:
+            if emitter.name not in self._emitters:
+                return
+            self._finalize(emitter)
+            del self._emitters[emitter.name]
+            if len(self._emitters) == 0:
+                self.flush(blocking=False)
+
     @abstractmethod
-    def emit(
+    def _emit(
         self,
+        emitter: BaseLogStoreEmitter,
         record: logging.LogRecord,
-        log_model: LogsResponse,
         metadata: Dict[str, Any],
     ) -> None:
         """Process a log record from the logging system.
 
         Args:
+            emitter: The emitter to emit the log record to.
             record: The Python logging.LogRecord to process.
-            log_model: The log model to emit the log record to.
             metadata: Additional metadata to attach to the log entry.
         """
 
     @abstractmethod
-    def finalize(
+    def _finalize(
         self,
-        log_model: LogsResponse,
+        emitter: BaseLogStoreEmitter,
     ) -> None:
-        """Finalize the stream of log records associated with a log model.
+        """Finalize the stream of log records associated with an emitter.
 
         This is used to announce the end of the stream of log records associated
-        with a log model and that no more log records will be emitted.
+        with an emitter and that no more log records will be emitted.
 
         The implementation should ensure that all log records associated with
-        the log model are flushed to the backend and any resources (clients,
+        the emitter are flushed to the backend and any resources (clients,
         connections, file descriptors, etc.) are released.
 
         Args:
-            log_model: The log model to finalize.
+            emitter: The emitter to finalize.
         """
 
-    def register_emitter(self) -> None:
-        """Register an emitter for the log store."""
-        with self._lock:
-            self._emitter_counter += 1
-
-    def deregister_emitter(self) -> None:
-        """Deregister an emitter for the log store."""
-        with self._lock:
-            self._emitter_counter -= 1
-            if self._emitter_counter == 0:
-                self.flush()
-
     @abstractmethod
-    def flush(self) -> None:
+    def flush(self, blocking: bool = True) -> None:
         """Flush the log store.
 
         This method is called to ensure that all logs are flushed to the backend.
+
+        Args:
+            blocking: Whether to block until the flush is complete.
         """
 
     @abstractmethod

src/zenml/log_stores/datadog/datadog_flavor.py

@@ -13,9 +13,9 @@
 #  permissions and limitations under the License.
 """Datadog log store flavor."""
 
-from typing import Type
+from typing import Any, Dict, Type
 
-from pydantic import Field, field_validator
+from pydantic import Field, field_validator, model_validator
 
 from zenml.enums import StackComponentType
 from zenml.log_stores import BaseLogStore, BaseLogStoreConfig
@@ -73,6 +73,22 @@ def validate_max_export_batch_size(cls, v: int) -> int:
             )
         return v
 
+    @model_validator(mode="before")
+    @classmethod
+    def set_default_endpoint(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Set the endpoint based on site if not provided.
+
+        Args:
+            data: The input data dictionary.
+
+        Returns:
+            The data dictionary with the endpoint set if not provided.
+        """
+        if isinstance(data, dict) and not data.get("endpoint"):
+            site = data.get("site", "datadoghq.com")
+            data["endpoint"] = f"https://http-intake.logs.{site}/api/v2/logs"
+        return data
+
 
 class DatadogLogStoreFlavor(Flavor):
     """Datadog log store flavor."""