Repository: Project-MONAI/MONAI

Issue:
NVTX handlers to probe workflow
**Is your feature request related to a problem? Please describe.**
NVTX transforms provides insights into the MONAI transforms; however, in order to probe other part of the MONAI workflow, we need NVTX handlers that can be attached to Ignite events.


Discussion:

Resolving patch:
diff --git a/monai/handlers/__init__.py b/monai/handlers/__init__.py
--- a/monai/handlers/__init__.py
+++ b/monai/handlers/__init__.py
@@ -22,6 +22,7 @@
 from .mean_dice import MeanDice
 from .metric_logger import MetricLogger, MetricLoggerKeys
 from .metrics_saver import MetricsSaver
+from .nvtx_handlers import MarkHandler, RangeHandler, RangePopHandler, RangePushHandler
 from .parameter_scheduler import ParamSchedulerHandler
 from .postprocessing import PostProcessing
 from .regression_metrics import MeanAbsoluteError, MeanSquaredError, PeakSignalToNoiseRatio, RootMeanSquaredError
diff --git a/monai/handlers/nvtx_handlers.py b/monai/handlers/nvtx_handlers.py
new file mode 100644
--- /dev/null
+++ b/monai/handlers/nvtx_handlers.py
@@ -0,0 +1,201 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Wrapper around NVIDIA Tools Extension for profiling MONAI ignite workflow
+"""
+
+from typing import TYPE_CHECKING, Optional, Tuple, Union
+
+from monai.config import IgniteInfo
+from monai.utils import ensure_tuple, min_version, optional_import
+
+_nvtx, _ = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
+if TYPE_CHECKING:
+    from ignite.engine import Engine, Events
+else:
+    Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine")
+    Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events")
+
+
+__all__ = ["RangeHandler", "RangePushHandler", "RangePopHandler", "MarkHandler"]
+
+
+class RangeHandler:
+    """
+    Attach a NVTX range to a pair of Ignite events.
+    It pushes an NVTX range at the first event and pops it at the second event.
+    Stores zero-based depth of the range that is started.
+
+    Args:
+        events: a string, pair of Ignite events, pair of Ignite event literals, or pair of Ignite events and literals.
+            If a single string is provided, it should  describe the base name of a pair of default Ignite events
+            with _STARTED and _COMPLETED postfix (like "EPOCH" for Events.EPOCH_STARTED and Events.EPOCH_COMPLETED).
+            The accepted events are: BATCH, ITERATION, EPOCH, and ENGINE.
+            If pair of literals, each should be the literal equivalent of an Ignite event, fo instance:
+            ("EPOCH_STARTED" and "EPOCH_COMPLETED").
+            One can combine events and literals, like (Events.EPOCH_STARTED and "EPOCH_COMPLETED").
+            For the complete list of Events,
+            check https://pytorch.org/ignite/generated/ignite.engine.events.Events.html.
+
+        msg: ASCII message to associate with range.
+            If not provided, the name of first event will be assigned to the NVTX range.
+    """
+
+    def __init__(
+        self,
+        events: Union[str, Tuple[Union[str, Events], Union[str, Events]]],
+        msg: Optional[str] = None,
+    ) -> None:
+        self.events = self.resolve_events(events)
+        if msg is None:
+            if isinstance(events, str):
+                # assign the prefix of the events
+                msg = events
+            else:
+                # combine events' names
+                msg = "/".join([e.name for e in self.events])
+        self.msg = msg
+        self.depth = None
+
+    def resolve_events(self, events: Union[str, Tuple]) -> Tuple[Events, Events]:
+        """
+        Resolve the input events to create a pair of Ignite events
+        """
+        events = ensure_tuple(events)
+        if len(events) == 1:
+            return self.create_paired_events(events[0])
+        if len(events) == 2:
+            return (
+                self.get_event(events[0]),
+                self.get_event(events[1]),
+            )
+        raise ValueError(f"Exactly two Ignite events should be provided [received {len(events)}].")
+
+    def create_paired_events(self, event: str) -> Tuple[Events, Events]:
+        """
+        Create pair of Ignite events from a event prefix name
+        """
+        event = event.upper()
+        event_prefix = {
+            "": "",
+            "ENGINE": "",
+            "EPOCH": "EPOCH_",
+            "ITERATION": "ITERATION_",
+            "BATCH": "GET_BATCH_",
+        }
+        return (
+            self.get_event(event_prefix[event] + "STARTED"),
+            self.get_event(event_prefix[event] + "COMPLETED"),
+        )
+
+    def get_event(self, event: Union[str, Events]) -> Events:
+        if isinstance(event, str):
+            event = event.upper()
+        return Events[event]
+
+    def attach(self, engine: Engine) -> None:
+        """
+        Attach an NVTX Range to specific Ignite events
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+        engine.add_event_handler(self.events[0], self.range_push)
+        engine.add_event_handler(self.events[1], self.range_pop)
+
+    def range_push(self):
+        self.depth = _nvtx.rangePushA(self.msg)
+
+    def range_pop(self):
+        _nvtx.rangePop()
+
+
+class RangePushHandler:
+    """
+    At a specific event, pushes a range onto a stack of nested range span.
+    Stores zero-based depth of the range that is started.
+
+    Args:
+        msg: ASCII message to associate with range
+    """
+
+    def __init__(self, event: Events, msg: Optional[str] = None) -> None:
+        if isinstance(event, str):
+            event = event.upper()
+        self.event = Events[event]
+        if msg is None:
+            msg = self.event.name
+        self.msg = msg
+        self.depth = None
+
+    def attach(self, engine: Engine) -> None:
+        """
+        Push an NVTX range at a specific Ignite event
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+        engine.add_event_handler(self.event, self.range_push)
+
+    def range_push(self):
+        self.depth = _nvtx.rangePushA(self.msg)
+
+
+class RangePopHandler:
+    """
+    At a specific event, pop a previously pushed range.
+    Stores zero-based depth of the range that is started.
+
+    Args:
+        msg: ASCII message to associate with range
+    """
+
+    def __init__(self, event: Events) -> None:
+        if isinstance(event, str):
+            event = event.upper()
+        self.event = Events[event]
+
+    def attach(self, engine: Engine) -> None:
+        """
+        Pop an NVTX range at a specific Ignite event
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+        engine.add_event_handler(self.event, self.range_pop)
+
+    def range_pop(self):
+        _nvtx.rangePop()
+
+
+class MarkHandler:
+    """
+    Mark an instantaneous event that occurred at some point.
+
+    Args:
+        msg: ASCII message to associate with range
+    """
+
+    def __init__(self, event: Events, msg: Optional[str] = None) -> None:
+        if isinstance(event, str):
+            event = event.upper()
+        self.event = Events[event]
+        if msg is None:
+            msg = self.event.name
+        self.msg = msg
+
+    def attach(self, engine: Engine) -> None:
+        """
+        Add an NVTX mark to a specific Ignite event
+        Args:
+            engine: Ignite Engine, it can be a trainer, validator or evaluator.
+        """
+        engine.add_event_handler(self.event, self.mark)
+
+    def mark(self):
+        _nvtx.markA(self.msg)


