Turn Stream into protocol and improve typing in torch/_C/__init__.pyi.in #145239
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/145239
Note: Links to docs will display an error until the docs builds have been completed. ❌ 9 New Failures, 1 Unrelated FailureAs of commit 27b061e with merge base caf6039 ( NEW FAILURES - The following jobs have failed:
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "topic: not user facing" |
torch/_C/__init__.pyi.in
Outdated
| def __exit__(self, exc_type, exc_val, exc_tb) -> None: ... | ||
| def __reduce__(self) -> Tuple[Any, ...]: ... # THPDevice_reduce | ||
|
|
||
| # Defined in torch/csrc/Stream.cpp | ||
| class Stream: | ||
| class Stream(Protocol): |
There was a problem hiding this comment.
We expect device-specific stream classes (for example, torch.cuda.Stream) to be explicitly derived from this class. What will happen if we remove Protocol here.
There was a problem hiding this comment.
Not all streams use this class as base, for example cpu stream in torch/cpu/__init__.py. The Protocol change should not affect the current behaviour of the class hierarchy in principle. However, turning it into a protocol makes it possible to type all kinds of stream implementation as Stream. Otherwise, we have to write ugly type annotation such as cpuStream | cudaStream | xxxStream for device agnostic code.
There was a problem hiding this comment.
If we change torch.Stream to Protocol, we must add the same methods defined in torch.Stream to torch.cpu.Stream.
What do you think we change torch.cpu.Stream like this:
class Stream(torch.Stream):
def __new__(cls, priority=-1):
return super().__new__(cls, device='cpu', priority=priority)It seems to look clearer.
There was a problem hiding this comment.
I prefer this if CI passes.
|
This will require some careful review and I won't have time before I go on leave |
|
@pytorchbot rebase |
|
@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here |
|
Successfully rebased |
f580b2f to
27b061e
Compare
| @overload | ||
| def __new__(self, stream_id: _int, device_index: _int, device_type: _int, *, priority: _int = 0) -> Stream: ... | ||
| @runtime_checkable | ||
| class StreamBase(Protocol): |
There was a problem hiding this comment.
Why do we need to split the method into an XxxBase class?
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Fixes #ISSUE_NUMBER
cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel