Print repro command when test_core_vs_torch_consistency fails with sample index specified

[crcrpar]

What does this PR do?

inspired by PyTorch's PYTORCH_OPINFO_SAMPLE_INPUT_INDEX, this gives test_core_vs_torch_consistency a mechanism to add repro command to error message.

An example is as follows:

            try:
                result = run_snippet(
                    snippet_torch_consistency,
                    op,
                    device,
                    dtype,
                    tfn,
                    op.torch_reference,
                    sample,
                    lambda a, b, **kwargs: sample_comp(a, b, equal_nan=True, **kwargs),
                )
            except Exception as e:
                if repro_cmd not in str(e):
>                   raise type(e)(f"{str(e)}\n{repro_cmd}") from e
E                   RuntimeError: Raise to get a repro command
E                   command to reproduce the error: THUNDER_OPTEST_SAMPLE_INDEX=0 pytest thunder/tests/test_ops.py -k "test_core_vs_torch_consistency and cpu and thunder.dtypes.float32 and copysign"

cc @mruberry

[crcrpar]

An example output is as follows:

============================================================================================================================================================================================================ test session starts =============================================================================================================================================================================================================
platform linux -- Python 3.11.12, pytest-8.3.5, pluggy-1.6.0
benchmark: 5.1.0 (defaults: timer=torch.utils.benchmark.utils.timer.timer disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=True warmup_iterations=100000)
Test order randomisation NOT enabled. Enable with --random-order or --random-order-bucket=<bucket_type>
rootdir: /home/mkozuki/ghq/github.com/crcrpar/lightning-thunder
configfile: pyproject.toml
plugins: xdist-3.6.1, timeout-2.3.1, timestamper-0.0.10, cov-6.1.1, benchmark-5.1.0, random-order-1.1.1, hypothesis-6.131.21, anyio-4.9.0
timeout: 900.0s
timeout method: signal
timeout func_only: False
collected 7403 items / 7402 deselected / 1 selected

thunder/tests/test_ops.py F                                                                                                                                                                                                                                                                                                                                                                                                            [100%]

================================================================================================================================================================================================================== FAILURES ==================================================================================================================================================================================================================
__________________________________________________________________________________________________________________________________________________________________________________ test_core_vs_torch_consistency_copysign_torch_cpu_thunder.dtypes.float32 __________________________________________________________________________________________________________________________________________________________________________________

snippet = <function snippet_torch_consistency at 0x734a27cc25c0>, opinfo = <thunder.tests.opinfos.OpInfo object at 0x734a27d7fd10>, devicetype = 'cpu', dtype = thunder.dtypes.float32
args = (<function copysign at 0x734a265e63e0>, <built-in method copysign of type object at 0x734b237470e0>, [SampleInput args...-2.5624e+00,  4.7711e-03]])) kwargs={}], <function test_core_vs_torch_consistency.<locals>.<lambda> at 0x734a265e6200>), kwargs = {}, exc_info = (<class 'RuntimeError'>, RuntimeError('Raise to get a repro command'), <traceback object at 0x734a264e9100>)

    def run_snippet(snippet, opinfo, devicetype, dtype, *args, **kwargs):
        try:
            snippet(*args, **kwargs)
        except Exception as ex:
            exc_info = sys.exc_info()

            # Raises exceptions that occur with pytest, and returns debug information when
            # called otherwise
            # NOTE: PYTEST_CURRENT_TEST is set by pytest
            if "PYTEST_CURRENT_TEST" in os.environ:
>               raise ex

thunder/tests/framework.py:568:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
thunder/tests/framework.py:560: in run_snippet
    snippet(*args, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

op = <function copysign at 0x734a265e63e0>, torch_op = <built-in method copysign of type object at 0x734b237470e0>, sample = [SampleInput args=(tensor([[ 1.5524,  1.0376, -4.0738, -5.4875],
        [-1.3917, -5.9171,  4.0513,  4.1672],
       ...+00,  6.2009e+00,  4.6911e+00,  8.9126e+00],
        [ 5.9076e+00,  6.7902e+00, -2.5624e+00,  4.7711e-03]])) kwargs={}]
comp = <function test_core_vs_torch_consistency.<locals>.<lambda> at 0x734a265e6200>

    def snippet_torch_consistency(op: OpInfo, torch_op, sample: SampleInput, comp: Callable):
        args, kwargs = sample.args, sample.kwargs

        thunder_result = op(*args, **kwargs)
        torch_result = torch_op(*args, **kwargs)

        # TODO Review how thunder.jit returns Exception information
        if isinstance(thunder_result, Exception):
            raise thunder_result

>       raise RuntimeError("Raise to get a repro command")
E       RuntimeError: Raise to get a repro command

thunder/tests/test_ops.py:48: RuntimeError

The above exception was the direct cause of the following exception:

    def test():
>       result = template(opinfo, device_str, dtype, executor, comp)

thunder/tests/framework.py:361:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

op = <thunder.tests.opinfos.OpInfo object at 0x734a27d7fd10>, device = 'cpu', dtype = thunder.dtypes.float32, executor = <thunder.tests.framework.TorchTestExecutor object at 0x734a282ed550>, comp = <function assert_close at 0x734acc391e40>

    @ops(tuple(op for op in opinfos if op.torch_reference is not None))
    def test_core_vs_torch_consistency(op, device: str, dtype: dtypes.dtype, executor, comp):
        if dtypes.is_complex_dtype(dtype) and not op.instantiate_complex_tests:
            pytest.skip("Skipping complex operator tests in CI for speed")
        if (
            torch.device(device).type == "cuda"
            and dtype is dtypes.bfloat16
            and not thunder.tests.bf16.device_supports_bf16(device)
        ):
            pytest.skip("Your CUDA device does not support bfloat16")

        # Check if a specific sample index is requested via environment variable
        target_sample_idx = int(os.environ.get("THUNDER_OPTEST_SAMPLE_INDEX", -1))

        for i, sample in enumerate(op.sample_inputs(device, dtype)):
            # Skip samples that don't match the requested index
            if target_sample_idx >= 0 and i != target_sample_idx:
                continue

            sample_comp = sample.comp if sample.comp is not None else comp

            tfn = thunder.jit(
                op.op,
                executors=executor.executors_list(),
                cache="no caching",
                disable_torch_autograd=True,
            )

            path = os.path.relpath(__file__)
            repro_cmd = (
                f"command to reproduce the error: THUNDER_OPTEST_SAMPLE_INDEX={i} "
                f'pytest {path} -k "test_core_vs_torch_consistency and {device} and {dtype} and {op.name}"'
            )

            try:
                result = run_snippet(
                    snippet_torch_consistency,
                    op,
                    device,
                    dtype,
                    tfn,
                    op.torch_reference,
                    sample,
                    lambda a, b, **kwargs: sample_comp(a, b, equal_nan=True, **kwargs),
                )
            except Exception as e:
                if repro_cmd not in str(e):
>                   raise type(e)(f"{str(e)}\n{repro_cmd}") from e
E                   RuntimeError: Raise to get a repro command
E                   command to reproduce the error: THUNDER_OPTEST_SAMPLE_INDEX=0 pytest thunder/tests/test_ops.py -k "test_core_vs_torch_consistency and cpu and thunder.dtypes.float32 and copysign"

thunder/tests/test_ops.py:122: RuntimeError
========================================================================================================================================================================================================== short test summary info ===========================================================================================================================================================================================================
FAILED thunder/tests/test_ops.py::test_core_vs_torch_consistency_copysign_torch_cpu_thunder.dtypes.float32 - RuntimeError: Raise to get a repro command

[IvanYashchuk]

Awesome! Makes debugging failing samples so much easier.

[t-vi]

Stamping. Thank you @crcrpar @IvanYashchuk