training-scripts/modal-training-script.py

"""
GRPO training run for CrisisOps + Unsloth on Modal cloud GPU infra.

This is a Modal-native port of simple-training-script.py. It keeps the exact
same Unsloth, TRL, and TrackIO implementation, but runs on Modal's cloud GPUs
instead of Hugging Face Jobs.

Usage:
    modal run training-scripts/modal-training-script.py

    # With overrides:
    modal run training-scripts/modal-training-script.py \
        --grpo-max-steps 200 \
        --task-id multi_zone_triage

Required Modal setup:
    - A Modal Secret named "huggingface-secret" (type: Hugging Face) so that
      HF_TOKEN is available in the container.
    - The Modal CLI authenticated (modal token set).
"""

from __future__ import annotations

# =============================================================================
# Modal scaffolding
# =============================================================================

import modal

image = (
    modal.Image.debian_slim(python_version="3.11")
    .uv_pip_install(
        "trl==0.19.1",
        "unsloth",
        "transformers",
        "accelerate",
        "bitsandbytes",
        "peft",
        "torch",
        "datasets",
        "trackio",
    )
)

app = modal.App("crisisops-grpo-modal", image=image)


# =============================================================================
# Original simple-training-script.py implementation (preserved verbatim)
# =============================================================================

import json
import os
import platform
import re
import subprocess
import sys
import textwrap
import threading
import time
import uuid
import warnings
from typing import Any, Dict, List, Mapping, Optional, Tuple
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


# --------------------------------------------------------------------------- #
# Config (env-var driven so Modal can override without code changes)
# --------------------------------------------------------------------------- #

ENV_URL = os.getenv("ENV_URL", "https://mrinaalarora-crisisops.hf.space").rstrip("/")
MODEL_ID = os.getenv("MODEL_ID", "unsloth/Qwen2.5-Coder-3B-Instruct-bnb-4bit").strip()
TASK_ID = os.getenv("TASK_ID", "single_zone_response").strip()

# Training scale. 100 GRPO steps with the sizing below = 100 weight updates.
# With grad_accum=4, per-device batch=3, and num_generations=4, Unsloth sees
# 3 unique prompts per generation cycle, satisfying its >2 prompt guidance.
GRPO_MAX_STEPS = int(os.getenv("GRPO_MAX_STEPS", "100"))
GRPO_NUM_GENERATIONS = int(os.getenv("GRPO_NUM_GENERATIONS", "4"))
GRPO_PER_DEVICE_BATCH = int(os.getenv("GRPO_PER_DEVICE_BATCH", "3"))
GRPO_GRAD_ACCUM = int(os.getenv("GRPO_GRAD_ACCUM", "4"))
GRPO_LEARNING_RATE = float(os.getenv("GRPO_LEARNING_RATE", "5e-6"))
GRPO_BETA = float(os.getenv("GRPO_BETA", "0.0"))
GRPO_TEMPERATURE = float(os.getenv("GRPO_TEMPERATURE", "0.9"))
GRPO_TOP_P = float(os.getenv("GRPO_TOP_P", "0.95"))
GRPO_MAX_PROMPT_LENGTH = int(os.getenv("GRPO_MAX_PROMPT_LENGTH", "3072"))
GRPO_MAX_COMPLETION_LENGTH = int(os.getenv("GRPO_MAX_COMPLETION_LENGTH", "512"))
GRPO_WARMUP_STEPS = int(os.getenv("GRPO_WARMUP_STEPS", "5"))
GRPO_LOGGING_STEPS = int(os.getenv("GRPO_LOGGING_STEPS", "1"))
GRPO_SAVE_STEPS = int(os.getenv("GRPO_SAVE_STEPS", "25"))
GRPO_LORA_RANK = int(os.getenv("GRPO_LORA_RANK", "16"))
GRPO_MAX_GRAD_NORM = float(os.getenv("GRPO_MAX_GRAD_NORM", "0.1"))
MAX_SEQ_LENGTH = int(os.getenv("MAX_SEQ_LENGTH", "4096"))
GRPO_TORCH_DTYPE = os.getenv("GRPO_TORCH_DTYPE", "float16").strip().lower()

OUTPUT_DIR = os.getenv("OUTPUT_DIR", "crisisops-grpo-easy-lora").strip()
HF_REPO_ID = os.getenv("HF_REPO_ID", "").strip() or None
TRACKIO_SPACE_ID = os.getenv("TRACKIO_SPACE_ID", "").strip() or None
HF_TOKEN = os.getenv("HF_TOKEN", "").strip() or None
RUN_NAME = os.getenv("RUN_NAME", f"crisisops-{TASK_ID}-grpo-one-step").strip()

# Prompt seeds. We repeat a compact set of deterministic reset observations to
# keep the run cheap and low-variance.
SEED_POOL_SIZE = int(os.getenv("SEED_POOL_SIZE", "16"))

# How many step-request rows to put in the dataset. GRPOTrainer iterates over
# these; with num_train_epochs=1 and max_steps capped, we just need enough to
# cover max_steps * per_device_batch * grad_accum. We oversize so the trainer
# never runs out before max_steps.
DATASET_ROWS = int(
    os.getenv(
        "DATASET_ROWS",
        str(max(256, GRPO_MAX_STEPS * GRPO_PER_DEVICE_BATCH * GRPO_GRAD_ACCUM * 2)),
    )
)


warnings.filterwarnings(
    "ignore",
    message=r".*AttentionMaskConverter.*deprecated.*",
    category=FutureWarning,
)

ACTION_JSON_RE = re.compile(r"\{[\s\S]*\}")


