wehub-resource-sync
30 lines
1.2 KiB
Python
30 lines
1.2 KiB
Python
# ABOUTME: Builds nsys (NVIDIA Nsight Systems) runtime_env configuration for Ray workers.
|
|
# ABOUTME: Configures nsys to wrap worker processes and use cudaProfilerApi for capture range control.
|
|
|
|
|
|
def runtime_env(outdir):
|
|
"""Build a runtime_env dict that wraps Ray workers with nsys profile.
|
|
|
|
Args:
|
|
outdir: Shared storage directory for .nsys-rep output files.
|
|
|
|
Returns:
|
|
Dict suitable for passing as runtime_env to map_batches() etc.
|
|
"""
|
|
# capture-range-end=stop: when cudaProfilerStop() is called, nsys finalizes
|
|
# the .nsys-rep file synchronously and detaches. Without this, nsys only
|
|
# finalizes at process exit; if Ray Data hard-kills the Infer actor at end
|
|
# of pipeline (which it sometimes does), the file is left empty/missing.
|
|
# cuda_profiler_fence() in profiling/nvtx.py is what triggers cudaProfilerStop
|
|
# after active_batches batches, so active_batches must be set lower than the
|
|
# number of batches each actor actually processes.
|
|
return {
|
|
"nsight": {
|
|
"t": "cuda,cudnn,cublas,nvtx",
|
|
"capture-range": "cudaProfilerApi",
|
|
"capture-range-end": "stop",
|
|
"stop-on-exit": "true",
|
|
"o": f"{outdir}/nsys_%h_%p",
|
|
}
|
|
}
|