Skip to content

Commit

Permalink
Add IProfiler protocol
Browse files Browse the repository at this point in the history
Summary: Add a protocol to define a Profiler.

Differential Revision: D50765964
  • Loading branch information
Danielle Pintz authored and facebook-github-bot committed Oct 29, 2023
1 parent d3fdc28 commit dfa3833
Showing 2 changed files with 42 additions and 0 deletions.
2 changes: 2 additions & 0 deletions torchtnt/utils/__init__.py
Original file line number Diff line number Diff line change
@@ -53,6 +53,7 @@
prepare_ddp,
prepare_fsdp,
)
from .profiler import IProfiler
from .progress import Progress
from .rank_zero_log import (
rank_zero_critical,
@@ -121,6 +122,7 @@
"NOOPStrategy",
"prepare_ddp",
"prepare_fsdp",
"IProfiler",
"Progress",
"rank_zero_critical",
"rank_zero_debug",
40 changes: 40 additions & 0 deletions torchtnt/utils/profiler.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import logging
from types import TracebackType
from typing import Optional, Protocol, Type

logger: logging.Logger = logging.getLogger(__name__)


class IProfiler(Protocol):
"""Protocol for profilers. Can be used as a context manager."""

def __enter__(self) -> None:
"""Enters the context manager."""
pass

Check warning on line 19 in torchtnt/utils/profiler.py

Codecov / codecov/patch

torchtnt/utils/profiler.py#L19

Added line #L19 was not covered by tests

def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_value: Optional[BaseException],
tb: Optional[TracebackType],
) -> Optional[bool]:
"""Exits the context manager."""
pass

Check warning on line 28 in torchtnt/utils/profiler.py

Codecov / codecov/patch

torchtnt/utils/profiler.py#L28

Added line #L28 was not covered by tests

def start(self) -> None:
"""Starts the profiler."""
pass

Check warning on line 32 in torchtnt/utils/profiler.py

Codecov / codecov/patch

torchtnt/utils/profiler.py#L32

Added line #L32 was not covered by tests

def stop(self) -> None:
"""Stops the profiler."""
pass

Check warning on line 36 in torchtnt/utils/profiler.py

Codecov / codecov/patch

torchtnt/utils/profiler.py#L36

Added line #L36 was not covered by tests

def step(self) -> None:
"""Signals to the profiler that the next profiling step has started."""
pass

Check warning on line 40 in torchtnt/utils/profiler.py

Codecov / codecov/patch

torchtnt/utils/profiler.py#L40

Added line #L40 was not covered by tests

0 comments on commit dfa3833

Please sign in to comment.