# --------------------------------------------------------------------------- #
# Domain constants (mirror baseline_smoke_test_unsloth_crisisops.py so the
# training prompt has the same shape as the baseline that produced 0.760
# on easy / 0.270 on medium).
# --------------------------------------------------------------------------- #

TASK_TIERS: Dict[str, str] = {
    "single_zone_response": "easy",
    "multi_zone_triage": "medium",
    "cascading_crisis": "hard",
    "multi_district_coordination": "expert",
}

TASK_CONFIGS: Dict[str, Dict[str, int]] = {
    "single_zone_response": {"episode_cap": 8, "expected_reports": 3, "stream_done_step": 0},
    "multi_zone_triage": {"episode_cap": 15, "expected_reports": 6, "stream_done_step": 0},
    "cascading_crisis": {"episode_cap": 25, "expected_reports": 10, "stream_done_step": 12},
    "multi_district_coordination": {"episode_cap": 40, "expected_reports": 16, "stream_done_step": 20},
}

INCIDENT_UNIT_TYPES: Dict[str, set[str]] = {
    "flood": {"rescue_team", "evac_bus"},
    "collapse": {"rescue_team", "medical_unit"},
    "medical_surge": {"medical_unit"},
    "fire": {"rescue_team"},
    "contamination": {"supply_truck", "medical_unit"},
    "power_outage": {"supply_truck"},
}

TASK_BY_UNIT_TYPE = {
    "rescue_team": "rescue",
    "medical_unit": "medical",
    "supply_truck": "supply_delivery",
    "evac_bus": "evacuation",
    "recon_drone": "recon",
}

SYSTEM_PROMPT = textwrap.dedent(
    """
    You are an emergency operations commander for CrisisOps. Return exactly one
    JSON object for the next environment action.

    Use only IDs that appear in the visible observation. Return JSON only. No
    markdown fences. No prose outside the JSON.
    """
).strip()

TASK_BRIEFS: Dict[str, str] = {
    "single_zone_response": (
        "EASY tier. Verify the report stream, allocate the matching unit type "
        "to the true zone before its deadline, then publish a sitrep."
    ),
    "multi_zone_triage": "MEDIUM tier with multiple concurrent incident zones.",
    "cascading_crisis": "HARD tier where incidents stream in mid-episode.",
    "multi_district_coordination": "EXPERT tier with mutual aid and comms degradation.",
}

ACTION_FORMAT_PROMPT = textwrap.dedent(
    """
    Valid action JSON shapes:
    {"type":"verify_report","report_id":"report-1","verification_method":"cross_check","rationale":"..."}
    {"type":"request_recon","zone_id":"zone-1","objective":"...","priority":"normal","report_id":null}
    {"type":"allocate_unit","unit_id":"unit-1","zone_id":"zone-1","task":"rescue","priority":"high","report_ids":["report-1"]}
    {"type":"reroute_unit","unit_id":"unit-1","route":{"route_id":"...","from_zone_id":"...","to_zone_id":"...","status":"open","travel_time_minutes":10,"hazards":[]},"reason":"..."}
    {"type":"issue_evacuation","zone_id":"zone-1","urgency":"critical","message":"...","route_id":null,"destination_shelter_id":null}
    {"type":"open_shelter","shelter":{"shelter_id":"...","zone_id":"...","name":"...","status":"open","capacity_total":100,"capacity_available":50,"supplies":{}},"reason":"..."}
    {"type":"dispatch_supplies","supplies":{"water":100},"destination_zone_id":"zone-1","priority":"high","unit_id":null,"destination_shelter_id":null}
    {"type":"flag_false_alarm","report_id":"report-1","rationale":"...","evidence":["..."]}
    {"type":"publish_sitrep","payload":{"incidents_confirmed":["report-1"],"incidents_resolved":["zone-1"],"unresolved_risks":[],"false_alarms_detected":[],"summary_text":"..."}}
    {"type":"noop","reason":"..."}
    """
).strip()


# --------------------------------------------------------------------------- #
# Training metrics — surfaced into Trackio and the final summary.
# --------------------------------------------------------------------------- #


class TrainingMetrics:
    """Small thread-safe metric bag used by the reward function."""

    def __init__(self) -> None:
        self.lock = threading.RLock()
        self.step_total_reward = 0.0
        self.step_total_count = 0
        self.valid_action_count = 0
        self.invalid_action_count = 0
        self.action_type_counts: Dict[str, int] = {}

    def record_invalid(self) -> None:
        with self.lock:
            self.invalid_action_count += 1

    def record_step(self, action: Mapping[str, Any], reward: float) -> None:
        with self.lock:
            self.valid_action_count += 1
            self.step_total_reward += float(reward)
            self.step_total_count += 1
            action_type = str(action.get("type", "unknown"))
            self.action_type_counts[action_type] = (
                self.action_type_counts.get(action_type, 0) + 1
            )

    def snapshot(self) -> Dict[str, Any]:
        with self.lock:
            total_actions = self.valid_action_count + self.invalid_action_count
            return {
                "average_step_reward": (
                    self.step_total_reward / self.step_total_count
                    if self.step_total_count
                    else 0.0
                ),
                "valid_action_fraction": (
                    self.valid_action_count / total_actions if total_actions else 0.0
                ),
                "valid_action_count": self.valid_action_count,
                "invalid_action_count": self.invalid_action_count,
                "action_type_counts": dict(self.action_type_counts),
            }


# --------------------------------------------------------------------------- #
# HTTP helpers (mirror the smoke test client so behaviour is identical)
# --------------------------------------------------------------------------- #


