# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import inspect
import os
import atexit
import traceback as _tb
import torch
from typing import Optional, Union, Callable, Dict, Any
from torch.utils.data import Dataset
from tqdm import tqdm
from itertools import islice
from torch.utils.data import DataLoader, DistributedSampler, Sampler
from datasets import concatenate_datasets


from cosmos_rl.dispatcher.data.packer.base import BaseDataPacker
from cosmos_rl.utils.logging import logger
from cosmos_rl.utils.parallelism import ParallelDims
from cosmos_rl.policy.config import Config as CosmosConfig
from cosmos_rl.policy.trainer.base import TrainerRegistry
from cosmos_rl.utils import util
from cosmos_rl.utils.distributed import destroy_distributed
import cosmos_rl.utils.distributed as dist_utils
import torch.distributed as dist
from cosmos_rl.utils.report.wandb_logger import (
    init_wandb,
    is_wandb_available,
    log_wandb,
)
from cosmos_rl.policy.config import (
    SFTDataConfig,
    config_hash,
)

from cosmos_rl.policy.trainer.sampler import SkippingSampler
import cosmos_rl.utils.cache as cache
from cosmos_rl.policy.trainer.llm_trainer.sft_trainer import SFTTrainer
from cosmos_rl.policy.worker.base import PolicyWorkerBase
from cosmos_rl.dispatcher.data.load_balanced_dataset import LoadBalancedDataset


class SFTDataset(Dataset):
    def __init__(
        self,
        config: SFTDataConfig,
        dataset: Dataset,
        data_packer: BaseDataPacker,
        is_user_dataset: bool = False,
    ):
        self.config = config
        self.column_name = config.conversation_column_name
        self.dataset = dataset
        self.data_packer = data_packer
        self.is_user_dataset = is_user_dataset
        self.cache = None
        if self.config.enable_dataset_cache:
            # TODO(zjx): can we reuse the cache between different training jobs?
            # It's not stable yet, we only checked if the config is the same
            # If there are any problems, it is recommended that the user clears the cache folder
            cache_folder = os.path.join(
                os.environ.get(
                    "COSMOS_CACHE",
                    os.path.join(os.path.expanduser("~"), ".cache/cosmos/"),
                ),
                "datasets_cache",
                f"{self.config.dataset.name}-{config_hash(config)}",
            )
            logger.info(f"SFTDataset Cache folder: {cache_folder}")
            self.cache = cache.DiskCache(cache_folder)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        # we only cache on_the_fly result
        if self.cache is not None:
            cache_obj = self.cache.get(idx)
            if cache_obj is not None:
                return cache_obj

        max_retries = 50
        for attempt in range(max_retries):
            raw_item = (
                self.dataset[idx][self.column_name]
                if not self.is_user_dataset and self.column_name
                else self.dataset[idx]
            )

            try:
                if isinstance(idx, list):  # a batch of items
                    item = [self.data_packer.sft_process_sample(x) for x in raw_item]
                else:
                    item: Dict[str, Any] = self.data_packer.sft_process_sample(raw_item)
                break
            except Exception as e:
                msg_info = []
                try:
                    msgs = raw_item if isinstance(raw_item, list) else []
                    if isinstance(raw_item, dict) and "messages" in raw_item:
                        msgs = raw_item["messages"]
                    for msg in msgs:
                        if not isinstance(msg, dict):
                            msg_info.append(f"non-dict:{type(msg).__name__}")
                            continue
                        role = msg.get("role", "?")
                        content = msg.get("content")
                        if isinstance(content, list):
                            items = []
                            for c in content:
                                if isinstance(c, dict):
                                    item_d = {}
                                    for k, v in c.items():
                                        if isinstance(
                                            v, (str, int, float, bool, type(None))
                                        ):
                                            item_d[k] = str(v)[:100]
                                        else:
                                            item_d[k] = type(v).__name__
                                    items.append(item_d)
                            msg_info.append(f"{role}:{items}")
                        elif isinstance(content, str):
                            msg_info.append(f"{role}:str({len(content)})")
                        else:
                            msg_info.append(
                                f"{role}:content={type(content).__name__}({repr(content)[:200]})"
                                if content is not None
                                else f"{role}:content=None"
                            )
                except Exception as dbg_e:
                    msg_info.append(f"debug_err:{dbg_e}")
                logger.warning(
                    f"sft_process_sample failed (attempt {attempt + 1}/{max_retries}): {e}"
                    f"\n  traceback: {_tb.format_exc().splitlines()[-3:]}"
                    f"\n  messages: {msg_info}"
                )
                if attempt >= max_retries - 1:
                    raise
                continue

        if self.cache is not None:
            self.cache.set(idx, item)
        return item


