from __future__ import annotations

from typing import Any, Optional

from paddle.fluid.framework import Program as Program
from paddle.fluid.incubate.fleet.base.fleet_base import DistributedOptimizer, Fleet

from .optimizer_factory import *

class PSLib(Fleet):
    def __init__(self) -> None: ...
    def init(self, role_maker: Optional[Any] = ...) -> None: ...
    def set_pull_local_thread_num(self, thread_num: Any) -> None: ...
    all_ips_: Any = ...
    def init_worker(self) -> None: ...
    def init_server(self, model_dir: Optional[Any] = ..., **kwargs: Any) -> None: ...
    def run_server(self) -> None: ...
    def end_pass(self, scope: Any) -> None: ...
    def train_from_dataset(
        self,
        executor: Any,
        program: Optional[Any] = ...,
        dataset: Optional[Any] = ...,
        scope: Optional[Any] = ...,
        thread: int = ...,
        debug: bool = ...,
        fetch_list: Optional[Any] = ...,
        fetch_info: Optional[Any] = ...,
        print_period: int = ...,
        fetch_handler: Optional[Any] = ...,
    ) -> None: ...
    def start_heter_trainer(
        self,
        executor: Any,
        program: Optional[Any] = ...,
        scope: Optional[Any] = ...,
        debug: bool = ...,
        fetch_list: Optional[Any] = ...,
        fetch_info: Optional[Any] = ...,
        print_period: int = ...,
        fetch_handler: Optional[Any] = ...,
    ) -> None: ...
    def stop_worker(self) -> None: ...
    def distributed_optimizer(self, optimizer: Any, strategy: Any = ...): ...
    def save_inference_model(
        self,
        executor: Any,
        dirname: Any,
        feeded_var_names: Optional[Any] = ...,
        target_vars: Optional[Any] = ...,
        main_program: Optional[Any] = ...,
        export_for_deployment: bool = ...,
    ) -> None: ...
    def print_table_stat(self, table_id: Any) -> None: ...
    def set_file_num_one_shard(self, table_id: Any, file_num: Any) -> None: ...
    def save_persistables(
        self, executor: Any, dirname: Any, main_program: Optional[Any] = ..., **kwargs: Any
    ) -> None: ...
    def save_model_with_whitelist(
        self, executor: Any, dirname: Any, whitelist_path: Any, main_program: Optional[Any] = ..., **kwargs: Any
    ) -> None: ...
    def save_multi_table_one_path(self, table_ids: Any, model_dir: Any, **kwargs: Any) -> None: ...
    def save_cache_model(self, executor: Any, dirname: Any, main_program: Optional[Any] = ..., **kwargs: Any): ...
    def shrink_sparse_table(self) -> None: ...
    def shrink_dense_table(
        self, decay: Any, emb_dim: int = ..., scope: Optional[Any] = ..., table_id: Optional[Any] = ...
    ) -> None: ...
    def clear_one_table(self, table_id: Any) -> None: ...
    def clear_model(self) -> None: ...
    def load_pslib_whitelist(self, table_id: Any, model_path: Any, **kwargs: Any) -> None: ...
    def load_one_table(self, table_id: Any, model_path: Any, **kwargs: Any) -> None: ...
    def confirm(self) -> None: ...
    def revert(self) -> None: ...
    def load_model(self, model_dir: Optional[Any] = ..., **kwargs: Any) -> None: ...
    def save_model(self, model_dir: Optional[Any] = ..., **kwargs: Any) -> None: ...
    def save_one_table(self, table_id: Any, model_dir: Any, **kwargs: Any) -> None: ...
    def set_date(self, table_id: Any, date: Any) -> None: ...

fleet: Any

class fleet_embedding:
    origin_emb: Any = ...
    origin_emb_v2: Any = ...
    click_name: Any = ...
    scale_sparse_grad: Any = ...
    accessor: str = ...
    def __init__(self, click_name: Any, scale_sparse_grad: bool = ...) -> None: ...
    def __enter__(self) -> None: ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: ...

class DownpourOptimizer(DistributedOptimizer):
    def __init__(self, optimizer: Any, strategy: Optional[Any] = ...) -> None: ...
    def backward(
        self,
        loss: Any,
        startup_program: Optional[Any] = ...,
        parameter_list: Optional[Any] = ...,
        no_grad_set: Optional[Any] = ...,
        callbacks: Optional[Any] = ...,
    ) -> None: ...
    def apply_gradients(self, params_grads: Any) -> None: ...
    def get_dist_env(self): ...
    def minimize(
        self,
        losses: Any,
        scopes: Optional[Any] = ...,
        startup_programs: Optional[Any] = ...,
        parameter_list: Optional[Any] = ...,
        no_grad_set: Optional[Any] = ...,
        program_mode: str = ...,
    ): ...
