Merge pull request #203 from Point72/tkp/logfire

Logfire integration

Author: timkpaine

csp_gateway/server/config/gateway/omnibus.yaml

@@ -43,6 +43,7 @@ gateway:
     UI: True
     API_KEY: "12345"
   modules:
+    - /modules/logfire
     - /modules/example_module
     - /modules/example_module_feedback
     - /modules/example_custom_table

csp_gateway/server/demo/config/logfire.yaml

@@ -0,0 +1,44 @@
+# @package _global_
+defaults:
+  - /gateway: omnibus
+  - _self_
+
+logfire:
+  _target_: csp_gateway.server.modules.logging.Logfire
+  project_name: "csp-gateway-poc"
+  service_name: "CSP Gateway"
+  instrument_fastapi: True
+  capture_logging: True
+  log_level: DEBUG
+  
+publish_logfire:
+  _target_: csp_gateway.server.modules.logging.PublishLogfire
+  selection:
+    include:
+      - example
+  service_name: "CSP Gateway"
+  include_metadata: True
+  use_spans: True
+  log_level: DEBUG
+
+gateway:
+  _target_: csp_gateway.Gateway
+  settings:
+    PORT: ${port}
+    AUTHENTICATE: ${authenticate}
+    UI: True
+    API_KEY: "12345"
+  modules:
+    - /modules/example_module
+    - /modules/mount_outputs
+    - /modules/mount_perspective_tables
+    - /modules/mount_rest_routes
+    - /logfire
+    - /publish_logfire
+  channels:
+    _target_: csp_gateway.server.demo.ExampleGatewayChannels
+
+# csp-gateway-start --config-dir=csp_gateway/server/demo +config=logfire
+
+authenticate: False
+port: 8000

csp_gateway/server/modules/controls/controls.py

@@ -4,8 +4,13 @@
 import threading
 from datetime import datetime, timezone
 
+try:
+    import psutil
+except ImportError:
+    # Hold and raise in model validator
+    psutil = None
+
 import csp
-import psutil
 from csp import ts
 
 from csp_gateway.server import GatewayChannels, GatewayModule
@@ -53,16 +58,16 @@ def manage_controls(self, data: ts[Controls]):
                 stats = {}
 
                 # Machine information
-                stats["cpu"] = psutil.cpu_percent()
-                stats["memory"] = psutil.virtual_memory().percent
+                stats["cpu"] = psutil.cpu_percent() if psutil else None
+                stats["memory"] = psutil.virtual_memory().percent if psutil else None
                 stats["memory-total"] = round(
-                    psutil.virtual_memory().available * 100 / psutil.virtual_memory().total,
+                    psutil.virtual_memory().available * 100 / psutil.virtual_memory().total if psutil else 0,
                     2,
                 )
 
                 # Process and thread information
-                current_process = psutil.Process()
-                stats["pid"] = current_process.pid
+                current_process = psutil.Process() if psutil else None
+                stats["pid"] = current_process.pid if current_process else None
                 stats["active_threads"] = threading.active_count()
 
                 # Get max threads from ulimit

csp_gateway/server/modules/filedrop/__init__.py

@@ -1,10 +1,2 @@
-import sys
-
-if sys.platform.startswith("linux"):
-    try:
-        from .adapter import *
-        from .filedrop import *
-    except ImportError:
-        pass
-else:
-    pass
+from .adapter import *
+from .filedrop import *

csp_gateway/server/modules/filedrop/adapter.py

@@ -5,15 +5,23 @@
 from enum import Enum, auto
 from typing import Any, Callable, Dict, List, Optional, TypeVar
 
-import csp
 import orjson
 import pyarrow.parquet as pq
 from csp.impl.pushadapter import PushInputAdapter
 from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
 from csp.impl.wiring import py_push_adapter_def
 from pydantic import TypeAdapter