class DPODataset(Dataset):
    """Dataset wrapper for DPO: yields processed {'chosen': dict, 'rejected': dict}."""

    def __init__(
        self,
        config: SFTDataConfig,
        dataset: Dataset,
        data_packer: BaseDataPacker,
        is_user_dataset: bool = False,
    ):
        self.config = config
        self.column_name = config.conversation_column_name
        self.dataset = dataset
        self.data_packer = data_packer
        self.is_user_dataset = is_user_dataset
        self.cache = None
        if self.config.enable_dataset_cache:
            cache_folder = os.path.join(
                os.environ.get(
                    "COSMOS_CACHE",
                    os.path.join(os.path.expanduser("~"), ".cache/cosmos/"),
                ),
                "datasets_cache",
                f"dpo-{self.config.dataset.name}-{config_hash(config)}",
            )
            logger.info(f"DPODataset Cache folder: {cache_folder}")
            self.cache = cache.DiskCache(cache_folder)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        if self.cache is not None:
            cache_obj = self.cache.get(idx)
            if cache_obj is not None:
                return cache_obj

        max_retries = 50
        for attempt in range(max_retries):
            raw_item = (
                self.dataset[idx][self.column_name]
                if not self.is_user_dataset and self.column_name
                else self.dataset[idx]
            )
            try:
                if hasattr(self.data_packer, "dpo_process_sample"):
                    item = self.data_packer.dpo_process_sample(raw_item)
                else:
                    raise ValueError("DPO requires data_packer with dpo_process_sample")
                break
            except Exception as e:
                if attempt >= max_retries - 1:
                    raise
                print(
                    f"WARNING: dpo_process_sample failed (attempt {attempt + 1}/{max_retries}): {e}"
                )
                continue

        if self.cache is not None:
            self.cache.set(idx, item)
        return item


def construct_dataset(
    cosmos_config: CosmosConfig,
    data_packer: BaseDataPacker,
    user_provided_dataset: Optional[Dataset] = None,
    val_data_packer: Optional[BaseDataPacker] = None,
    user_provided_val_dataset: Optional[Dataset] = None,
):
    config = cosmos_config.train.train_policy
    if user_provided_dataset is not None:
        dataset = None
        train_dataset = user_provided_dataset
        logger.info("Using user-provided dataset, which will skip split processing.")
    else:
        dataset = util.load_data_from_disk_or_hf(
            config.dataset.name,
            config.dataset.subset,
            config.dataset.revision or None,
        )
        dataset_list = []
        for split_name in config.dataset.split:
            logger.info(
                f"Appending split {split_name}, dataset size = {len(dataset[split_name])}"
            )
            dataset_list.append(dataset[split_name])
        train_dataset = concatenate_datasets(dataset_list)
    logger.info(f"Final dataset size = {len(train_dataset)}")

    if cosmos_config.validation.enable:
        if user_provided_val_dataset is not None:
            test_dataset = user_provided_val_dataset
            logger.info(
                "Using user-provided validation dataset, which will skip split processing."
            )
        elif cosmos_config.validation.dataset.name:
            dataset = util.load_data_from_disk_or_hf(
                cosmos_config.validation.dataset.name,
                cosmos_config.validation.dataset.subset,
                cosmos_config.validation.dataset.revision or None,
            )
            dataset_list = []
            for split_name in cosmos_config.validation.dataset.split:
                logger.info(
                    f"Appending validation split {split_name}, validation dataset size = {len(dataset[split_name])}"
                )
                dataset_list.append(dataset[split_name])
            test_dataset = concatenate_datasets(dataset_list)
        else:
            # Split train/val from training dataset if no val dataset and no val dataset name provided
            train_dataset, test_dataset = util.split_train_n_val_dataset(
                train_dataset, cosmos_config
            )
    else:

        class EmptyDataset(Dataset):
            def __len__(self):
                return 0

            def __getitem__(self, idx):
                raise IndexError("EmptyDataset has no items")

        test_dataset = EmptyDataset()

    trainer_type = cosmos_config.train.train_policy.trainer_type
    DatasetCls = DPODataset if trainer_type == "dpo" else SFTDataset
    train_sft_dataset = DatasetCls(
        config,
        dataset=train_dataset,
        data_packer=data_packer,
        is_user_dataset=user_provided_dataset is not None,
    )
    test_sft_dataset = DatasetCls(
        config,
        dataset=test_dataset,
        data_packer=val_data_packer,
        is_user_dataset=user_provided_dataset is not None,
    )

    return train_sft_dataset, test_sft_dataset


def collate_fn(
    batch,
):
    return batch


