relationalai 0.13.0__py3-none-any.whl → 0.13.0.dev0__py3-none-any.whl

v0/relationalai/util/span_tracker.py

@@ -0,0 +1,207 @@
+import json
+from logging import LogRecord
+import os
+import subprocess
+from typing import List, Tuple, Union, Dict, Any
+from uuid import UUID, uuid4
+
+from v0.relationalai.debugging import Span, encode_log_message, filter_span_attrs, logger
+from v0.relationalai.util.constants import SPAN_ATTR_ALLOW_LIST, SPAN_TYPES_KEYS
+from v0.relationalai.util.format import default_serialize
+from v0.relationalai.clients.config import Config
+
+TRACE_ID = uuid4()
+
+class SpanTracker:
+    """
+    Abstract base class for tracking spans produced as log messages. 
+    Provides methods to handle the start and end of spans, format span messages, and manage span attributes.
+    """
+
+    def __init__(self, trace_id: UUID, span_types_to_skip=None, log_span_attrs_as_str=True):
+        self.trace_id = trace_id
+        self._set_span_types_to_skip(span_types_to_skip)
+        self._log_span_attrs_as_str = log_span_attrs_as_str
+        self.open_spans = {}
+        self.skipped_spans = set()
+
+    def handle_span_start(self, msg: dict) -> Union[None,Span]:
+        span: Span = msg["span"]
+        assert_key_exists_for_span_type(span)
+        if span.type in self.span_types_to_skip:
+            self.skipped_spans.add(str(span.id))
+            return None
+        self.open_spans[str(span.id)] = span
+        return span
+
+    def handle_span_end(self, msg: dict) -> Union[Tuple[dict, str, dict], None]:
+        if str(msg["id"]) in self.skipped_spans:
+            return None
+        span = self.open_spans.pop(msg["id"])
+        joined_key_path = get_span_path_as_str(span)
+        end_attrs = get_span_allowed_attributes_values(span, msg["end_attrs"])
+        combined_attributes = {
+            **span.attrs, # span_start attrs
+            **end_attrs, # span_end attrs
+            "key_path": joined_key_path,
+            "elapsed_s": (span.end_timestamp - span.start_timestamp).total_seconds(),
+        }
+        filtered_attrs = filter_span_attrs(combined_attributes)
+        formatted_span_msg = format_span_msg(self.trace_id, span, filtered_attrs, attrs_as_str=self._log_span_attrs_as_str)
+        return formatted_span_msg, joined_key_path, filtered_attrs
+
+    def _set_span_types_to_skip(self, span_types_to_skip):
+        if span_types_to_skip is None:
+            self.span_types_to_skip = set()
+        elif isinstance(span_types_to_skip, str):
+            self.span_types_to_skip = {span_types_to_skip}
+        elif isinstance(span_types_to_skip, (list, tuple)):
+            self.span_types_to_skip = set(span_types_to_skip)
+        elif isinstance(span_types_to_skip, set):
+            self.span_types_to_skip = span_types_to_skip
+        else:
+            raise ValueError(f"Invalid span_types_to_skip: {span_types_to_skip}, type: {type(span_types_to_skip)}")
+
+def format_span_json(trace_id: UUID, span: Span, attrs: dict, attrs_as_str:bool=True) -> dict:
+    span_key = get_span_value_from_key(span)
+    if attrs_as_str:
+        a = json.dumps(attrs, default=default_serialize)
+    else:
+        a = encode_log_message(attrs)
+
+    span_json = span.to_json()
+    extra_attrs = {
+        "trace_id": trace_id,
+        "key": str(span_key),
+        "attrs": a
+    }
+    formatted_span_msg = {**span_json, **extra_attrs}
+    return formatted_span_msg
+
+#region Span formatting and filtering functions
+def format_span_msg(trace_id: UUID, span: Span, attrs: dict, attrs_as_str:bool=True) -> dict:
+    span_key = get_span_value_from_key(span)
+    if attrs_as_str:
+        a = json.dumps(attrs, default=default_serialize)
+    else:
+        a = encode_log_message(attrs)
+
+    span_json = span.to_json()
+    extra_attrs = {
+        "trace_id": str(trace_id) if attrs_as_str else trace_id,
+        "key": str(span_key) if attrs_as_str else span_key,
+        "attrs": a
+    }
+    formatted_span_msg = {**span_json, **extra_attrs}
+    return formatted_span_msg
+
+def get_span_allowed_attributes_values(span: Span, attrs=None) -> dict:
+    # For an input span with attrs (or instead for an `attrs` dict), return a dict with only the allowed attributes for the span type
+    attrs = attrs or span.attrs
+    out = {}
+    allowed_attributes = SPAN_ATTR_ALLOW_LIST.get(span.type, set())
+    for attr in allowed_attributes:
+        if attr in attrs:
+            out[attr] = attrs[attr]
+    return out
+
+def assert_key_exists_for_span_type(span: Span):
+    if span.type in SPAN_TYPES_KEYS:
+        key = SPAN_TYPES_KEYS[span.type]
+        if key not in span.attrs:
+            raise ValueError(f"Attribute {key} is required for spans of type {span.type}, but wasn't provided. Provided attributes:\n{span.attrs}")
+
+def get_span_key_path(span: Span) -> str:
+    # Return the key path of the current span
+    value = get_span_value_from_key(span)
+    if value is None:
+        return span.type
+    return f"{span.type}({value})"
+
+def get_span_path_as_list(span: Span) -> List[Span]:
+    # Return the path from the root to the current span
+    path = []
+    cur = span
+    while cur is not None:
+        path.append(cur)
+        cur = cur.parent
+    path.reverse()
+    return path
+
+def get_span_path_as_str(span: Span) -> str:
+    # Return a string representation of the path from the root to the current span
+    # formerly format_key_path function
+    return '.'.join([get_span_key_path(span) for span in get_span_path_as_list(span)])
+
+def get_span_value_from_key(span: Span):
+    if span.type in SPAN_TYPES_KEYS:
+        key = SPAN_TYPES_KEYS[span.type]
+        if key in span.attrs:
+            return span.attrs[key]
+        raise ValueError(f"Key {key} not found in span of type {span.type}\n{span.attrs}")
+    return None
+
+def record_is_span(record: LogRecord) -> bool:
+    return (
+        isinstance(record.msg, dict) and
+        "event" in record.msg and
+        record.msg["event"] in ("span_start", "span_end")
+    )
+
+def _get_local_git_info() -> Dict[str, str]:
+    try:
+        branch = subprocess.check_output(['git', 'rev-parse', '--abbrev-ref', 'HEAD']).strip().decode('utf-8')
+        sha = subprocess.check_output(['git', 'rev-parse', 'HEAD']).strip().decode('utf-8')
+        repo = subprocess.check_output(['git', 'config', '--get', 'remote.origin.url']).strip().decode('utf-8').rsplit('/', maxsplit=1)[-1].replace('.git', '')
+        return {
+            'commit': sha,
+            'branch': branch,
+            'repo': repo,
+        }
+    except subprocess.CalledProcessError as e:
+        logger.warning(f"Could not get git info: {e}")
+        return {
+            'commit': 'local',
+            'branch': 'local',
+            'repo': 'local',
+        }
+
+def _get_git_info() -> Dict[str, str]:
+    if os.getenv("GITHUB_REF_NAME") is not None:
+        return {
+            'commit': os.getenv("GITHUB_SHA", ""),
+            'branch': os.getenv("GITHUB_REF_NAME", ""),
+            'repo': os.getenv("GITHUB_REPOSITORY", ""),
+        }
+    return _get_local_git_info()
+
+def _get_masked_config() -> Dict[str, Any]:
+    """Return a config with all sensitive keys (secrets and passwords) removed, such that it can be safely logged"""
+    config = Config()
+    masked_props = {
+        k: v for k, v in config.props.items()
+        if not isinstance(v, str)
+        or not any(prop_key_pattern in k for prop_key_pattern in ["secret", "password"])
+    }
+    return masked_props
+
+def _get_cloud_provider() -> Dict[str, str]:
+    provider = os.getenv("RAI_CLOUD_PROVIDER")
+    if not provider:
+        config = Config()
+        provider = config.get("platform")
+    assert provider, "Could not retrieve cloud provider from environment or config"
+    return {"platform": provider}
+
+def get_root_span_attrs(get_full_config: bool = False) -> Dict[str, Any]:
+    git_info = _get_git_info()
+    cloud_provider = _get_cloud_provider()
+    masked_config = _get_masked_config() if get_full_config else {}
+    attrs = {
+        **git_info,
+        **cloud_provider,
+        **masked_config
+    }
+    return attrs
+
+#endregion

