from __future__ import annotations

from typing import Any, Optional

logger: Any

class DistributeMode:
    COLLECTIVE: int = ...
    PS: int = ...
    PS_HETER: int = ...

class DeviceMode:
    UNKNOWN: int = ...
    CPU: int = ...
    GPU: int = ...
    KUNLUN: int = ...
    XPU: int = ...
    ASCEND_NPU: int = ...
    MLU: int = ...

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 world_device_ids(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:
    accelerators: Any = ...
    endpoint: Any = ...
    stage: Any = ...
    def __init__(self) -> None: ...
    def __eq__(self, t: Any) -> Any: ...
    def __ne__(self, t: Any) -> Any: ...
    def rank(self): ...

class Pod:
    id: Any = ...
    addr: Any = ...
    port: Any = ...
    trainers: Any = ...
    servers: Any = ...
    workers: Any = ...
    heter_workers: Any = ...
    accelerators: Any = ...
    device_mode: 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 rank(self): ...
    def get_visible_accelerators(self): ...

def get_logger(log_level: int = ..., name: str = ...): ...
def get_cluster(node_ips: Any, node_ip: Any, trainer_endpoints: Any, device_mode: Any, devices_per_proc: 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): ...
def get_ports(num: Any, offset: Any): ...
def pretty_print_envs(envs: Any, header: Optional[Any] = ...): ...

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

def run_with_coverage(*args: Any): ...
def start_local_trainers(
    cluster: Any,
    pod: Any,
    training_script: Any,
    training_script_args: Any,
    log_dir: Optional[Any] = ...,
    envs: Optional[Any] = ...,
): ...
def pull_worker_log(tp: Any) -> None: ...
def watch_local_trainers(procs: Any, nranks: Any): ...
def get_gpus(gpus: Any): ...
def get_xpus(xpus: Any): ...
def get_npus(npus: Any): ...
def get_mlus(mlus: Any): ...
def get_device_mode(backend: Any): ...
def get_device_proc_info(args: Any): ...
def direct_start(args: Any) -> None: ...
def get_custom_endpoints(origin_endpoints: Any, offset: int = ...): ...
def get_mapped_cluster_without_rank_mapping(
    node_ips: Any, node_ip: Any, trainer_endpoints: Any, device_mode: Any, node_ranks: Any
): ...
def get_mapped_cluster_from_args_without_rank_mapping(args: Any, device_mode: Any): ...
def get_mapped_cluster_with_rank_mapping(
    node_ips: Any, node_ip: Any, trainer_endpoints: Any, device_mode: Any, node_ranks: Any, node_rank_mappings: Any
): ...
def get_mapped_cluster_from_args_with_rank_mapping(args: Any, device_mode: Any): ...

class ParameterServerLauncher:
    args: Any = ...
    distribute_mode: Any = ...
    server_num: int = ...
    worker_num: int = ...
    heter_worker_num: int = ...
    server_endpoints: str = ...
    server_endpoints_ips: Any = ...
    server_endpoints_port: Any = ...
    worker_endpoints: str = ...
    worker_endpoints_ips: Any = ...
    worker_endpoints_port: Any = ...
    heter_worker_endpoints: str = ...
    heter_worker_endpoints_ips: Any = ...
    heter_worker_endpoints_port: Any = ...
    is_local: bool = ...
    current_node_ip: str = ...
    stage_trainer_num: Any = ...
    stage_heter_map: Any = ...
    stage_list: Any = ...
    stage_device_map: Any = ...
    stage_num: int = ...
    def __init__(self, args: Any, distribute_mode: Any) -> None: ...
    stage_heter_trainer_num: Any = ...
    http_port: Any = ...
    node_ips: Any = ...
    node_rank: Any = ...
    def get_role_endpoints(self, args: Any) -> None: ...
    gloo_rendezvous_dir: Any = ...
    procs: Any = ...
    cmds: Any = ...
    log_fns: Any = ...
    def start_ps(self) -> None: ...
    def start_pod_server(self, args: Any, pod: Any) -> None: ...
    def start_pod_worker(self, args: Any, pod: Any) -> None: ...
    def start_pod_heter_worker(self, args: Any, pod: Any) -> None: ...

def check_backend(backend: Any) -> None: ...
def block_windows_and_macos(backend: Any) -> None: ...
def get_backend_by_compile_flag(): ...