def post_json(path: str, payload: Mapping[str, Any]) -> dict:
    request = Request(
        f"{ENV_URL}{path}",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urlopen(request, timeout=60) as response:
            return json.loads(response.read().decode("utf-8"))
    except HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} from {path}: {body}") from exc
    except URLError as exc:
        raise RuntimeError(f"Could not reach {ENV_URL}{path}: {exc}") from exc


# --------------------------------------------------------------------------- #
# Action parsing & sanitization (copied from baseline so training and eval
# share exactly one schema-handling implementation)
# --------------------------------------------------------------------------- #


def parse_action_json(response_text: str) -> dict:
    match = ACTION_JSON_RE.search(response_text or "")
    if not match:
        raise ValueError("model response did not contain a JSON object")
    parsed = json.loads(match.group(0))
    if not isinstance(parsed, dict) or "type" not in parsed:
        raise ValueError("model response was not an action object")
    return parsed


def _required_str(payload: Mapping[str, Any], key: str) -> str:
    value = payload.get(key)
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"action missing required string field: {key}")
    return value.strip()


def _text_value(value: Any, default: str) -> str:
    if isinstance(value, str) and value.strip():
        return value.strip()
    return default


def _enum_value(value: Any, allowed: set[str], default: str) -> str:
    if isinstance(value, str) and value in allowed:
        return value
    return default


def _priority_value(value: Any, default: str) -> str:
    return _enum_value(
        str(value) if value is not None else value,
        {"low", "normal", "high", "critical"},
        default,
    )


def _string_list(value: Any) -> List[str]:
    if not isinstance(value, list):
        return []
    return [str(item) for item in value if isinstance(item, (str, int))]


def sanitize_model_action(parsed: Mapping[str, Any]) -> dict:
    action_type = str(parsed.get("type", ""))
    if not action_type:
        raise ValueError("model response missing type")

    if action_type == "verify_report":
        return {
            "type": "verify_report",
            "report_id": _required_str(parsed, "report_id"),
            "verification_method": _enum_value(
                parsed.get("verification_method"),
                {
                    "cross_check",
                    "contact_source",
                    "field_recon",
                    "sensor_review",
                    "official_confirmation",
                },
                "cross_check",
            ),
            "rationale": _text_value(parsed.get("rationale"), "Verify report."),
        }
    if action_type == "request_recon":
        return {
            "type": "request_recon",
            "zone_id": _required_str(parsed, "zone_id"),
            "objective": _text_value(parsed.get("objective"), "Clarify incident status."),
            "priority": _priority_value(parsed.get("priority"), "normal"),
            "report_id": parsed.get("report_id"),
        }
    if action_type == "allocate_unit":
        return {
            "type": "allocate_unit",
            "unit_id": _required_str(parsed, "unit_id"),
            "zone_id": _required_str(parsed, "zone_id"),
            "task": _enum_value(
                parsed.get("task"),
                {
                    "rescue",
                    "medical",
                    "evacuation",
                    "fire_suppression",
                    "supply_delivery",
                    "recon",
                    "route_clearance",
                },
                "recon",
            ),
            "priority": _priority_value(parsed.get("priority"), "normal"),
            "report_ids": _string_list(parsed.get("report_ids")),
        }
    if action_type == "reroute_unit":
        route = parsed.get("route")
        if not isinstance(route, Mapping):
            raise ValueError("reroute_unit requires route object")
        return {
            "type": "reroute_unit",
            "unit_id": _required_str(parsed, "unit_id"),
            "route": dict(route),
            "reason": _text_value(parsed.get("reason"), "Use safer route."),
        }
    if action_type == "issue_evacuation":
        return {
            "type": "issue_evacuation",
            "zone_id": _required_str(parsed, "zone_id"),
            "urgency": _priority_value(parsed.get("urgency"), "high"),
            "message": _text_value(parsed.get("message"), "Evacuate immediately."),
            "route_id": parsed.get("route_id"),
            "destination_shelter_id": parsed.get("destination_shelter_id"),
        }
    if action_type == "open_shelter":
        shelter = parsed.get("shelter")
        if not isinstance(shelter, Mapping):
            raise ValueError("open_shelter requires shelter object")
        return {
            "type": "open_shelter",
            "shelter": dict(shelter),
            "reason": _text_value(parsed.get("reason"), "Open shelter capacity."),
        }
    if action_type == "dispatch_supplies":
        supplies = parsed.get("supplies")
        if not isinstance(supplies, Mapping) or not supplies:
            raise ValueError("dispatch_supplies requires non-empty supplies")
        sanitized_supplies = {
            str(key): int(value)
            for key, value in supplies.items()
            if isinstance(value, (int, float)) and value > 0
        }
        if not sanitized_supplies:
            raise ValueError("dispatch_supplies requires positive supply amounts")
        return {
            "type": "dispatch_supplies",
            "supplies": sanitized_supplies,
            "destination_zone_id": _required_str(parsed, "destination_zone_id"),
            "priority": _priority_value(parsed.get("priority"), "normal"),
            "unit_id": parsed.get("unit_id"),
            "destination_shelter_id": parsed.get("destination_shelter_id"),
        }
    if action_type == "flag_false_alarm":
        return {
            "type": "flag_false_alarm",
            "report_id": _required_str(parsed, "report_id"),
            "rationale": _text_value(parsed.get("rationale"), "Report is disputed."),
            "evidence": _string_list(parsed.get("evidence")),
        }
    if action_type == "publish_sitrep":
        payload = parsed.get("payload")
        if not isinstance(payload, Mapping):
            raise ValueError("publish_sitrep requires payload object")
        return {
            "type": "publish_sitrep",
            "payload": {
                "incidents_confirmed": _string_list(payload.get("incidents_confirmed")),
                "incidents_resolved": _string_list(payload.get("incidents_resolved")),
                "unresolved_risks": _string_list(payload.get("unresolved_risks")),
                "false_alarms_detected": _string_list(payload.get("false_alarms_detected")),
                "summary_text": _text_value(
                    payload.get("summary_text"), "Situation report published."
                )[:800],
            },
        }
    if action_type == "noop":
        return {
            "type": "noop",
            "reason": _text_value(parsed.get("reason"), "Waiting for more evidence."),
        }
    raise ValueError(f"unknown action type: {action_type}")


