nicetrace
1from .tracing import ( 2 Metadata, 3 Tag, 4 TracingNode, 5 trace, 6 trace_instant, 7 TracingNodeState, 8 current_tracing_node, 9 with_trace, 10) 11from .serialization import ( 12 register_custom_serializer, 13 unregister_custom_serializer, 14 serialize_with_type, 15) 16from .data.html import Html 17from .data.blob import DataWithMime 18from .writer.base import current_writer, TraceWriter 19from .writer.filewriter import DirWriter, FileWriter 20from .reader.filereader import DirReader, TraceReader 21from .html.statichtml import get_full_html, write_html 22 23__all__ = [ 24 "trace", 25 "trace_instant", 26 "with_trace", 27 "TracingNode", 28 "Metadata", 29 "TracingNodeState", 30 "Tag", 31 "current_tracing_node", 32 "current_writer", 33 "register_custom_serializer", 34 "unregister_custom_serializer", 35 "serialize_with_type", 36 "Html", 37 "DataWithMime", 38 "TraceWriter", 39 "DirWriter", 40 "FileWriter", 41 "TraceReader", 42 "DirReader", 43 "get_full_html", 44 "write_html", 45]
288@contextmanager 289def trace( 290 name: str, 291 kind: str | None = None, 292 inputs: dict[str, Any] | None = None, 293 meta: Metadata | None = None, 294 writer: Optional["TraceWriter"] = None, 295): 296 """ 297 The main function that creates a tracing context manager. Returns an instance of `TracingNode`. 298 ```python 299 with trace("my node", inputs={"z": 42}) as c: 300 c.add_input("x", 1) 301 y = do_some_computation(x=1) 302 # The tracing node would also note any exceptions raised here 303 # (letting it propagate upwards), but an output needs to be set manually: 304 c.add_output("", y) 305 # <- Here the tracing node is already closed. 306 ``` 307 """ 308 node, token = start_trace_block(name, kind, inputs, meta, writer) 309 try: 310 yield node 311 except BaseException as e: 312 end_trace_block(node, token, e, writer) 313 raise e 314 end_trace_block(node, token, None, writer)
The main function that creates a tracing context manager. Returns an instance of TracingNode
.
with trace("my node", inputs={"z": 42}) as c:
c.add_input("x", 1)
y = do_some_computation(x=1)
# The tracing node would also note any exceptions raised here
# (letting it propagate upwards), but an output needs to be set manually:
c.add_output("", y)
# <- Here the tracing node is already closed.
317def trace_instant( 318 name: str, 319 kind: Optional[str] = None, 320 inputs: Optional[dict[str, Any]] = None, 321 output: Optional[Any] = None, 322 meta: Optional[Metadata] = None, 323): 324 """ 325 Trace an instant event that does not have a duration. 326 """ 327 return current_tracing_node().add_instant(name, kind, inputs, meta)
Trace an instant event that does not have a duration.
330def with_trace( 331 fn: Callable = None, *, name=None, kind=None, meta: Optional[Metadata] = None 332): 333 """ 334 A decorator wrapping every execution of the function in a new `TracingNode`. 335 336 The `inputs`, `output`, and `error` (if any) are set automatically. 337 Note that you can access the created tracing in your function using `current_tracing_node`. 338 339 *Usage:* 340 341 ```python 342 @with_trace 343 def func(): 344 pass 345 346 @with_trace(name="custom_name", kind="custom_kind", tags=['tag1', 'tag2']) 347 def func(): 348 pass 349 ``` 350 """ 351 if isinstance(fn, str): 352 raise TypeError("use `with_tracing()` with explicit `name=...` parameter") 353 354 def helper(func): 355 signature = inspect.signature(func) 356 357 @functools.wraps(func) 358 def wrapper(*a, **kw): 359 binding = signature.bind(*a, **kw) 360 with trace( 361 name=name or func.__name__, 362 kind=kind or "call", 363 inputs=binding.arguments, 364 meta=meta, 365 ) as node: 366 output = func(*a, **kw) 367 node.add_output("", output) 368 return output 369 370 async def async_wrapper(*a, **kw): 371 binding = signature.bind(*a, **kw) 372 with trace( 373 name=name or func.__name__, 374 kind=kind or "acall", 375 inputs=binding.arguments, 376 ) as node: 377 output = await func(*a, **kw) 378 node.add_output("", output) 379 return output 380 381 if inspect.iscoroutinefunction(func): 382 return async_wrapper 383 else: 384 return wrapper 385 386 if fn is not None: 387 assert callable(fn) 388 return helper(fn) 389 else: 390 return helper
A decorator wrapping every execution of the function in a new TracingNode
.
The inputs
, output
, and error
(if any) are set automatically.
Note that you can access the created tracing in your function using current_tracing_node
.
Usage:
@with_trace
def func():
pass
@with_trace(name="custom_name", kind="custom_kind", tags=['tag1', 'tag2'])
def func():
pass
59class TracingNode: 60 """ 61 A tracing object that represents a single request or (sub)task in a nested hierarchy. 62 """ 63 64 def __init__( 65 self, 66 name: str, 67 kind: Optional[str] = None, 68 meta: Optional[Metadata] = None, 69 lock=None, 70 is_instant=False, 71 ): 72 """ 73 - `name` - A description or name for the tracing node. 74 - `kind` - Indicates category of the tracing node, may e.g. influence display of the tracing node. 75 - `inputs` - A dictionary of inputs for the tracing node. 76 - `meta` - A dictionary of any metadata for the tracing node, e.g. UI style data. 77 This allows you to split the stored data across multiple files. 78 - `output` - The output value of the tracing node, if it has already been computed. 79 """ 80 81 if meta: 82 assert isinstance(meta, Metadata) 83 84 self.name = name 85 self.kind = kind 86 self.uid = generate_uid() 87 self.entries: list | None = None 88 self.children: list[TracingNode] | None = None 89 if is_instant: 90 self.start_time = None 91 self.end_time = datetime.now() 92 self.state = TracingNodeState.FINISHED 93 else: 94 self.start_time = datetime.now() 95 self.end_time = None 96 self.state = TracingNodeState.OPEN 97 self.meta = meta 98 self._lock = lock 99 100 def _to_dict(self): 101 result = {"name": self.name, "uid": self.uid} 102 if self.state != TracingNodeState.FINISHED: 103 result["state"] = self.state.value 104 for name in "kind", "entries": 105 value = getattr(self, name) 106 if value is not None: 107 result[name] = value 108 if self.kind: 109 result["kind"] = self.kind 110 if self.entries: 111 result["entries"] = self.entries 112 if self.children: 113 result["children"] = [c._to_dict() for c in self.children] 114 if self.start_time: 115 result["start_time"] = self.start_time.isoformat() 116 if self.end_time: 117 result["end_time"] = self.end_time.isoformat() 118 if self.meta is not None: 119 result["meta"] = serialize_with_type(self.meta) 120 return result 121 122 def to_dict(self): 123 """ 124 Serialize `TracingNode` object into JSON structure. 125 """ 126 with self._lock: 127 result = self._to_dict() 128 result["version"] = TRACING_FORMAT_VERSION 129 return result 130 131 def add_tag(self, tag: str | Tag): 132 """ 133 Add a tag to the tracing node. 134 """ 135 with self._lock: 136 if self.meta is None: 137 self.meta = Metadata() 138 if self.meta.tags is None: 139 self.meta.tags = [] 140 self.meta.tags.append(tag) 141 142 def add_instant( 143 self, 144 name: str, 145 kind: Optional[str] = None, 146 inputs: Optional[dict[str, Any]] = None, 147 meta: Optional[Metadata] = None, 148 ) -> "TracingNode": 149 node = TracingNode( 150 name=name, 151 kind=kind, 152 meta=meta, 153 is_instant=True, 154 ) 155 if inputs: 156 for key, value in inputs.items(): 157 node._add_entry("input", key, value) 158 with self._lock: 159 if self.children is None: 160 self.children = [] 161 self.children.append(node) 162 return node 163 164 def add_entry(self, kind: str, name: str, value: object): 165 """ 166 Add a entry into the node. 167 """ 168 with self._lock: 169 self._add_entry(kind, name, value) 170 171 def _add_entry(self, kind: str, name: str, value: object): 172 if self.entries is None: 173 self.entries = [] 174 entry = {"kind": kind, "value": serialize_with_type(value)} 175 if name: 176 entry["name"] = name 177 self.entries.append(entry) 178 179 def add_input(self, name: str, value: object): 180 """ 181 A shortcut for .add_entry(entry_type="input") 182 """ 183 self.add_entry("input", name, value) 184 185 def add_inputs(self, inputs: dict[str, object]): 186 """ 187 A shortcut for calling multiple .add_entry(entry_type="input") 188 """ 189 with self._lock: 190 for key, value in inputs.items(): 191 self._add_entry("input", key, value) 192 193 def add_output(self, name: str, value: Any): 194 """ 195 A shortcut for .add_entry(entry_type="output") 196 """ 197 self.add_entry("output", name, value) 198 199 def set_error(self, exc: Any): 200 """ 201 Set the error value of the tracing node (usually an `Exception` instance). 202 """ 203 with self._lock: 204 self.state = TracingNodeState.ERROR 205 self._add_entry("error", "", exc) 206 207 def find_nodes(self, predicate: Callable) -> list["TracingNode"]: 208 """ 209 Find all nodes matching the given callable `predicate`. 210 211 The predicate is called with a single argument, the `TracingNode` to check, and should return `bool`. 212 """ 213 214 def _helper(node: TracingNode): 215 if predicate(node): 216 result.append(node) 217 if node.children: 218 for child in node.children: 219 _helper(child) 220 221 result = [] 222 with self._lock: 223 _helper(self) 224 return result 225 226 def _repr_html_(self): 227 from .html.statichtml import get_inline_html 228 229 return get_inline_html(self)
A tracing object that represents a single request or (sub)task in a nested hierarchy.
64 def __init__( 65 self, 66 name: str, 67 kind: Optional[str] = None, 68 meta: Optional[Metadata] = None, 69 lock=None, 70 is_instant=False, 71 ): 72 """ 73 - `name` - A description or name for the tracing node. 74 - `kind` - Indicates category of the tracing node, may e.g. influence display of the tracing node. 75 - `inputs` - A dictionary of inputs for the tracing node. 76 - `meta` - A dictionary of any metadata for the tracing node, e.g. UI style data. 77 This allows you to split the stored data across multiple files. 78 - `output` - The output value of the tracing node, if it has already been computed. 79 """ 80 81 if meta: 82 assert isinstance(meta, Metadata) 83 84 self.name = name 85 self.kind = kind 86 self.uid = generate_uid() 87 self.entries: list | None = None 88 self.children: list[TracingNode] | None = None 89 if is_instant: 90 self.start_time = None 91 self.end_time = datetime.now() 92 self.state = TracingNodeState.FINISHED 93 else: 94 self.start_time = datetime.now() 95 self.end_time = None 96 self.state = TracingNodeState.OPEN 97 self.meta = meta 98 self._lock = lock
name
- A description or name for the tracing node.kind
- Indicates category of the tracing node, may e.g. influence display of the tracing node.inputs
- A dictionary of inputs for the tracing node.meta
- A dictionary of any metadata for the tracing node, e.g. UI style data. This allows you to split the stored data across multiple files.output
- The output value of the tracing node, if it has already been computed.
122 def to_dict(self): 123 """ 124 Serialize `TracingNode` object into JSON structure. 125 """ 126 with self._lock: 127 result = self._to_dict() 128 result["version"] = TRACING_FORMAT_VERSION 129 return result
Serialize TracingNode
object into JSON structure.
131 def add_tag(self, tag: str | Tag): 132 """ 133 Add a tag to the tracing node. 134 """ 135 with self._lock: 136 if self.meta is None: 137 self.meta = Metadata() 138 if self.meta.tags is None: 139 self.meta.tags = [] 140 self.meta.tags.append(tag)
Add a tag to the tracing node.
164 def add_entry(self, kind: str, name: str, value: object): 165 """ 166 Add a entry into the node. 167 """ 168 with self._lock: 169 self._add_entry(kind, name, value)
Add a entry into the node.
179 def add_input(self, name: str, value: object): 180 """ 181 A shortcut for .add_entry(entry_type="input") 182 """ 183 self.add_entry("input", name, value)
A shortcut for .add_entry(entry_type="input")
185 def add_inputs(self, inputs: dict[str, object]): 186 """ 187 A shortcut for calling multiple .add_entry(entry_type="input") 188 """ 189 with self._lock: 190 for key, value in inputs.items(): 191 self._add_entry("input", key, value)
A shortcut for calling multiple .add_entry(entry_type="input")
193 def add_output(self, name: str, value: Any): 194 """ 195 A shortcut for .add_entry(entry_type="output") 196 """ 197 self.add_entry("output", name, value)
A shortcut for .add_entry(entry_type="output")
199 def set_error(self, exc: Any): 200 """ 201 Set the error value of the tracing node (usually an `Exception` instance). 202 """ 203 with self._lock: 204 self.state = TracingNodeState.ERROR 205 self._add_entry("error", "", exc)
Set the error value of the tracing node (usually an Exception
instance).
207 def find_nodes(self, predicate: Callable) -> list["TracingNode"]: 208 """ 209 Find all nodes matching the given callable `predicate`. 210 211 The predicate is called with a single argument, the `TracingNode` to check, and should return `bool`. 212 """ 213 214 def _helper(node: TracingNode): 215 if predicate(node): 216 result.append(node) 217 if node.children: 218 for child in node.children: 219 _helper(child) 220 221 result = [] 222 with self._lock: 223 _helper(self) 224 return result
Find all nodes matching the given callable predicate
.
The predicate is called with a single argument, the TracingNode
to check, and should return bool
.
45@dataclass 46class Metadata: 47 """ 48 Metadata of tracing, allows to set colors and icons when TracingNode is visualized 49 """ 50 51 icon: str | None = None 52 color: str | None = None 53 tags: list[Tag] | None = None 54 counters: dict[str, int] | None = None 55 collapse: str | None = None 56 custom: Any = None
Metadata of tracing, allows to set colors and icons when TracingNode is visualized
20class TracingNodeState(Enum): 21 """ 22 An enumeration representing the state of a tracing node. 23 """ 24 25 OPEN = "open" 26 """The tracing node is currently running.""" 27 FINISHED = "finished" 28 """The tracing node has successfully finished execution.""" 29 ERROR = "error" 30 """The tracing node finished with an exception."""
An enumeration representing the state of a tracing node.
The tracing node has successfully finished execution.
Inherited Members
- enum.Enum
- name
- value
33@dataclass 34class Tag: 35 """ 36 A simple class representing a tag that can be applied to a tracing node. Optionally with style information. 37 """ 38 39 name: str 40 """The name of the tag; any short string.""" 41 color: Optional[str] = None 42 """HTML color code, e.g. `#ff0000`."""
A simple class representing a tag that can be applied to a tracing node. Optionally with style information.
393def current_tracing_node(check: bool = True) -> Optional[TracingNode]: 394 """ 395 Returns the inner-most open tracing node, if any. 396 397 Throws an error if `check` is `True` and there is no current tracing node. If `check` is `False` and there is 398 no current tracing node, it returns `None`. 399 """ 400 stack = _TRACING_STACK.get() 401 if not stack: 402 if check: 403 raise Exception("No current tracing") 404 return None 405 return stack[-1]
Returns the inner-most open tracing node, if any.
Throws an error if check
is True
and there is no current tracing node. If check
is False
and there is
no current tracing node, it returns None
.
44def current_writer() -> TraceWriter | None: 45 """ 46 Get the current global writer. 47 """ 48 return _TRACE_WRITER.get()
Get the current global writer.
101def register_custom_serializer(cls: type[T], serialize_fn: Callable[[T], Data]): 102 """ 103 Register a custom serializer for a given type 104 """ 105 CUSTOM_SERIALIZERS[cls] = serialize_fn
Register a custom serializer for a given type
108def unregister_custom_serializer(cls): 109 """ 110 Unregister a custom serializer for a given type 111 """ 112 CUSTOM_SERIALIZERS.pop(cls, None)
Unregister a custom serializer for a given type
2class Html: 3 """ 4 Wrapper around HTML code that is directly rendered in Data Browser 5 """ 6 7 def __init__(self, html: str): 8 self.html = html 9 10 def __trace_to_node__(self): 11 return {"_type": "$html", "html": self.html}
Wrapper around HTML code that is directly rendered in Data Browser
9class DataWithMime: 10 """ 11 Wrapper around bytes that are serialized by base64. 12 Data Browser renders some MIME types in a specific way 13 (e.g. images are rendered directly into browser). 14 """ 15 16 def __init__(self, data: bytes, mime_type: str = MIME_OCTET_STREAM): 17 self.data = data 18 self.mime_type = mime_type 19 20 def __trace_to_node__(self): 21 return { 22 "_type": "$blob", 23 "data": base64.b64encode(self.data).decode(), 24 "mime_type": self.mime_type, 25 }
Wrapper around bytes that are serialized by base64. Data Browser renders some MIME types in a specific way (e.g. images are rendered directly into browser).
13class TraceWriter(ABC): 14 """ 15 Abstract base class for all trace writers. 16 """ 17 18 @abstractmethod 19 def write_node(self, node: TracingNode, final: bool): 20 raise NotImplementedError() 21 22 @abstractmethod 23 def sync(self): 24 pass 25 26 @abstractmethod 27 def start(self): 28 pass 29 30 @abstractmethod 31 def stop(self): 32 pass 33 34 def __enter__(self): 35 self.__token = _TRACE_WRITER.set(self) 36 self.start() 37 return self 38 39 def __exit__(self, exc_type, exc_val, exc_tb): 40 self.stop() 41 _TRACE_WRITER.reset(self.__token)
Abstract base class for all trace writers.
92class DirWriter(DelayedWriter): 93 """ 94 Writes JSON serialized trace into a given directory. 95 Trace is saved under filename trace-<ID>.json. 96 It allows to write multiple traces at once. 97 """ 98 99 def __init__( 100 self, path: str, min_write_delay: timedelta = timedelta(milliseconds=300) 101 ): 102 super().__init__(min_write_delay) 103 Path(path).mkdir(parents=True, exist_ok=True) 104 self.path = path 105 106 def _write_node_to_file(self, node): 107 json_data = json.dumps(node.to_dict()) 108 write_file(os.path.join(self.path, f"trace-{node.uid}.json"), json_data) 109 110 def write_node(self, node: TracingNode, final: bool): 111 with self.lock: 112 self._write_node(node, final)
Writes JSON serialized trace into a given directory.
Trace is saved under filename trace-
115class FileWriter(DelayedWriter): 116 """ 117 Write JSON serialized trace into given file. 118 It allows to write only one trace at time. 119 It throws an error if more then one top-level trace node is created at once. 120 """ 121 122 def __init__( 123 self, filename: str, min_write_delay: timedelta = timedelta(milliseconds=300) 124 ): 125 super().__init__(min_write_delay) 126 127 filename = os.path.abspath(filename) 128 path = os.path.dirname(filename) 129 Path(path).mkdir(parents=True, exist_ok=True) 130 if not filename.endswith(".json"): 131 filename = filename + ".json" 132 133 self.filename = filename 134 self.current_node = None 135 136 def _write_node_to_file(self, node): 137 json_data = json.dumps(node.to_dict()) 138 write_file(self.filename, json_data) 139 140 def write_node(self, node: TracingNode, final: bool): 141 with self.lock: 142 if self.current_node is None: 143 self.current_node = node 144 elif node != self.current_node: 145 raise Exception( 146 "FileWriter allows to write only one root node at once," 147 "use DirWriter if you need " 148 ) 149 self._write_node(node, final) 150 if final: 151 self.current_node = None
Write JSON serialized trace into given file. It allows to write only one trace at time. It throws an error if more then one top-level trace node is created at once.
5class TraceReader(ABC): 6 """Abstract base class for reading traces""" 7 8 @abstractmethod 9 def list_summaries(self) -> list[dict]: 10 """Get summaries of traces in storage""" 11 raise NotImplementedError() 12 13 @abstractmethod 14 def read_trace(self, uid: str) -> dict: 15 """Read a trace from storage""" 16 raise NotImplementedError()
Abstract base class for reading traces
9class DirReader(TraceReader): 10 """ 11 Reads a traces from a given directory. 12 """ 13 14 def __init__(self, path: str): 15 if not os.path.isdir(path): 16 raise Exception(f"Path '{path}' does not exists") 17 self.path = path 18 self.finished_paths = {} 19 self.uids_to_filenames = {} 20 self.lock = Lock() 21 22 def list_summaries(self) -> list[dict]: 23 summaries = [] 24 with self.lock: 25 finished_paths = self.finished_paths 26 for filename in os.listdir(self.path): 27 if filename.endswith(".json"): 28 summary = finished_paths.get(filename) 29 if summary: 30 summaries.append(summary) 31 continue 32 with open(os.path.join(self.path, filename)) as f: 33 snode = json.loads(f.read()) 34 state = snode.get("state", "finished") 35 summary = { 36 "storage_id": filename[: -len(".json")], 37 "uid": snode["uid"], 38 "name": snode["name"], 39 "state": state, 40 "start_time": snode["start_time"], 41 "end_time": snode.get("end_time"), 42 } 43 if state != "open": 44 finished_paths[filename] = summary 45 summaries.append(summary) 46 return summaries 47 48 def read_trace(self, storage_id: str) -> dict: 49 assert "/" not in storage_id 50 assert not storage_id.startswith(".") 51 with open(os.path.join(self.path, f"{storage_id}.json")) as f: 52 return json.loads(f.read())
Reads a traces from a given directory.
22 def list_summaries(self) -> list[dict]: 23 summaries = [] 24 with self.lock: 25 finished_paths = self.finished_paths 26 for filename in os.listdir(self.path): 27 if filename.endswith(".json"): 28 summary = finished_paths.get(filename) 29 if summary: 30 summaries.append(summary) 31 continue 32 with open(os.path.join(self.path, filename)) as f: 33 snode = json.loads(f.read()) 34 state = snode.get("state", "finished") 35 summary = { 36 "storage_id": filename[: -len(".json")], 37 "uid": snode["uid"], 38 "name": snode["name"], 39 "state": state, 40 "start_time": snode["start_time"], 41 "end_time": snode.get("end_time"), 42 } 43 if state != "open": 44 finished_paths[filename] = summary 45 summaries.append(summary) 46 return summaries
Get summaries of traces in storage
62def write_html(node: TracingNode, filename: str | os.PathLike): 63 """Write a `TracingNode` as static HTML file""" 64 write_file(filename, get_full_html(node))
Write a TracingNode
as static HTML file