cli_args.py

import argparse
import json
import os
from dataclasses import MISSING as dataclass_missing
from dataclasses import asdict, dataclass, field, fields
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar

import uvloop
import yaml
from hydra import compose as hydra_compose
from hydra import initialize as hydra_init
from hydra.core.global_hydra import GlobalHydra
from omegaconf import MISSING, DictConfig, OmegaConf

from areal.utils import logging, name_resolve, pkg_version
from areal.utils.constants import (
    PROX_LOGP_METHOD_RECOMPUTE,
    PROX_LOGP_METHODS_ALL,
)
from areal.utils.pkg_version import is_version_less

if TYPE_CHECKING:
    from transformers import PreTrainedTokenizerFast

uvloop.install()

logger = logging.getLogger("CLIArgs")


@dataclass
class NormConfig:
    """Configuration for reward/advantage normalization."""

    mean_level: str | None = field(
        default="batch",
        metadata={
            "help": "Mean level for normalization. None for no mean normalization.",
            "choices": ["batch", "group", None],
        },
    )
    mean_leave1out: bool = field(
        default=False,
        metadata={"help": "Whether to use leave-one-out average."},
    )
    std_level: str | None = field(
        default="batch",
        metadata={
            "help": "Standard deviation level for normalization. None for no std normalization.",
            "choices": ["batch", "group", None],
        },
    )
    std_unbiased: bool = field(
        default=True,
        metadata={
            "help": "Whether to use unbiased standard deviation computation. Defaults to True (changed from False in v0.3.4)."
        },
    )
    eps: float = field(
        default=1e-5,
        metadata={
            "help": "The eps when dividing by standard deviation to avoid numerical issues."
        },
    )
    group_size: int = field(
        default=1, metadata={"help": "Group size for group-level normalization"}
    )


@dataclass
class MicroBatchSpec:
    """Specification for splitting micro-batches during training."""

    n_mbs: int | None = field(
        default=1,
        metadata={
            "help": "Number of micro-batches (or minimum number if max_tokens_per_mb is set). Used when max_tokens_per_mb is None or as minimum count",
        },
    )
    granularity: int = field(
        default=1,
        metadata={
            "help": "Granularity of each micro-batch. Adjacent sequences are grouped by this size when dividing microbatches.",
        },
    )
    max_tokens_per_mb: int | None = field(
        default=None,
        metadata={
            "help": "Maximum tokens per micro-batch for each forward pass. When set, n_mbs becomes the minimum number of micro-batches.",
        },
    )
    n_mbs_divisor: int = field(
        default=1,
        metadata={
            "help": "Divisor for the number of micro-batches. The final number of micro-batches will be adjusted to be divisible by this value.",
        },
    )

    @classmethod
    def new(cls, mb_spec: "MicroBatchSpec", **kwargs):
        """Create new spec with updated fields while maintaining Omegaconf compatibility."""
        fields = dict(
            n_mbs=mb_spec.n_mbs,
            granularity=mb_spec.granularity,
            max_tokens_per_mb=mb_spec.max_tokens_per_mb,
            n_mbs_divisor=mb_spec.n_mbs_divisor,
        )
        fields.update(kwargs)
        return cls(**fields)


@dataclass
class GenerationHyperparameters:
    """Controls text generation behavior for rollout."""

    n_samples: int = field(
        default=1, metadata={"help": "Number of sequences to generate per prompt."}
    )
    max_new_tokens: int = field(
        default=16384, metadata={"help": "Maximum number of tokens to generate."}
    )
    min_new_tokens: int = field(
        default=0, metadata={"help": "Minimum number of tokens to generate."}
    )
    max_tokens: int = field(
        default=32768,
        metadata={
            "help": "Maximum number of tokens including prompt and generated tokens."
        },
    )
    greedy: bool = field(
        default=False,
        metadata={"help": "Whether to use greedy decoding (max probability)."},
    )
    top_p: float = field(
        default=1.0,
        metadata={"help": "Nucleus sampling probability threshold (0.0, 1.0]."},
    )
    top_k: int = field(
        default=int(1e8),
        metadata={"help": "Number of highest probability tokens to consider."},
    )
    temperature: float = field(
        default=1.0,
        metadata={"help": "Sampling temperature. Higher values increase diversity."},
    )
    stop_token_ids: list[int] = field(
        default_factory=list,
        metadata={"help": "Stop generation when encountering these token IDs."},
    )
    ignore_eos: bool = field(
        default=False,
        metadata={"help": "Do not stop generation when EOS is encountered."},
    )
    skip_special_tokens: bool = field(
        default=True,
        metadata={"help": "Skip special tokens when decoding/displaying outputs."},
    )
    stop: list[str] | None = field(
        default=None,
        metadata={
            "help": "One or multiple stop words. Generation will stop if one of these words is sampled."
        },
    )
    frequency_penalty: float = field(
        default=0.0,
        metadata={
            "help": (
                "Penalizes tokens based on their frequency in generation so far. "
                "Must be between -2 and 2 where negative numbers encourage repetition."
            )
        },
    )
    lora_name: str = field(
        default="",
        metadata={"help": "Lora name to be used for this generation."},
    )
    use_beam_search: bool = field(
        default=False,
        metadata={
            "help": "Enable beam search in the vLLM engine. When enabled, sampling parameters like temperature, top-p, and top-k are auto ignored."
        },
    )
    # NOTE: to add new parameters, please correctly handle them in the `to_openai_args_dict` method.

    def new(self, **kwargs):
        args = asdict(self)
        args.update(kwargs)
        return GenerationHyperparameters(**args)

    def new_with_stop_and_pad_token_ids(self, tokenizer: "PreTrainedTokenizerFast"):
        """Create a new generation hyperparameters with stop and pad token ids added."""
        new_stop_token_ids = self.stop_token_ids.copy()
        if tokenizer.pad_token_id not in new_stop_token_ids:
            new_stop_token_ids.append(tokenizer.pad_token_id)
        if tokenizer.eos_token_id not in new_stop_token_ids:
            new_stop_token_ids.append(tokenizer.eos_token_id)
        return self.new(stop_token_ids=new_stop_token_ids)

    def to_openai_completions_args_dict(
        self, exclude_args: list[str] | None = None
    ) -> dict[str, Any]:
        return self.to_openai_args_dict(
            exclude_args=exclude_args, api_format="completions"
        )

    def to_openai_responses_args_dict(
        self, exclude_args: list[str] | None = None
    ) -> dict[str, Any]:
        return self.to_openai_args_dict(
            exclude_args=exclude_args, api_format="responses"
        )

    def to_openai_agents_model_settings_dict(
        self, exclude_args: list[str] | None = None
    ) -> dict[str, Any]:
        return self.to_openai_args_dict(
            exclude_args=exclude_args, api_format="openai-agents"
        )

    _OPENAI_UNSUPPORTED_ARGS: ClassVar[set[str]] = {
        "min_new_tokens",  # Not supported by OpenAI
        "greedy",  # Not directly supported by OpenAI
        "top_k",  # Not supported by OpenAI
        "stop_token_ids",  # Not supported by OpenAI
        "ignore_eos",  # Not supported by OpenAI
        "skip_special_tokens",  # Not supported by OpenAI
        "lora_name",  # Not supported by OpenAI
        "use_beam_search",  # Not supported by OpenAI
        "max_tokens",  # deprecated by "completions", not used in "responses", should be `max_new_tokens` in "openai-agents"
    }

    def to_openai_args_dict(
        self, exclude_args: list[str] | None = None, api_format: str = "completions"
    ) -> dict[str, Any]:
        """Convert the generation hyperparameters to a dictionary of arguments for OpenAI client."""
        final_exclude_args = set(exclude_args) if exclude_args is not None else set()
        final_exclude_args.update(self._OPENAI_UNSUPPORTED_ARGS)
        # TODO: move the excluded args into extra body, so they can be passed through the client request

        mapping = {"n_samples": "n"}
        if api_format == "completions":
            mapping["max_new_tokens"] = "max_completion_tokens"
        elif api_format == "responses":
            mapping["max_new_tokens"] = "max_output_tokens"
        elif api_format == "openai-agents":
            # NOTE: max_tokens in openai-agents means `max_new_tokens` in sglang/vllm. This is not a bug
            mapping["max_new_tokens"] = "max_tokens"
        else:
            raise ValueError(f"Unsupported API format: {api_format}")

        res = {}
        for k, v in asdict(self).items():
            if k in final_exclude_args:
                should_warn = False

                current_value = getattr(self, k)
                f = next(_field for _field in fields(self) if _field.name == k)

                # Check if equal to the default value
                if f.default is not dataclass_missing:
                    if current_value != f.default:
                        should_warn = True
                elif f.default_factory is not dataclass_missing:
                    if current_value != f.default_factory():
                        should_warn = True

                if should_warn:
                    logger.warning(
                        f"Unsupported arg for openai format: `{k}` with value {current_value}"
                    )
                continue
            key = mapping.get(k, k)
            if key in res:
                logger.warning(f"Overriding key: {key} from {k} with value: {v}")
            res[key] = v

        return res