# --------------------------------------------------------------------------- #
# Prompt rendering — same shape as the baseline (no recommended-action block).
# We render fresh for every step request because the env state is mutating.
# --------------------------------------------------------------------------- #


def _compact_observation(obs: Mapping[str, Any]) -> dict:
    return {
        "time_step": obs.get("time_step"),
        "metadata": obs.get("metadata", {}),
        "visible_zones": [
            {
                "zone_id": zone.get("zone_id"),
                "name": zone.get("name"),
                "incident_type": zone.get("incident_type"),
                "severity": zone.get("severity"),
                "population_at_risk": zone.get("population_at_risk"),
                "deadline_steps": zone.get("deadline_steps"),
                "access_status": zone.get("access_status"),
                "district_id": zone.get("district_id"),
                "required_unit_types": sorted(zone.get("required_unit_types") or []),
            }
            for zone in (obs.get("visible_zones") or [])
        ],
        "reports": [
            {
                "report_id": report.get("report_id"),
                "zone_id": report.get("zone_id"),
                "source": report.get("source"),
                "report_type": report.get("report_type"),
                "severity": report.get("severity"),
                "description": (str(report.get("description") or ""))[:200],
                "verified_status": report.get("verified_status"),
                "confidence": report.get("confidence"),
                "reveal_at_step": report.get("reveal_at_step"),
            }
            for report in (obs.get("reports") or [])
        ],
        "resources": [
            {
                "unit_id": unit.get("unit_id"),
                "unit_type": unit.get("unit_type"),
                "status": unit.get("status"),
                "current_zone_id": unit.get("current_zone_id"),
                "capacity": unit.get("capacity"),
                "capabilities": unit.get("capabilities"),
                "fatigue": unit.get("fatigue"),
                "district_id": unit.get("district_id"),
                "mutual_aid_unlock_step": unit.get("mutual_aid_unlock_step"),
            }
            for unit in (obs.get("resources") or [])
        ],
        "incident_log": (obs.get("incident_log") or [])[-6:],
    }


def render_user_prompt(task_id: str, obs: Mapping[str, Any]) -> str:
    brief = TASK_BRIEFS.get(task_id, "(no brief)")
    episode_cap = (obs.get("metadata") or {}).get(
        "episode_cap", TASK_CONFIGS.get(task_id, {}).get("episode_cap", 8)
    )
    return (
        f"Task: {task_id}\n"
        f"Brief: {brief}\n"
        f"Time step: {obs.get('time_step')}\n"
        f"Episode cap: {episode_cap}\n\n"
        f"Action contract:\n{ACTION_FORMAT_PROMPT}\n\n"
        f"Current observation:\n"
        f"{json.dumps(_compact_observation(obs), sort_keys=True)}\n\n"
        "Return exactly one JSON object and nothing else."
    )


def build_prompt_messages(task_id: str, obs: Mapping[str, Any]) -> List[Dict[str, str]]:
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": render_user_prompt(task_id, obs)},
    ]


# --------------------------------------------------------------------------- #
# Reward function — this is where the env interaction happens. GRPO will call
# this with `completions` (a list of model-generated strings) for each row
# in the dataset batch. We turn each completion into an action, hit /step,
# and return the env's reward.
# --------------------------------------------------------------------------- #

# Metrics created in main() and bound here so the reward function can update it.
_METRICS: Optional[TrainingMetrics] = None


def resolve_torch_dtype(torch_module):
    if GRPO_TORCH_DTYPE in {"float16", "fp16", "half"}:
        return torch_module.float16
    if GRPO_TORCH_DTYPE in {"bfloat16", "bf16"}:
        return torch_module.bfloat16
    if GRPO_TORCH_DTYPE in {"float32", "fp32"}:
        return torch_module.float32
    if GRPO_TORCH_DTYPE in {"auto", "none", ""}:
        return None
    raise ValueError(
        "GRPO_TORCH_DTYPE must be one of float16, bfloat16, float32, or auto; "
        f"got {GRPO_TORCH_DTYPE!r}"
    )


def align_trainable_parameter_dtype(model: Any, dtype: Any) -> None:
    if dtype is None:
        return
    converted = 0
    for parameter in model.parameters():
        if parameter.requires_grad and parameter.dtype != dtype:
            parameter.data = parameter.data.to(dtype=dtype)
            converted += parameter.numel()
    print(
        f"[TRAIN] trainable_parameter_dtype={dtype} converted_params={converted}",
        flush=True,
    )


def _completion_text(completion: Any) -> str:
    """Normalize a completion to a string. TRL passes either str or [{role, content}]."""
    if isinstance(completion, str):
        return completion
    if isinstance(completion, list) and completion:
        last = completion[-1]
        if isinstance(last, Mapping) and "content" in last:
            return str(last.get("content", ""))
    if isinstance(completion, Mapping) and "content" in completion:
        return str(completion.get("content", ""))
    return str(completion)


