src/model_training/training_modules/training_loop.py

import logging
import os
from typing import Optional

import torch
import torch.distributed as dist
from accelerate import Accelerator
from tqdm import tqdm

from src.model_training.transformers_compat import patch_transformers_hybrid_cache

patch_transformers_hybrid_cache()

from diffsynth.trainers.utils import DiffusionTrainingModule
from src.model_training.fov_retrieval import FOVMemoryRetriever
from src.model_training.fov_retrieval import retrieve_context_frames_advanced, retrieve_fov_context_frames
from src.model_training.training_modules.model_logger import ModelLogger

logger = logging.getLogger(__name__)


def launch_training_task(
    dataset: torch.utils.data.Dataset,
    model: DiffusionTrainingModule,
    model_logger: ModelLogger,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    num_epochs: int = 1,
    gradient_accumulation_steps: int = 1,
    per_device_train_batch_size: int = 1,
    seed: int = 42,
    spike_threshold: float = 5.0,
    resume_step_count: int = 0,
    enable_fov_retrieval: bool = False,
    retrieval_method: str = "fov",  # fov | latent_sim
    latent_retrieval_dir: Optional[str] = None,
    dataset_base_path: str = None,
    fov_retriever: Optional[FOVMemoryRetriever] = None,
    context_memory_frames: int = 5,
    prev_chunk_frames: int = 81,
    fov_top_k: int = 4,  # Number of overlap frames to retrieve. GT frame 0 will be added automatically.
    use_rt_relative: bool = False,  # Experiment 1_4_2: Use RT relative conversion (aligned with Context-as-Memory)
    strict_overlap_context: bool = False,
    dataset_repeat: int = 1,  # Add dataset_repeat parameter for step calculation
    use_camera_encoder: bool = False,  # exp1_4_3: use CameraEncoder (action_mlp unused -> need find_unused_parameters)
    num_workers: int = 0,  # DataLoader workers: 0=main process, >0=parallel preload (recommend 4 for video)
    context_source: str = "fov",
    max_train_steps: int = 0,
    progress_total_steps: int = 0,
):
    prev_chunk_frames = int(prev_chunk_frames)
    # VideoDataset can return None when file loading fails; keep distributed batches aligned.
    def collate_fn(batch):
        valid_batch = [item for item in batch if item is not None]
        return valid_batch or None
    
    num_workers = max(0, int(num_workers))
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=per_device_train_batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=num_workers,
        drop_last=True,
        persistent_workers=(num_workers > 0),
        pin_memory=(num_workers > 0 and torch.cuda.is_available()),
    )
    if num_workers > 0:
        logger.info(f"[DataLoader] num_workers={num_workers}, persistent_workers=True, pin_memory={torch.cuda.is_available()} (data preload parallel to GPU)")
    
    timeout_seconds = int(os.environ.get('TORCH_DISTRIBUTED_DEFAULT_TIMEOUT', 2400))
    os.environ['TORCH_DISTRIBUTED_DEFAULT_TIMEOUT'] = str(timeout_seconds)
    logger.info(f"[Timeout Config] Setting TORCH_DISTRIBUTED_DEFAULT_TIMEOUT={timeout_seconds} seconds ({timeout_seconds/60:.1f} minutes)")
    
    # Conditional context paths can leave parameters unused on some iterations.
    need_find_unused = bool(use_camera_encoder) or model_logger.context_drop_prob > 0.0
    if need_find_unused:
        from accelerate import DistributedDataParallelKwargs
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
        accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, kwargs_handlers=[ddp_kwargs])
        logger.info("[DDP] find_unused_parameters=True (conditional modules / context_drop_prob enabled)")
    else:
        accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
    
    if model_logger.enable_video_sampling and model_logger.total_steps is not None:
        dataset_size = len(dataset)
        num_processes = accelerator.num_processes
        effective_dataset_size = dataset_size * dataset_repeat
        total_steps_per_gpu = (effective_dataset_size * num_epochs) // (gradient_accumulation_steps * num_processes * per_device_train_batch_size)
        total_steps_global = total_steps_per_gpu * num_processes
        model_logger.total_steps = total_steps_global
        
        if accelerator.is_main_process:
            logger.info("="*80)
            logger.info("[Step Calculation] Corrected total_steps after accelerator.init")
            logger.info("="*80)
            logger.info(f"  Dataset size (unique samples): {dataset_size}")
            logger.info(f"  Dataset repeat: {dataset_repeat}")
            logger.info(f"  Effective dataset size: {effective_dataset_size} (unique * repeat)")
            logger.info(f"  Number of epochs: {num_epochs}")
            logger.info(f"  Number of GPUs: {num_processes}")
            logger.info(f"  Gradient accumulation steps: {gradient_accumulation_steps}")
            logger.info(f"  Per-device batch size: {per_device_train_batch_size}")
            logger.info(f"  Total samples to process: {effective_dataset_size * num_epochs}")
            logger.info(f"  Steps per GPU: ~{total_steps_per_gpu}")
            logger.info(f"  Total steps (global): {total_steps_global}")
            logger.info("")
            logger.info(f"  ✓ Each GPU will process ~{total_steps_per_gpu} steps")
            logger.info(f"  ✓ This ensures traversal of all {effective_dataset_size} samples")
            logger.info(f"    ({dataset_size} unique samples × {dataset_repeat} repeats)")
            logger.info(f"  ✓ Over {num_epochs} epoch(s)")
            logger.info("="*80)
    
    step = resume_step_count
    traj_loss = 0.0
    if resume_step_count > 0:
        adaptation_steps = max(200, resume_step_count // 100)
        spike_detection_start_step = resume_step_count + adaptation_steps
        logger.info(f"Resuming from step {resume_step_count}, spike detection will start at step {spike_detection_start_step} (after {adaptation_steps} adaptation steps)")
    else:
        spike_detection_start_step = 100

    for epoch_id in range(num_epochs):
        epoch_seed = seed + epoch_id
        torch.manual_seed(epoch_seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(epoch_seed)
            torch.cuda.manual_seed_all(epoch_seed)
        
        if resume_step_count > 0 and epoch_id == 0:
            estimated_skip = resume_step_count // gradient_accumulation_steps
            if estimated_skip > 0:
                logger.info(f"Skipping {estimated_skip} data samples to resume from step {resume_step_count}...")
                dataloader_iter = iter(dataloader)
                for _ in tqdm(range(estimated_skip), desc="Skipping data", unit="samples", leave=False):
                    try:
                        next(dataloader_iter)
                    except StopIteration:
                        break
                dataloader = dataloader_iter
                logger.info(f"Successfully skipped {estimated_skip} data samples, resuming training...")
        
        # Track consecutive None data to detect if we're stuck in a loop
        consecutive_none_count = 0
        max_consecutive_none = 100  # If we get 100 consecutive None values, something is wrong
        
        progress_total = int(progress_total_steps)
        if progress_total <= 0:
            progress_total = len(dataloader)
        progress_bar = tqdm(
            dataloader,
            total=progress_total,
            initial=resume_step_count if progress_total_steps else 0,
            desc="Training steps",
            unit="step",
        )
        for data_idx, data in enumerate(progress_bar):
            # Handle None data (can happen if all files in batch fail to load)
            if data is None:
                consecutive_none_count += 1
                if consecutive_none_count >= max_consecutive_none:
                    logger.error(f"Received {max_consecutive_none} consecutive None data samples. This suggests a serious dataset issue. Stopping training.")
                    raise ValueError(f"Too many consecutive None data samples ({max_consecutive_none}). Check dataset files.")
                
                # Log warning but continue (will skip this step)
                if consecutive_none_count <= 10 or consecutive_none_count % 10 == 0:
                    logger.warning(f"Received None data at index {data_idx} (consecutive: {consecutive_none_count}). This may indicate missing or corrupted files. Skipping...")
                
                # Still increment step to keep step_count synchronized
                step += 1
                dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False)
                model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples)
                continue
            
            # Reset consecutive None counter when we get valid data
            consecutive_none_count = 0
            
            # Normalize to list of samples for batch processing (per_device_train_batch_size > 1)
            samples = data if isinstance(data, list) else [data]
            
            # Simplified context-based retrieval OR replay/prev_chunk_tail (aligned with multichunk eval)
            context_retrieval_success = True  # Set False if any sample fails (for strict mode)
            _umodel = accelerator.unwrap_model(model)
            _cm_frames = int(_umodel.context_memory_frames)
            _cs = context_source.strip().lower()
            if _cs not in ("fov", "replay", "prev_chunk_tail"):
                _cs = "fov"

            if _cs == "replay" and dataset_base_path:
                from src.model_training.multichunk_sample_utils import (
                    replay_context_actions_from_segment_actions,
                    replay_context_global_indices,
                    synthetic_replay_context_from_segment,
                )
                for d in samples:
                    vf = d.get("video") or []
                    n_seg = min(int(prev_chunk_frames), len(vf)) if vf else 0
                    ctx_pil = synthetic_replay_context_from_segment(vf, n_seg, _cm_frames) if n_seg > 0 else None
                    if not ctx_pil:
                        context_retrieval_success = False
                        break
                    d["context_frames"] = ctx_pil
                    d["context_source"] = "replay_synthetic"
                    acts = d.get("actions")
                    if isinstance(acts, list) and len(acts) >= n_seg:
                        ra = replay_context_actions_from_segment_actions(acts[:n_seg], n_seg, _cm_frames)
                        if ra is not None:
                            d["context_actions"] = ra
                    sf = int(d.get("start_frame", 0) or 0)
                    idxs = replay_context_global_indices(n_seg, _cm_frames)
                    d["context_frame_indices"] = [sf + int(i) for i in idxs]

            elif _cs == "prev_chunk_tail" and dataset_base_path:
                from src.model_training.multichunk_sample_utils import load_prev_chunk_tail_from_disk, load_prev_chunk_tail_rt_actions
                _ctx_pos = os.environ.get("CONTEXT_POSITION", "suffix").strip().lower()
                _nearest_first = (_ctx_pos == "suffix")
                for d in samples:
                    sf = int(d.get("start_frame", 0) or 0)
                    vn = d.get("video_name", "")
                    pil_list, idxs = load_prev_chunk_tail_from_disk(
                        dataset_base_path, str(vn), sf, _cm_frames, nearest_first=_nearest_first
                    )
                    if not pil_list:
                        context_retrieval_success = False
                        break
                    d["context_frames"] = pil_list
                    d["context_frame_indices"] = list(idxs) if idxs else []
                    d["context_source"] = "prev_chunk_tail"
                    ra, _ = load_prev_chunk_tail_rt_actions(
                        dataset_base_path,
                        str(vn),
                        sf,
                        _cm_frames,
                        use_rt_relative=use_rt_relative,
                        nearest_first=_nearest_first,
                    )
                    if ra:
                        d["context_actions"] = ra

            elif enable_fov_retrieval and dataset_base_path:
                for d in samples:
                    if retrieval_method == "latent_sim":
                        (
                            context_frames,
                            context_actions,
                            context_indices,
                            ref_frame_idx,
                            video_name,
                            source,
                        ) = retrieve_context_frames_advanced(
                            data=d,
                            dataset_base_path=dataset_base_path,
                            top_k=fov_top_k,
                            drop_overlap_probability=0.1,
                            use_rt_relative=use_rt_relative,
                            retrieval_method="latent_sim",
                            latent_retrieval_dir=latent_retrieval_dir,
                            strict_overlap_labels=strict_overlap_context,
                        )
                    else:
                        (
                            context_frames,
                            context_actions,
                            context_indices,
                            ref_frame_idx,
                            video_name,
                            source,
                        ) = retrieve_fov_context_frames(
                            data=d,
                            dataset_base_path=dataset_base_path,
                            fov_retriever=fov_retriever,  # unused in simplified retrieval, kept for compat
                            top_k=fov_top_k,  # fov_top_k is number of overlap frames (4), GT frame 0 will be added automatically
                            use_precomputed_overlaps=True,
                            strict_overlap_labels=strict_overlap_context,
                            allow_realtime_fallback=(not strict_overlap_context),
                            allow_segment_fallback=(not strict_overlap_context),
                        )

                    if context_frames and len(context_frames) > 0:
                        # Use retrieved frames as context
                        d["context_frames"] = context_frames
                        if context_actions:
                            d["context_actions"] = context_actions
                        # Store retrieval metadata for visualization/debugging
                        d["context_frame_indices"] = context_indices
                        d["context_ref_frame_idx"] = ref_frame_idx
                        d["context_video_name"] = video_name
                        d["context_source"] = source
                    else:
                        context_retrieval_success = False
                        break
                        
            # Strict mode: if we require context but retrieval failed, skip this step
            _need_ctx_strict = (
                strict_overlap_context
                and (not context_retrieval_success)
                and (
                    enable_fov_retrieval
                    or context_source.strip().lower() in ("replay", "prev_chunk_tail")
                )
            )
            if _need_ctx_strict:
                if step % 50 == 0 and accelerator.is_main_process:
                    logger.warning(f"[CONTEXT][STRICT] No context at step={step}, skipping this training sample.")
                step += 1
                dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False)
                model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples)
                continue

            with accelerator.accumulate(model):
                optimizer.zero_grad()
                # One forward over full batch: data is list of B dicts when per_device_train_batch_size > 1
                # Main loss on current batch
                loss = model(data)

                step += 1
                if traj_loss == 0.0:
                    traj_loss = loss.item()
                else:
                    alpha = 0.01
                    traj_loss = (1 - alpha) * traj_loss + alpha * loss.item()
                
                if step >= spike_detection_start_step and traj_loss > 0:
                    relative_loss = loss.item() / traj_loss
                    if resume_step_count > 0 and step < resume_step_count + 500:
                        effective_threshold = spike_threshold * 1.5
                    else:
                        effective_threshold = spike_threshold
                    
                    should_skip = relative_loss > effective_threshold
                    # Keep the skip decision identical across ranks to avoid DDP hangs.
                    skip_t = torch.tensor(1.0 if should_skip else 0.0, device=accelerator.device, dtype=torch.float32)
                    if accelerator.num_processes > 1:
                        dist.all_reduce(skip_t, op=dist.ReduceOp.MAX)
                    skip_global = skip_t.item() > 0.5
                    
                    if skip_global:
                        if accelerator.is_main_process:
                            logger.warning(f"Spike detected at step {step} (loss={loss.item():.4f}, traj_loss={traj_loss:.4f}, ratio={relative_loss:.2f}), sync skip across all ranks")
                        dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False)
                        model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples)
                        del loss
                        torch.cuda.empty_cache()
                        continue
                accelerator.backward(loss)
                optimizer.step()
                model_logger.on_step_end(loss, accelerator, model, current_batch=samples)
                scheduler.step()

                if max_train_steps and step >= max_train_steps:
                    if progress_total_steps:
                        progress_bar.n = min(step, progress_bar.total) if progress_bar.total is not None else step
                        progress_bar.refresh()
                    if accelerator.is_main_process:
                        logger.info(f"[TRAIN] Reached max_train_steps={max_train_steps}; stopping without epoch checkpoint.")
                    accelerator.wait_for_everyone()
                    return
            
        model_logger.on_epoch_end(accelerator, model, epoch_id)