from __future__ import annotations

from typing import Any, Optional

class TaskNode:
    previous: Any = ...
    node: Any = ...
    program: Any = ...
    def __init__(
        self,
        cur_rank: Any,
        max_run_times: Any,
        max_slot_times: Any,
        role: Optional[Any] = ...,
        node_type: str = ...,
        task_id: Optional[Any] = ...,
        ops: Optional[Any] = ...,
        program: Optional[Any] = ...,
    ) -> None: ...
    def set_type(self, interceptor_type: Any) -> None: ...
    def task_node(self): ...
    def set_program(self, program: Any) -> None: ...
    def get_program(self): ...
    def set_run_pre_steps(self, steps: Any) -> None: ...
    def set_run_at_offset(self, offset: Any) -> None: ...
    def add_upstream_task(self, upstream: Any, buffer_size: Any) -> None: ...
    def add_downstream_task(self, downstream: Any, buffer_size: Any) -> None: ...
    def role(self): ...
    def task_id(self): ...

class CoordSys:
    dp_degree: Any = ...
    pp_degree: Any = ...
    sharding_degree: Any = ...
    mp_degree: Any = ...
    def __init__(self, dist_opt: Any) -> None: ...
    def coord_to_rank(self, coord: Any): ...
    def rank_to_coord(self, rank: Any): ...

def is_optimizer_op(op_role: Any): ...
def is_lr_sched_op(op_role: Any): ...
def is_forward_op(op_role: Any): ...
def is_backward_op(op_role: Any): ...
def run1f1b(program: Any, cur_rank: Any, max_run_times: Any, dist_opt: Any, nrank: Any): ...
def origin(program: Any, cur_rank: Any): ...