def crisisops_step_reward(prompts, completions, **kwargs) -> List[float]:
    """One reward per completion: env step reward (action accepted) or penalty."""
    metrics = _METRICS
    seeds = kwargs.get("seed") or [42] * len(completions)
    rewards: List[float] = []
    for index, completion in enumerate(completions):
        text = _completion_text(completion)
        try:
            parsed = parse_action_json(text)
            action = sanitize_model_action(parsed)
        except Exception:
            if metrics is not None:
                metrics.record_invalid()
            rewards.append(-0.5)  # invalid JSON penalty
            continue

        seed = int(seeds[index % len(seeds)])
        episode_id = f"grpo-{seed}-{uuid.uuid4().hex[:8]}"

        try:
            reset_response = post_json(
                "/reset",
                {"task_id": TASK_ID, "seed": seed, "episode_id": episode_id},
            )
        except Exception as exc:
            print(f"[REWARD] /reset error seed={seed} error={exc}", flush=True)
            rewards.append(-0.2)
            continue
        session_id = str(reset_response.get("session_id", ""))
        if not session_id:
            rewards.append(-0.5)
            continue

        try:
            response = post_json(
                "/step", {"session_id": session_id, "action": action}
            )
        except Exception as exc:
            print(f"[REWARD] /step error session={session_id} error={exc}", flush=True)
            rewards.append(-0.2)
            continue

        step_reward = float(response.get("reward") or 0.0)
        reward = step_reward + action_shaping_bonus(action)
        if metrics is not None:
            metrics.record_step(action, reward)

        rewards.append(reward)
    return rewards


def action_shaping_bonus(action: Mapping[str, Any]) -> float:
    """Small bias toward actions that can make progress on the easy task."""
    action_type = str(action.get("type", ""))
    if action_type == "verify_report":
        return 0.05
    if action_type in {"flag_false_alarm", "allocate_unit"}:
        return 0.02
    if action_type in {"request_recon", "noop"}:
        return -0.03
    return 0.0


# --------------------------------------------------------------------------- #
# Trackio metric callback — flush episode-level summary stats every N steps so
# the live dashboard shows the curve we care about (avg episode score).
# --------------------------------------------------------------------------- #


def make_pool_metric_callback():
    from transformers import TrainerCallback

    class PoolMetricCallback(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):
            metrics = _METRICS
            if metrics is None or logs is None:
                return
            snapshot = metrics.snapshot()
            logs["crisisops/avg_step_reward"] = snapshot["average_step_reward"]
            logs["crisisops/valid_action_fraction"] = snapshot["valid_action_fraction"]
            logs["crisisops/valid_action_count"] = snapshot["valid_action_count"]

    return PoolMetricCallback


def patch_text_only_unsloth_grpo_trainer(trainer: Any) -> None:
    """Patch text-only compatibility gaps in Unsloth's generated GRPO trainer."""
    for attr in (
        "image_token",
        "image_token_id",
        "vision_start_token_id",
        "vision_end_token_id",
    ):
        if not hasattr(trainer, attr):
            setattr(trainer, attr, None)
    processing_class = getattr(trainer, "processing_class", None)
    for attr in ("pad_token", "pad_token_id", "eos_token", "eos_token_id"):
        if not hasattr(trainer, attr) and hasattr(processing_class, attr):
            setattr(trainer, attr, getattr(processing_class, attr))
    args = getattr(trainer, "args", None)
    compatibility_defaults = {
        "importance_sampling_level": "token",
        "top_entropy_quantile": 1.0,
        "vllm_importance_sampling_correction": False,
        "vllm_importance_sampling_cap": 2.0,
        "num_iterations": getattr(args, "num_iterations", 1),
        "epsilon_low": getattr(args, "epsilon", 0.2),
        "epsilon_high": (
            getattr(args, "epsilon_high", None)
            if getattr(args, "epsilon_high", None) is not None
            else getattr(args, "epsilon", 0.2)
        ),
        "loss_type": getattr(args, "loss_type", "bnpo"),
        "mask_truncated_completions": getattr(
            args, "mask_truncated_completions", False
        ),
    }
    for attr, value in compatibility_defaults.items():
        if not hasattr(trainer, attr):
            setattr(trainer, attr, value)
    if getattr(trainer, "scale_rewards", None) is True:
        trainer.scale_rewards = "group"
    elif getattr(trainer, "scale_rewards", None) is False:
        trainer.scale_rewards = "none"

    def truncate_with_protected_tokens(input_ids, attention_mask, max_length, protected):
        del protected
        if max_length is None:
            return input_ids, attention_mask
        if input_ids.shape[-1] <= max_length:
            return input_ids, attention_mask
        return input_ids[..., -max_length:], attention_mask[..., -max_length:]

    patched_scopes = 0

    for module_name, module in list(sys.modules.items()):
        if module_name.endswith("UnslothGRPOTrainer"):
            module.__dict__.setdefault(
                "truncate_with_protected_tokens", truncate_with_protected_tokens
            )
            module.__dict__.setdefault("has_images", False)
            module.__dict__.setdefault("images", None)
            patched_scopes += 1

    method = getattr(trainer, "_generate_and_score_completions", None)
    candidates = [method, getattr(method, "__func__", None)]
    for candidate in candidates:
        while candidate is not None:
            globals_dict = getattr(candidate, "__globals__", None)
            if isinstance(globals_dict, dict):
                globals_dict["truncate_with_protected_tokens"] = (
                    truncate_with_protected_tokens
                )
                globals_dict.setdefault("has_images", False)
                globals_dict.setdefault("images", None)
                patched_scopes += 1
            candidate = getattr(candidate, "__wrapped__", None)

    print(
        f"[TRAIN] patched_text_only_unsloth_grpo_trainer scopes={patched_scopes}",
        flush=True,
    )


