Source code for composer.profiler.marker

# Copyright 2022 MosaicML Composer authors
# SPDX-License-Identifier: Apache-2.0

"""Profiler Marker."""

from __future__ import annotations

import functools
import time
from types import TracebackType
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union

from composer.profiler.trace_handler import TraceHandler

if TYPE_CHECKING:
    from composer.core import State, Timestamp

__all__ = ['Marker']


[docs]class Marker: """Profiler Marker. Used by the :class:`.Engine` to measure the duration of :class:`.Event` during training. .. note:: :class:`.Marker` should not be instantiated directly; instead use :meth:`.Profiler.marker`. Markers can record the following types of events: #. Duration: Records the start and stop time of an event of interest (:meth:`.Marker.start()`, :meth:`.Marker.finish()`). #. Instant: Record time a particular event occurs, but not the full duration (:meth:`.Marker.instant()`). #. Counter: The value of a variable at given time (:meth:`.Marker.counter()`). A :class:`.Marker` can also be used as a context manager or decorator to record a duration: #. Use a :class:`.Marker` with a context manager: .. testsetup:: from composer.profiler import Profiler, cyclic_schedule profiler = Profiler(schedule=cyclic_schedule(), trace_handlers=[], torch_prof_memory_filename=None) profiler.bind_to_state(state) .. doctest:: >>> def something_to_measure(): ... print("something_to_measure") >>> marker = profiler.marker("foo") >>> with marker: ... something_to_measure() something_to_measure #. Use a :class:`.Marker` as a decorator: .. testsetup:: from composer.profiler import Profiler, cyclic_schedule profiler = Profiler(schedule=cyclic_schedule(), trace_handlers=[], torch_prof_memory_filename=None) profiler.bind_to_state(state) .. doctest:: >>> marker = profiler.marker("foo") >>> @marker ... def something_to_measure(): ... print("something_to_measure") >>> something_to_measure() something_to_measure """ def __init__(self, state: State, should_record: Callable[[State], bool], trace_handlers: Sequence[TraceHandler], name: str, record_instant_on_start: bool, record_instant_on_finish: bool, categories: Union[List[str], Tuple[str, ...]]) -> None: self.state = state self.trace_handlers = trace_handlers self.name = name self.categories = categories self.record_instant_on_start = record_instant_on_start self.record_instant_on_finish = record_instant_on_finish self.should_record = should_record self._started = False self._recorded_start = False def _record_duration_event(self, is_start: bool, wall_clock_time_ns: int, timestamp: Timestamp): """Record a duration event.""" for handler in self.trace_handlers: handler.process_duration_event( name=self.name, categories=self.categories, timestamp=timestamp, is_start=is_start, wall_clock_time_ns=wall_clock_time_ns, ) def _record_instant_event(self, wall_clock_time_ns: int, timestamp: Timestamp): """Record an instant event.""" for handler in self.trace_handlers: handler.process_instant_event( name=self.name, categories=self.categories, timestamp=timestamp, wall_clock_time_ns=wall_clock_time_ns, ) def _record_counter_event(self, wall_clock_time_ns: int, timestamp: Timestamp, values: Dict[str, Union[int, float]]) -> None: """Record a counter invent.""" for handler in self.trace_handlers: handler.process_counter_event( name=self.name, categories=self.categories, wall_clock_time_ns=wall_clock_time_ns, timestamp=timestamp, values=values, )
[docs] def start(self) -> None: """Record the start of a duration event. To record the duration of an event, invoke :meth:`.Marker.start` followed by :meth:`.Marker.finish`. .. testsetup:: from composer.profiler import Profiler, cyclic_schedule profiler = Profiler(schedule=cyclic_schedule(), trace_handlers=[], torch_prof_memory_filename=None) profiler.bind_to_state(state) .. doctest:: >>> def something_to_measure(): ... print("something_to_measure") >>> marker = profiler.marker("foo") >>> marker.start() >>> something_to_measure() something_to_measure >>> marker.finish() """ if self._started: raise RuntimeError( f'Attempted to start profiler event {self.name}; however, this marker is already started') self._recorded_start = self.should_record(self.state) if self._recorded_start: wall_clock_time = time.time_ns() self._record_duration_event( is_start=True, wall_clock_time_ns=wall_clock_time, timestamp=self.state.timestamp, ) if self.record_instant_on_start: self._record_instant_event( timestamp=self.state.timestamp, wall_clock_time_ns=wall_clock_time, ) self._started = True
[docs] def finish(self) -> None: """Record the end of a duration event. See :meth:`.Marker.start()` for a usage example. """ if not self._started: raise RuntimeError( f'Attempted to finish profiler event {self.name}; however, this profiler event is not yet started') wall_clock_time = time.time_ns() self._record_duration_event( is_start=False, timestamp=self.state.timestamp, wall_clock_time_ns=wall_clock_time, ) if self.record_instant_on_finish: self._record_instant_event( wall_clock_time_ns=wall_clock_time, timestamp=self.state.timestamp, ) self._started = False
[docs] def instant(self) -> None: """Record an instant event. To record an instant event: .. testsetup:: from composer.profiler import Profiler, cyclic_schedule profiler = Profiler(schedule=cyclic_schedule(), trace_handlers=[], torch_prof_memory_filename=None) profiler.bind_to_state(state) .. doctest:: >>> def something_to_measure(): ... print("something_to_measure") >>> marker = profiler.marker("instant") >>> marker.instant() >>> something_to_measure() something_to_measure """ if self.should_record(self.state): self._record_instant_event( wall_clock_time_ns=time.time_ns(), timestamp=self.state.timestamp, )
[docs] def counter(self, values: Dict[str, Union[float, int]]) -> None: """Record a counter event. To record a counter event: .. testsetup:: from composer.profiler import Profiler, cyclic_schedule profiler = Profiler(schedule=cyclic_schedule(), trace_handlers=[], torch_prof_memory_filename=None) profiler.bind_to_state(state) .. doctest:: >>> marker = profiler.marker("foo") >>> counter_event = 5 >>> marker.counter({"counter_event": counter_event}) >>> counter_event = 10 >>> marker.counter({"counter_event": counter_event}) """ if self.should_record(self.state): self._record_counter_event( wall_clock_time_ns=time.time_ns(), values=values, timestamp=self.state.timestamp, )
def __enter__(self) -> Marker: self.start() return self def __exit__(self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], traceback: Optional[TracebackType]) -> None: del exc_type, exc, traceback # unused self.finish() def __call__(self, func: Optional[Callable[..., Any]] = None) -> Callable[..., Any]: if func is None: # for decorators of the style @Marker(), # return self so it's equivalent to @Marker return self @functools.wraps(func) def wrapped(*args: Any, **kwargs: Any): with self: func(*args, **kwargs) return wrapped