/rdagent/scenarios/finetune/train/eval.py

import json
from typing import Any, Dict, List, Optional

from rdagent.app.finetune.llm.conf import FT_RD_SETTING
from rdagent.components.coder.CoSTEER.evaluators import (
    CoSTEEREvaluator,
    CoSTEERSingleFeedback,
)
from rdagent.components.coder.finetune.conf import (
    FT_DATA_FILE_NAME,
    FT_DATA_SCRIPT_NAME,
    FT_YAML_FILE_NAME,
    clear_workspace,
    get_data_processing_cache_key,
    get_data_processing_env,
    get_ft_env,
    get_workspace_prefix,
    inject_data_stats,
)
from rdagent.components.coder.finetune.exp import FTTask
from rdagent.components.coder.finetune.unified_validator import LLMConfigValidator
from rdagent.core.evolving_framework import QueriedKnowledge
from rdagent.core.experiment import FBWorkspace
from rdagent.log import rdagent_logger as logger
from rdagent.scenarios.finetune.benchmark import get_benchmark_ranges, run_benchmark
from rdagent.utils.agent.tpl import T
from rdagent.utils.agent.workflow import build_cls_from_json_with_retry


def extract_loss_history(output_path) -> Dict[str, List[Dict[str, Any]]]:
    """
    Extract training and evaluation loss history from LlamaFactory's trainer_state.json.

    Args:
        output_path: Path to the training output directory

    Returns:
        Dict with 'train' and 'eval' keys, each containing a list of loss entries.
    """
    trainer_state_path = output_path / "trainer_state.json"
    result = {"train": [], "eval": []}

    if not trainer_state_path.exists():
        logger.warning(f"trainer_state.json not found at {trainer_state_path}")
        return result

    try:
        with open(trainer_state_path) as f:
            trainer_state = json.load(f)

        log_history = trainer_state.get("log_history", [])
        for entry in log_history:
            if "loss" in entry:
                result["train"].append(
                    {
                        "step": entry.get("step"),
                        "epoch": entry.get("epoch"),
                        "loss": entry.get("loss"),
                    }
                )
            if "eval_loss" in entry:
                result["eval"].append(
                    {
                        "step": entry.get("step"),
                        "epoch": entry.get("epoch"),
                        "eval_loss": entry.get("eval_loss"),
                    }
                )

        logger.info(f"Extracted {len(result['train'])} train + {len(result['eval'])} eval entries")

    except (json.JSONDecodeError, OSError) as e:
        logger.warning(f"Failed to parse trainer_state.json: {e}")

    return result