# --------------------------------------------------------------------------- #
# CUDA / GPU sanity (carry over from the smoke test so logs are useful)
# --------------------------------------------------------------------------- #


def log_gpu_preflight() -> None:
    try:
        result = subprocess.run(
            ["nvidia-smi", "-L"], check=False, capture_output=True, text=True
        )
        output = (result.stdout or result.stderr or "").strip()
        if output:
            print(f"[TRAIN] nvidia_smi={output}", flush=True)
    except Exception as exc:
        print(f"[TRAIN] nvidia_smi_unavailable={exc}", flush=True)

    try:
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=name,driver_version",
                "--format=csv,noheader",
            ],
            check=False,
            capture_output=True,
            text=True,
        )
        output = (result.stdout or result.stderr or "").strip()
        if output:
            print(f"[TRAIN] nvidia_smi_detail={output}", flush=True)
    except Exception as exc:
        print(f"[TRAIN] nvidia_smi_detail_unavailable={exc}", flush=True)


def wait_for_cuda_runtime() -> None:
    """Wait for CUDA in child processes so failed probes do not poison training."""
    retries = int(os.getenv("CUDA_WAIT_RETRIES", "60"))
    sleep_seconds = int(os.getenv("CUDA_WAIT_SLEEP_SECONDS", "10"))
    probe_timeout = int(os.getenv("CUDA_PROBE_TIMEOUT", "30"))
    last_error: Optional[str] = None
    for attempt in range(1, retries + 1):
        try:
            probe = subprocess.run(
                [
                    sys.executable,
                    "-c",
                    (
                        "import os, torch; "
                        "print('torch_version=' + torch.__version__); "
                        "print('cuda_visible_devices=' + str(os.environ.get('CUDA_VISIBLE_DEVICES'))); "
                        "torch.cuda.init(); "
                        "print('device_count=' + str(torch.cuda.device_count())); "
                        "print('device_name=' + torch.cuda.get_device_name(0))"
                    ),
                ],
                check=False,
                capture_output=True,
                text=True,
                timeout=probe_timeout,
            )
        except subprocess.TimeoutExpired:
            last_error = f"probe timed out after {probe_timeout}s"
            print(
                f"[TRAIN] cuda_not_ready attempt={attempt}/{retries} "
                f"sleep_seconds={sleep_seconds} last_error={last_error}",
                flush=True,
            )
            if attempt < retries:
                time.sleep(sleep_seconds)
            continue
        if probe.returncode == 0:
            output = (probe.stdout or "").strip()
            print(f"[TRAIN] cuda_ready probe={output}", flush=True)
            return
        last_error = ((probe.stderr or "") + (probe.stdout or "")).strip()
        print(
            f"[TRAIN] cuda_not_ready attempt={attempt}/{retries} "
            f"sleep_seconds={sleep_seconds} last_error={last_error or 'none'}",
            flush=True,
        )
        if attempt < retries:
            time.sleep(sleep_seconds)
    hint = (
        " Error 802/system not yet initialized means nvidia-smi can see the "
        "GPU but the CUDA runtime cannot initialize compute in this job "
        "container yet. Increase CUDA_WAIT_RETRIES if it is transient; if it "
        "persists for the full wait, the failure is below the GRPO/Unsloth "
        "training code."
    )
    raise RuntimeError(
        "CUDA did not become ready inside the container. "
        f"Last error: {last_error or 'torch.cuda probe failed'}." + hint
    )


# --------------------------------------------------------------------------- #
# Original main() — renamed so the Modal entrypoint can call it.
# --------------------------------------------------------------------------- #