# Train Engine Configs


@dataclass
class OptimizerConfig:
    """Configuration for model optimization during training."""

    type: str = field(
        default="adam",
        metadata={
            "help": "Optimizer type. Adam_bf16 currently only supported FSDP Engine.",
            "choices": ["adam", "sgd", "adam_bf16"],
        },
    )
    lr: float = field(default=1e-3, metadata={"help": "Learning rate"})
    weight_decay: float = field(default=0.01, metadata={"help": "Weight decay"})
    beta1: float = field(
        default=0.9,
        metadata={
            "help": "Adam beta1 parameter. Only effective when optimizer_type is adam/adam_bf16"
        },
    )
    beta2: float = field(
        default=0.999,
        metadata={
            "help": "Adam beta2 parameter. Only effective when optimizer_type is adam/adam_bf16"
        },
    )
    eps: float = field(
        default=1e-8,
        metadata={
            "help": "Adam epsilon parameter. Only effective when optimizer_type is adam/adam_bf16"
        },
    )
    min_lr_ratio: float = field(
        default=0.0,
        metadata={
            "help": "Minimum learning rate ratio after annealing",
        },
    )
    lr_scheduler_type: str = field(
        default="constant",
        metadata={
            "help": "Learning rate scheduler type",
            "choices": ["linear", "cosine", "constant"],
        },
    )
    warmup_steps_proportion: float = field(
        default=0.001,
        metadata={
            "help": "Proportion of training steps for warmup",
        },
    )
    offload: bool = field(
        default=False, metadata={"help": "Enable optimizer state offloading"}
    )
    initial_loss_scale: float = field(
        default=2**32, metadata={"help": "Initial loss scaling factor"}
    )
    min_loss_scale: float = field(
        default=1.0, metadata={"help": "Minimum loss scaling factor"}
    )
    loss_scale_window: float = field(
        default=5, metadata={"help": "Window size for loss scaling adjustment"}
    )
    hysteresis: int = field(
        default=2, metadata={"help": "Hysteresis (scaling factor) for loss scaling"}
    )
    gradient_clipping: float = field(
        default=1.0, metadata={"help": "Gradient clipping threshold"}
    )


@dataclass
class FSDPWrapPolicy:
    """Policy configuration for FSDP model layer wrapping. None defaults to wrapping transformer decoder layers defined by transformers."""

    transformer_layer_cls_to_wrap: list[str] | None = field(
        default=None,
        metadata={"help": "A list of transformer layer names for FSDP to wrap."},
    )


@dataclass
class FSDPEngineConfig:
    """Configuration for Fully Sharded Data Parallel (FSDP) training backend."""

    wrap_policy: FSDPWrapPolicy | None = field(
        default=None,
        metadata={"help": "FSDP wrap policy, specifying model layers to wrap."},
    )
    offload_params: bool = field(
        default=False,
        metadata={"help": "Whether to offload FSDP parameters to CPU."},
    )
    memory_efficient_load: bool = field(
        default=False,
        metadata={
            "help": "Enable memory-efficient model loading. When enabled, model weights "
            "are initialized on CPU and only rank 0 loads pretrained weights, which are "
            "then broadcast to all ranks after FSDP sharding. This reduces peak GPU memory "
            "during initialization for large models. Note: For VLMs, rank 0 broadcast is "
            "not used; each rank loads weights independently on CPU."
        },
    )


@dataclass
class ArchonEngineConfig:
    """Configuration for Archon Engine training backend."""

    # Attention backend
    attn_type: str = field(
        default="varlen",
        metadata={
            "help": "Attention backend type. Use 'tree' for tree training.",
            "choices": ["varlen", "sdpa", "tree"],
        },
    )

    # CPU offloading for FSDP
    offload_params: bool = field(
        default=False,
        metadata={"help": "Whether to offload FSDP parameters to CPU."},
    )

    # Whether to enable torch.compile
    enable_compile: bool = field(
        default=True,
        metadata={"help": "Enable torch.compile for TransformerBlocks."},
    )

    # Activation Checkpointing (enabled when gradient_checkpointing=True)
    ac_mode: str = field(
        default="selective",
        metadata={
            "help": "Activation checkpointing mode. "
            "'memory_budget' requires enable_compile=True.",
            "choices": ["none", "full", "selective", "memory_budget"],
        },
    )
    selective_ac_option: str = field(
        default="op",
        metadata={
            "help": "Selective AC option: 'op' for op-level, "
            "or integer string (e.g., '2') for every Nth layer."
        },
    )
    ac_memory_budget: float = field(
        default=0.5,
        metadata={
            "help": "Memory budget for 'memory_budget' AC mode. "
            "0.0 = minimum memory (max recompute), 1.0 = default behavior (no recompute)."
        },
    )
    ac_preserve_rng_state: bool = field(
        default=False,
        metadata={
            "help": "Preserve RNG state during checkpointing for deterministic output. "
            "Enabling this may slow down training."
        },
    )
    ac_debug: bool = field(
        default=False,
        metadata={
            "help": "(Testing only) Capture AC debug information. Will be slower."
        },
    )

    # Pipeline Parallel Schedule
    pp_schedule: str = field(
        default="Interleaved1F1B",
        metadata={
            "help": "Pipeline parallel schedule type.",
            "choices": ["1F1B", "Interleaved1F1B", "ZBVZeroBubble"],
        },
    )
    # NOTE: The following three PP layer distribution parameters are advanced options
    # that most users do not need to configure. The defaults work well for typical cases.
    # TODO: Consider simplifying or refactoring these parameters in the future.
    # Currently kept for consistency with Megatron's pipeline parallel configuration.
    pp_layers_per_stage: int | None = field(
        default=None,
        metadata={
            "help": "Number of transformer layers per (virtual) pipeline stage. "
            "If set, num_virtual_stages is calculated from num_layers. "
            "If None, stages are inferred from schedule type "
            "(1 stage/rank for 1F1B, 2 stages/rank for Interleaved1F1B/ZBVZeroBubble).",
        },
    )
    pp_first_stage_less_layers: int = field(
        default=1,
        metadata={
            "help": "Number of layers to reduce in the first pipeline stage. "
            "Accounts for embedding layer overhead.",
        },
    )
    pp_last_stage_less_layers: int = field(
        default=1,
        metadata={
            "help": "Number of layers to reduce in the last pipeline stage. "
            "Accounts for output layer overhead.",
        },
    )

    def __post_init__(self):
        if self.pp_layers_per_stage is not None and self.pp_layers_per_stage < 1:
            raise ValueError(
                f"pp_layers_per_stage must be >= 1, got {self.pp_layers_per_stage}"
            )
        if self.pp_first_stage_less_layers < 0:
            raise ValueError(
                f"pp_first_stage_less_layers must be >= 0, "
                f"got {self.pp_first_stage_less_layers}"
            )
        if self.pp_last_stage_less_layers < 0:
            raise ValueError(
                f"pp_last_stage_less_layers must be >= 0, "
                f"got {self.pp_last_stage_less_layers}"
            )


