__init__.pyi 1.32 KB
Newer Older
Tom Birch's avatar
Tom Birch committed
1
2
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

3
from typing import Union, Callable, Optional, Any
4
from torch.futures import Future
Tom Birch's avatar
Tom Birch committed
5

6
7
class RRef:
    def __init__(self, t: Any) -> None: ...
8
    def local_value(self) -> Any: ...
9
10
11
12
    def owner(self) -> WorkerInfo: ...
    def remote(self) -> Any: ...
    def rpc_sync(self) -> Any: ...
    def to_here(self) -> Any: ...
13
class WorkerInfo: ...
Tom Birch's avatar
Tom Birch committed
14

15
16
17
18
class BackendType:
    TENSORPIPE: Any
    PROCESS_GROUP: Any

19
20
def TensorPipeRpcBackendOptions(init_method: str) -> Any: ...
def ProcessGroupRpcBackendOptions(init_method: str) -> Any: ...
21
22
23
24
25
26
27
def remote(
    to: Union[str, WorkerInfo],
    func: Callable,
    args: Optional[tuple] = None,
    kwargs: Optional[dict] = None,
    timeout=-1.0,
) -> RRef: ...
Tom Birch's avatar
Tom Birch committed
28
29
30
31
32
33
def rpc_async(
    to: Union[str, WorkerInfo],
    func: Callable,
    args: Optional[tuple] = None,
    kwargs: Optional[dict] = None,
    timeout=-1.0,
34
) -> Future: ...
Tom Birch's avatar
Tom Birch committed
35
36
37
38
39
40
def rpc_sync(
    to: Union[str, WorkerInfo],
    func: Callable,
    args: Optional[tuple] = None,
    kwargs: Optional[dict] = None,
    timeout=-1.0,
41
) -> Any: ...
42
43
44
45
46
47
48
49
def init_rpc(
    name: str,
    backend: Optional[Any] = None,
    rank: int = -1,
    world_size: Optional[int] = None,
    rpc_backend_options: Optional[Any] = None,
) -> None: ...
def shutdown(graceful: Optional[bool] = True) -> None: ...