def training_main() -> Dict[str, Any]:
    global _METRICS

    if TASK_ID not in TASK_CONFIGS:
        raise ValueError(f"Unknown TASK_ID={TASK_ID!r}")

    print(
        f"[TRAIN] env_url={ENV_URL} task_id={TASK_ID} model_id={MODEL_ID} "
        f"max_steps={GRPO_MAX_STEPS} num_generations={GRPO_NUM_GENERATIONS} "
        f"per_device_batch={GRPO_PER_DEVICE_BATCH} grad_accum={GRPO_GRAD_ACCUM}",
        flush=True,
    )

    log_gpu_preflight()
    wait_for_cuda_runtime()

    import torch

    torch.cuda.init()

    print(
        f"[TRAIN] python={platform.python_version()} platform={platform.platform()}",
        flush=True,
    )
    print(
        f"[TRAIN] cuda_device_count={torch.cuda.device_count()} "
        f"device_name={torch.cuda.get_device_name(0)} "
        f"cuda_capability={torch.cuda.get_device_capability(0)}",
        flush=True,
    )

    # ---- Load model + LoRA via unsloth -----------------------------------
    from unsloth import FastLanguageModel

    train_dtype = resolve_torch_dtype(torch)
    print(
        f"[TRAIN] loading model={MODEL_ID} max_seq_length={MAX_SEQ_LENGTH} "
        f"torch_dtype={train_dtype}",
        flush=True,
    )
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_ID,
        max_seq_length=MAX_SEQ_LENGTH,
        dtype=train_dtype,
        load_in_4bit=True,
        fast_inference=False,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=GRPO_LORA_RANK,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_alpha=GRPO_LORA_RANK,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=3407,
        use_rslora=False,
    )
    align_trainable_parameter_dtype(model, train_dtype)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    print("[TRAIN] model loaded and LoRA adapter attached", flush=True)

    # ---- Build prompt pool from deterministic resets ---------------------
    seeds = [42 + offset for offset in range(SEED_POOL_SIZE)]
    _METRICS = TrainingMetrics()
    prompt_observations: Dict[int, Dict[str, Any]] = {}
    for seed in seeds:
        response = post_json(
            "/reset",
            {
                "task_id": TASK_ID,
                "seed": seed,
                "episode_id": f"prompt-{seed}-{uuid.uuid4().hex[:8]}",
            },
        )
        prompt_observations[seed] = response.get("observation") or {}
    print(f"[TRAIN] prompt pool warmed seeds={seeds}", flush=True)

    # ---- Build dataset -----------------------------------------------------
    # Each row carries the seed that produced its prompt. The reward function
    # resets that same seed before applying the completion, so prompt and
    # scored environment state stay aligned.
    from datasets import Dataset

    def make_row(index: int) -> Dict[str, Any]:
        seed = seeds[index % len(seeds)]
        messages = build_prompt_messages(TASK_ID, prompt_observations[seed])
        # GRPOTrainer accepts prompts as either string or list-of-messages;
        # list-of-messages avoids manual chat template formatting.
        return {
            "prompt": messages,
            "seed": seed,
            "task_id": TASK_ID,
        }

    rows = [make_row(i) for i in range(DATASET_ROWS)]
    train_dataset = Dataset.from_list(rows)
    print(f"[TRAIN] built dataset rows={len(train_dataset)}", flush=True)

    # ---- GRPO config -------------------------------------------------------
    from trl import GRPOConfig, GRPOTrainer

    if TRACKIO_SPACE_ID:
        os.environ["TRACKIO_SPACE_ID"] = TRACKIO_SPACE_ID
        os.environ.setdefault("TRACKIO_PROJECT", "crisisops-grpo")

    grpo_kwargs: Dict[str, Any] = dict(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=GRPO_MAX_STEPS,
        learning_rate=GRPO_LEARNING_RATE,
        beta=GRPO_BETA,
        optim="paged_adamw_8bit",
        max_grad_norm=GRPO_MAX_GRAD_NORM,
        per_device_train_batch_size=GRPO_PER_DEVICE_BATCH,
        gradient_accumulation_steps=GRPO_GRAD_ACCUM,
        num_generations=GRPO_NUM_GENERATIONS,
        warmup_steps=GRPO_WARMUP_STEPS,
        max_prompt_length=GRPO_MAX_PROMPT_LENGTH,
        max_completion_length=GRPO_MAX_COMPLETION_LENGTH,
        temperature=GRPO_TEMPERATURE,
        top_p=GRPO_TOP_P,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        logging_steps=GRPO_LOGGING_STEPS,
        save_steps=GRPO_SAVE_STEPS,
        save_total_limit=3,
        push_to_hub=bool(HF_REPO_ID),
        hub_token=HF_TOKEN,
        hub_model_id=HF_REPO_ID,
        hub_strategy="every_save" if HF_REPO_ID else "end",
        report_to="trackio" if TRACKIO_SPACE_ID else "none",
        run_name=RUN_NAME,
        fp16=False,
        bf16=train_dtype is torch.bfloat16,
    )
    grpo_config = GRPOConfig(**grpo_kwargs)

    PoolMetricCallback = make_pool_metric_callback()

    trainer = GRPOTrainer(
        model=model,
        processing_class=tokenizer,
        reward_funcs=[crisisops_step_reward],
        args=grpo_config,
        train_dataset=train_dataset,
        callbacks=[PoolMetricCallback()],
    )
    patch_text_only_unsloth_grpo_trainer(trainer)

    # ---- Train -------------------------------------------------------------
    start_mem = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
    total_mem = round(torch.cuda.get_device_properties(0).total_memory / 1024 / 1024 / 1024, 3)
    print(
        f"[TRAIN] memory_before reserved_gb={start_mem} total_gb={total_mem}",
        flush=True,
    )

    train_start = time.time()
    trainer.train()
    train_end = time.time()

    peak_mem = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
    print(
        f"[TRAIN] memory_after peak_gb={peak_mem} duration_minutes={(train_end - train_start) / 60:.2f}",
        flush=True,
    )

    # ---- Save adapter ------------------------------------------------------
    save_dir = OUTPUT_DIR
    trainer.save_model(save_dir)
    if HF_REPO_ID:
        try:
            trainer.push_to_hub()
            print(f"[TRAIN] pushed adapter to {HF_REPO_ID}", flush=True)
        except Exception as exc:
            print(f"[TRAIN] push_to_hub failed error={exc}", flush=True)

    # ---- Final pool snapshot ---------------------------------------------
    final = _METRICS.snapshot()
    summary_path = os.path.join(save_dir, "crisisops_training_summary.json")
    try:
        os.makedirs(save_dir, exist_ok=True)
        with open(summary_path, "w", encoding="utf-8") as handle:
            json.dump(
                {
                    "task_id": TASK_ID,
                    "env_url": ENV_URL,
                    "model_id": MODEL_ID,
                    "grpo_max_steps": GRPO_MAX_STEPS,
                    "seeds": seeds,
                    "training_minutes": (train_end - train_start) / 60,
                    **final,
                },
                handle,
                indent=2,
            )
        print(f"[TRAIN] wrote summary to {summary_path}", flush=True)
    except Exception as exc:
        print(f"[TRAIN] summary write failed error={exc}", flush=True)

    print(
        f"[TRAIN] DONE avg_step_reward={final['average_step_reward']:.3f} "
        f"valid_action_fraction={final['valid_action_fraction']:.3f}",
        flush=True,
    )

    return final


