Expose StreamBase and EventBase

[cyyever]

So that we can use them in type hinting to avoid device-specific types like CUDAStream and XPUStream.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/126257

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 15 Unrelated Failures

As of commit fb7e8d8 with merge base bed1c60 ():

NEW FAILURE - The following job has failed:

• pull / linux-docs / build-docs-python-false (gh)
  Process completed with exit code 1.

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-focal-py3.11-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-focal-py3.11-clang10 / test (default, 2, 3, linux.2xlarge) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-focal-py3.12-clang10 / test (default, 2, 3, linux.2xlarge) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-focal-py3.8-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-focal-py3.8-clang10 / test (default, 2, 3, linux.2xlarge) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-jammy-py3.10-clang15-asan / test (default, 2, 6, linux.4xlarge) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names
• pull / linux-jammy-py3.8-gcc11 / test (default, 2, 3, linux.2xlarge) (gh) (similar failure)
  test_public_bindings.py::TestPublicBindings::test_correct_module_names

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  hf_BigBird

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) ()
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) ()
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) ()
  hf_BigBird

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ezyang]

I did some check in fbcode and the blast radius does not seem too big

[cyyever]

I did some check in fbcode and the blast radius does not seem too big

If we are careful not to break the existing extensions relying on the old class names, it should work in general. On the other hand, such reliance indicates that they should indeed be exposed as public classes.

[albanD]

I don't think you need this. The shared class here are torch.Stream and torch.Event you can use these to do proper typing and subclass checking.
Unless I missed something in that (quite complex) collection of classes.

[cyyever]

I don't think you need this. The shared class here are torch.Stream and torch.Event you can use these to do proper typing and subclass checking. Unless I missed something in that (quite complex) collection of classes.

So you also think they should be moved to torch.types? If both of you think it is a better way to organize code, I will do it.

[albanD]

I'm afraid I don't know what torch.types is ?
Also what you mean by "move" here?

[ezyang]

@cyyever is talking about the module, torch/types.py

[albanD]

Ho I missed that one, thanks!

My expectation is that you can use torch.Stream and torch.Event as type hints. Since they are already available in our public API, we don't need any special entry for them in torch.types.
So I'm not sure to understand why is this PR needed and what is the problem it is helping solve.

[cyyever]

@albanD Their class names contains an underscore prefix before, which causing pylint and mypy complaining because that is the convention of inner classes. The main fix here is just to delete the underscore prefix.

[albanD]

I am not sure what you mean? torch.Stream and torch.Event don't have any underscore in their names?

[albanD]

To be clear, these two APIs were added recently as the way to support device-generic stream/event.
The old _StreamBase/_EventBase classes don't really have a purpose anymore and can be deleted when we get around to do it.

You can check the mro for any stream/event-like object we have and they should all properly contain torch.Stream / torch.Event.