U |Ph@sdZddlmZddlmZddlmZddlmZm Z m Z e ddd\Z Z erbdd l mZmZn0e d eje d d d \ZZ e d eje dd d \ZZ ddddgZGdddZGdddZGdddZGdddZdS)zK Wrapper around NVIDIA Tools Extension for profiling MONAI ignite workflow ) annotations) TYPE_CHECKING) IgniteInfo) ensure_tuple min_versionoptional_importztorch._C._nvtxz:NVTX is not installed. Are you sure you have a CUDA build?) descriptor)EngineEventsz ignite.enginer decorator)as_typer RangeHandlerRangePushHandlerRangePopHandler MarkHandlerc@steZdZdZdddddddZd d d d d Zdd dddZdddddZdddddZddZ ddZ dS)r a 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. Nz'str | tuple[str | Events, str | Events] str | NoneNone)eventsmsgreturncCsJ|||_|dkr:t|tr$|}nddd|jD}||_d|_dS)N/cSsg|] }|jqS)name).0errQ/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/handlers/nvtx_handlers.py Asz)RangeHandler.__init__..)resolve_eventsr isinstancestrjoinrdepth)selfrrrrr__init__9s  zRangeHandler.__init__z str | tupleztuple[Events, Events])rrcCsbt|}t|dkr"||dSt|dkrJ||d||dfStdt|ddS)zL Resolve the input events to create a pair of Ignite events rz7Exactly two Ignite events should be provided [received z].N)rlencreate_paired_events get_event ValueError)r"rrrrrEs   zRangeHandler.resolve_eventsreventrcCs<|}dddddd}|||d|||dfS)zG Create pair of Ignite events from a event prefix name ZEPOCH_Z ITERATION_Z GET_BATCH_)r,ZENGINEEPOCHZ ITERATIONBATCHSTARTED COMPLETED)upperr()r"r+Z event_prefixrrrr'Psz!RangeHandler.create_paired_events str | Eventsr cCst|trt|S|SN)rrr r1r"r+rrrr(XszRangeHandler.get_eventr enginercCs,||jd|j||jd|jdS)z Attach an NVTX Range to specific Ignite events Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. rr$N)add_event_handlerr range_push range_popr"r6rrrattach[szRangeHandler.attachcCst|j|_dSr3_nvtx rangePushArr!r"rrrr8dszRangeHandler.range_pushcCs tdSr3r=rangePopr?rrrr9gszRangeHandler.range_pop)N) __name__ __module__ __qualname____doc__r#rr'r(r;r8r9rrrrr $s   c@s<eZdZdZdddddddZd dd d d Zd dZdS)rz 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 Nr2rrr+rrcCs<t|trt|n||_|dkr,|jj}||_d|_dSr3)rrr r1r+rrr!r"r+rrrrr#ts zRangePushHandler.__init__r r5cCs||j|jdS)z Push an NVTX range at a specific Ignite event Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. N)r7r+r8r:rrrr;{szRangePushHandler.attachcCst|j|_dSr3r<r?rrrr8szRangePushHandler.range_push)N)rBrCrDrEr#r;r8rrrrrksc@s8eZdZdZdddddZdddd d Zd d Zd S)rz 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 r2rr*cCs t|trt|n||_dSr3)rrr r1r+r4rrrr#szRangePopHandler.__init__r r5cCs||j|jdS)z Pop an NVTX range at a specific Ignite event Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. N)r7r+r9r:rrrr;szRangePopHandler.attachcCs tdSr3r@r?rrrr9szRangePopHandler.range_popN)rBrCrDrEr#r;r9rrrrrsc@s<eZdZdZdddddddZd dd d d Zd dZdS)rz Mark an instantaneous event that occurred at some point. Args: msg: ASCII message to associate with range Nr2rrrFcCs6t|trt|n||_|dkr,|jj}||_dSr3)rrr r1r+rrrGrrrr#szMarkHandler.__init__r r5cCs||j|jdS)z Add an NVTX mark to a specific Ignite event Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. N)r7r+markr:rrrr;szMarkHandler.attachcCst|jdSr3)r=markArr?rrrrHszMarkHandler.mark)N)rBrCrDrEr#r;rHrrrrrsN)rE __future__rtypingr monai.configr monai.utilsrrrr=_ ignite.enginer r OPT_IMPORT_VERSION__all__r rrrrrrr s4      G