# These configurations are used by Megatron Bridge to build Megatron models.
@dataclass
class DistributedDataParallelConfig:
    """Configuration for Megatron's DistributedDataParallel.
    Refer to Megatron-LM documentation for details.
    """

    grad_reduce_in_fp32: bool = True
    overlap_grad_reduce: bool = False
    overlap_param_gather: bool = False
    align_param_gather: bool = False
    use_distributed_optimizer: bool = True
    check_for_nan_in_grad: bool = False
    bucket_size: int | None = None
    average_in_collective: bool = False
    fp8_param_gather: bool = False


@dataclass
class FP8EngineConfig:
    """Configuration for FP8 (8-bit floating point) training.

    This configuration encapsulates all FP8-related parameters and can be reused
    across different engines (e.g., Megatron, FSDP). When None in the parent config,
    FP8 training is disabled.
    """

    mode: str = field(
        default="e4m3",
        metadata={
            "help": "FP8 precision mode. Options: "
            "'e4m3' (uniform e4m3), "
            "'hybrid' (e4m3 for activations/weights, e5m2 for output activation gradients)."
        },
    )

    recipe: str = field(
        default="delayed",
        metadata={
            "help": "FP8 scaling recipe. Options: 'tensorwise', 'delayed', 'mxfp8' (Blackwell only), 'blockwise'."
        },
    )

    param: bool = field(
        default=False,
        metadata={
            "help": "Keep parameters in FP8 precision to save memory. "
            "Not all parameters will be converted to fp8; for example, biases will remain unchanged."
        },
    )

    margin: int = field(
        default=0,
        metadata={"help": "Margin for FP8 scaling factor computation."},
    )

    amax_history_len: int = field(
        default=1,
        metadata={
            "help": "Length of amax history window for scaling factor computation."
        },
    )

    amax_compute_algo: str = field(
        default="most_recent",
        metadata={
            "help": "Algorithm for choosing amax value. Options: 'max' (largest in history window), 'most_recent'."
        },
    )

    wgrad: bool = field(
        default=True,
        metadata={
            "help": "When False, override FP8 config and compute weight gradients in higher precision."
        },
    )

    dot_product_attention: bool = field(
        default=False,
        metadata={"help": "Use FP8 implementation of Dot Product Attention."},
    )

    multi_head_attention: bool = field(
        default=False,
        metadata={"help": "Use FP8 implementation of Multi Head Attention."},
    )

    tp_only_amax_red: bool = field(
        default=False,
        metadata={"help": "Reduce FP8 AMAX only in TP or TP-CP domain."},
    )

    first_last_layers_bf16: bool = field(
        default=False,
        metadata={
            "help": "Retain first and last N TransformerBlocks in BF16 instead of FP8."
        },
    )

    num_layers_at_start_in_bf16: int = field(
        default=1,
        metadata={
            "help": "Number of layers at start to keep in BF16 when first_last_layers_bf16 is True."
        },
    )

    num_layers_at_end_in_bf16: int = field(
        default=1,
        metadata={
            "help": "Number of layers at end to keep in BF16 when first_last_layers_bf16 is True."
        },
    )

    direct_convert: bool = field(
        default=True,
        metadata={
            "help": "Whether to use direct FP8 conversion during weight updates and save/load. "
            "When True, FP8 parameters are directly converted between TE FP8 and PyTorch FP8 "
            "without intermediate dequantization/quantization."
        },
    )


@dataclass
class MegatronEngineConfig:
    """Configuration for Megatron-LM training framework.
    Refer to Megatron-LM documentation for implementation details.
    """

    # Distributed Training Configuration
    wrap_with_ddp: bool = True
    use_torch_fsdp2: bool = False  # TODO: pending test
    use_custom_fsdp: bool = False  # TODO: pending test
    ddp: DistributedDataParallelConfig = field(
        default_factory=DistributedDataParallelConfig
    )
    virtual_pipeline_parallel_size: int = field(
        default=1,
        metadata={
            "help": (
                "Virtual pipeline parallel size for Megatron interleaved schedule. "
                "Set to >1 to enable VPP. Default is 1 (disabled)."
            )
        },
    )
    # Don't use MegatronOptimizerConfig here because OmegaConf
    # does not recognize the annotation "torch.dtype"
    overlap_param_gather_with_optimizer_step: bool = False

    # Precision Configuration
    use_precision_aware_optimizer: bool = False
    main_grads_dtype: str = "float32"
    main_params_dtype: str = "float32"
    exp_avg_dtype: str = "float32"
    exp_avg_sq_dtype: str = "float32"

    # Checkpointing Configuration
    async_save: bool = False
    use_checkpoint_opt_param_scheduler: bool = True

    # Deterministic Option
    # NOTE: This option forces torch to use deterministic algorithms,
    # which makes sure that two forward passes with the same input
    # will produce the same output. However, it may have a performance impact.
    # It is recommended to set this option to True for RL training on MoE models for stability.
    use_deterministic_algorithms: bool = False

    # Gradient checkpointing options, only effective when gradient_checkpointing=True
    recompute_granularity: str | None = "full"
    recompute_method: str | None = "uniform"
    recompute_num_layers: int | None = 1
    distribute_saved_activations: bool | None = None
    recompute_modules: list[str] | None = None

    # MoE
    moe_router_dtype: str | None = "fp32"
    moe_shared_expert_overlap: bool = field(
        default=False,
        metadata={
            "help": "Enable overlapping between shared expert computations and dispatcher communications. "
            "Without this, the shared experts execute after the routed experts."
        },
    )
    moe_enable_deepep: bool = False
    moe_token_dispatcher_type: str = field(
        default="alltoall",
        metadata={
            "help": "Type of token dispatcher. Options: 'allgather','alltoall' and 'flex'."
        },
    )
    moe_permute_fusion: bool = field(
        default=False,
        metadata={"help": "Fuse token rearrangement ops during token dispatching."},
    )

    # FP8 Training Configuration
    fp8_config: FP8EngineConfig | None = None


class SchedulingStrategyType(str, Enum):
    separation = "separation"
    colocation = "colocation"


@dataclass
class SchedulingStrategy:
    type: str = field(
        default="separation",
        metadata={"choices": ["separation", "colocation"]},
    )
    target: str | None = field(
        default=None, metadata={"help": "The target role to be colocated with"}
    )
    fork: bool = field(
        default=True,
        metadata={
            "help": "When True with colocation, the target worker spawns a new "
            "process on the same node/GPUs instead of sharing its process. "
            "Provides process isolation while sharing GPU resources."
        },
    )


@dataclass
class SchedulingSpec:
    cpu: int = field(
        default=8, metadata={"help": "Number of CPU cores required per GPU"}
    )
    gpu: int = field(
        default=0,
        metadata={
            "help": "Number of GPU units required. Used only when allocating pods."
        },
    )
    mem: int = field(
        default=32, metadata={"help": "Amount of memory (GB) required per GPU"}
    )
    port_count: int = field(default=2, metadata={"help": "Number of ports to expose"})
    image: str = field(
        default="/storage/openpsi/images/areal-latest.sif",
        metadata={
            "help": "Docker/Singularity container image to use. "
            "Currently only used by Slurm. Will be potentially used by Kubernetes in the future."
        },
    )
    task_type: str = field(
        default="worker",
        metadata={
            "help": "Task type (e.g., worker, engine)",
            "choices": ["worker", "engine"],
        },
    )
    env_vars: dict[str, str] = field(
        default_factory=dict,
        metadata={"help": "Environment variables for the container"},
    )
    cmd: str | None = field(
        default=None,
        metadata={
            "help": "Command to execute inside the container. Defaults to AReaL's RPC server."
        },
    )
    # Slurm specific options
    srun_additional_args: str = field(
        default="--unbuffered --mpi=pmi2 -K --chdir $PWD",
        metadata={
            "help": "Additional arguments to pass to the srun command. Only used by slurm."
        },
    )
    additional_bash_cmds: list[str] | None = field(
        default=None,
        metadata={
            "help": "Additional bash commands to setup the container before running "
            "the torchrun command. Only used by slurm."
        },
    )
    container_type: str = field(
        default="apptainer",
        metadata={
            "help": "Type of containers used in slurm",
            "choices": ["apptainer", "none"],
        },
    )
    mount: str = field(
        default="/storage:/storage", metadata={"help": "Mount path for slurm."}
    )
    nodelist: str | None = field(
        default=None, metadata={"help": "sbatch/srun's `--nodelist` option for slurm."}
    )
    exclude: str | None = field(
        default=None, metadata={"help": "sbatch/srun's `--exclude` option for slurm."}
    )


