-
Notifications
You must be signed in to change notification settings - Fork 400
/
trace_handler.py
102 lines (83 loc) · 3.94 KB
/
trace_handler.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0
"""Profiler Trace Handler."""
from __future__ import annotations
import abc
import pathlib
from typing import Dict, List, Tuple, Union
from composer.core.callback import Callback
from composer.core.time import Timestamp
__all__ = ['TraceHandler']
class TraceHandler(Callback, abc.ABC):
"""Base class for Composer Profiler trace handlers.
Subclasses should implement :meth:`process_duration_event`, :meth:`process_instant_event`,
:meth:`process_counter_event`, and :meth:`process_chrome_json_trace_file` to record trace events.
Since :class:`TraceHandler` subclasses :class:`~composer.core.callback.Callback`, a trace handler can run on any
:class:`.Event` (such as on :attr:`.Event.INIT` to open files or on :attr:`.Event.BATCH_END` to periodically dump
data to files) and use :meth:`.Callback.close` to perform any cleanup.
"""
def process_duration_event(
self,
name: str,
categories: Union[List[str], Tuple[str, ...]],
is_start: bool,
timestamp: Timestamp,
wall_clock_time_ns: int,
) -> None:
"""Invoked whenever there is a duration event to record.
This method is called twice for each duration event -- once with ``is_start = True``,
and then again with ``is_start = False``. Interleaving events are not permitted.
Specifically, for each event (identified by the ``name``), a call with ``is_start = True`` will be followed
by a call with ``is_start = False`` before another call with ``is_start = True``.
Args:
name (str): The name of the event.
categories (Union[List[str], Tuple[str, ...]]): The categories for the event.
is_start (bool): Whether the event is a start event or end event.
timestamp (Timestamp): Snapshot of the training time.
wall_clock_time_ns (int): The :py:func:`time.time_ns` corresponding to the event.
"""
del name, categories, is_start, timestamp, wall_clock_time_ns # unused
pass
def process_instant_event(
self,
name: str,
categories: Union[List[str], Tuple[str, ...]],
timestamp: Timestamp,
wall_clock_time_ns: int,
) -> None:
"""Invoked whenever there is an instant event to record.
Args:
name (str): The name of the event.
categories (List[str] | Tuple[str, ...]): The categories for the event.
timestamp (Timestamp): Snapshot of current training time.
wall_clock_time_ns (int): The :py:func:`time.time_ns` corresponding to the event.
"""
del name, categories, timestamp, wall_clock_time_ns # unused
pass
def process_counter_event(
self,
name: str,
categories: Union[List[str], Tuple[str, ...]],
timestamp: Timestamp,
wall_clock_time_ns: int,
values: Dict[str, Union[int, float]],
) -> None:
"""Invoked whenever there is an counter event to record.
Args:
name (str): The name of the event.
categories (List[str] | Tuple[str, ...]): The categories for the event.
timestamp (Timestamp): The timestamp.
wall_clock_time_ns (int): The :py:func:`time.time_ns` corresponding to the event.
values (Dict[str, int | float]): The values corresponding to this counter event.
"""
del name, categories, timestamp, wall_clock_time_ns, values # unused
pass
def process_chrome_json_trace_file(self, filepath: pathlib.Path) -> None:
"""Invoked when there are events in Chrome JSON format to record.
See `this document <https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview>`_
for more information.
Args:
filepath (pathlib.Path): The filepath to a Chrome JSON trace file.
"""
del filepath # unused
pass