o  i@sdZddlmZddlmZddlmZmZmZm Z e ddd\Z Z er-ddl m Z mZne d ejed d d \Z Z e d ejed d d \ZZ gdZGdddZGdddZGdddZGdddZdS)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@sTeZdZdZdd d d Zd!d dZd"ddZd#ddZd$ddZddZ 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. Nevents'str | tuple[str | Events, str | Events]msg str | NonereturnNonecCsJ|||_|durt|tr|}n ddd|jD}||_d|_dS)N/cSsg|]}|jqS)name).0err^/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/handlers/nvtx_handlers.py @sz)RangeHandler.__init__..)resolve_eventsr isinstancestrjoinrdepth)selfrrrrr__init__8s   zRangeHandler.__init__ str | tupletuple[Events, Events]cCs^t|}t|dkr||dSt|dkr%||d||dfStdt|d)zL Resolve the input events to create a pair of Ignite events rz7Exactly two Ignite events should be provided [received z].)rlencreate_paired_events get_event ValueError)r#rrrrrDs   zRangeHandler.resolve_eventseventr cCs<|}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*Os$z!RangeHandler.create_paired_events str | Eventsr cCst|tr t|S|SN)rr r r3r#r-rrrr+WszRangeHandler.get_eventenginer cCs,||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#r7rrrattachZszRangeHandler.attachcCt|j|_dSr5_nvtx rangePushArr"r#rrrr9czRangeHandler.range_pushcC tdSr5r?rangePoprArrrr:f zRangeHandler.range_popr5)rrrrrr)rr%rr&)r-r rr&)r-r4rr r7r rr) __name__ __module__ __qualname____doc__r$rr*r+r<r9r:rrrrr #s    r c@.eZdZdZddd d Zdd dZddZdS)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 Nr-r4rrrrcCs<t|tr t|n||_|dur|jj}||_d|_dSr5)rr r r3r-rrr"r#r-rrrrr$ss  zRangePushHandler.__init__r7r cC||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)r8r-r9r;rrrr<zzRangePushHandler.attachcCr=r5r>rArrrr9rBzRangePushHandler.range_pushr5r-r4rrrrrG)rHrIrJrKr$r<r9rrrrrjs    rc@s,eZdZdZdddZdd d Zd d ZdS)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 r-r4rrcCs&t|trt||_dS||_dSr5)rr r r3r-r6rrrr$s&zRangePopHandler.__init__r7r cCrN)z Pop an NVTX range at a specific Ignite event Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. N)r8r-r:r;rrrr<rOzRangePopHandler.attachcCrCr5rDrArrrr:rFzRangePopHandler.range_popN)r-r4rrrG)rHrIrJrKr$r<r:rrrrrs    rc@rL)rz Mark an instantaneous event that occurred at some point. Args: msg: ASCII message to associate with range Nr-r4rrrrcCs6t|tr t|n||_|dur|jj}||_dSr5)rr r r3r-rrrMrrrr$s zMarkHandler.__init__r7r cCrN)z Add an NVTX mark to a specific Ignite event Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. N)r8r-markr;rrrr<rOzMarkHandler.attachcCst|jdSr5)r?markArrArrrrQszMarkHandler.markr5rPrG)rHrIrJrKr$r<rQrrrrrs    rN)rK __future__rtypingr monai.utilsrrrrr?_ ignite.enginer r OPT_IMPORT_VERSION__all__r rrrrrrrs$      G