from __future__ import annotations

from typing import Any, Optional

def global_scatter(
    x: Any, local_count: Any, global_count: Any, group: Optional[Any] = ..., use_calc_stream: bool = ...
): ...
def global_gather(
    x: Any, local_count: Any, global_count: Any, group: Optional[Any] = ..., use_calc_stream: bool = ...
): ...

class Hdfs:
    hdfs_ugi: Any = ...
    hdfs_name: Any = ...
    hdfs_path: Any = ...
    def __init__(self) -> None: ...
    def is_valid(self): ...
    def __eq__(self, n: Any) -> Any: ...
    def __ne__(self, n: Any) -> Any: ...

class Cluster:
    job_server: Any = ...
    pods: Any = ...
    hdfs: Any = ...
    job_stage_flag: Any = ...
    def __init__(self, hdfs: Any) -> None: ...
    def __eq__(self, cluster: Any) -> Any: ...
    def __ne__(self, cluster: Any) -> Any: ...
    def update_pods(self, cluster: Any) -> None: ...
    def trainers_nranks(self): ...
    def pods_nranks(self): ...
    def trainers_endpoints(self): ...
    def pods_endpoints(self): ...
    def get_pod_by_id(self, pod_id: Any): ...

class JobServer:
    endpoint: Any = ...
    def __init__(self) -> None: ...
    def __eq__(self, j: Any) -> Any: ...
    def __ne__(self, j: Any) -> Any: ...

class Trainer:
    gpus: Any = ...
    endpoint: Any = ...
    rank: Any = ...
    def __init__(self) -> None: ...
    def __eq__(self, t: Any) -> Any: ...
    def __ne__(self, t: Any) -> Any: ...
    def get_rank(self): ...

class Pod:
    rank: Any = ...
    id: Any = ...
    addr: Any = ...
    port: Any = ...
    trainers: Any = ...
    gpus: Any = ...
    def __init__(self) -> None: ...
    def __eq__(self, pod: Any) -> Any: ...
    def __ne__(self, pod: Any) -> Any: ...
    def parse_response(self, res_pods: Any) -> None: ...
    def get_visible_gpus(self): ...

def get_logger(log_level: Any, name: str = ...): ...
def get_cluster(node_ips: Any, node_ip: Any, trainer_endpoints: Any, selected_gpus: Any): ...
def terminate_local_procs(procs: Any) -> None: ...
def get_host_name_ip(): ...
def add_arguments(argname: Any, type: Any, default: Any, help: Any, argparser: Any, **kwargs: Any) -> None: ...
def find_free_ports(num: Any): ...

class TrainerProc:
    proc: Any = ...
    log_fn: Any = ...
    log_offset: Any = ...
    rank: Any = ...
    local_rank: Any = ...
    cmd: Any = ...
    def __init__(self) -> None: ...

def start_local_trainers(
    cluster: Any, pod: Any, training_script: Any, training_script_args: Any, log_dir: Optional[Any] = ...
): ...
def pull_worker_log(tp: Any) -> None: ...
def watch_local_trainers(procs: Any, nranks: Any): ...
