from __future__ import annotations

from typing import Any, Optional

from paddle.fluid.incubate.fleet.base.fleet_base import DistributedOptimizer, Fleet
from paddle.fluid.incubate.fleet.parameter_server.distribute_transpiler.distributed_strategy import (
    TrainerRuntimeConfig as TrainerRuntimeConfig,
)

class FleetTranspiler(Fleet):
    startup_program: Any = ...
    main_program: Any = ...
    def __init__(self) -> None: ...
    def init(self, role_maker: Optional[Any] = ...) -> None: ...
    def init_worker(self) -> None: ...
    def init_server(self, model_dir: Optional[Any] = ..., **kwargs: Any) -> None: ...
    def run_server(self) -> None: ...
    def stop_worker(self) -> None: ...
    def distributed_optimizer(self, optimizer: Any, strategy: Optional[Any] = ...): ...
    def save_inference_model(
        self,
        executor: Any,
        dirname: Any,
        feeded_var_names: Any,
        target_vars: Any,
        main_program: Optional[Any] = ...,
        export_for_deployment: bool = ...,
    ) -> None: ...
    def save_persistables(
        self, executor: Any, dirname: Any, main_program: Optional[Any] = ..., **kwargs: Any
    ) -> None: ...

fleet: Any

class ParameterServerOptimizer(DistributedOptimizer):
    type: str = ...
    data_norm_name: Any = ...
    def __init__(self, optimizer: Any, strategy: Any, mode: 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 minimize(
        self,
        losses: Any,
        scopes: Optional[Any] = ...,
        startup_programs: Optional[Any] = ...,
        parameter_list: Optional[Any] = ...,
        no_grad_set: Optional[Any] = ...,
    ) -> None: ...