v0/relationalai/util/spans_file_handler.py

@@ -0,0 +1,72 @@
+from copy import deepcopy
+from logging import FileHandler
+import os
+import sys
+from uuid import UUID
+from typing import Dict, Union
+
+from v0.relationalai.debugging import JsonFormatter
+from v0.relationalai.util import get_timestamp
+from v0.relationalai.util.format import str_to_bool
+from v0.relationalai.util.span_tracker import SpanTracker, record_is_span
+
+# Span types that are automatically generated by PyRel but are not relevant for performance analysis
+SKIP_SPAN_TYPES = ['check_status']
+
+
+class SpansFileHandler(FileHandler):
+    """
+    This logger writes spans to a local .jsonl file inside the `spans_output` folder as they finish.
+    The primary goal is to ingest these spans into Snowflake for analytics.
+    """
+
+    def __init__(self, trace_id: UUID, filename: Union[str,None] = None):
+        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=SKIP_SPAN_TYPES, log_span_attrs_as_str=False)
+        if filename is None:
+            filename = self._get_default_log_file_path()
+            if not os.path.exists(os.path.dirname(filename)):
+                os.makedirs(os.path.dirname(filename))
+        FileHandler.__init__(self, filename=filename)
+        self.setFormatter(JsonFormatter())
+
+    def emit(self, record):
+        if record_is_span(record):
+            msg: Dict = record.msg # type: ignore
+            if msg["event"] == "span_start":
+                self.span_tracker.handle_span_start(msg)
+            if msg["event"] == "span_end":
+                span_end_result = self.span_tracker.handle_span_end(msg)
+                if span_end_result:
+                    formatted_span_msg = span_end_result[0]
+                    if formatted_span_msg is None: # This will happen for span_types_to_skip
+                        return
+                    record_copy = deepcopy(record)  # Create a deep copy so as to not interfere with other handlers
+                    record_copy.msg = formatted_span_msg
+                    FileHandler.emit(self, record_copy)
+
+    def _get_default_log_file_path(self) -> str:
+        now_str = get_timestamp().strftime("%Y%m%d_%H%M%S")
+        trace_id_trunc = str(self.span_tracker.trace_id)[:8]
+        spans_subfolder = os.getenv("SPANS_SUBFOLDER", "")
+        return os.path.join(
+            os.getcwd(), "output_spans", spans_subfolder, f"{now_str}__{trace_id_trunc}.jsonl"
+        )
+
+
+def is_collect_only():
+    return "pytest" in sys.modules and "--collect-only" in sys.argv
+
+
+def maybe_add_spans_file_handler(logger, trace_id:UUID, if_not_set: bool = False):
+    """
+    Add a SpansFileHandler to the logger if the 'ENABLE_SPANS_FILE_HANDLER' environment variable is set to True,
+    or if the 'if_not_set' parameter is True.
+    """
+    if str_to_bool(os.getenv("ENABLE_SPANS_FILE_HANDLER")) or if_not_set:
+        if is_collect_only():
+            print("Running in collect-only mode, skipping SpansFileHandler")
+        elif any(isinstance(handler, SpansFileHandler) for handler in logger.handlers):
+            print("logger already has a SpansFileHandler instance, won't add another one")
+        else:
+            logger.debug("Adding SpansFileHandler to logger")
+            logger.addHandler(SpansFileHandler(trace_id))