# =============================================================================
# Modal remote function + local entrypoint
# =============================================================================


def _set_env(name: str, value: Any) -> None:
    if value is not None:
        os.environ[name] = str(value)


@app.function(
    gpu="H200",
    secrets=[modal.Secret.from_name("huggingface-secret")],
    cpu=8,
    memory=32768,
    ephemeral_disk=524288,
    timeout=4 * 60 * 60,
)
def run_training(
    env_url: str = "https://mrinaalarora-crisisops.hf.space",
    task_id: str = "single_zone_response",
    model_id: str = "unsloth/Qwen2.5-Coder-3B-Instruct-bnb-4bit",
    grpo_max_steps: int = 100,
    grpo_num_generations: int = 4,
    grpo_per_device_batch: int = 3,
    grpo_grad_accum: int = 4,
    grpo_learning_rate: float = 5e-6,
    grpo_beta: float = 0.0,
    grpo_temperature: float = 0.9,
    grpo_top_p: float = 0.95,
    grpo_max_prompt_length: int = 3072,
    grpo_max_completion_length: int = 512,
    grpo_warmup_steps: int = 5,
    grpo_logging_steps: int = 1,
    grpo_save_steps: int = 25,
    grpo_lora_rank: int = 16,
    grpo_max_grad_norm: float = 0.1,
    grpo_torch_dtype: str = "float16",
    max_seq_length: int = 4096,
    output_dir: str = "crisisops-grpo-easy-lora",
    hf_repo_id: Optional[str] = "mrinaalarora/crisisops-grpo-easy-lora",
    trackio_space_id: Optional[str] = "mrinaalarora/crisisops-grpo-trackio",
    run_name: str = "",
    seed_pool_size: int = 16,
) -> Dict[str, Any]:
    """
    Run the CrisisOps GRPO training loop on Modal.

    Parameters mirror the env vars used by simple-training-script.py so you
    can override any setting from the CLI via the local_entrypoint.
    """
    # Push parameters into os.environ so the original config block picks them up.
    _set_env("ENV_URL", env_url)
    _set_env("TASK_ID", task_id)
    _set_env("MODEL_ID", model_id)
    _set_env("GRPO_MAX_STEPS", grpo_max_steps)
    _set_env("GRPO_NUM_GENERATIONS", grpo_num_generations)
    _set_env("GRPO_PER_DEVICE_BATCH", grpo_per_device_batch)
    _set_env("GRPO_GRAD_ACCUM", grpo_grad_accum)
    _set_env("GRPO_LEARNING_RATE", grpo_learning_rate)
    _set_env("GRPO_BETA", grpo_beta)
    _set_env("GRPO_TEMPERATURE", grpo_temperature)
    _set_env("GRPO_TOP_P", grpo_top_p)
    _set_env("GRPO_MAX_PROMPT_LENGTH", grpo_max_prompt_length)
    _set_env("GRPO_MAX_COMPLETION_LENGTH", grpo_max_completion_length)
    _set_env("GRPO_WARMUP_STEPS", grpo_warmup_steps)
    _set_env("GRPO_LOGGING_STEPS", grpo_logging_steps)
    _set_env("GRPO_SAVE_STEPS", grpo_save_steps)
    _set_env("GRPO_LORA_RANK", grpo_lora_rank)
    _set_env("GRPO_MAX_GRAD_NORM", grpo_max_grad_norm)
    _set_env("GRPO_TORCH_DTYPE", grpo_torch_dtype)
    _set_env("MAX_SEQ_LENGTH", max_seq_length)
    _set_env("OUTPUT_DIR", output_dir)
    _set_env("HF_REPO_ID", hf_repo_id)
    _set_env("TRACKIO_SPACE_ID", trackio_space_id)
    _set_env("RUN_NAME", run_name or f"crisisops-{task_id}-grpo-one-step")
    _set_env("SEED_POOL_SIZE", seed_pool_size)

    # HF_TOKEN is injected automatically by the Modal Secret.
    return training_main()


@app.local_entrypoint()
def main(
    env_url: str = "https://mrinaalarora-crisisops.hf.space",
    task_id: str = "single_zone_response",
    model_id: str = "unsloth/Qwen2.5-Coder-3B-Instruct-bnb-4bit",
    grpo_max_steps: int = 100,
    grpo_num_generations: int = 4,
    grpo_per_device_batch: int = 3,
    grpo_grad_accum: int = 4,
    grpo_torch_dtype: str = "float16",
    hf_repo_id: str = "mrinaalarora/crisisops-grpo-easy-lora",
    trackio_space_id: str = "mrinaalarora/crisisops-grpo-trackio",
    run_name: str = "",
):
    """
    Launch the CrisisOps GRPO trainer on Modal's cloud GPUs.

    Example:
        modal run training-scripts/modal-training-script.py

        modal run training-scripts/modal-training-script.py \
            --grpo-max-steps 200 \
            --task-id multi_zone_triage
    """
    result = run_training.remote(
        env_url=env_url,
        task_id=task_id,
        model_id=model_id,
        grpo_max_steps=grpo_max_steps,
        grpo_num_generations=grpo_num_generations,
        grpo_per_device_batch=grpo_per_device_batch,
        grpo_grad_accum=grpo_grad_accum,
        grpo_torch_dtype=grpo_torch_dtype,
        hf_repo_id=hf_repo_id,
        trackio_space_id=trackio_space_id,
        run_name=run_name,
    )
    print("\n=== Training Summary ===")
    for key, value in result.items():
        print(f"  {key}: {value}")