@dataclass
class TrainEngineConfig:
    """Core configuration for model training, including optimization and backend settings."""

    experiment_name: str = MISSING
    trial_name: str = MISSING
    path: str = field(default="", metadata={"help": "Path to HuggingFace checkpoint"})
    attn_impl: str = field(
        default="flash_attention_2",
        metadata={
            "help": "Attention implementation for huggingface transformers model.",
            "choices": ["flash_attention_2"],
        },
    )
    init_from_scratch: bool = field(
        default=False, metadata={"help": "Initialize model weights randomly"}
    )
    is_critic: bool = field(
        default=False,
        metadata={"help": "Whether to use a critic/reward model"},
    )
    temperature: float = field(
        default=1.0, metadata={"help": "Temperature during generation."}
    )
    # Runtime microbatch limit
    mb_spec: MicroBatchSpec = field(default_factory=MicroBatchSpec)
    pad_to_maximum: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to pad each microbatch to the length upper bound specified by mb_spec. "
                "Can reduce memory fragmentation but slows down training."
            )
        },
    )

    # Training Backend Configuration
    disable_dropout: bool = field(
        default=False, metadata={"help": "Disable dropout layers during training"}
    )
    gradient_checkpointing: bool = field(
        default=False, metadata={"help": "Enable gradient checkpointing"}
    )
    dtype: str = field(default="bfloat16", metadata={"help": "Parameter data type."})
    grad_reduce_dtype: str = field(
        default="float32", metadata={"help": "Gradient reduction data type."}
    )
    optimizer: OptimizerConfig | None = field(
        default=None,
        metadata={"help": "Optimizer configuration. None means no training."},
    )

    weight_update_mode: str = field(
        default="xccl",
        metadata={"help": "Weight update backend type.", "choices": ["disk", "xccl"]},
    )
    fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig)
    archon: ArchonEngineConfig = field(default_factory=ArchonEngineConfig)
    megatron: MegatronEngineConfig = field(default_factory=MegatronEngineConfig)

    # Lora
    use_lora: bool = field(
        default=False,
        metadata={
            "help": "Whether to use LoRA. Only support FSDP. Note that should be enabled together with vLLM/SGLang."
        },
    )
    lora_rank: int = field(default=32, metadata={"help": "lora rank"})
    lora_alpha: int = field(default=16, metadata={"help": "lora alpha"})
    target_modules: list[str] = field(
        default_factory=list,
        metadata={"help": "lora target_modules."},
    )
    peft_type: str = field(
        default="lora",
        metadata={"help": "peft method type. Only LoRA is supported for now."},
    )

    # Tree training
    enable_tree_training: bool = field(
        default=False,
        metadata={"help": "Enable tree training with flex attention module."},
    )

    # Scheduling
    scheduling_spec: tuple[SchedulingSpec, ...] = field(
        default_factory=lambda: (
            SchedulingSpec(cmd="python -m areal.scheduler.rpc.rpc_server"),
        ),
        metadata={
            "help": "Train engine schedule specs. Can accept 1 or 2 SchedulingSpec: "
            "if 1 spec provided, it's used for both worker and engine, engine is embedded in the worker; "
            "if 2 specs provided, first one is for worker, second one is for engine. "
            "Currently only used by the TrainController."
        },
    )
    scheduling_strategy: SchedulingStrategy = field(
        default_factory=SchedulingStrategy,
        metadata={
            "help": "The scheduling strategy of this TrainEngine, either separation or colocation. "
            "Currently only used by the TrainController."
        },
    )

    def __post_init__(self):
        """Validate scheduling_spec length and config combinations."""
        if len(self.scheduling_spec) not in (1, 2):
            raise ValueError(
                f"scheduling_spec must contain 1 or 2 SchedulingSpec, "
                f"got {len(self.scheduling_spec)}"
            )
        if self.fsdp.memory_efficient_load and self.init_from_scratch:
            raise ValueError(
                "memory_efficient_load cannot be used with init_from_scratch=True. "
                "memory_efficient_load is for loading pretrained weights on CPU, "
                "but init_from_scratch creates a model without loading any weights."
            )


@dataclass
class PPOActorConfig(TrainEngineConfig):
    """Configuration for PPO actor model, a subclass of a TrainEngine."""

    # Core PPO/GRPO Parameters
    ppo_n_minibatches: int = field(
        default=4, metadata={"help": "Number of minibatches for each PPO update"}
    )
    eps_clip: float = field(
        default=0.2, metadata={"help": "Clipping factor for policy ratio"}
    )
    eps_clip_higher: float | None = field(
        default=None,
        metadata={
            "help": "Clipping factor (higher value) for policy ratio. Default is None. When eps_clip_higher is set (decoupled), eps_clip will be used as the lower value."
        },
    )
    c_clip: float | None = field(
        default=None,
        metadata={
            "help": "Dual clipping factor for policy ratio, must be > 1.0. None disables dual clipping."
        },
    )
    # M2PO
    m2_threshold: float | None = field(
        default=None, metadata={"help": "The second momentum threshold for M2PO."}
    )
    # Reward
    reward_norm: NormConfig | None = field(
        default=None,
        metadata={"help": "Normalization configuration for rewards"},
    )
    reward_scaling: float = field(
        default=1.0, metadata={"help": "Reward scaling factor"}
    )
    reward_bias: float = field(default=0.0, metadata={"help": "Reward bias"})
    reward_clip: float = field(
        default=20.0, metadata={"help": "Maximum absolute value for reward clipping"}
    )
    overlong_reward_penalty: bool = field(
        default=False,
        metadata={"help": "Penalty for overlong sequences. Used within DAPO."},
    )
    overlong_tokens: int | None = field(
        default=None,
        metadata={"help": "Number of tokens in the tail that will receive a penalty"},
    )
    overlong_penalty_factor: float | None = field(
        default=None,
        metadata={"help": "Penalty factor for tokens in the tail"},
    )
    mask_no_eos_with_zero: bool = field(
        default=False,
        metadata={
            "help": "Mask truncated generations (no EOS token) and exclude from training"
        },
    )

    # Advantage Estimation
    discount: float = field(
        default=1.0, metadata={"help": "Discount factor for future rewards"}
    )
    gae_lambda: float = field(
        default=1.0, metadata={"help": "Lambda parameter for GAE"}
    )
    adv_norm: NormConfig | None = field(
        default=None, metadata={"help": "Normalization configuration for advantages."}
    )

    # KL Control
    kl_ctl: float = field(default=0.1, metadata={"help": "KL divergence coefficient"})
    kl_estimator: str = field(
        default="k1",
        metadata={"help": "KL divergence estimator", "choices": ["k1", "k2", "k3"]},
    )

    # SAPO (Soft Adaptive Policy Optimization) - https://arxiv.org/abs/2511.20347
    use_sapo_loss: bool = field(
        default=False,
        metadata={"help": "Use SAPO loss (mutually exclusive with PPO clipping)"},
    )
    sapo_tau_pos: float = field(
        default=1.0,
        metadata={"help": "SAPO temperature for positive advantages"},
    )
    sapo_tau_neg: float = field(
        default=1.05,
        metadata={"help": "SAPO temperature for negative advantages"},
    )

    # Asynchronous RL
    recompute_logprob: bool = field(
        default=False,
        metadata={
            "help": "Recompute log probability and replace the log probability returned by inference."
        },
    )
    use_decoupled_loss: bool = field(
        default=False,
        metadata={
            "help": "Use the decoupled loss. Implicitly enables recompute_logprob."
        },
    )
    behav_imp_weight_cap: float | None = field(
        default=None,
        metadata={
            "help": "Filter out tokens where behav_imp_weight exceeds behav_imp_weight_cap when computing loss. Must be > 1.0. use_decoupled_loss must be true."
        },
    )
    importance_sampling_level: str = field(
        default="token",
        metadata={
            "help": "Level at which to compute importance sampling ratios. 'token': per-token ratios (standard PPO). 'sequence': sequence-level geometric mean of per-token ratios (GSPO).",
            "choices": ["token", "sequence"],
        },
    )
    # Proximal Log-Probability Computation Method
    prox_logp_method: str = field(
        default=PROX_LOGP_METHOD_RECOMPUTE,
        metadata={
            "help": "Method for computing proximal policy log-probabilities in decoupled PPO. "
            "Only effective when use_decoupled_loss=True. Options: "
            "'recompute' (default): Standard decoupled PPO, recompute proximal policy via forward pass. "
            "'loglinear': Use log-linear interpolation to approximate proximal policy (skip forward pass). "
            "'metrics': Like 'recompute', but also compute approximation metrics for evaluation.",
            "choices": PROX_LOGP_METHODS_ALL,
        },
    )

    # Logging Agent Trajectories
    log_agent_stats: bool = field(
        default=False,
        metadata={"help": "Log statistics for agent trajectories"},
    )
    log_agent_stats_keys: list[str] = field(
        default_factory=lambda: [],
        metadata={"help": "Keys for logging agent trajectory statistics"},
    )
    # Others
    max_new_tokens: int = field(
        default=1024,
        metadata={"help": "Maximum number of new tokens to generate"},
    )

    def should_compute_prox_logp(self) -> bool:
        """Determine if forward pass is needed for proximal log-probabilities.

        Returns:
            True if compute_logp() should be called, False to skip.
        """
        from areal.utils.constants import ProxLogpMethod

        method = ProxLogpMethod(self.prox_logp_method)
        return (self.use_decoupled_loss and not method.skips_forward_pass()) or (
            not self.use_decoupled_loss and self.recompute_logprob
        )


