from __future__ import annotations

from typing import Any, Optional

from ..common import OP_ROLE_KEY as OP_ROLE_KEY
from ..common import OpRole as OpRole
from ..common import is_optimizer_op as is_optimizer_op
from ..common import is_update_op as is_update_op
from .shard import Shard as Shard

class PlaceType:
    CPU: int = ...
    CUDA: int = ...
    CUDA_PINNED: int = ...
    XPU: int = ...
    NPU: int = ...
    NPU_PINNED: int = ...
    @staticmethod
    def default_device(): ...
    @staticmethod
    def default_pinned(): ...

class OffloadHelper:
    cpu_place_type: int = ...
    cuda_place_type: Any = ...
    cuda_pinned_place_type: Any = ...
    mp_ring_id: Any = ...
    dp_ring_id: Any = ...
    def __init__(self, mp_ring_id: Optional[Any] = ..., dp_ring_id: Optional[Any] = ...) -> None: ...
    def offload_fp32param(self, block: Any, startup_block: Any, offload: bool = ...) -> None: ...
    def cast_fp32param_in_optimize(self, block: Any, startup_block: Any) -> None: ...
    def offload(self, block: Any, startup_block: Any) -> None: ...
    def opt_sharding_cast_fp32param(self, block: Any, startup_block: Any, params: Any, offload: bool = ...) -> None: ...