Test patch:
diff --git a/tests/min_tests.py b/tests/min_tests.py
--- a/tests/min_tests.py
+++ b/tests/min_tests.py
@@ -60,6 +60,7 @@ def run_testsuit():
         "test_handler_mean_dice",
         "test_handler_metrics_saver",
         "test_handler_metrics_saver_dist",
+        "test_handler_nvtx",
         "test_handler_parameter_scheduler",
         "test_handler_post_processing",
         "test_handler_prob_map_producer",
diff --git a/tests/test_handler_nvtx.py b/tests/test_handler_nvtx.py
new file mode 100644
--- /dev/null
+++ b/tests/test_handler_nvtx.py
@@ -0,0 +1,114 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+from ignite.engine import Events
+from parameterized import parameterized
+
+from monai.engines import SupervisedEvaluator
+from monai.handlers import StatsHandler, from_engine
+from monai.handlers.nvtx_handlers import MarkHandler, RangeHandler, RangePopHandler, RangePushHandler
+from monai.utils import optional_import
+
+_, has_nvtx = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")
+
+TENSOR_0 = torch.tensor(
+    [
+        [
+            [[1.0], [2.0]],
+            [[3.0], [4.0]],
+        ]
+    ]
+)
+
+TENSOR_1 = torch.tensor(
+    [
+        [
+            [[0.0], [-2.0]],
+            [[-3.0], [4.0]],
+        ]
+    ]
+)
+
+TENSOR_1_EXPECTED = torch.tensor(
+    [
+        [[1.0], [0.5]],
+        [[0.25], [5.0]],
+    ]
+)
+
+TEST_CASE_0 = [[{"image": TENSOR_0}], TENSOR_0[0] + 1.0]
+TEST_CASE_1 = [[{"image": TENSOR_1}], TENSOR_1_EXPECTED]
+
+
+class TestHandlerDecollateBatch(unittest.TestCase):
+    @parameterized.expand(
+        [
+            TEST_CASE_0,
+            TEST_CASE_1,
+        ]
+    )
+    @unittest.skipUnless(has_nvtx, "CUDA is required for NVTX!")
+    def test_compute(self, data, expected):
+        # Set up handlers
+        handlers = [
+            # Mark with Ignite Event
+            MarkHandler(Events.STARTED),
+            # Mark with literal
+            MarkHandler("EPOCH_STARTED"),
+            # Mark with literal and providing the message
+            MarkHandler("EPOCH_STARTED", "Start of the epoch"),
+            # Define a range using one prefix (between BATCH_STARTED and BATCH_COMPLETED)
+            RangeHandler("Batch"),
+            # Define a range using a pair of events
+            RangeHandler((Events.STARTED, Events.COMPLETED)),
+            # Define a range using a pair of literals
+            RangeHandler(("GET_BATCH_STARTED", "GET_BATCH_COMPLETED"), msg="Batching!"),
+            # Define a range using a pair of literal and events
+            RangeHandler(("GET_BATCH_STARTED", Events.COMPLETED)),
+            # Define the start of range using literal
+            RangePushHandler("ITERATION_STARTED"),
+            # Define the start of range using event
+            RangePushHandler(Events.ITERATION_STARTED, "Iteration 2"),
+            # Define the start of range using literals and providing message
+            RangePushHandler("EPOCH_STARTED", "Epoch 2"),
+            # Define the end of range using Ignite Event
+            RangePopHandler(Events.ITERATION_COMPLETED),
+            RangePopHandler(Events.EPOCH_COMPLETED),
+            # Define the end of range using literal
+            RangePopHandler("ITERATION_COMPLETED"),
+            # Other handlers
+            StatsHandler(tag_name="train", output_transform=from_engine(["label"], first=True)),
+        ]
+
+        # Set up an engine
+        engine = SupervisedEvaluator(
+            device=torch.device("cpu:0"),
+            val_data_loader=data,
+            epoch_length=1,
+            network=torch.nn.PReLU(),
+            postprocessing=lambda x: dict(pred=x["pred"] + 1.0),
+            decollate=True,
+            val_handlers=handlers,
+        )
+        # Run the engine
+        engine.run()
+
+        # Get the output from the engine
+        output = engine.state.output[0]
+
+        torch.testing.assert_allclose(output["pred"], expected)
+
+
+if __name__ == "__main__":
+    unittest.main()