-from watchdog.events import FileSystemEvent, FileSystemEventHandler
-from watchdog.observers import Observer
+
+try:
+    from watchdog.events import FileSystemEvent, FileSystemEventHandler
+    from watchdog.observers import Observer
+except ImportError:
+    # Hold and raise in model validator
+    FileSystemEvent = object
+    FileSystemEventHandler = object
+    Observer = None
+
+import csp
 
 __all__ = (
     "FileDropType",

csp_gateway/server/modules/filedrop/filedrop.py

@@ -16,6 +16,12 @@
 
 from .adapter import FileDropAdapterConfiguration, FileDropType, filedrop_adapter_def
 
+try:
+    import watchdog
+except ImportError:
+    # Hold and raise in model validator
+    watchdog = None
+
 __all__ = (
     "ReadFileDropConfiguration",
     "ReadFileDrop",
@@ -44,6 +50,12 @@ class ReadFileDropConfiguration(BaseModel):
         description=("If False, replaces the timestamp field on the GatewayStruct with a timestamp autogenerated by the current Gateway."),
     )
 
+    @model_validator(mode="before")
+    def check_import(cls, values):
+        if watchdog is None:
+            raise ImportError("watchdog is required for ReadFileDrop. Install it with: pip install watchdog")
+        return values
+
     @model_validator(mode="after")
     def check_filedrop_type(self) -> "ReadFileDropConfiguration":
         if isinstance(self.filedrop_type, str):
@@ -65,6 +77,12 @@ class ReadFileDrop(GatewayModule):
 
     configs: List[ReadFileDropConfiguration] = Field(description="List of configs for the directories to monitor")
 
+    @model_validator(mode="before")
+    def check_import(cls, values):
+        if watchdog is None:
+            raise ImportError("watchdog is required for ReadFileDrop. Install it with: pip install watchdog")
+        return values
+
     @csp.node
     def handle_list_basket(self, data: csp.ts[List[T]], list_size: int) -> csp.OutputBasket(List[csp.ts[T]], shape="list_size"):
         if csp.ticked(data):

csp_gateway/server/modules/logging/__init__.py

@@ -1,17 +1,11 @@
-from .logging import LogChannels
+from .datadog import PublishDatadog
+from .logfire import (
+    Logfire,
+    PublishLogfire,
+    configure_logfire_early,
+    is_logfire_configured,
+)
+from .opsgenie import PublishOpsGenie
 from .printing import PrintChannels
-
-try:
-    from .symphony import PublishSymphony
-except ImportError:
-    pass
-
-try:
-    from .datadog import PublishDatadog
-except ImportError:
-    pass
-
-try:
-    from .opsgenie import PublishOpsGenie
-except ImportError:
-    pass
+from .stdlib import LogChannels
+from .symphony import PublishSymphony

csp_gateway/server/modules/logging/datadog.py

@@ -6,7 +6,7 @@
 
 import csp
 from csp import ts
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from csp_gateway.server import ChannelSelection, GatewayModule
 from csp_gateway.server.modules.logging.util import (
@@ -15,6 +15,12 @@
 )
 from csp_gateway.utils import get_thread
 
+try:
+    from datadog import api
+except ImportError:
+    # Hold and raise in model validator
+    api = None
+
 log = logging.getLogger(__name__)
 logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
 logging.getLogger("datadog.api").setLevel(logging.ERROR)
@@ -54,8 +60,6 @@ def _log_result(
 
 
 def _send_to_datadog(queue: Queue, dd_latency_log_threshold_seconds: int) -> None:
-    from datadog import api
-
     while True:
         data = queue.get()
         if data is None:
@@ -98,6 +102,12 @@ class PublishDatadog(GatewayModule):
 
     dd_tags: Optional[Dict[str, str]] = Field(default=None, description="Tags to be included with Datadog submissions.")
 
+    @model_validator(mode="before")
+    def check_import(cls, values):
+        if api is None:
+            raise ImportError("datadog is required for PublishDatadog. Install it with: pip install datadog")
+        return values
+
     def connect(self, channels):
         """
         Channels to be connected to graph