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

v0/relationalai/util/span_tracker.py

@@ -1,207 +0,0 @@
-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

@@ -1,72 +0,0 @@
-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

@@ -1,34 +0,0 @@
-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)