mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
c855f8632e
Adds suppressions to pyrefly will typecheck clean: https://github.com/pytorch/pytorch/issues/163283 Almost there! Test plan: dmypy restart && python3 scripts/lintrunner.py -a pyrefly check step 1: delete lines in the pyrefly.toml file from the project-excludes field step 2: run pyrefly check step 3: add suppressions, clean up unused suppressions before: https://gist.github.com/maggiemoss/4b3bf2037014e116bc00706a16aef199 after: INFO 0 errors (6,884 ignored) Pull Request resolved: https://github.com/pytorch/pytorch/pull/164913 Approved by: https://github.com/oulgen
324 lines
12 KiB
Python
324 lines
12 KiB
Python
"""Common utilities for testing Dynamo's minifier functionality.
|
|
|
|
This module provides the base infrastructure for running minification tests in Dynamo.
|
|
It includes:
|
|
- MinifierTestResult: A dataclass for storing and processing minifier test results
|
|
- MinifierTestBase: A base test class with utilities for:
|
|
- Running tests in isolated environments
|
|
- Managing temporary directories and configurations
|
|
- Executing minifier launcher scripts
|
|
- Running and validating reproduction scripts
|
|
- Supporting both compile-time and runtime error testing
|
|
|
|
The minifier helps reduce failing Dynamo compilations to minimal reproductions.
|
|
"""
|
|
|
|
import dataclasses
|
|
import io
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import traceback
|
|
from collections.abc import Sequence
|
|
from typing import Any, Optional, Union
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
import torch._dynamo
|
|
import torch._dynamo.test_case
|
|
from torch._dynamo.trace_rules import _as_posix_path
|
|
from torch.utils._traceback import report_compile_source_on_error
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MinifierTestResult:
|
|
minifier_code: str
|
|
repro_code: str
|
|
|
|
def _get_module(self, t: str) -> str:
|
|
match = re.search(r"class Repro\(torch\.nn\.Module\):\s+([ ].*\n| *\n)+", t)
|
|
assert match is not None, "failed to find module"
|
|
r = match.group(0)
|
|
r = re.sub(r"\s+$", "\n", r, flags=re.MULTILINE)
|
|
r = re.sub(r"\n{3,}", "\n\n", r)
|
|
return r.strip()
|
|
|
|
def get_exported_program_path(self) -> Optional[str]:
|
|
# Extract the exported program file path from AOTI minifier's repro.py
|
|
# Regular expression pattern to match the file path
|
|
pattern = r'torch\.export\.load\(\s*["\'](.*?)["\']\s*\)'
|
|
# Search for the pattern in the text
|
|
match = re.search(pattern, self.repro_code)
|
|
# Extract and print the file path if a match is found
|
|
if match:
|
|
file_path = match.group(1)
|
|
return file_path
|
|
return None
|
|
|
|
def minifier_module(self) -> str:
|
|
return self._get_module(self.minifier_code)
|
|
|
|
def repro_module(self) -> str:
|
|
return self._get_module(self.repro_code)
|
|
|
|
|
|
class MinifierTestBase(torch._dynamo.test_case.TestCase):
|
|
DEBUG_DIR = tempfile.mkdtemp()
|
|
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
super().setUpClass()
|
|
if not os.path.exists(cls.DEBUG_DIR):
|
|
cls.DEBUG_DIR = tempfile.mkdtemp()
|
|
cls._exit_stack.enter_context( # type: ignore[attr-defined]
|
|
torch._dynamo.config.patch(debug_dir_root=cls.DEBUG_DIR)
|
|
)
|
|
# These configurations make new process startup slower. Disable them
|
|
# for the minification tests to speed them up.
|
|
cls._exit_stack.enter_context( # type: ignore[attr-defined]
|
|
torch._inductor.config.patch(
|
|
{
|
|
# https://github.com/pytorch/pytorch/issues/100376
|
|
"pattern_matcher": False,
|
|
# multiprocess compilation takes a long time to warmup
|
|
"compile_threads": 1,
|
|
# https://github.com/pytorch/pytorch/issues/100378
|
|
"cpp.vec_isa_ok": False,
|
|
}
|
|
)
|
|
)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
if os.getenv("PYTORCH_KEEP_TMPDIR", "0") != "1":
|
|
shutil.rmtree(cls.DEBUG_DIR)
|
|
else:
|
|
print(f"test_minifier_common tmpdir kept at: {cls.DEBUG_DIR}")
|
|
cls._exit_stack.close() # type: ignore[attr-defined]
|
|
|
|
def _gen_codegen_fn_patch_code(self, device: str, bug_type: str) -> str:
|
|
assert bug_type in ("compile_error", "runtime_error", "accuracy")
|
|
return f"""\
|
|
{torch._dynamo.config.codegen_config()}
|
|
{torch._inductor.config.codegen_config()}
|
|
torch._inductor.config.{"cpp" if device == "cpu" else "triton"}.inject_relu_bug_TESTING_ONLY = {bug_type!r}
|
|
"""
|
|
|
|
def _maybe_subprocess_run(
|
|
self, args: Sequence[Any], *, isolate: bool, cwd: Optional[str] = None
|
|
) -> subprocess.CompletedProcess[bytes]:
|
|
from torch._inductor.cpp_builder import normalize_path_separator
|
|
|
|
if not isolate:
|
|
assert len(args) >= 2, args
|
|
assert args[0] == "python3", args
|
|
if args[1] == "-c":
|
|
assert len(args) == 3, args
|
|
code = args[2]
|
|
args = ["-c"]
|
|
else:
|
|
assert len(args) >= 2, args
|
|
with open(args[1]) as f:
|
|
# Need normalize path of the code.
|
|
code = normalize_path_separator(f.read())
|
|
args = args[1:]
|
|
|
|
# WARNING: This is not a perfect simulation of running
|
|
# the program out of tree. We only interpose on things we KNOW we
|
|
# need to handle for tests. If you need more stuff, you will
|
|
# need to augment this appropriately.
|
|
|
|
# NB: Can't use save_config because that will omit some fields,
|
|
# but we must save and reset ALL fields
|
|
dynamo_config = torch._dynamo.config.get_config_copy()
|
|
inductor_config = torch._inductor.config.get_config_copy()
|
|
try:
|
|
stderr = io.StringIO()
|
|
log_handler = logging.StreamHandler(stderr)
|
|
log = logging.getLogger("torch._dynamo")
|
|
log.addHandler(log_handler)
|
|
try:
|
|
prev_cwd = _as_posix_path(os.getcwd())
|
|
if cwd is not None:
|
|
cwd = _as_posix_path(cwd)
|
|
os.chdir(cwd)
|
|
with patch("sys.argv", args), report_compile_source_on_error():
|
|
exec(code, {"__name__": "__main__", "__compile_source__": code})
|
|
rc = 0
|
|
except Exception:
|
|
rc = 1
|
|
traceback.print_exc(file=stderr)
|
|
finally:
|
|
log.removeHandler(log_handler)
|
|
if cwd is not None:
|
|
os.chdir(prev_cwd) # type: ignore[possibly-undefined]
|
|
# Make sure we don't leave buggy compiled frames lying
|
|
# around
|
|
torch._dynamo.reset()
|
|
finally:
|
|
torch._dynamo.config.load_config(dynamo_config)
|
|
torch._inductor.config.load_config(inductor_config)
|
|
|
|
# TODO: return a more appropriate data structure here
|
|
return subprocess.CompletedProcess(
|
|
args,
|
|
rc,
|
|
b"",
|
|
stderr.getvalue().encode("utf-8"),
|
|
)
|
|
else:
|
|
if cwd is not None:
|
|
cwd = _as_posix_path(cwd)
|
|
return subprocess.run(args, capture_output=True, cwd=cwd, check=False)
|
|
|
|
# Run `code` in a separate python process.
|
|
# Returns the completed process state and the directory containing the
|
|
# minifier launcher script, if `code` outputted it.
|
|
def _run_test_code(
|
|
self, code: str, *, isolate: bool
|
|
) -> tuple[subprocess.CompletedProcess[bytes], Union[str, Any]]:
|
|
proc = self._maybe_subprocess_run(
|
|
["python3", "-c", code], isolate=isolate, cwd=self.DEBUG_DIR
|
|
)
|
|
|
|
print("test stdout:", proc.stdout.decode("utf-8"))
|
|
print("test stderr:", proc.stderr.decode("utf-8"))
|
|
repro_dir_match = re.search(
|
|
r"(\S+)minifier_launcher.py", proc.stderr.decode("utf-8")
|
|
)
|
|
if repro_dir_match is not None:
|
|
return proc, repro_dir_match.group(1)
|
|
return proc, None
|
|
|
|
# Runs the minifier launcher script in `repro_dir`
|
|
def _run_minifier_launcher(
|
|
self,
|
|
repro_dir: str,
|
|
isolate: bool,
|
|
*,
|
|
minifier_args: Sequence[Any] = (),
|
|
repro_after: Optional[str] = None,
|
|
) -> tuple[subprocess.CompletedProcess[bytes], str]:
|
|
self.assertIsNotNone(repro_dir)
|
|
launch_file = _as_posix_path(os.path.join(repro_dir, "minifier_launcher.py"))
|
|
with open(launch_file) as f:
|
|
launch_code = f.read()
|
|
|
|
self.assertTrue(os.path.exists(launch_file))
|
|
|
|
args = ["python3", launch_file, "minify", *minifier_args]
|
|
if not isolate and repro_after != "aot_inductor":
|
|
# AOTI minifier doesn't have --no-isolate flag.
|
|
# Everything in AOTI minifier is in no-isolate mode.
|
|
args.append("--no-isolate")
|
|
launch_proc = self._maybe_subprocess_run(args, isolate=isolate, cwd=repro_dir)
|
|
print("minifier stdout:", launch_proc.stdout.decode("utf-8"))
|
|
stderr = launch_proc.stderr.decode("utf-8")
|
|
print("minifier stderr:", stderr)
|
|
|
|
self.assertNotIn("Input graph did not fail the tester", stderr)
|
|
|
|
return launch_proc, launch_code
|
|
|
|
# Runs the repro script in `repro_dir`
|
|
def _run_repro(
|
|
self, repro_dir: str, *, isolate: bool = True
|
|
) -> tuple[subprocess.CompletedProcess[bytes], str]:
|
|
self.assertIsNotNone(repro_dir)
|
|
repro_file = _as_posix_path(os.path.join(repro_dir, "repro.py"))
|
|
with open(repro_file) as f:
|
|
repro_code = f.read()
|
|
|
|
self.assertTrue(os.path.exists(repro_file))
|
|
|
|
repro_proc = self._maybe_subprocess_run(
|
|
["python3", repro_file], isolate=isolate, cwd=repro_dir
|
|
)
|
|
print("repro stdout:", repro_proc.stdout.decode("utf-8"))
|
|
print("repro stderr:", repro_proc.stderr.decode("utf-8"))
|
|
return repro_proc, repro_code
|
|
|
|
# Template for testing code.
|
|
# `run_code` is the code to run for the test case.
|
|
# `patch_code` is the code to be patched in every generated file; usually
|
|
# just use this to turn on bugs via the config
|
|
def _gen_test_code(self, run_code: str, repro_after: str, repro_level: int) -> str:
|
|
repro_after_line = ""
|
|
if repro_after == "aot_inductor":
|
|
repro_after_line = (
|
|
"torch._inductor.config.aot_inductor.dump_aoti_minifier = True"
|
|
)
|
|
elif repro_after:
|
|
repro_after_line = f"""\
|
|
torch._dynamo.config.repro_after = "{repro_after}"
|
|
"""
|
|
return f"""\
|
|
import torch
|
|
import torch._dynamo
|
|
import torch._inductor
|
|
{_as_posix_path(torch._dynamo.config.codegen_config())}
|
|
{_as_posix_path(torch._inductor.config.codegen_config())}
|
|
{repro_after_line}
|
|
torch._dynamo.config.repro_level = {repro_level}
|
|
torch._inductor.config.aot_inductor.repro_level = {repro_level}
|
|
torch._dynamo.config.debug_dir_root = "{_as_posix_path(self.DEBUG_DIR)}"
|
|
{run_code}
|
|
"""
|
|
|
|
# Runs a full minifier test.
|
|
# Minifier tests generally consist of 3 stages:
|
|
# 1. Run the problematic code
|
|
# 2. Run the generated minifier launcher script
|
|
# 3. Run the generated repro script
|
|
#
|
|
# If possible, you should run the test with isolate=False; use
|
|
# isolate=True only if the bug you're testing would otherwise
|
|
# crash the process
|
|
def _run_full_test(
|
|
self,
|
|
run_code: str,
|
|
repro_after: str,
|
|
expected_error: Optional[str],
|
|
*,
|
|
isolate: bool,
|
|
minifier_args: Sequence[Any] = (),
|
|
) -> Optional[MinifierTestResult]:
|
|
if isolate:
|
|
repro_level = 3
|
|
elif expected_error is None or expected_error == "AccuracyError":
|
|
repro_level = 4
|
|
else:
|
|
repro_level = 2
|
|
test_code = self._gen_test_code(run_code, repro_after, repro_level)
|
|
print("running test", file=sys.stderr)
|
|
test_proc, repro_dir = self._run_test_code(test_code, isolate=isolate)
|
|
if expected_error is None:
|
|
# Just check that there was no error
|
|
self.assertEqual(test_proc.returncode, 0)
|
|
|
|
self.assertIsNone(repro_dir)
|
|
return None
|
|
# NB: Intentionally do not test return code; we only care about
|
|
# actually generating the repro, we don't have to crash
|
|
|
|
self.assertIn(expected_error, test_proc.stderr.decode("utf-8"))
|
|
|
|
self.assertIsNotNone(repro_dir)
|
|
print("running minifier", file=sys.stderr)
|
|
_minifier_proc, minifier_code = self._run_minifier_launcher(
|
|
repro_dir,
|
|
isolate=isolate,
|
|
minifier_args=minifier_args,
|
|
repro_after=repro_after,
|
|
)
|
|
print("running repro", file=sys.stderr)
|
|
repro_proc, repro_code = self._run_repro(repro_dir, isolate=isolate)
|
|
|
|
self.assertIn(expected_error, repro_proc.stderr.decode("utf-8"))
|
|
self.assertNotEqual(repro_proc.returncode, 0)
|
|
return MinifierTestResult(minifier_code=minifier_code, repro_code=repro_code)
|