nemo_rl.utils.nsys

Module Contents

Classes

ProfilablePolicy

Functions

maybe_gpu_profile_step
wrap_with_nvtx_name	A decorator to wrap a function with an NVTX range with the given name.

Data

NRL_NSYS_WORKER_PATTERNS
NRL_NSYS_PROFILE_STEP_RANGE

API

nemo_rl.utils.nsys.NRL_NSYS_WORKER_PATTERNS

‘get(…)’

nemo_rl.utils.nsys.NRL_NSYS_PROFILE_STEP_RANGE

‘get(…)’

class nemo_rl.utils.nsys.ProfilablePolicy

Bases: typing.Protocol

start_gpu_profiling() → None

stop_gpu_profiling() → None

nemo_rl.utils.nsys.maybe_gpu_profile_step(

policy: nemo_rl.utils.nsys.ProfilablePolicy,

step: int,

)

nemo_rl.utils.nsys.wrap_with_nvtx_name(name: str)

A decorator to wrap a function with an NVTX range with the given name.