@dataclass
class PPOCriticConfig(TrainEngineConfig):
    """Configuration for PPO critic model, a subclass of a TrainEngine."""

    ppo_n_minibatches: int = field(
        default=4, metadata={"help": "Number of minibatches for each PPO update"}
    )
    eps_clip: float = field(
        default=0.5, metadata={"help": "Clipping factor for value loss"}
    )
    mask_no_eos_with_zero: bool = field(
        default=False,
        metadata={
            "help": "Mask truncated generations (no EOS token) and exclude from training"
        },
    )


def get_py_cmd(module: str, args: dict[str, Any]):
    # convert to flags
    cmd = ["python3", "-m", module]
    for k, v in args.items():
        if v is None or v is False or v == "" or (isinstance(v, list) and not v):
            continue
        flag = f"--{k.replace('_', '-')}"
        if v is True:
            cmd.append(flag)
        elif isinstance(v, list):
            cmd.append(flag)
            cmd.extend(map(str, v))
        else:
            cmd.append(flag)
            cmd.append(str(v))
    return cmd


@dataclass
class vLLMConfig:
    """Configuration for vLLM runtime. Refer to:
    https://docs.vllm.ai/en/stable/api/index.html for detailed documentation.
    """

    model: str = ""
    seed: int = 1
    skip_tokenizer_init: bool = False
    enforce_eager: bool = False
    dtype: str = "bfloat16"
    distributed_executor_backend: str = "mp"
    # original
    max_num_seqs: int = 256
    # kv_cache_type: str = "auto"
    block_size: int = 16
    swap_space: int = 4
    cpu_offload_gb: float = 0
    disable_sliding_window: bool = True
    max_model_len: int | None = 32768
    # NOTE: We use no_enable_* prefix (instead of enable_*) because get_py_cmd()
    # ignores parameters with False values. Setting enable_chunked_prefill=False
    # or enable_prefix_caching=False has NO effect - vLLM will use its default
    # values (True). Using no_enable_*=True correctly passes --no-enable-* flags
    # to vLLM, achieving enable_*=False behavior.
    #
    # IMPORTANT: vLLM V1 engine forces enable_chunked_prefill=True by default
    # for non-pooling tasks (generation tasks). And no_enable_chunked_prefill=True
    # has NO effect for generation tasks in vLLM v0.11.0.
    #
    # TODO(vllm-v0.11.0): vLLM v0.11.0 has inference quality issues when
    # temperature=1.0 - multiple sampling runs produce garbled/corrupted outputs.
    # This affects generation quality in RL training workflows.
    no_enable_chunked_prefill: bool = False
    # NOTE: Disables prefix caching (vLLM default is enabled) because it will
    # make RL training corrupted in single controller mode.
    no_enable_prefix_caching: bool = True
    gpu_memory_utilization: float = 0.9
    worker_extension_cls: str = (
        "areal.engine.vllm_ext.vllm_worker_extension.VLLMWorkerExtension"
    )
    enable_sleep_mode: bool = False
    uvicorn_log_level: str = "warning"
    enable_lora: bool = False
    lora_modules: str = ""

    @staticmethod
    def build_args(
        vllm_config: "vLLMConfig",
        tp_size: int,
        pp_size: int,
        host: str | None = None,
        port: int | None = None,
        dist_init_addr: str | None = None,
    ):
        args: dict = conf_as_dict(vllm_config)
        args = dict(
            # Model and tokenizer
            tokenizer=vllm_config.model,
            load_format="auto",
            trust_remote_code=True,
            tensor_parallel_size=tp_size,
            pipeline_parallel_size=pp_size,
            **args,
        )
        if port is not None:
            args["port"] = port
        if host is not None:
            args["host"] = host
        # handle lora modules separately
        lm = args.get("lora_modules")
        if lm:
            if isinstance(lm, str):
                lm = [lm]
            if isinstance(lm, (list, tuple)):
                try:
                    args["lora_modules"] = [
                        json.dumps(json.loads(s), separators=(",", ":")) for s in lm
                    ]
                except json.JSONDecodeError as e:
                    raise ValueError(f"Invalid JSON string in lora_modules: {e}") from e
        return args

    @staticmethod
    def build_cmd_from_args(args: dict[str, Any]):
        return get_py_cmd("areal.engine.vllm_ext.areal_vllm_server", args)

    @staticmethod
    def build_cmd(
        vllm_config: "vLLMConfig",
        tp_size: int,
        pp_size: int,
        host: str | None = None,
        port: int | None = None,
        dist_init_addr: str | None = None,
    ):
        args = vLLMConfig.build_args(
            vllm_config=vllm_config,
            tp_size=tp_size,
            pp_size=pp_size,
            host=host,
            port=port,
            dist_init_addr=dist_init_addr,
        )
        return vLLMConfig.build_cmd_from_args(args)


