Configuration

@dataclass
class SFTPlacementConfig(BaseConfig):
    """Placement configuration for SFT training"""

    num_nodes: int = 1
    num_gpus_per_node: int = 4

@dataclass
class SFTConfig(BaseConfig):
    """Configuration for SFT training.

    Usage::

        cfg = SFTConfig(
            strategy="megatron",
            placement=SFTPlacementConfig(num_gpus_per_node=4),
            megatron_config=MegatronConfig(tensor_model_parallel_size=2,
                                    pipeline_model_parallel_size=2),
        )

    Or from CLI::

        cfg = SFTConfig.from_cli_overrides(sys.argv[1:])
    """

    @classmethod
    def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig":
        """Construct an SFTConfig from CLI arguments or a dict of overrides.

        Parses CLI dotlist arguments via OmegaConf and builds a typed config.
        Dataclass field defaults are used for any values not specified.

        Args:
            args: Either a list of CLI arguments in 'key.path=value' format, or a dict
                  mapping dot-notation keys to values.
                  Example list: ['strategy=megatron', 'model.path=Qwen/Qwen3-0.6B']
                  Example dict: {'strategy': 'megatron', 'model.path': 'Qwen/Qwen3-0.6B'}
                  Dict values are serialized as JSON, so ``None``, bools, strings,
                  lists and nested dicts keep their types.

        Returns:
            A fully constructed SFTConfig with CLI overrides applied.

        Raises:
            ValueError: If both ``num_epochs`` and ``num_steps`` are explicitly provided.
        """
        if isinstance(args, dict):
            args = overrides_dict_to_dotlist(args)

        overrides = OmegaConf.from_cli(args)
        # Check for mutual exclusion before constructing the full config
        if "num_epochs" in overrides and "num_steps" in overrides:
            raise ValueError("Cannot specify both num_epochs and num_steps")
        # Accept the deprecated ``use_sample_packing`` key as an alias for
        # ``remove_microbatch_padding``. Remap it before construction so the
        # strict key validation does not reject the old name.
        if "use_sample_packing" in overrides:
            if "remove_microbatch_padding" in overrides:
                raise ValueError(
                    "Specify only one of use_sample_packing (deprecated) and remove_microbatch_padding, not both."
                )
            import warnings

            warnings.warn(
                "use_sample_packing has been renamed to remove_microbatch_padding; "
                "use remove_microbatch_padding instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            overrides["remove_microbatch_padding"] = overrides["use_sample_packing"]
            del overrides["use_sample_packing"]
        return cls.from_dict_config(overrides)

    # ---- Reused SkyRL config objects ----
    model: ModelConfig = field(default_factory=lambda: ModelConfig(path="Qwen/Qwen3-0.6B"))
    optimizer_config: OptimizerConfig = field(default_factory=OptimizerConfig)
    placement: SFTPlacementConfig = field(default_factory=SFTPlacementConfig)
    megatron_config: MegatronConfig = field(
        default_factory=lambda: MegatronConfig(
            tensor_model_parallel_size=2,
            pipeline_model_parallel_size=2,
        )
    )
    fsdp_config: FSDPConfig = field(default_factory=FSDPConfig)

    # Ulysses sequence parallelism
    sequence_parallel_size: int = 1
    """Ulysses sequence parallelism size"""

    model_config_kwargs: dict = field(default_factory=dict)
    """Pass-through kwargs for the HuggingFace model config (FSDP backends).
    For Megatron, use ``megatron_config.transformer_config_kwargs`` instead."""
    use_torch_compile: bool = False
    """Apply torch.compile to logits calculation."""
    record_memory: bool = False
    """Save memory snapshots to ``{ckpt_path}/memory_snapshots/``.
    Visualize by dragging pickle files to https://docs.pytorch.org/memory_viz."""
    torch_profiler_config: TorchProfilerConfig = field(default_factory=TorchProfilerConfig)
    """torch.profiler config for policy training steps."""

    # ---- SFT-specific flat fields ----
    strategy: str = "megatron"  # "megatron" or "fsdp"
    dataset_name: Optional[str] = None
    """Deprecated: use ``train_datasets`` instead. Translated to ``train_datasets=[dataset_name]``
    with a DeprecationWarning. Cannot be combined with ``train_datasets``."""
    dataset_split: Optional[str] = None
    """Deprecated: use ``train_dataset_splits`` instead."""
    train_datasets: Optional[List[str]] = None
    """HuggingFace dataset names (or paths) to train on. With multiple datasets, batches are
    mixed per-source by :class:`~skyrl.train.dataset.samplers.DataMixingSampler` according to
    ``train_dataset_weights``. Defaults to ``["yahma/alpaca-cleaned"]``. All datasets must share
    the same ``messages_key``/``tools_key``/``system_key`` columns and modality."""
    train_dataset_splits: Optional[List[str]] = None
    """Split to load for each entry of ``train_datasets`` (e.g. ``"train[:50000]"``). Must match
    ``train_datasets`` in length. Defaults to ``["train[:100]"]``."""
    train_dataset_weights: Optional[List[float]] = None
    """Per-dataset sampling weights: the approximate per-batch ratio of samples drawn from each
    dataset, independent of dataset sizes. Only supported with ``sampler="random"`` (custom
    samplers receive ratios via ``sampler_kwargs``). Defaults to equal mixing (``1/N`` each)."""
    pretokenized_dataset_paths: Optional[List[str]] = None
    """Local paths to *pretokenized* training datasets, each a file or
    directory holding parquet/JSONL/arrow files or a HF
    ``Dataset.save_to_disk`` directory. Rows must carry unpadded
    ``input_ids`` and a full-sequence 0/1 ``loss_mask`` (``num_actions`` is
    inferred); VLM rows additionally carry ``pixel_values`` /
    ``image_grid_thw``. See ``skyrl.train.dataset.pretokenized``. When set,
    online tokenization is skipped; cannot be combined with ``train_datasets``.
    Multiple stores are concatenated and mixed per ``train_dataset_weights``
    (like ``train_datasets``)."""
    messages_key: str = "messages"  # column name for chat-format datasets
    tools_key: str = "tools"
    """Column name holding per-row tool/function schemas for tool-calling datasets
    (e.g. APIGen-MT, xLAM, ToolACE). May be a list[dict] or a JSON-encoded string.
    Ignored if the column is absent from the dataset."""
    system_key: str = "system"
    """Column name holding a per-row system prompt to prepend when ``messages``
    does not already start with a system turn. Ignored if absent."""

    # ---- Evaluation datasets ----
    eval_dataset_name: Optional[str] = None
    """Deprecated: use ``eval_datasets`` instead. Translated to ``eval_datasets=[eval_dataset_name]``
    with a DeprecationWarning. Cannot be combined with ``eval_datasets``."""
    eval_dataset_split: Optional[str] = None
    """Deprecated: use ``eval_dataset_splits`` instead."""
    eval_datasets: Optional[List[str]] = None
    """HuggingFace dataset names (or paths) used to compute eval loss during training.
    When ``None`` (default), eval is disabled. Metrics are logged per dataset under
    ``eval/{name}/`` (nested even with a single eval dataset)."""
    eval_dataset_splits: Optional[List[str]] = None
    """Split to load for each entry of ``eval_datasets`` (e.g. ``"validation"``, ``"test[:500]"``).
    Must match ``eval_datasets`` in length. Defaults to ``["validation"]`` on the deprecated path."""
    eval_dataset_names: Optional[List[str]] = None
    """Optional shorthand names used only for logging (``eval/{name}/loss``). Must be unique and
    match ``eval_datasets`` (or ``eval_pretokenized_dataset_paths``) in length. Defaults to each
    dataset name with ``/`` replaced by ``_`` (path basenames for pretokenized stores)."""
    eval_pretokenized_dataset_paths: Optional[List[str]] = None
    """Paths to *pretokenized* eval datasets (same formats and schema as
    ``pretokenized_dataset_paths``). Cannot be combined with ``eval_datasets``.
    Metrics are logged under ``eval/{name}/`` where the names come from
    ``eval_dataset_names`` when set, defaulting to each path's basename."""
    eval_interval: int = 0
    """Run eval every N training steps. Eval also runs once at the end of training
    when an eval dataset is configured. ``0`` disables periodic eval."""
    eval_before_train: bool = False
    """If True, run a baseline eval pass before training begins (logged at step 0)."""
    max_length: Optional[int] = None
    """Maximum length of tokenized sequences. If specified, all sequences will be truncated to this value
    By default, no truncation is performed"""
    num_steps: Optional[int] = None
    """Number of training steps. If None, num_epochs is used to derive the step count."""
    num_epochs: Optional[int] = 1
    """Number of training epochs. Used when num_steps is None. Default: 1 epoch."""
    batch_size: int = 4
    micro_train_batch_size_per_gpu: int = 2
    logger: str = "console"  # "console" or "wandb"
    project_name: str = "skyrl_sft"
    run_name: str = "skyrl_sft_run"
    tags: Optional[List[str]] = None
    """Optional list of tags to apply to the W&B run. Has no effect on other backends."""
    ckpt_path: str = ""
    ckpt_interval: int = 0  # <= 0 -> no checkpointing
    enable_ray_gpu_monitor: bool = True
    """Enable background Ray GPU/RAM metrics collection and logging to wandb."""
    max_ckpts_to_keep: int = -1
    """-1 to keep all checkpoints, N to keep only the last N."""
    resume_from: str = ""  # "" = no resume, "latest" = latest checkpoint, or path to global_step_N dir

    # ---- HF export ----
    hf_save_interval: int = 0
    """Save HuggingFace-format weights every N steps. 0 = disabled."""
    export_path: str = ""
    """Directory for HF-format exports. Defaults to ckpt_path/hf_exports if empty."""

    seed: int = 42

    # ---- Data loading ----
    num_workers: int = 8
    """Number of worker processes for parallel tokenization during dataset loading. Set to 0 for single-threaded."""
    async_batch_collation: bool = True
    """Overlap the next stateful-dataloader batch with the current GPU step.

    Checkpoint state remains pinned after the current batch. Set to False for
    serial data loading."""

    # ---- Dataloader / sampler ----
    dataloader_num_workers: int = 0
    """Number of worker processes for the training/eval ``StatefulDataLoader``. ``0`` loads in the main process."""
    dataloader_persistent_workers: bool = False
    """Keep dataloader workers alive across epochs. Only takes effect when ``dataloader_num_workers > 0``."""
    sampler: str = "random"
    """Training sampler: ``"random"`` (shuffle each epoch), ``"sequential"`` (in-order), or ``"custom"``
    (load from ``sampler_class_path``)."""
    sampler_class_path: Optional[str] = None
    """Import path (``"module.path.ClassName"``) to a custom stateful sampler. Required when ``sampler='custom'``.
    Instantiated as ``ClassName(tokenized, **sampler_kwargs)``."""
    sampler_kwargs: dict = field(default_factory=dict)
    """Keyword arguments forwarded to the custom sampler constructor."""

    # ---- Tokenized dataset caching ----
    cache_dir: str = os.path.join(
        os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")), "skyrl", "tokenized_datasets"
    )
    """Directory to cache tokenized datasets. For multi-node training, set this to an NFS-mounted path so all nodes can
    share the cache."""
    force_recache: bool = False
    """If True, ignore existing cache and re-tokenize the dataset."""
    disable_cache: bool = False
    """If True, disable cache completely (always tokenize from scratch)."""

    # ---- Training target ----
    train_on_what: TrainOnWhat = TrainOnWhat.LAST_ASSISTANT_MESSAGE
    """Which tokens to compute loss on. See :class:`TrainOnWhat` for options."""

    # ---- Packing ----
    remove_microbatch_padding: bool = True  # Pack multiple sequences per microbatch (requires flash_attn)
    use_sequence_packing: bool = False
    """Enable controller-level FFD bin-packing across the global mini-batch.
    Requires ``remove_microbatch_padding=True`` and the Megatron backend. When
    enabled, ``SFTTrainer`` uses ``PackedDataCollator`` instead of
    ``DefaultCollator``. Each bin row becomes one row in the dispatched batch
    and one worker micro-batch.
    """
    max_tokens_per_microbatch: Optional[int] = None
    """FFD bin capacity (max tokens per bin) when ``use_sequence_packing=True``.
    Each bin row becomes one worker micro-batch, so this is the token budget for
    one micro-batch. Must be ``>= max_length`` so any single sequence fits in a
    bin. ``None`` (default) resolves to ``max_length`` (each bin holds one
    sequence)."""

    # ---- Dummy run / benchmarking ----
    dummy_run_full_ctx: bool = False  # Skip real data; fabricate full-context sequences
    dummy_run_max_steps: int = 5  # Number of steps to run in dummy mode

    # ---- CI / smoke test support ----
    max_training_steps: Optional[int] = None
    """If set, stop training after this many steps regardless of num_steps or num_epochs.
    Useful for CI smoke tests and quick validation runs."""

    def resolved_bin_capacity(self) -> int:
        """FFD bin capacity (max tokens per bin) when sequence packing is enabled.

        Resolves ``max_tokens_per_microbatch`` against ``max_length``: when the
        token budget is ``None`` it falls back to ``max_length`` (each bin holds
        one sequence). Requires ``max_length`` to be set and the resolved budget
        to be ``>= max_length`` so any single sequence fits in a bin.
        """
        if self.max_length is None:
            raise ValueError("max_tokens_per_microbatch requires max_length to be set.")
        max_tokens = self.max_tokens_per_microbatch
        if max_tokens is None:
            max_tokens = self.max_length
        if max_tokens < self.max_length:
            raise ValueError(
                f"max_tokens_per_microbatch ({max_tokens}) must be >= max_length "
                f"({self.max_length}) so any single sequence fits in a bin."
            )
        return max_tokens

    @classmethod
    def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig":
        """Construct an SFTConfig from CLI arguments or a dict of overrides.

        Parses CLI dotlist arguments via OmegaConf and builds a typed config.
        Dataclass field defaults are used for any values not specified.

        Args:
            args: Either a list of CLI arguments in 'key.path=value' format, or a dict
                  mapping dot-notation keys to values.
                  Example list: ['strategy=megatron', 'model.path=Qwen/Qwen3-0.6B']
                  Example dict: {'strategy': 'megatron', 'model.path': 'Qwen/Qwen3-0.6B'}
                  Dict values are serialized as JSON, so ``None``, bools, strings,
                  lists and nested dicts keep their types.

        Returns:
            A fully constructed SFTConfig with CLI overrides applied.

        Raises:
            ValueError: If both ``num_epochs`` and ``num_steps`` are explicitly provided.
        """
        if isinstance(args, dict):
            args = overrides_dict_to_dotlist(args)

        overrides = OmegaConf.from_cli(args)
        # Check for mutual exclusion before constructing the full config
        if "num_epochs" in overrides and "num_steps" in overrides:
            raise ValueError("Cannot specify both num_epochs and num_steps")
        # Accept the deprecated ``use_sample_packing`` key as an alias for
        # ``remove_microbatch_padding``. Remap it before construction so the
        # strict key validation does not reject the old name.
        if "use_sample_packing" in overrides:
            if "remove_microbatch_padding" in overrides:
                raise ValueError(
                    "Specify only one of use_sample_packing (deprecated) and remove_microbatch_padding, not both."
                )
            import warnings

            warnings.warn(
                "use_sample_packing has been renamed to remove_microbatch_padding; "
                "use remove_microbatch_padding instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            overrides["remove_microbatch_padding"] = overrides["use_sample_packing"]
            del overrides["use_sample_packing"]
        return cls.from_dict_config(overrides)

    def resolved_bin_capacity(self) -> int:
        """FFD bin capacity (max tokens per bin) when sequence packing is enabled.

        Resolves ``max_tokens_per_microbatch`` against ``max_length``: when the
        token budget is ``None`` it falls back to ``max_length`` (each bin holds
        one sequence). Requires ``max_length`` to be set and the resolved budget
        to be ``>= max_length`` so any single sequence fits in a bin.
        """
        if self.max_length is None:
            raise ValueError("max_tokens_per_microbatch requires max_length to be set.")
        max_tokens = self.max_tokens_per_microbatch
        if max_tokens is None:
            max_tokens = self.max_length
        if max_tokens < self.max_length:
            raise ValueError(
                f"max_tokens_per_microbatch ({max_tokens}) must be >= max_length "
                f"({self.max_length}) so any single sequence fits in a bin."
            )
        return max_tokens

class TrainOnWhat(StrEnum):
    """Enum controlling which parts of the sequence to compute loss on.

    Members:
        LAST_ASSISTANT_MESSAGE: Train only on the final assistant message.
        ALL_ASSISTANT_MESSAGES: Train on every assistant message in the conversation.
    """

    LAST_ASSISTANT_MESSAGE = "last_assistant_message"
    ALL_ASSISTANT_MESSAGES = "all_assistant_messages"