()
| 89 | |
| 90 | |
| 91 | def main(): |
| 92 | args = sys.argv[1:] |
| 93 | sep = args.index("--") |
| 94 | wrapper_args = args[:sep] |
| 95 | child_args = args[sep + 1 :] |
| 96 | |
| 97 | report_dir = wrapper_args[wrapper_args.index("--report-dir") + 1] |
| 98 | name = wrapper_args[wrapper_args.index("--name") + 1] |
| 99 | rank = os.environ["LOCAL_RANK"] |
| 100 | |
| 101 | _numa_bind(int(rank)) |
| 102 | _print_numa_info(rank) |
| 103 | report_path = f"{report_dir}/{name}-rank{rank}" |
| 104 | |
| 105 | cmd = [ |
| 106 | "nsys", |
| 107 | "profile", |
| 108 | "-o", |
| 109 | report_path, |
| 110 | "--force-overwrite=true", |
| 111 | "--capture-range=cudaProfilerApi", |
| 112 | "--cuda-memory-usage=true", |
| 113 | sys.executable, |
| 114 | *child_args, |
| 115 | ] |
| 116 | print(f"[rank {rank}] Running: {' '.join(cmd)}", flush=True) |
| 117 | result = subprocess.run(cmd) |
| 118 | # nsys exits with 143 (SIGTERM) after cudaProfilerStop, which is expected. |
| 119 | # Report 0 to torchrun so it doesn't kill the other rank mid-write. |
| 120 | rc = 0 if result.returncode in (0, 143, -15) else result.returncode |
| 121 | sys.exit(rc) |
| 122 | |
| 123 | |
| 124 | if __name__ == "__main__": |
no test coverage detected