class FTRunnerEvaluator(CoSTEEREvaluator):
    """LLM Fine-tuning specific evaluator that uses LLM Docker environment."""

    def evaluate(
        self,
        target_task: FTTask,
        implementation: FBWorkspace,
        gt_implementation: FBWorkspace,
        queried_knowledge: Optional[QueriedKnowledge] = None,
        **kwargs,
    ) -> CoSTEERSingleFeedback:
        """Evaluate LLM fine-tuning implementation using dedicated LLM environment.

        This evaluator performs three stages:
        0. Clean workspace (remove old training outputs)
        1. Full data processing (without --debug flag) to generate complete data.json
        2. Full training with the complete dataset
        """

        # Check if FT_YAML_FILE_NAME exists
        if FT_YAML_FILE_NAME not in implementation.file_dict:
            fb = CoSTEERSingleFeedback(
                execution=f"No {FT_YAML_FILE_NAME} found in workspace",
                return_checking="Config file missing",
                code="No valid configuration file",
                final_decision=False,
            )
            implementation.feedback = fb
            logger.log_object(fb, tag="evaluator_feedback.FTRunnerEvaluator")
            return fb

        # Use LLM-specific environment with appropriate timeout for training
        env = get_ft_env(operation="full_training")

        # ========== Stage 0: Clean Workspace ==========
        # Clean old training outputs before data processing and training
        clear_workspace(implementation, env)

        # ========== Stage 1: Full Data Processing ==========
        # Execute data processing WITHOUT --debug flag to generate complete data.json
        data_result = self._run_full_data_processing(implementation)
        data_stdout = data_result.stdout or ""

        if data_result.exit_code != 0:
            # Data processing failed, return feedback to enter next loop
            logger.error(f"Full data processing failed with exit_code={data_result.exit_code}")
            return self._generate_llm_feedback(
                target_task=target_task,
                implementation=implementation,
                raw_stdout=data_stdout,
                exit_code=data_result.exit_code,
                model_files_exist=False,
                benchmark_result=None,
                loss_history=None,
                failed_stage="data_processing",
            )

        logger.info("Full data processing completed successfully")

        # Update data_stats.json with full dataset statistics
        # This ensures feedback sees the correct sample count, not debug mode count
        data_json_path = implementation.workspace_path / FT_DATA_FILE_NAME
        if data_json_path.exists():
            with open(data_json_path, "r", encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, list) and len(data) > 0:
                inject_data_stats(implementation, data, data_stdout)

        # ========== Stage 2: Full Training ==========

        # Execute LlamaFactory training
        train_result = implementation.run(
            env=env,
            entry=f"llamafactory-cli train {FT_YAML_FILE_NAME}",
        )
        # Combine data processing and training stdout for comprehensive feedback
        combined_stdout = (
            f"=== DATA PROCESSING OUTPUT ===\n{data_stdout}\n\n=== TRAINING OUTPUT ===\n{train_result.stdout or ''}"
        )
        implementation.running_info.running_time = train_result.running_time
        # NOTE: Docker execution is logged by FTWorkspace.run() automatically

        # Simple success check: exit code
        training_success = train_result.exit_code == 0

        # Check for model output files
        workspace_path = implementation.workspace_path
        output_path = workspace_path / "output"
        model_output_files = (
            list(output_path.glob("*.safetensors"))
            + list(output_path.glob("*.bin"))
            + list(output_path.glob("adapter_*"))
            if output_path.exists()
            else []
        )

        # Early return if training failed
        if not training_success or len(model_output_files) == 0:
            return self._generate_llm_feedback(
                target_task=target_task,
                implementation=implementation,
                raw_stdout=combined_stdout,
                exit_code=train_result.exit_code,
                model_files_exist=len(model_output_files) > 0,
                benchmark_result=None,
                loss_history=None,
                failed_stage="training",
            )

        # Extract loss history from training output
        loss_history = extract_loss_history(output_path)

        val_range, test_range = get_benchmark_ranges()

        # Validation set - used for SOTA judgment, visible to agent
        validation_result = run_benchmark(
            workspace_path=str(workspace_path),
            model_path=output_path,
            model_name=target_task.base_model,
            benchmark_name=target_task.benchmark,
            gpu_count=self.scen.gpu_count,
            test_range=val_range,
            result_subdir="validation",
        )

        # Test set - only for frontend display, not visible to agent
        test_result = run_benchmark(
            workspace_path=str(workspace_path),
            model_path=output_path,
            model_name=target_task.base_model,
            benchmark_name=target_task.benchmark,
            gpu_count=self.scen.gpu_count,
            test_range=test_range,
            result_subdir="test",
        )

        # Build comprehensive result with training metrics and benchmark results
        # Note: "benchmark" is for agent (SOTA judgment), "benchmark_test" is for frontend only
        train_history = loss_history.get("train", []) if loss_history else []
        implementation.running_info.result = {
            "benchmark": validation_result,  # Agent visible - used for SOTA judgment
            "benchmark_test": test_result,  # Agent invisible - frontend display only
            "training_metrics": {
                "loss_history": loss_history,
                "final_loss": train_history[-1]["loss"] if train_history else None,
                "initial_loss": train_history[0]["loss"] if train_history else None,
            },
        }
        benchmark_result = validation_result  # For backward compatibility with feedback

        # Call LLM for feedback analysis - LLM will determine final_decision
        return self._generate_llm_feedback(
            target_task=target_task,
            implementation=implementation,
            raw_stdout=combined_stdout,
            exit_code=train_result.exit_code,
            model_files_exist=len(model_output_files) > 0,
            benchmark_result=benchmark_result,
            loss_history=loss_history,
        )

    def _generate_llm_feedback(
        self,
        target_task: FTTask,
        implementation: FBWorkspace,
        raw_stdout: str,
        exit_code: int,
        model_files_exist: bool,
        benchmark_result: Optional[Dict] = None,
        loss_history: Optional[Dict[str, List[Dict]]] = None,
        failed_stage: Optional[str] = None,
    ) -> CoSTEERSingleFeedback:
        """Generate LLM-based feedback for runner evaluation.

        LLM will determine final_decision based on all provided information.

        Args:
            failed_stage: Which stage failed - "data_processing" or "training"
        """
        # Parse execution log to extract structured info (reuse unified_validator's method)
        # Reduces ~36k tokens to ~500 tokens by extracting: status, errors, metrics, warnings
        parsed_stdout = LLMConfigValidator()._parse_execution_log(raw_stdout, exit_code, failed_stage)

        # Get timeout config for the failed stage
        timeout_seconds = None
        if failed_stage == "data_processing":
            timeout_seconds = FT_RD_SETTING.data_processing_timeout
        elif failed_stage == "training":
            timeout_seconds = FT_RD_SETTING.full_timeout

        # Pass loss_history directly (simpler and preserves full information)
        # Sample train entries if too many to avoid token bloat
        if loss_history and len(loss_history.get("train", [])) > 60:
            loss_history["train"] = loss_history["train"][:30] + loss_history["train"][-30:]

        system_prompt = T("rdagent.components.coder.finetune.prompts:runner_eval.system").r()
        user_prompt = T("rdagent.components.coder.finetune.prompts:runner_eval.user").r(
            task_desc=target_task.get_task_information(),
            config_yaml=implementation.file_dict.get(FT_YAML_FILE_NAME, ""),
            exit_code=exit_code,
            model_files_status="Found" if model_files_exist else "Not found",
            stdout=parsed_stdout,  # Structured JSON instead of raw truncated log
            benchmark_result=(
                json.dumps(benchmark_result, indent=2) if benchmark_result else "N/A (not executed or failed)"
            ),
            loss_history=(
                json.dumps(loss_history, indent=2)
                if (loss_history and (loss_history.get("train") or loss_history.get("eval")))
                else "N/A"
            ),
            failed_stage=failed_stage,
            timeout_seconds=timeout_seconds,
        )

        feedback = build_cls_from_json_with_retry(
            CoSTEERSingleFeedback,
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            init_kwargs_update_func=CoSTEERSingleFeedback.val_and_update_init_dict,
        )
        feedback.raw_execution = raw_stdout
        implementation.feedback = feedback
        logger.log_object(feedback, tag="evaluator_feedback.FTRunnerEvaluator")
        return feedback

    def _run_full_data_processing(self, implementation: FBWorkspace):
        """Execute full data processing (without --debug flag) to generate complete data.json.

        This is called at the beginning of the running stage to regenerate data.json
        with all samples instead of the debug subset created during coding stage.

        Args:
            implementation: The workspace containing process_data.py

        Returns:
            EnvResult with exit_code, stdout, etc.
        """
        # Get data processing environment with LLM API access
        env, env_vars = get_data_processing_env()
        ws_prefix = get_workspace_prefix(env)

        logger.info("Starting full data processing (without --debug flag)")

        # Execute WITHOUT --debug flag to generate all samples
        result = implementation.run(
            env=env,
            entry=f"python {ws_prefix}/{FT_DATA_SCRIPT_NAME}",  # No --debug flag
            env_vars=env_vars,
            cache_key_extra_func=get_data_processing_cache_key,
            cache_files_to_extract=[FT_DATA_FILE_NAME],
        )

        return result