server/_heuristic.py

"""Rule-based heuristic baseline for the /baseline endpoint.

Extracted from app.py to keep route definitions clean.
"""

from __future__ import annotations

from ml_training_debugger.models import MLTrainingAction
from server.environment import MLTrainingEnvironment

ALL_TASK_IDS = [
    "task_001",
    "task_002",
    "task_003",
    "task_004",
    "task_005",
    "task_006",
    "task_007",
]


def run_baseline_all_tasks() -> dict[str, float]:
    """Run the rule-based baseline on all tasks. Returns {task_id: score}."""
    scores: dict[str, float] = {}
    for task_id in ALL_TASK_IDS:
        env = MLTrainingEnvironment()
        env.reset(seed=42, episode_id=f"baseline_{task_id}", task_id=task_id)
        scores[task_id] = round(_run_heuristic_episode(env), 4)
    return scores


def _run_heuristic_episode(
    env: MLTrainingEnvironment, task_id: str = "",
) -> float:
    """Run one heuristic baseline episode. Returns grader score."""
    # Step 1: inspect_gradients
    obs = env.step(MLTrainingAction(action_type="inspect_gradients"))

    if obs.gradient_stats:
        if any(g.is_exploding for g in obs.gradient_stats):
            env.step(MLTrainingAction(
                action_type="modify_config", target="learning_rate", value=0.001,
            ))
            env.step(MLTrainingAction(action_type="restart_run"))
            env.step(MLTrainingAction(
                action_type="mark_diagnosed", diagnosis="lr_too_high",
            ))
            return _get_score(env)

        if any(g.is_vanishing for g in obs.gradient_stats):
            env.step(MLTrainingAction(
                action_type="modify_config", target="learning_rate", value=0.01,
            ))
            env.step(MLTrainingAction(action_type="restart_run"))
            env.step(MLTrainingAction(
                action_type="mark_diagnosed", diagnosis="vanishing_gradients",
            ))
            return _get_score(env)

    # Step 2: inspect_data_batch
    obs = env.step(MLTrainingAction(action_type="inspect_data_batch"))
    if obs.data_batch_stats and obs.data_batch_stats.class_overlap_score > 0.5:
        env.step(MLTrainingAction(action_type="patch_data_loader"))
        env.step(MLTrainingAction(action_type="restart_run"))
        env.step(MLTrainingAction(
            action_type="mark_diagnosed", diagnosis="data_leakage",
        ))
        return _get_score(env)

    # Detect overfitting pattern
    looks_like_overfitting = _detect_overfitting(obs)

    # Step 3: inspect_model_modes
    obs = env.step(MLTrainingAction(action_type="inspect_model_modes"))
    if obs.model_mode_info:
        if any(v == "eval" for v in obs.model_mode_info.values()):
            env.step(MLTrainingAction(action_type="fix_model_mode"))
            env.step(MLTrainingAction(action_type="restart_run"))
            env.step(MLTrainingAction(
                action_type="mark_diagnosed", diagnosis="batchnorm_eval_mode",
            ))
            return _get_score(env)

    # Step 4: inspect_code (for Task 6)
    obs = env.step(MLTrainingAction(action_type="inspect_code"))
    if obs.code_snippet:
        code = obs.code_snippet.code
        _try_code_fix(env, code)

        session = env._get_session()
        if session and session.state.fix_action_taken:
            env.step(MLTrainingAction(action_type="restart_run"))

        env.step(MLTrainingAction(
            action_type="mark_diagnosed", diagnosis="code_bug",
        ))
        return _get_score(env)

    # Step 5: scheduler issue (loss stagnates after initial progress)
    if _detect_scheduler_issue(obs):
        env.step(MLTrainingAction(
            action_type="modify_config", target="learning_rate", value=0.001,
        ))
        env.step(MLTrainingAction(action_type="restart_run"))
        env.step(MLTrainingAction(
            action_type="mark_diagnosed", diagnosis="scheduler_misconfigured",
        ))
        return _get_score(env)

    # Overfitting fallback
    if looks_like_overfitting:
        env.step(MLTrainingAction(
            action_type="modify_config", target="weight_decay", value=0.01,
        ))
        env.step(MLTrainingAction(action_type="restart_run"))
        env.step(MLTrainingAction(
            action_type="mark_diagnosed", diagnosis="overfitting",
        ))
        return _get_score(env)

    # Final fallback
    env.step(MLTrainingAction(
        action_type="mark_diagnosed", diagnosis="overfitting",
    ))
    return _get_score(env)


def _try_code_fix(env: MLTrainingEnvironment, code: str) -> None:
    """Attempt to fix a detected code bug."""
    if "model.eval()" in code and "model.train()" not in code:
        env.step(MLTrainingAction(
            action_type="fix_code", line=5, replacement="model.train()",
        ))
    elif ".detach()" in code:
        env.step(MLTrainingAction(
            action_type="fix_code", line=14,
            replacement="        loss = criterion(output, batch_y)",
        ))
    elif "inplace=True" in code:
        env.step(MLTrainingAction(
            action_type="fix_code", line=15,
            replacement="        output = F.relu(output)",
        ))
    elif "optimizer.zero_grad()" not in code and "optimizer.step()" in code:
        env.step(MLTrainingAction(
            action_type="fix_code", line=11,
            replacement="        optimizer.zero_grad()",
        ))


def _detect_overfitting(obs: object) -> bool:
    """Detect overfitting pattern from observation."""
    if not (obs.val_loss_history and obs.training_loss_history
            and len(obs.val_loss_history) >= 10):
        return False
    early_train = sum(obs.training_loss_history[:5]) / 5
    late_train = sum(obs.training_loss_history[-5:]) / 5
    early_val = sum(obs.val_loss_history[:5]) / 5
    late_val = sum(obs.val_loss_history[-5:]) / 5
    train_dropped = late_train < early_train * 0.5
    train_loss_low = late_train < 0.15
    val_not_improving = late_val >= early_val * 0.95
    gap_widening = (late_val - late_train) > (early_val - early_train)
    return (
        (train_dropped or train_loss_low)
        and (val_not_improving or gap_widening)
        and obs.data_batch_stats
        and obs.data_batch_stats.class_overlap_score < 0.3
    )


def _detect_scheduler_issue(obs: object) -> bool:
    """Detect scheduler misconfiguration from loss history."""
    if not (obs.training_loss_history and len(obs.training_loss_history) >= 10):
        return False
    early_loss = sum(obs.training_loss_history[:3]) / 3
    mid_loss = sum(obs.training_loss_history[5:8]) / 3
    finite_late = [v for v in obs.training_loss_history[-3:] if v != float("inf")]
    late_loss = sum(finite_late) / max(len(finite_late), 1)
    return early_loss > mid_loss and abs(late_loss - mid_loss) < 0.3


def _get_score(env: MLTrainingEnvironment) -> float:
    """Extract the grader score from the environment."""
    session = env._get_session()
    if session and session.last_score is not None:
        return session.last_score
    return 0.0