@dataclass
class SGLangConfig:
    """Configuration for SGLang runtime. Refer to:
    https://github.com/sgl-project/sglang for detailed documentation.
    """

    model_path: str = ""
    random_seed: int = 1
    skip_tokenizer_init: bool = False
    disable_cuda_graph: bool = False
    disable_radix_cache: bool = True
    disable_cuda_graph_padding: bool = False
    enable_nccl_nvls: bool = False
    disable_outlines_disk_cache: bool = False
    disable_custom_all_reduce: bool = False
    disable_overlap_schedule: bool = False
    enable_mixed_chunk: bool = False
    enable_dp_attention: bool = False
    enable_ep_moe: bool = False
    enable_torch_compile: bool = False
    torch_compile_max_bs: int = 32
    cuda_graph_max_bs: int | None = None
    cuda_graph_bs: list[int] | None = None
    torchao_config: str = ""
    enable_nan_detection: bool = False
    enable_p2p_check: bool = False
    triton_attention_reduce_in_fp32: bool = False
    triton_attention_num_kv_splits: int = 8
    num_continuous_decode_steps: int = 1
    enable_memory_saver: bool = False
    allow_auto_truncate: bool = False
    attention_backend: str | None = "fa3"
    enable_multimodal: bool = False
    sampling_backend: str | None = None
    context_length: int | None = 32768
    mem_fraction_static: float | None = 0.9
    max_running_requests: int | None = None
    # NOTE: chunked_prefill_size is by default 8192 on GPUs with 80GB mem in SGLang,
    # but we disable it to avoid precision issues
    chunked_prefill_size: int | None = -1
    max_prefill_tokens: int = 32768
    schedule_policy: str = "lpm"
    schedule_conservativeness: float = 1.0
    cpu_offload_gb: int = 0
    dtype: str = "bfloat16"
    kv_cache_dtype: str = "auto"
    dp_size: int = 1  # only used for dp attention
    ep_size: int = 1
    # lora
    enable_lora: bool | None = None
    max_lora_rank: int | None = None
    lora_target_modules: list[str] | None = None
    lora_paths: list[str] | None = None
    max_loaded_loras: int = 1
    max_loras_per_batch: int = 1
    lora_backend: str = "triton"
    # logging
    log_level: str = "warning"
    log_level_http: str | None = "warning"
    log_requests: bool = False
    log_requests_level: int = 0
    show_time_cost: bool = False
    enable_metrics: bool = True  # Exports Prometheus-like metrics
    # The interval (in decoding iterations) to log throughput
    # and update prometheus metrics
    decode_log_interval: int = 1
    # Extra loader arguments
    # NOTE: These arguments will be parsed into a dict json-string
    # and passed as `model_loader_extra_config` to SGLang.
    enable_multithread_load: bool = False

    # Use staticmethod to make OmegaConf happy.
    @staticmethod
    def build_cmd(
        sglang_config: "SGLangConfig",
        tp_size,
        base_gpu_id,
        host: str | None = None,
        port: int | None = None,
        dist_init_addr: str | None = None,
        n_nodes: int = 1,
        node_rank: int = 0,
    ):
        args = SGLangConfig.build_args(
            sglang_config=sglang_config,
            tp_size=tp_size,
            base_gpu_id=base_gpu_id,
            host=host,
            port=port,
            dist_init_addr=dist_init_addr,
            n_nodes=n_nodes,
            node_rank=node_rank,
        )

        return SGLangConfig.build_cmd_from_args(args)

    @staticmethod
    def build_cmd_from_args(args: dict[str, Any]):
        return get_py_cmd("sglang.launch_server", args)

    @staticmethod
    def build_args(
        sglang_config: "SGLangConfig",
        tp_size: int,
        base_gpu_id: int,
        host: str | None = None,
        port: int | None = None,
        dist_init_addr: str | None = None,
        n_nodes: int = 1,
        node_rank: int = 0,
    ):
        # Map "all-linear" to "all"
        args: dict = conf_as_dict(sglang_config)
        if sglang_config.enable_multithread_load:
            model_loader_extra_config = dict(
                enable_multithread_load=sglang_config.enable_multithread_load,
            )
            args["model_loader_extra_config"] = json.dumps(
                model_loader_extra_config, separators=(",", ":")
            )
        args.pop("enable_multithread_load", None)
        # Map "all-linear" to "all"
        if "lora_target_modules" in args and args["lora_target_modules"]:
            args["lora_target_modules"] = [
                x.replace("-linear", "") for x in args["lora_target_modules"]
            ]

        args = dict(
            # Model and tokenizer
            tokenizer_path=sglang_config.model_path,
            tokenizer_mode="auto",
            load_format="auto",
            trust_remote_code=True,
            is_embedding=False,
            # Other runtime options
            tp_size=tp_size,
            # Because we have set CUDA_VISIBLE_DEVICES to a single GPU in each process
            base_gpu_id=base_gpu_id,
            nnodes=n_nodes,
            node_rank=node_rank,
            # initialization addresses and ports
            dist_init_addr=dist_init_addr,
            **args,
        )
        if host is not None:
            args["host"] = host
        if port is not None:
            args["port"] = port
        if not pkg_version.is_version_greater_or_equal("sglang", "0.4.9.post2"):
            raise RuntimeError("Needs sglang>=0.4.9.post2 to run the code.")
        if is_version_less("sglang", "0.4.10.post2"):
            args.pop("max_loaded_loras", None)
        return args


@dataclass
class OpenAIProxyConfig:
    """Configuration for OpenAI proxy when using agent workflows."""

    mode: str = field(
        default="inline",
        metadata={
            "help": (
                "OpenAI proxy mode: 'inline' (in-process) or 'subproc' (subprocess). "
                "`inline` mode runs the provided agent workflow directly in the same process. "
                "It can use the provided `base_url` and `http_client` to reduce overhead. "
                "`subproc` mode launches a separate process to run the agent with `OPENAI_BASE_URL` environment variable, "
                "which offers more flexible deployment options at the cost of larger overhead."
            ),
            "choices": ["inline", "subproc"],
        },
    )
    tool_call_parser: str = field(
        default="qwen3",
        metadata={"help": "Parser for tool calls in model output."},
    )
    reasoning_parser: str = field(
        default="qwen3",
        metadata={"help": "Parser for reasoning content (<think> tags)."},
    )
    chat_template_type: str = field(
        default="hf",
        metadata={
            "help": "Chat template type: 'hf' (standard) or 'concat' (multi-turn concatenation).",
            "choices": ["hf", "concat"],
        },
    )
    engine_max_tokens: int | None = field(
        default=None,
        metadata={"help": "Maximum total tokens for the engine (prompt + completion)."},
    )
    turn_discount: float = field(
        default=1.0,
        metadata={"help": "Discount factor for multi-turn reward propagation."},
    )
    export_style: str = field(
        default="individual",
        metadata={
            "help": "Export style: 'individual' (all interactions) or 'concat' (leaf nodes only). "
            "The 'individual' style exports each interaction (input-output-reward) step separately, "
            "and treats them as independent samples to train the model. "
            "The 'concat' style exports only the final concatenated trajectory from the root. "
            "It is only suitable for linear conversation histories without token mismatching (whether valid depends on the tokenizer).",
            "choices": ["individual", "concat"],
        },
    )
    subproc_max_workers: int = field(
        default=4,
        metadata={
            "help": "Maximum number of worker processes for subprocess mode execution pool."
        },
    )
    session_timeout_seconds: int = field(
        default=3600,
        metadata={
            "help": "Session timeout in seconds. Sessions inactive longer than this will be garbage collected."
        },
    )


@dataclass
class InferenceEngineConfig:
    """Configuration for inference servers, including offpolicyness control."""

    experiment_name: str | None = None
    trial_name: str | None = None
    fileroot: str | None = field(
        default=None,
        metadata={"help": "Root directory for logs and trajectory dumps."},
    )
    max_concurrent_rollouts: None | int = field(
        default=None,
        metadata={
            "help": "Maximum number of concurrent rollouts to "
            "the inference engine. Defaults to consumer_batch_size."
        },
    )
    queue_size: None | int = field(
        default=None,
        metadata={"help": "Input/Output queue size for async rollout."},
    )
    consumer_batch_size: int = field(
        default=1,
        metadata={"help": "Batch size for consuming rollouts from the queue."},
    )
    max_head_offpolicyness: int = field(
        default=0,
        metadata={
            "help": "Maximum off-policyness for the head. "
            "If the current version is more than this many versions behind, "
            "the request will not be accepted.",
        },
    )
    enable_rollout_tracing: bool = field(
        default=False,
        metadata={
            "help": "Whether to output verbose tracing messages for each generation request."
        },
    )
    check_trajectory_format: bool = field(
        default=False,
        metadata={
            "help": "Whether to check the format of produced trajectories of a customized workflow. Useful when debugging the workflow in isolation. Should be False during RL training."
        },
    )
    schedule_policy: str = field(
        default="round_robin",
        metadata={"help": "Request scheduling policy", "choices": ["round_robin"]},
    )
    tokenizer_path: str = field(
        default="",
        metadata={"help": "Path to tokenizer for trajectory text decoding."},
    )
    dump_to_file: bool = field(
        default=False,
        metadata={"help": "Whether to dump the trajectories to files under fileroot."},
    )
    setup_timeout: float = field(
        default=300.0,
        metadata={
            "help": "Timeout in seconds of connecting to remote servers or launching local servers."
        },
    )
    request_timeout: float = field(
        default=3600, metadata={"help": "Timeout for HTTP requests."}
    )
    request_retries: int = field(
        default=3, metadata={"help": "Number of retries for failed requests."}
    )
    pause_grace_period: float = field(
        default=0.0,
        metadata={
            "help": "The grace period after calling /pause_generation. Wait until all requests have been dropped."
        },
    )
    scheduling_spec: tuple[SchedulingSpec, ...] = field(
        default_factory=lambda: (
            SchedulingSpec(cmd="python -m areal.scheduler.rpc.rpc_server"),
        ),
        metadata={
            "help": "inference engine schedule specs. Can accept 1 or 2 SchedulingSpec: "
            "if 1 spec provided, it's used for both worker and engine, engine is embedded in the worker; "
            "if 2 specs provided, first one is for worker, second one is for engine. "
            "Currently only used by the RolloutController."
        },
    )
    scheduling_strategy: SchedulingStrategy = field(
        default_factory=SchedulingStrategy,
        metadata={
            "help": "The scheduling strategy of this TrainEngine, either separation or colocation. "
            "Currently only used by the RolloutController."
        },
    )
    use_lora: bool = field(
        default=False,
        metadata={"help": "Whether to use LoRA. Should be same as actors LORA option."},
    )
    openai: OpenAIProxyConfig | None = field(
        default=None,
        metadata={
            "help": "OpenAI proxy configuration (used when workflow is an agent workflow)."
        },
    )

    def __post_init__(self):
        """Validate scheduling_spec length."""
        if len(self.scheduling_spec) not in (1, 2):
            raise ValueError(
                f"scheduling_spec must contain 1 or 2 SchedulingSpec, "
                f"got {len(self.scheduling_spec)}"
            )