v0/relationalai/util/tracing_handler.py

@@ -0,0 +1,34 @@
+from logging import Handler
+from typing import Dict
+from uuid import UUID
+
+from v0.relationalai.util.span_tracker import SpanTracker, get_span_allowed_attributes_values, get_span_path_as_str, record_is_span
+
+
+class TracingHandler(Handler):
+    """
+    This logger prints spans to stdout as they start and finish.
+    """
+
+    def __init__(self, trace_id: UUID):
+        super().__init__()
+        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=None)
+
+    def emit(self, record):
+        if record_is_span(record):
+            msg: Dict = record.msg # type: ignore
+            if msg["event"] == "span_start":
+                span_to_print = self.span_tracker.handle_span_start(msg)
+                if span_to_print is not None:
+                    attrs = get_span_allowed_attributes_values(span_to_print)
+                    print("start", get_span_path_as_str(span_to_print), attrs)
+            elif msg["event"] == "span_end":
+                span_end_result = self.span_tracker.handle_span_end(msg)
+                if span_end_result:
+                    joined_key_path, filtered_attrs = span_end_result[1], span_end_result[2]
+                    if joined_key_path is not None and filtered_attrs is not None:
+                        # key_path is already printed separately, and name is already in the key_path
+                        filtered_attrs = {k: v for k, v in filtered_attrs.items() if k not in ["name", "key_path"]}
+                        # Exclude class instances, which can have a large nested structure to print
+                        filtered_attrs = {k: v for k, v in filtered_attrs.items() if not hasattr(v, "__dict__")}
+                        print("  end", joined_key_path, filtered_attrs)

frontend/debugger/dist/.gitignore

@@ -1,2 +0,0 @@
-*
-!.gitignore