| 101 | |
| 102 | @functools.wraps(test_fn) |
| 103 | def wrapper(*args, **kwargs): |
| 104 | if os.environ.get("SKIP_COMPUTE_SANITIZER") == "1": |
| 105 | test_fn(*args, **kwargs) |
| 106 | return |
| 107 | |
| 108 | import psutil |
| 109 | |
| 110 | if target_kwargs.pop("clear_torch_cache", False): |
| 111 | # If we don't pop clear_torch_cache, it won't pass |
| 112 | # target_kwargs.items() <= kwargs.items() condition below. |
| 113 | torch.cuda.empty_cache() |
| 114 | tools_to_check = target_kwargs.pop("tools_to_check", [ComputeSanitizerTool.MEMCHECK]) |
| 115 | assert isinstance(tools_to_check, list), f"{tools_to_check=}" |
| 116 | assert all(tool in ComputeSanitizerTool for tool in tools_to_check), ( |
| 117 | f"{(tool for tool in tools_to_check if tool not in ComputeSanitizerTool)=}") |
| 118 | |
| 119 | ppid_name = psutil.Process(os.getppid()).exe() |
| 120 | run_compute_sanitizer = target_kwargs.items() <= kwargs.items() |
| 121 | if "run_sanitizer" in kwargs: |
| 122 | run_compute_sanitizer &= kwargs["run_sanitizer"] |
| 123 | if run_compute_sanitizer and "compute-sanitizer" not in ppid_name: |
| 124 | for tool in tools_to_check: |
| 125 | path = os.path.realpath(test_fn.__globals__["__file__"]) |
| 126 | # get path of current file |
| 127 | env = { |
| 128 | "PATH": os.environ["PATH"], |
| 129 | "PYTORCH_NO_CUDA_MEMORY_CACHING": "1", |
| 130 | "TORCH_SHOW_CPP_STACKTRACES": "1", |
| 131 | "CUDA_LAUNCH_BLOCKING": "1", |
| 132 | } |
| 133 | if "CUDA_VISIBLE_DEVICES" in os.environ: |
| 134 | env["CUDA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] |
| 135 | assert "request_fixture" in kwargs, ( |
| 136 | "memcheck'ed test must have a (possibly unused) `request` fixture") |
| 137 | test_id = kwargs["request_fixture"].node.callspec.id |
| 138 | cmd = f"{path}::{test_fn.__name__}[{test_id}]" |
| 139 | cmd = [ |
| 140 | "compute-sanitizer", |
| 141 | "--target-processes=application-only", |
| 142 | "--destroy-on-device-error=context", |
| 143 | f"--tool={tool.value}", |
| 144 | sys.executable, |
| 145 | "-m", |
| 146 | "pytest", |
| 147 | "-vsx", |
| 148 | cmd, |
| 149 | ] |
| 150 | for opt in ["--update_checksum", "--ignore_checksum_error"]: |
| 151 | if opt in sys.argv: |
| 152 | cmd.append(opt) |
| 153 | out = subprocess.run( |
| 154 | cmd, |
| 155 | stdout=subprocess.PIPE, |
| 156 | stderr=subprocess.STDOUT, |
| 157 | env=env, |
| 158 | ) |
| 159 | sanitizer_ok = "ERROR SUMMARY: 0 errors" in str( |
| 160 | out.stdout) or "RACECHECK SUMMARY: 0 hazards displayed" in str(out.stdout) |