@dataclass
class _Timer:
    experiment_name: str = MISSING
    trial_name: str = MISSING
    fileroot: str = MISSING
    freq_epochs: int | None = field(
        default=None,
        metadata={
            "help": "Trigger frequency in epochs. None disables epoch-based saving."
        },
    )
    freq_steps: int | None = field(
        default=None,
        metadata={
            "help": "Trigger frequency in steps. None disables step-based saving."
        },
    )
    freq_secs: int | None = field(
        default=None,
        metadata={
            "help": "Trigger frequency in seconds. None disables time-based saving."
        },
    )


@dataclass
class EvaluatorConfig(_Timer):
    """Configuration for model evaluation scheduling and timing."""


@dataclass
class SaverConfig(_Timer):
    """Configuration for model checkpoint saving scheduling and timing."""


@dataclass
class RecoverConfig(_Timer):
    """Configuration for experiment recovery and fault tolerance."""

    mode: str = field(
        default="disabled",
        metadata={
            "help": "Recovery mode for the launcher. "
            "Options: "
            "'on' or 'auto': Automatically recover from previous runs if recover info and checkpoints are available. "
            "'off' or 'disabled': Never recover from previous runs."
        },
    )
    retries: int = field(
        default=3,
        metadata={"help": "Number of recovery retries when recovery is enabled."},
    )

    def __post_init__(self):
        valid_modes = {"on", "off", "auto", "disabled"}
        if self.mode not in valid_modes:
            raise ValueError(
                f"Invalid recover mode '{self.mode}'. "
                f"Valid options: {valid_modes}. "
                f"Note: 'fault' and 'resume' modes have been removed."
            )


@dataclass
class WandBConfig:
    """Configuration for Weights & Biases experiment tracking."""

    mode: str = "disabled"
    wandb_base_url: str = ""
    wandb_api_key: str = ""
    entity: str | None = None
    project: str | None = None
    name: str | None = None
    job_type: str | None = None
    group: str | None = None
    notes: str | None = None
    tags: list[str] | None = None
    config: dict | None = None
    id_suffix: str | None = "train"


@dataclass
class SwanlabConfig:
    """Configuration for SwanLab experiment tracking and monitoring."""

    project: str | None = None
    name: str | None = None
    config: dict | None = None
    logdir: str | None = None
    mode: str | None = "disabled"
    api_key: str | None = os.getenv("SWANLAB_API_KEY", None)


@dataclass
class TensorBoardConfig:
    """Configuration for TensorBoard logging and visualization."""

    path: str | None = None


@dataclass
class StatsLoggerConfig:
    """Configuration for experiment statistics logging and tracking services."""

    experiment_name: str = MISSING
    trial_name: str = MISSING
    fileroot: str = MISSING
    wandb: WandBConfig = field(
        default_factory=WandBConfig,
        metadata={"help": "Weights & Biases configuration."},
    )
    swanlab: SwanlabConfig = field(
        default_factory=SwanlabConfig,
        metadata={"help": "SwanLab configuration."},
    )
    tensorboard: TensorBoardConfig = field(
        default_factory=TensorBoardConfig,
        metadata={"help": "TensorBoard configuration. Only 'path' field required."},
    )


@dataclass
class SessionTracerConfig:
    """Configuration for per-session lifecycle tracing."""

    enabled: bool = field(
        default=False,
        metadata={
            "help": (
                "Enable per-session lifecycle tracing alongside perf events. "
                "When true, session metadata is captured to sessions.jsonl."
            )
        },
    )
    flush_threshold: int = field(
        default=256,
        metadata={
            "help": (
                "Flush session trace records once this many entries are ready. "
                "Values <= 0 fall back to 1."
            )
        },
    )


@dataclass
class PerfTracerConfig:
    """Configuration for perf tracer emission."""

    experiment_name: str = MISSING
    trial_name: str = MISSING
    fileroot: str = MISSING
    enabled: bool = field(
        default=False,
        metadata={
            "help": (
                "Explicitly enable or disable perf tracing. Set to true to capture perf traces."
            )
        },
    )
    save_interval: int = field(
        default=1,
        metadata={
            "help": (
                "Flush trace events to disk every N calls to save(step=...). "
                "A value of 1 writes on every step; values <= 0 fall back to 1."
            )
        },
    )
    profile_steps: list[int] | None = field(
        default=None,
        metadata={
            "help": (
                "List of step numbers at which to capture detailed profiling traces. "
                "If None, no detailed profiling traces are captured."
            )
        },
    )
    session_tracer: SessionTracerConfig | None = field(
        default=None,
        metadata={"help": "Session tracing configuration."},
    )


@dataclass
class NameResolveConfig:
    """Configuration for distributed name resolution and service discovery."""

    type: str = field(
        default="nfs",
        metadata={
            "help": "Type of the distributed KV store for name resolving.",
            "choices": ["nfs", "etcd3", "ray"],
        },
    )
    nfs_record_root: str = field(
        default="/tmp/areal/name_resolve",
        metadata={
            "help": "Record root for NFS name resolving. Should be available on all nodes."
        },
    )
    etcd3_addr: str = field(
        default="localhost:2379", metadata={"help": "Address of the ETCD3 server."}
    )
    ray_actor_name: str = field(
        default="ray_kv_store",
        metadata={"help": "Name of the distributed Ray KV store."},
    )


@dataclass
class ClusterSpecConfig:
    """Configuration for cluster specification and distributed computing setup."""

    name_resolve: NameResolveConfig = field(
        default_factory=NameResolveConfig,
        metadata={"help": "Name resolving configuration."},
    )
    cluster_name: str = field(
        default="local",
        metadata={"help": "Name of the cluster. Used to set specific environs."},
    )
    fileroot: str = field(
        default="/tmp/areal/",
        metadata={
            "help": "Root for logs and checkpoints. Should be available on all nodes."
        },
    )
    n_nodes: int = field(
        default=32,
        metadata={
            "help": "The size of the cluster. Used to decide slurm hostname suffix."
        },
    )
    n_gpus_per_node: int = field(
        default=8,
        metadata={"help": "Number of GPUs per node (physical)."},
    )


@dataclass
class SchedulerConfig:
    """Configuration for worker scheduling. Used in the single-controller mode. Experimental."""

    type: str | None = field(default=None)
    endpoint: str = field(default="http://localhost:8081")
    deploy_mode: str = field(default="separation")
    functioncall_service_domain: str = field(default="http://localhost:8080")
    reward_functioncall_config: dict = field(default_factory=dict)
    reward_model_path: str = field(default="")
    reward_model_service_url: str = field(default="http://localhost:30000/classify")


