from __future__ import annotations

from typing import Any

from paddle.distributed.fleet.runtime.runtime_base import RuntimeBase
from paddle.distributed.ps.utils.public import *

class Service:
    def __init__(self) -> None: ...

class GpuService(Service):
    def __init__(self) -> None: ...

class Accessor:
    accessor_class: str = ...
    optimizer: Any = ...
    feature_dim: int = ...
    embedding_dim: int = ...
    def __init__(self) -> None: ...

class CommonAccessor(Accessor):
    table_name: str = ...
    entry: str = ...
    attrs: Any = ...
    params: Any = ...
    dims: Any = ...
    trainer_num: int = ...
    sync: bool = ...
    initializers: Any = ...
    opt_input_map: Any = ...
    opt_attr_map: Any = ...
    opt_init_map: Any = ...
    def __init__(self) -> None: ...
    def define_optimize_map(self) -> None: ...
    def parse_entry(self, varname: Any, program_id: Any, context: Any) -> None: ...
    def get_shard(self, total_dim: Any, shard_num: Any, pserver_id: Any): ...
    def get_initializer_attr(self, value_name: Any, o_startup_program: Any): ...
    table_num: Any = ...
    table_dim: Any = ...
    accessor_class: str = ...
    def parse_by_optimizer(self, ctx: Any, context: Any) -> None: ...

class Tensor:
    tensor_dict: Any = ...
    def __init__(self, tesnor_dcit: Any) -> None: ...

class Table:
    table_class: Any = ...
    shard_num: int = ...
    type: Any = ...
    accessor: Any = ...
    common: Any = ...
    tensor: Any = ...
    def __init__(self) -> None: ...

class BarrierTable(Table):
    type: Any = ...
    shard_num: int = ...
    is_heter_ps_mode: Any = ...
    role_maker: Any = ...
    idx: Any = ...
    is_sync: Any = ...
    def __init__(self, context: Any, idx: Any) -> None: ...

class TensorTable(Table):
    idx: Any = ...
    tensor_dict: Any = ...
    role_maker: Any = ...
    def __init__(self, idx: Any, tensor_dict: Any, role_maker: Any) -> None: ...

class SparseTable(Table):
    context: Any = ...
    ctx: Any = ...
    type: Any = ...
    table_class: str = ...
    accessor: Any = ...
    def __init__(self, context: Any, send_ctx: Any) -> None: ...

class GeoSparseTable(SparseTable):
    table_class: str = ...
    def __init__(self, context: Any, send_ctx: Any) -> None: ...

class DenseTable(Table):
    context: Any = ...
    ctx: Any = ...
    accessor: Any = ...
    def __init__(self, context: Any, send_ctx: Any) -> None: ...

class Server:
    def __init__(self) -> None: ...

class DownpourServer(Server):
    def __init__(self) -> None: ...

class Worker:
    def __init__(self) -> None: ...

class DownpourWorker(Worker):
    def __init__(self) -> None: ...

class fsClient:
    fs_client_param: Any = ...
    def __init__(self, fs_client_param: Any) -> None: ...

class PsDescBuilder:
    context: Any = ...
    is_sync: Any = ...
    ps_mode: Any = ...
    is_heter_ps_mode: Any = ...
    use_ps_gpu: Any = ...
    barrier_table_id: Any = ...
    send_ctx: Any = ...
    tensor_table_dict: Any = ...
    tables: Any = ...
    service: Any = ...
    fs_client: Any = ...
    ps_desc: Any = ...
    def __init__(self, context: Any) -> None: ...
    def build_worker_desc(self): ...
    sparse_table_maps: Any = ...
    def build_server_desc(self): ...

class TheOnePSRuntime(RuntimeBase):
    def __init__(self) -> None: ...
    def load_model(self, path: Any, mode: Any) -> None: ...
