graphlearn_torch/python/distributed/__init__.py

# Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from .dist_client import (
  init_client, shutdown_client, async_request_server, request_server
)
from .dist_context import DistRole, DistContext, get_context, init_worker_group
from .dist_dataset import DistDataset
from .dist_feature import PartialFeature, DistFeature
from .dist_graph import DistGraph
from .dist_link_neighbor_loader import DistLinkNeighborLoader
from .dist_loader import DistLoader
from .dist_neighbor_loader import DistNeighborLoader
from .dist_neighbor_sampler import DistNeighborSampler
from .dist_options import (
  CollocatedDistSamplingWorkerOptions,
  MpDistSamplingWorkerOptions,
  RemoteDistSamplingWorkerOptions
)
from .dist_random_partitioner import DistRandomPartitioner
from .dist_sampling_producer import (
  DistMpSamplingProducer, DistCollocatedSamplingProducer
)
from .dist_server import (
  DistServer, get_server, init_server, wait_and_shutdown_server
)
from .dist_subgraph_loader import DistSubGraphLoader
from .dist_table_dataset import DistTableDataset, DistTableRandomPartitioner
from .event_loop import ConcurrentEventLoop
from .rpc import (
  init_rpc, shutdown_rpc, rpc_is_initialized,
  get_rpc_master_addr, get_rpc_master_port,
  all_gather, barrier, global_all_gather, global_barrier,
  RpcDataPartitionRouter, rpc_sync_data_partitions,
  RpcCalleeBase, rpc_register, rpc_request_async, rpc_request,
  rpc_global_request_async, rpc_global_request
)