@dataclass
class _DatasetConfig:
    """Configuration for dataset loading and preprocessing."""

    path: str = field(
        default=MISSING,
        metadata={
            "help": "Path to the dataset. Can be a local path or a HuggingFace dataset name."
        },
    )
    type: str = field(
        default=MISSING,
        metadata={"help": "Type of training method, e.g., 'sft', 'rl', etc."},
    )
    batch_size: int = field(
        default=1, metadata={"help": "Batch size for the dataloader"}
    )
    shuffle: bool = field(
        default=True, metadata={"help": "Whether to shuffle the dataset"}
    )
    pin_memory: bool = field(
        default=False,
        metadata={
            "help": "Pin memory for faster data loading (set True for GPU training)"
        },
    )
    num_workers: int = field(
        default=0, metadata={"help": "Number of worker processes for data loading"}
    )
    drop_last: bool = field(
        default=True, metadata={"help": "Drop the last incomplete batch"}
    )
    max_length: int | None = field(
        default=None,
        metadata={
            "help": "Maximum token length of sequences in dataset. Longer sequences are filtered out."
        },
    )


@dataclass
class TrainDatasetConfig(_DatasetConfig):
    """Configuration for training dataset loading and preprocessing."""


@dataclass
class ValidDatasetConfig(_DatasetConfig):
    """Configuration for validation dataset loading and preprocessing.

    It has different default values with `TrainDatasetConfig`.
    `shuffle` and `drop_last` default to False.
    """

    shuffle: bool = field(
        default=False, metadata={"help": "Whether to shuffle the dataset"}
    )
    drop_last: bool = field(
        default=False, metadata={"help": "Drop the last incomplete batch"}
    )


@dataclass
class BaseExperimentConfig:
    """Base configuration class for all experiment types with common settings."""

    # NOTE: we need this unified config class because different experiments
    # have different config structures, e.g., GRPO has two engine configs,
    # but SFT only has a single one. We use subclasses to represent these structures.
    experiment_name: str = field(
        default=MISSING,
        metadata={"help": "Name of the experiment (no '_' or '/'). Required."},
    )
    trial_name: str = field(
        default=MISSING,
        metadata={"help": "Name of the trial (no '-' or '/'). Required."},
    )
    cluster: ClusterSpecConfig = field(
        default_factory=ClusterSpecConfig,
        metadata={"help": "Cluster specification. Mainly used by slurm."},
    )
    allocation_mode: str = field(
        default="",
        metadata={"help": "Pattern-based GPU parallel strategy allocation mode. "},
    )
    seed: int = field(default=1, metadata={"help": "Random seed for reproducibility."})
    enable_offload: bool = field(
        default=False,
        metadata={
            "help": "Whether to enable training offload using torch_memory_saver. "
            "This requires setting up the environment for TMS (e.g., via LD_PRELOAD)."
        },
    )
    total_train_epochs: int = field(
        default=1, metadata={"help": "Total number of epochs to train the model."}
    )
    total_train_steps: int | None = field(
        default=None,
        metadata={
            "help": "Terminate training after this number of steps. "
            "For benchmarking purposes only. None indicates normal training."
        },
    )
    total_train_n_seqs: int | None = field(
        default=None,
        metadata={
            "help": "Terminate training after consuming this number of samples. "
            "For benchmarking purposes only. None indicates normal training."
        },
    )
    tokenizer_path: str = field(
        default="",
        metadata={"help": "Path to the tokenizer."},
    )

    train_dataset: TrainDatasetConfig = field(default_factory=TrainDatasetConfig)
    valid_dataset: ValidDatasetConfig | None = field(default=None)

    saver: SaverConfig = field(default_factory=SaverConfig)
    evaluator: EvaluatorConfig = field(default_factory=EvaluatorConfig)
    stats_logger: StatsLoggerConfig = field(default_factory=StatsLoggerConfig)
    perf_tracer: PerfTracerConfig | None = field(
        default=None,
        metadata={"help": "Performance tracer configuration. None means disabled."},
    )
    recover: RecoverConfig = field(default_factory=RecoverConfig)

    sglang: SGLangConfig = field(default_factory=SGLangConfig)
    vllm: vLLMConfig = field(default_factory=vLLMConfig)

    scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)


@dataclass
class SFTConfig(BaseExperimentConfig):
    """Configuration for Supervised Fine-Tuning (SFT) experiments."""

    actor: TrainEngineConfig = field(default_factory=TrainEngineConfig)


@dataclass
class RWConfig(BaseExperimentConfig):
    """Configuration for Reward Model (RW) training experiments."""

    actor: TrainEngineConfig = field(default_factory=TrainEngineConfig)


@dataclass
class PPOConfig(BaseExperimentConfig):
    """Configuration for Proximal Policy Optimization (PPO) reinforcement learning experiments."""

    gconfig: GenerationHyperparameters = field(
        default_factory=GenerationHyperparameters
    )
    eval_gconfig: GenerationHyperparameters | None = field(
        default=None,
        metadata={
            "help": "Generation hyperparameters for evaluation. If None, use gconfig."
        },
    )
    rollout: InferenceEngineConfig = field(default_factory=InferenceEngineConfig)
    actor: PPOActorConfig = field(default_factory=PPOActorConfig)
    ref: PPOActorConfig | None = field(default=None)
    critic: PPOCriticConfig | None = field(default=None)
    dynamic_bs: bool = field(
        default=False,
        metadata={
            "help": "Enable dynamic batch sizing in prepare_batch. When True, batch collection "
            "stops when (accepted + rejected) >= batch_size, returning only accepted results. "
            "This results in variable-sized batches of valid data."
        },
    )

    def __post_init__(self):
        """Validate the eval generation config."""
        if self.eval_gconfig is None:
            self.eval_gconfig = self.gconfig.new()


@dataclass
class GRPOConfig(PPOConfig):
    """A dummy place holder of GRPO config for backward compatibility."""

    pass


def parse_cli_args(argv: list[str]):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config", help="Path to the main configuration file", required=True
    )
    # The first argument might be the path to a training script,
    # which should be ignored by the argument parser.
    if argv and argv[0].endswith(".py"):
        argv = argv[1:]
    args, overrides = parser.parse_known_args(argv)
    # Initialize hydra config
    config_file = Path(args.config).absolute()
    assert config_file.exists(), f"Config file {config_file} does not exist."
    # hydra only recognize relative paths
    relpath = Path(os.path.relpath(str(config_file), Path(__file__).parent.absolute()))
    if GlobalHydra.instance().is_initialized():
        GlobalHydra.instance().clear()
    hydra_init(config_path=str(relpath.parent), job_name="app", version_base=None)
    cfg = hydra_compose(
        config_name=str(relpath.name).split(".yaml")[0],
        overrides=overrides,
    )
    return cfg, config_file


def to_structured_cfg(cfg, config_cls):
    # Merge with the default configuration.
    # The yaml and commandline can omit some default values defined in python dataclasses.
    default_cfg = OmegaConf.structured(config_cls)
    cfg = OmegaConf.merge(default_cfg, cfg)
    return cfg


def load_expr_config[ConfigT](
    argv: list[str], config_cls: type[ConfigT]
) -> tuple[ConfigT, str]:
    cfg, config_file = parse_cli_args(argv)
    cfg = to_structured_cfg(cfg, config_cls=config_cls)
    cfg = OmegaConf.to_object(cfg)
    assert isinstance(cfg, config_cls)
    # Setup environment

    name_resolve.reconfigure(cfg.cluster.name_resolve)

    from areal.utils.stats_logger import StatsLogger

    # Save configuration as yaml
    if os.getenv("RANK", "0") == "0":
        save_config(cfg, StatsLogger.get_log_path(cfg.stats_logger))

    return cfg, str(config_file)


def conf_as_dict(cfg):
    if isinstance(cfg, (OmegaConf, DictConfig)):
        return OmegaConf.to_container(cfg, resolve=True)
    return asdict(cfg)


def save_config(cfg, log_dir):
    os.makedirs(log_dir, exist_ok=True)
    config_save_path = os.path.join(log_dir, "config.yaml")
    with open(config_save_path, "w") as f:
        config_dict: dict = asdict(cfg)
        yaml.dump(
            config_dict,
            f,
            default_flow_style=False,
            sort_keys=False,
        )