Add Chrome trace support

Author: dmelikyan

graphsignal/__init__.py

@@ -38,20 +38,22 @@ def _parse_env_param(name: str, value: Any, expected_type: Type) -> Any:
     raise ValueError(f"Invalid type for {name}: expected {expected_type.__name__}, got {type(value).__name__}")
 
 
-def _read_config_param(name: str, expected_type: Type, provided_value: Optional[Any] = None, required: bool = False) -> Any:
+def _read_config_param(name: str, expected_type: Type, provided_value: Optional[Any] = None, default_value: Optional[Any] = None, required: bool = False) -> Any:
     # Check if the value was provided as an argument
     if provided_value is not None:
         return provided_value
 
     # Check if the value was provided as an environment variable
     env_value = os.getenv(f'GRAPHSIGNAL_{name.upper()}')
-    if env_value is None:
-        if required:
-            raise ValueError(f"Missing required argument: {name}")
-        return None
+    if env_value is not None:
+        parsed_env_value = _parse_env_param(name, env_value, expected_type)
+        if parsed_env_value is not None:
+            return parsed_env_value
+
+    if required:
+        raise ValueError(f"Missing required argument: {name}")
 
-    # Parse the environment variable
-    return _parse_env_param(name, env_value, expected_type)
+    return default_value
 
 
 def _read_config_tags(provided_value: Optional[dict] = None, prefix: str = "GRAPHSIGNAL_TAG_") -> Dict[str, str]:
@@ -68,10 +70,10 @@ def configure(
     api_url: Optional[str] = None,
     deployment: Optional[str] = None,
     tags: Optional[Dict[str, str]] = None,
-    auto_instrument: Optional[bool] = True,
-    record_payloads: Optional[bool] = True,
-    profiling_rate: Optional[float] = 0.1,
-    debug_mode: Optional[bool] = False
+    auto_instrument: Optional[bool] = None,
+    record_payloads: Optional[bool] = None,
+    profiling_rate: Optional[float] = None,
+    debug_mode: Optional[bool] = None
 ) -> None:
     global _tracer
 
@@ -82,10 +84,10 @@ def configure(
     api_key = _read_config_param("api_key", str, api_key, required=True)
     api_url = _read_config_param("api_url", str, api_url)
     tags = _read_config_tags(tags)
-    auto_instrument = _read_config_param("auto_instrument", bool, auto_instrument)
-    record_payloads = _read_config_param("record_payloads", bool, record_payloads)
-    profiling_rate = _read_config_param("profiling_rate", float, profiling_rate)
-    debug_mode = _read_config_param("debug_mode", bool, debug_mode)
+    auto_instrument = _read_config_param("auto_instrument", bool, auto_instrument, default_value=True)
+    record_payloads = _read_config_param("record_payloads", bool, record_payloads, default_value=True)
+    profiling_rate = _read_config_param("profiling_rate", float, profiling_rate, default_value=0.1)
+    debug_mode = _read_config_param("debug_mode", bool, debug_mode, default_value=False)
 
     # left for compatibility
     if deployment and isinstance(deployment, str):

graphsignal/recorders/profiler_utils.py

@@ -0,0 +1,39 @@
+import logging
+import os
+import tempfile
+import gzip
+import shutil
+import glob
+
+logger = logging.getLogger('graphsignal')
+
+def create_log_dir():
+    log_dir = tempfile.mkdtemp(prefix='graphsignal-')
+    logger.debug('Created temporary log directory %s', log_dir)
+    return log_dir
+
+def remove_log_dir(log_dir):
+    shutil.rmtree(log_dir)
+    logger.debug('Removed temporary log directory %s', log_dir)
+
+def find_and_read(log_dir, file_pattern, decompress=True, max_size=None):
+    file_paths = glob.glob(os.path.join(log_dir, file_pattern))
+    if len(file_paths) == 0:
+        logger.debug('Files are not found at %s', os.path.join(log_dir, file_pattern))
+        return None
+
+    found_path = file_paths[-1]
+
+    if max_size:
+        file_size = os.path.getsize(found_path)
+        if file_size > max_size:
+            raise Exception('File is too big: {0}'.format(file_size))
+
+    if decompress and found_path.endswith('.gz'):
+        last_file = gzip.open(found_path, "rb")
+    else:
+        last_file = open(found_path, "rb")
+    data = last_file.read()
+    last_file.close()
+
+    return data

graphsignal/recorders/pytorch_recorder.py

@@ -1,16 +1,20 @@
 import logging
-import torch
+import os
 import random
 import json
+import time
+import torch
 
 import graphsignal
 from graphsignal.recorders.base_recorder import BaseRecorder
+from graphsignal.recorders.profiler_utils import create_log_dir, remove_log_dir
 
 logger = logging.getLogger('graphsignal')
 
 class PyTorchRecorder(BaseRecorder):
     def __init__(self):
         self._torch_prof = None
+        self._log_dir = None
 
     def setup(self):
         pass
@@ -87,17 +91,43 @@ def on_span_read(self, span, context):
                             count = 1,
                             duration_ns = _ns(kernel.duration)
                         )
-
+            
             device_profile = kernel_index.values()
             if len(device_profile) > 0:
                 span.set_profile('device-profile', 'event-averages', json.dumps(device_profile))
                 span.set_tag('profile_type', 'device') # override cpu value
 
+            chrome_trace = self._export_chrome_trace()
+            if chrome_trace:
+                span.set_profile('event-timeline', 'chrome-trace', chrome_trace)
+
             if len(cpu_profile) > 0 or len(device_profile) > 0:
                 span.set_tag('profiler', f'pytorch-{torch.__version__}')
         finally:
             self._torch_prof = None
 
+    def _export_chrome_trace(self):
+        try:
+            read_start_time = time.time()
+
+            self._log_dir = create_log_dir()
+
+            trace_path = os.path.join(self._log_dir, 'trace.json')
+            self._torch_prof.export_chrome_trace(trace_path)
+
+            trace_file_size = os.path.getsize(trace_path)
+            logger.debug('Chrome trace size: %s', trace_file_size)
+            if trace_file_size > 50 * 1e6:
+                raise Exception('Trace file too big: {0}'.format(trace_file_size))
+
+            with open(trace_path, "r") as f:
+                return str(f.read())
+        finally:
+            remove_log_dir(self._log_dir)
+            logger.debug('Chrome trace export time: %s', time.time() - read_start_time)
+
+        return None
+
 def _ns(val):
     return int(max(val, 0) * 1e3)
 

graphsignal/version.py

@@ -1 +1 @@
-__version__ = '0.16.2'
+__version__ = '0.16.3'