class SFTPolicyWorker(PolicyWorkerBase):
    trainer: SFTTrainer

    def __init__(
        self,
        config: CosmosConfig,
        parallel_dims: ParallelDims,
        dataset: Optional[Union[Dataset, Callable[[CosmosConfig], Dataset]]] = None,
        data_packer: Optional[BaseDataPacker] = None,
        val_dataset: Optional[Union[Dataset, Callable[[CosmosConfig], Dataset]]] = None,
        val_data_packer: Optional[BaseDataPacker] = None,
        sampler: Optional[Callable] = None,
        batch_sampler: Optional[Callable] = None,
        val_sampler: Optional[Callable] = None,
        val_batch_sampler: Optional[Callable] = None,
        **kwargs,
    ):
        super(SFTPolicyWorker, self).__init__(config, parallel_dims, **kwargs)

        # Enlarge the compile cache size for validation
        if self.config.train.compile and self.config.validation.enable:
            torch._dynamo.config.cache_size_limit = 64

        self.hook_fns = self.hook_fns if self.hook_fns is not None else {}
        self.custom_logger_fns = (
            self.custom_logger_fns if self.custom_logger_fns is not None else []
        )

        # Prepare wandb
        if "wandb" in self.config.logging.logger and is_wandb_available():
            # Only initialize wandb on the first dp replicate coord and first rank for policy
            if self.parallel_dims.dp_replicate_coord[0] == 0 and self.global_rank == 0:
                init_wandb(self.config)
        else:
            logger.warning(
                "Wandb is not available. Please install it to use wandb logging features."
            )

        self.train_step = 0
        self.start_epoch = 0
        self.enable_dp_load_balancing = (
            self.config.train.train_policy.enable_dp_load_balancing
        )

        self.build_runner(
            data_packer=data_packer,
            val_data_packer=val_data_packer,
            sampler=sampler,
            batch_sampler=batch_sampler,
            val_sampler=val_sampler,
            val_batch_sampler=val_batch_sampler,
            dataset=dataset,
            val_dataset=val_dataset,
        )

        # Register the exit function to be called when the program exits
        atexit.register(self.handle_shutdown)

    def setup(
        self,
        data_packer: Optional[Union[BaseDataPacker, Callable]] = None,
        val_data_packer: Optional[Union[BaseDataPacker, Callable]] = None,
    ):
        # setup data packer first
        self.init_data_packer(
            data_packer=data_packer,
            val_data_packer=val_data_packer,
        )
        self.setup_hooks()

    def setup_hooks(self):
        if "pre_validation_hook" in self.hook_fns:
            self.pre_validation_hook = self.hook_fns["pre_validation_hook"]
        else:
            self.pre_validation_hook = None

        if "pre_per_step_validation_hook" in self.hook_fns:
            self.pre_per_step_validation_hook = self.hook_fns[
                "pre_per_step_validation_hook"
            ]
        else:
            self.pre_per_step_validation_hook = None

        if "post_per_step_validation_hook" in self.hook_fns:
            self.post_per_step_validation_hook = self.hook_fns[
                "post_per_step_validation_hook"
            ]
        else:
            self.post_per_step_validation_hook = None

        if "post_validation_hook" in self.hook_fns:
            self.post_validation_hook = self.hook_fns["post_validation_hook"]
        else:
            self.post_validation_hook = None

    def build_runner(
        self,
        data_packer: Optional[BaseDataPacker] = None,
        val_data_packer: Optional[BaseDataPacker] = None,
        sampler: Optional[Callable] = None,
        batch_sampler: Optional[Callable] = None,
        val_sampler: Optional[Callable] = None,
        val_batch_sampler: Optional[Callable] = None,
        dataset: Optional[Union[Dataset, Callable[[CosmosConfig], Dataset]]] = None,
        val_dataset: Optional[Union[Dataset, Callable[[CosmosConfig], Dataset]]] = None,
    ):
        self.setup(
            data_packer=data_packer,
            val_data_packer=val_data_packer,
        )
        trainer_type = self.config.train.train_policy.trainer_type
        if self.config.policy.is_diffusers:
            trainer_type = "diffusers_" + trainer_type
        self.trainer = TrainerRegistry.get_trainer_cls(trainer_type)(
            config=self.config,
            parallel_dims=self.parallel_dims,
            train_stream=self.train_stream,
            data_packer=self.data_packer,
            val_data_packer=self.val_data_packer,
            hook_fns=self.hook_fns,
        )
        self.ckpt_total_steps, self.train_step, _ = self.trainer.load_model()
        if isinstance(dataset, Callable):
            # Incase it is a factory function, we need to call it to get the dataset
            dataset = dataset(self.config)
            util.call_setup(dataset, self.config)

        if isinstance(val_dataset, Callable):
            val_dataset = val_dataset(self.config)
            util.call_setup(val_dataset, self.config)

        if not self.val_data_packer:
            self.val_data_packer = self.data_packer

        # Prepare dataset
        train_dataset, val_dataset = construct_dataset(
            self.config,
            data_packer=self.data_packer,
            user_provided_dataset=dataset,
            val_data_packer=self.val_data_packer,
            user_provided_val_dataset=val_dataset,
        )

        # Apply load-balanced dynamic batching if enabled
        if self.enable_dp_load_balancing:
            logger.info("Enabling load-balanced dynamic batching for training dataset.")
            # Determine max_tokens_for_batch if not specified
            max_tokens_for_batch = (
                self.config.train.train_policy.load_balanced_max_tokens_for_batch
            )
            if max_tokens_for_batch is None:
                # Default: model_max_length
                model_max_length = self.config.policy.model_max_length
                max_tokens_for_batch = model_max_length
                logger.info(
                    f"max_tokens_for_batch not specified, using default: "
                    f"{max_tokens_for_batch} = {model_max_length}"
                )

            accumulate_steps = (
                self.config.train.train_policy.load_balanced_batches_per_optimizer_step
            )
            train_dataset = LoadBalancedDataset(
                base_dataset=train_dataset,
                pool_size=self.config.train.train_policy.load_balanced_pool_size,
                max_tokens_for_batch=max_tokens_for_batch,
                length_key="input_ids",
                batching_strategy=self.config.train.train_policy.load_balanced_batching_strategy,
                max_tokens_len=self.config.policy.model_max_length,
                seq_packing_enabled=self.config.train.sequence_packing,
                seed=self.config.train.train_policy.dataloader_seed,
                dp_rank=self.dp_rank,
                dp_world_size=self.dp_world_size,
                accumulate_steps=accumulate_steps,
            )
            logger.info(
                f"Wrapped training dataset with LoadBalancedDataset: "
                f"pool_size={self.config.train.train_policy.load_balanced_pool_size}, "
                f"max_tokens_for_batch={max_tokens_for_batch}, "
                f"accumulate_steps={accumulate_steps}"
            )

        # For sampler, we won't drop data for un-even distribution DP.
        # Note: If enable_dp_load_balancing, we don't need a sampler
        # as the LoadBalancedDataset handles data distribution internally
        if self.enable_dp_load_balancing:
            train_sampler = None
            logger.info(
                "Skipping sampler setup for load-balanced batching (LoadBalancedDataset handles distribution)."
            )
        elif sampler is not None:
            logger.info("Using user-provided sampler for training dataset.")
            if isinstance(sampler, Callable):
                train_sampler = sampler(
                    train_dataset,
                    num_replicas=self.dp_world_size,
                    rank=self.dp_rank,
                    shuffle=self.config.train.train_policy.dataloader_shuffle,
                    drop_last=False,
                )
            else:
                train_sampler = sampler
        else:
            train_sampler = DistributedSampler(
                train_dataset,
                num_replicas=self.dp_world_size,
                rank=self.dp_rank,
                shuffle=self.config.train.train_policy.dataloader_shuffle,
                drop_last=False,
                seed=self.config.train.train_policy.dataloader_seed,
            )
        self.train_sampler = train_sampler

        if batch_sampler is not None and isinstance(batch_sampler, Callable):
            sig = inspect.signature(batch_sampler)
            kwargs = {
                "dataset": train_dataset.dataset,
                "num_replicas": self.dp_world_size,
                "rank": self.dp_rank,
                "num_workers": self.config.train.train_policy.dataloader_num_workers,
                "config": self.config,
                "sampler": self.train_sampler,
                "batch_size": self.config.train.train_batch_per_replica,
                "drop_last": False,
            }
            # Filter kwargs to only those the function accepts
            filtered = {k: v for k, v in kwargs.items() if k in sig.parameters}
            batch_sampler = batch_sampler(**filtered)
        self.train_batch_sampler = batch_sampler

        def get_train_data_loader(
            sampler: Union[Sampler[int], Sampler[list[int]]],
            sampler_in_batch: Optional[Sampler[list[int]]] = None,
        ):
            if self.enable_dp_load_balancing:
                # For IterableDataset with load-balanced batching, batches are already formed
                # We set batch_size=None and let the dataset yield batches directly
                data_loader = DataLoader(
                    train_dataset,
                    batch_size=None,  # Batches are already formed by IterableDataset
                    num_workers=self.config.train.train_policy.dataloader_num_workers,
                    prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                    collate_fn=collate_fn,  # Still need collate_fn for final batch formatting
                )
            elif sampler_in_batch is not None:
                logger.info(
                    "Using custom batch Sampler that yields list of indices for training dataset."
                )
                data_loader = DataLoader(
                    train_dataset,
                    num_workers=self.config.train.train_policy.dataloader_num_workers,
                    prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                    batch_sampler=sampler_in_batch,
                    collate_fn=collate_fn,
                )
            else:
                data_loader = DataLoader(
                    train_dataset,
                    batch_size=self.config.train.train_batch_per_replica,
                    shuffle=False,
                    num_workers=self.config.train.train_policy.dataloader_num_workers,
                    prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                    sampler=sampler,
                    collate_fn=collate_fn,
                    drop_last=self.config.train.train_policy.dataloader_drop_last,
                )
            return data_loader

        if self.config.train.resume and self.train_step > 0:
            """
            Note: Here both shuffle and no shuffle samplers are supported for deterministic resuming.
            Note: Resume logic for load-balanced batching is handled differently since IterableDataset
            manages its own iteration state.
            """
            if self.enable_dp_load_balancing:
                # For load-balanced batching, training is step-based, not epoch-based.
                # Since total_steps = max_num_steps, train_step is always < max_num_steps
                # when resuming (otherwise training would have already completed).
                # LoadBalancedDataset will automatically manage epoch increments when data loops.

                # Set initial epoch to 0 for deterministic data ordering
                # (epoch will be automatically incremented by LoadBalancedDataset when data loops)
                initial_epoch = 0
                # Skip batches equal to completed steps (use modulo for robustness)
                batches_to_skip = self.train_step % self.config.train.max_num_steps

                if hasattr(train_dataset, "set_epoch"):
                    train_dataset.set_epoch(initial_epoch)
                self.start_epoch = initial_epoch

                # Skip batches that have already been processed
                train_dataset.skip_batches(batches_to_skip)
                logger.info(
                    f"Resuming load-balanced training: initial_epoch={initial_epoch}, "
                    f"skipping {batches_to_skip} batches "
                    f"(train_step={self.train_step}, dp_rank={self.dp_rank})"
                )
            else:
                # Resume training from the last checkpoint if needed
                total_steps_per_epoch = len(
                    get_train_data_loader(self.train_sampler, self.train_batch_sampler)
                )
                data_loader_bias = self.train_step % total_steps_per_epoch
                data_loader_bias *= self.config.train.train_batch_per_replica
                logger.info(
                    f"Resuming training from step {self.train_step}/{self.ckpt_total_steps}"
                )
                if self.train_sampler is not None and hasattr(
                    self.train_sampler, "set_epoch"
                ):
                    self.train_sampler.set_epoch(
                        self.train_step // total_steps_per_epoch
                    )
                if self.train_sampler is not None:
                    self.train_sampler = SkippingSampler(
                        self.train_sampler,
                        skip_samples=data_loader_bias
                        // (
                            len(list(islice(iter(self.train_sampler), 1))[0])
                            if isinstance(
                                list(islice(iter(self.train_sampler), 1))[0], list
                            )
                            else 1
                        ),
                    )
                if self.train_batch_sampler is not None:
                    if hasattr(self.train_batch_sampler, "set_epoch"):
                        self.train_batch_sampler.set_epoch(
                            self.train_step // total_steps_per_epoch
                        )
                    batched_loader_bias = self.train_step % total_steps_per_epoch
                    self.train_batch_sampler = SkippingSampler(
                        self.train_batch_sampler,
                        skip_samples=batched_loader_bias,
                    )
                self.start_epoch = self.train_step // total_steps_per_epoch

        if val_sampler is not None:
            logger.info("Using user-provided sampler for validation dataset.")
            if isinstance(val_sampler, Callable):
                val_sampler = val_sampler(
                    val_dataset,
                    num_replicas=self.dp_world_size,
                    rank=self.dp_rank,
                    shuffle=False,
                    drop_last=False,
                )
        else:
            val_sampler = DistributedSampler(
                val_dataset,
                num_replicas=self.dp_world_size,
                rank=self.dp_rank,
                shuffle=False,
                drop_last=False,
            )
        self.epoch = self.config.train.epoch

        if hasattr(train_dataset, "dataset") and hasattr(
            train_dataset.dataset, "data_loader"
        ):
            # Use custom data loader if provided by dataset
            self.train_data_loader = train_dataset.dataset.data_loader
        else:
            self.train_data_loader = get_train_data_loader(
                self.train_sampler, self.train_batch_sampler
            )

        if hasattr(val_dataset.dataset, "data_loader"):
            # Use custom data loader if provided by dataset
            self.val_data_loader = val_dataset.dataset.data_loader
        elif val_batch_sampler is not None:
            logger.info(
                "Using custom batch Sampler that yields list of indices for validation dataset."
            )
            if isinstance(val_batch_sampler, Callable):
                sig = inspect.signature(val_batch_sampler)
                kwargs = {
                    "dataset": val_dataset.dataset,
                    "num_replicas": self.dp_world_size,
                    "rank": self.dp_rank,
                    "num_workers": self.config.train.train_policy.dataloader_num_workers,
                    "config": self.config,
                    "sampler": val_sampler,
                    "batch_size": self.config.validation.batch_size
                    or self.config.train.train_batch_per_replica,
                    "drop_last": False,
                }
                # Filter kwargs to only those the function accepts
                filtered = {k: v for k, v in kwargs.items() if k in sig.parameters}
                val_batch_sampler = val_batch_sampler(**filtered)
            self.val_data_loader = DataLoader(
                val_dataset,
                num_workers=self.config.train.train_policy.dataloader_num_workers,
                prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                batch_sampler=val_batch_sampler,
                collate_fn=collate_fn,
            )
        else:
            self.val_data_loader = DataLoader(
                val_dataset,
                batch_size=self.config.validation.batch_size
                or self.config.train.train_batch_per_replica,
                num_workers=self.config.train.train_policy.dataloader_num_workers,
                prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                sampler=val_sampler,
                collate_fn=collate_fn,
                drop_last=self.config.train.train_policy.dataloader_drop_last,
            )

        steps_by_dataset = (
            self.ckpt_total_steps
            if self.ckpt_total_steps > 0
            else len(self.train_data_loader) * self.epoch
        )

        if self.enable_dp_load_balancing:
            self.total_steps = self.config.train.max_num_steps
            logger.info(
                f"Total training steps set to max_num_steps={self.config.train.max_num_steps} for load-balanced dynamic batching"
            )
        elif self.config.train.max_num_steps is not None:
            self.total_steps = min(steps_by_dataset, self.config.train.max_num_steps)
        else:
            self.total_steps = steps_by_dataset

        # Calculate the step interval to save the checkpoint
        if self.config.train.ckpt.save_freq_in_epoch > 0:
            # Use save_freq_in_epoch to calculate the save frequency in priority
            self._save_freq = (
                self.config.train.ckpt.save_freq_in_epoch * len(self.train_data_loader)
            ) // self.dp_world_size
            logger.info(
                f"Checkpoint will be saved every {self._save_freq} steps, which is approximately every `train.ckpt.save_freq_in_epoch` {self.config.train.ckpt.save_freq_in_epoch} epochs. `train.ckpt.save_freq` will be ignored."
            )
        else:
            self._save_freq = self.config.train.ckpt.save_freq

    def validate(self, current_epoch: int, is_last_step: bool = False):
        if not self.config.validation.enable:
            return None

        if (
            (self.train_step == 0 and self.config.validation.val_before_train)
            or (
                self.train_step != 0
                and self.train_step % self.config.validation.freq == 0
            )
            or is_last_step
        ):
            pass
        else:
            return None

        # Call pre_validation_hook
        if self.pre_validation_hook is not None:
            report_data = {
                "current_epoch": current_epoch,
                "is_last_step": is_last_step,
            }
            self.pre_validation_hook(self, report_data=report_data)

        # validation
        logger.info(f"Validation at step {self.train_step}/{self.total_steps}...")
        val_total_loss = 0.0

        for batch_index, val_global_batch in enumerate(
            tqdm(
                self.get_batch_from_dataloader(self.val_data_loader), desc="Validation"
            )
        ):
            # Call pre_per_step_validation_hook
            if self.pre_per_step_validation_hook is not None:
                report_data = {
                    "current_epoch": current_epoch,
                    "batch_index": batch_index,
                }
                self.pre_per_step_validation_hook(self, report_data=report_data)

            val_score = self.trainer.step_validation(
                val_global_batch, self.train_step, self.total_steps
            )

            # Call post_per_step_validation_hook
            if self.post_per_step_validation_hook is not None:
                report_data = {
                    "current_epoch": current_epoch,
                    "batch_index": batch_index,
                    "val_score": val_score,
                }
                self.post_per_step_validation_hook(self, report_data=report_data)

            val_total_loss += val_score

        val_avg_loss = val_total_loss / len(self.val_data_loader.dataset)
        logger.info(
            f"[SFT] Validation loss: {val_avg_loss} for train step {self.train_step}/{self.total_steps}, epoch {current_epoch}"
        )

        # Call post_validation_hook
        if self.post_validation_hook is not None:
            report_data = {
                "current_epoch": current_epoch,
                "val_avg_loss": val_avg_loss,
            }
            self.post_validation_hook(self, report_data=report_data)

        # Call custom logger functions
        report_data = {
            "val/cur_epoch": current_epoch,
            "val/avg_loss": val_avg_loss,
            "val/train_epochs": self.epoch,
            "val/total_steps": self.total_steps,  # This total_steps is for training
            "val/train_step": self.train_step,
        }

        if util.is_master_rank(self.parallel_dims, self.global_rank):
            if "wandb" in self.config.logging.logger and is_wandb_available():
                log_wandb(
                    data=report_data,
                    step=self.train_step,
                )
            for custom_logger_fn in self.custom_logger_fns:
                try:
                    custom_logger_fn(report_data, self.train_step)
                except Exception as e:
                    logger.warning(f"[SFT] Error calling custom logger function: {e}")

        return val_avg_loss

    def collect_broadcast_info(self, item):
        if isinstance(item, list):
            return [self.collect_broadcast_info(x) for x in item]
        elif isinstance(item, dict):
            return {k: self.collect_broadcast_info(v) for k, v in item.items()}
        elif isinstance(item, torch.Tensor):
            return {"tensor_to_be_recv": (item.shape, item.dtype, item.device)}
        elif isinstance(item, (int, float, str)) or item is None:
            return item
        else:
            raise ValueError(
                f"Unsupported item type for broadcast info collection: {type(item)}"
            )

    def recv_tensor_from_info(self, item):
        if isinstance(item, list):
            return [self.recv_tensor_from_info(x) for x in item]
        elif isinstance(item, dict):
            if (
                "tensor_to_be_recv" in item
                and len(item) == 1
                and isinstance(item["tensor_to_be_recv"], tuple)
                and len(item["tensor_to_be_recv"]) == 3
            ):
                placeholder = torch.empty(
                    *item["tensor_to_be_recv"][0],
                    dtype=item["tensor_to_be_recv"][1],
                    device=self.device,
                )
                dist.broadcast(
                    placeholder,
                    group_src=0,
                    group=self.parallel_dims.mesh["pp_cp_tp"].get_group(),
                )
                return placeholder.to(item["tensor_to_be_recv"][2])
            return {k: self.recv_tensor_from_info(v) for k, v in sorted(item.items())}
        elif isinstance(item, (int, float, str)) or item is None:
            return item
        else:
            raise ValueError(
                f"Unsupported item type for broadcast info collection: {type(item)}"
            )

    def send_tensor_from_info(self, item):
        if isinstance(item, list):
            for x in item:
                self.send_tensor_from_info(x)
        elif isinstance(item, dict):
            for _, v in sorted(item.items()):
                self.send_tensor_from_info(v)
        elif isinstance(item, torch.Tensor):
            dist.broadcast(
                item.to(self.device),
                group_src=0,
                group=self.parallel_dims.mesh["pp_cp_tp"].get_group(),
            )
        elif isinstance(item, (int, float, str)) or item is None:
            pass
        else:
            raise ValueError(
                f"Unsupported item type for broadcast info collection: {type(item)}"
            )

    def get_batch_from_dataloader(self, data_loader):
        # self.iter = iter(self.train_data_loader)
        if self.config.train.train_policy.dataloader_broadcast and (
            self.parallel_dims.pp_enabled
            or self.parallel_dims.cp_enabled
            or self.parallel_dims.tp_enabled
        ):
            # Only the first rank of the pp/cp/tp mesh will read from the dataloader and broadcast to other ranks, to avoid redundant dataloader workers and potential data mismatches across ranks.
            if self.parallel_dims.mesh["pp_cp_tp"].get_local_rank() == 0:
                for batch in data_loader:
                    info = self.collect_broadcast_info(batch)
                    dist_utils.broadcast_object_cpu(
                        info,
                        group=self.parallel_dims.mesh["pp_cp_tp"].get_group(),
                        group_src=0,
                    )
                    self.send_tensor_from_info(batch)
                    yield batch
                dist_utils.broadcast_object_cpu(
                    "end",
                    group=self.parallel_dims.mesh["pp_cp_tp"].get_group(),
                    group_src=0,
                )
            else:
                while True:
                    info = dist_utils.broadcast_object_cpu(
                        None,
                        group=self.parallel_dims.mesh["pp_cp_tp"].get_group(),
                        group_src=0,
                    )
                    if info == "end":
                        break
                    batch = self.recv_tensor_from_info(info)
                    # Do check to verify that the broadcast batch matches the dataloader batch for non-zero ranks, to ensure correctness of broadcasting logic.
                    # ref = next(self.iter)
                    # assert util.recursive_check_equal(
                    #     batch, ref
                    # ), f"Broadcast batch does not match dataloader batch for non-zero rank {batch} {ref}"
                    yield batch
        else:
            # If dataloader_broadcast is enabled but no relevant parallelism is enabled, just yield batches without broadcasting
            for batch in data_loader:
                yield batch

    def main_loop(self):
        self.profiler.start()
        pp_last_stage = False

        if self.parallel_dims.pp_enabled:
            pp_last_stage = (
                self.parallel_dims.pp_coord[0] == self.parallel_dims.pp_coord[1] - 1
            )

        cur_epoch = self.start_epoch
        if self.enable_dp_load_balancing:
            logger.info(
                f"Epoch set to {cur_epoch + 1} for load-balanced dynamic batching"
            )
            self.epoch = cur_epoch + 1
        stop_training = False
        # For pre-train validation
        val_avg_loss = self.validate(current_epoch=cur_epoch, is_last_step=False)
        for _ in range(self.start_epoch, self.epoch):
            if hasattr(self.train_sampler, "set_epoch"):
                self.train_sampler.set_epoch(cur_epoch)
            if hasattr(self.train_batch_sampler, "set_epoch"):
                self.train_batch_sampler.set_epoch(cur_epoch)
            if hasattr(self.train_data_loader, "dataset") and hasattr(
                self.train_data_loader.dataset, "set_epoch"
            ):
                self.train_data_loader.dataset.set_epoch(cur_epoch)
            logger.info(f"Training epoch {cur_epoch + 1}/{self.epoch}")

            data_arrival_event = torch.cuda.Event(enable_timing=True)
            data_arrival_event.record()
            # global_batch is a list of items from `datapacker.sft_process_sample()`
            for global_batch in self.get_batch_from_dataloader(self.train_data_loader):
                # if [profiler.enable_nsys] is true, cudaProfilerStart() / cudaProfilerStop() are used to trigger nsys capture
                # settings from [profiler.sub_profiler_config] are reused
                if (
                    self.config.profiler.enable_nsys
                    and self.profiler.global_rank in self.profiler.rank_filter
                ):
                    if (
                        self.train_step
                        == self.profiler.wait_steps + self.profiler.warmup_steps
                    ):
                        torch.cuda.cudart().cudaProfilerStart()
                    elif (
                        self.train_step
                        == self.profiler.wait_steps
                        + self.profiler.warmup_steps
                        + self.profiler.active_steps
                    ):
                        torch.cuda.cudart().cudaProfilerStop()

                report_data = self.trainer.step_training(
                    global_batch=global_batch,
                    total_steps=self.total_steps,
                    train_step=self.train_step,
                    save_freq=self._save_freq,
                    data_arrival_event=data_arrival_event,
                )
                report_data["train/epoch"] = cur_epoch

                self.train_step += 1

                if report_data and util.is_master_rank(
                    self.parallel_dims, self.global_rank
                ):
                    if "wandb" in self.config.logging.logger and is_wandb_available():
                        log_wandb(
                            data=report_data,
                            step=self.train_step,
                        )
                    if "console" in self.config.logging.logger:
                        log_info = f"Step: {self.train_step}/{self.total_steps}, Loss: {report_data['train/loss_avg']:.5f}, Grad norm: {report_data['optimizer/grad_norm']:.5f}, Iteration time: {report_data['train/iteration_time']:.2f}s"
                        # Append learning rate of each optimizer to log_info
                        for key in report_data:
                            if key.startswith("optimizer/lr_"):
                                log_info += f", {key}: {report_data[key]:.5e}"

                        logger.info(log_info)

                    for custom_logger_fn in self.custom_logger_fns:
                        try:
                            custom_logger_fn(report_data, self.train_step)
                        except Exception as e:
                            logger.warning(
                                f"[SFT] Error calling custom logger function: {e}"
                            )

                if self.train_step >= self.total_steps:
                    stop_training = True
                    break  # break outer epoch loop

                val_avg_loss = self.validate(
                    current_epoch=cur_epoch, is_last_step=False
                )

                self.trainer.checkpointing(
                    total_steps=self.total_steps,
                    train_step=self.train_step,
                    save_freq=self._save_freq,
                    pp_last_stage=False,
                    is_last_step=False,
                    val_score=val_avg_loss,
                )

                self.profiler.step()
                data_arrival_event = torch.cuda.Event(enable_timing=True)
                data_arrival_event.record()

                if self.signal_handler is not None and any(
                    self.signal_handler.signals_received()
                ):
                    # If processes was killed by signal trapped, stop training and finish the main_loop.
                    stop_training = True
                    self.signal_handler.release()
                    break

            if stop_training:
                break
            cur_epoch += 1

        # Finally: validation and save checkpoint
        val_avg_loss = self.validate(current_epoch=cur_epoch, is_last_step=True)
        self.trainer.checkpointing(
            total_steps=self.total_steps,
            train_step=self.train_step,
            save_freq=self._save_freq,
            is_last_step=True,
            pp_last_stage=pp_last_stage,
            val_score=val_avg_loss,
        )

    def handle_shutdown(self):
        # handle the ckpt saving
        logger.info("Handling shutdown...")
        if (
            hasattr(self.trainer, "upload_thread")
            and self.trainer.upload_thread is not None
        ):
            logger.info("[Policy] Waiting for upload thread to finish...")
            self.trainer.upload_thread.join()
            logger.info("[Policy] Upload thread finished.")
            self.trainer.upload_thread = None

        self.unregister_from_controller()

    def destroy_worker(self):
        destroy_distributed()
        logger.info("[Policy] Process group destroyed.")
