scripts/gcp_finetune_vqa_entrypoint.py

"""GCP Vertex AI Custom Job entrypoint — **VQA finetune** from an existing run.

Take a fully-trained run (stage 2 done) and continue training Stage 2 with
**VQA-heavy task weights** + low LR + few epochs, into a NEW run_id (so the
original checkpoint stays untouched on HF Hub).

Continuation strategy:
  - run_1/stage2/best/   ← already-trained projection + LoRA on HF
       checkpoint_projection.pt
       checkpoint_lora/
       checkpoint_chexpert_classifier.pt        (if present)

  We download those and place them at:
       run_2/stage1_projection/stage1_final_projection.pt   ← renamed
       run_2/stage1_projection/stage1_final_lora/           ← renamed
       run_2/stage1_projection/stage1_final_chexpert_classifier.pt

  Then train.py:
    1. detect_resume_point sees stage1_final_projection.pt → ("stage2", None)
       → skips Stage 1 entirely
    2. Builds the model, calls load_checkpoint(stage1_final.pt) which loads
       BOTH the projection AND the LoRA from the renamed files (load_checkpoint
       derives the LoRA dir from the .pt stem).
    3. Runs Stage 2 fresh with the new task weights + LR + epochs, starting
       from the loaded weights. Optimizer state is reset — that's deliberate
       (different task mix; old momentum would point the wrong way).

Required env vars:
  HF_TOKEN              — HuggingFace token (read+write)
  DATASET_NAME          — 'IU-Xray' | 'MIMIC-CXR' | 'MIMIC-CXR_resized'
  SOURCE_RUN_ID         — run on HF_RUNS_REPO whose stage2/best is the seed
                          e.g. 'MIMIC-CXR_resized_run_1'
  TARGET_RUN_ID         — NEW run id to write into, e.g.
                          'MIMIC-CXR_resized_run_2'

Optional env vars (defaults shown — tuned for "vqa-heavy with rehearsal"):
  HF_USER               = hieu3636
  HF_RUNS_REPO          = hieu3636/cxr-vlm-runs
  SOURCE_CKPT_PICK      = best                  # 'best' | 'last'
  REPORT_MODE           =                       # blank → from source run's snapshot
  IMAGE_MODE            =                       # blank → from source run's snapshot
  W_FINDINGS            = 0.15                  # rehearsal weight
  W_IMPRESSION          = 0.10                  # rehearsal weight
  W_VQA                 = 0.75                  # focus
  S2_EPOCHS             = 3
  S2_LR                 = 5e-5                  # 4× lower than original 2e-4
  STRICT_VQA_REQUIRED   = 1                     # fail fast if VQA samples = 0
  WORK                  = /workspace
"""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
import tarfile
import zipfile
from pathlib import Path

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "warning")
os.environ.setdefault("PYTHONUNBUFFERED", "1")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")


def env(name: str, default: str | None = None, *, required: bool = False) -> str:
    val = os.environ.get(name, default)
    if required and not val:
        sys.exit(f"[gcp_finetune_vqa] ERROR: required env var {name} not set")
    return val or ""


# ── 1) Resolve config from env ────────────────────────────────────────────────
HF_TOKEN             = env("HF_TOKEN", required=True)
DATASET_NAME         = env("DATASET_NAME", required=True)
SOURCE_RUN_ID        = env("SOURCE_RUN_ID", required=True)
TARGET_RUN_ID        = env("TARGET_RUN_ID", required=True)
HF_USER              = env("HF_USER", "hieu3636")
HF_RUNS_REPO         = env("HF_RUNS_REPO", "hieu3636/cxr-vlm-runs")
SOURCE_CKPT_PICK     = env("SOURCE_CKPT_PICK", "best")
REPORT_MODE_OVERRIDE = env("REPORT_MODE", "")
IMAGE_MODE_OVERRIDE  = env("IMAGE_MODE", "")
W_FINDINGS           = float(env("W_FINDINGS",   "0.15"))
W_IMPRESSION         = float(env("W_IMPRESSION", "0.10"))
W_VQA                = float(env("W_VQA",        "0.75"))
S2_EPOCHS            = int(env("S2_EPOCHS", "3"))
S2_LR                = float(env("S2_LR",   "5e-5"))
STRICT_VQA_REQUIRED  = env("STRICT_VQA_REQUIRED", "1") in ("1", "true", "True")
WORK                 = Path(env("WORK", "/workspace"))

assert DATASET_NAME in ("IU-Xray", "MIMIC-CXR", "MIMIC-CXR_resized"), DATASET_NAME
assert SOURCE_CKPT_PICK in ("best", "last"), SOURCE_CKPT_PICK
assert TARGET_RUN_ID != SOURCE_RUN_ID, \
    f"TARGET_RUN_ID must differ from SOURCE_RUN_ID — refusing to overwrite the source run."

PROJECT       = Path(__file__).resolve().parent.parent
DATA_SRC      = WORK / "data"
RUN_PULL_ROOT = WORK / "run_pull"
CKPT_ROOT     = WORK / "ckpt"
for d in (DATA_SRC, RUN_PULL_ROOT, CKPT_ROOT):
    d.mkdir(parents=True, exist_ok=True)

# Make project modules (utils.*, data.*, model.*) importable from this script
# itself. Python's sys.path includes the SCRIPT's dir (/workspace/code/scripts)
# but NOT the project root (/workspace/code) when invoked directly. The
# subprocess we spawn later uses `python -m training.train` which resolves
# modules relative to cwd, so it's fine — but the pre-flight VQA check below
# needs `from utils.dataset_resolver import ...` to work in-kernel.
sys.path.insert(0, str(PROJECT))

print(f"[gcp_finetune_vqa] PROJECT       = {PROJECT}")
print(f"[gcp_finetune_vqa] WORK          = {WORK}")
print(f"[gcp_finetune_vqa] DATASET_NAME  = {DATASET_NAME}")
print(f"[gcp_finetune_vqa] SOURCE_RUN_ID = {SOURCE_RUN_ID}  (ckpt: stage2/{SOURCE_CKPT_PICK})")
print(f"[gcp_finetune_vqa] TARGET_RUN_ID = {TARGET_RUN_ID}  (NEW — original stays untouched)")
print(f"[gcp_finetune_vqa] Task mix      = findings:{W_FINDINGS}  impression:{W_IMPRESSION}  vqa:{W_VQA}")
print(f"[gcp_finetune_vqa] S2_EPOCHS={S2_EPOCHS}  S2_LR={S2_LR}")

# ── 2) Download dataset payload from HF Hub ───────────────────────────────────
from huggingface_hub import HfApi, hf_hub_download, snapshot_download  # noqa: E402

if DATASET_NAME == "MIMIC-CXR_resized":
    mr_dir = DATA_SRC / "MIMIC-CXR_resized"
    mr_dir.mkdir(parents=True, exist_ok=True)
    files_dir = mr_dir / "files"
    manifests_present = all(
        (mr_dir / f).is_file()
        for f in ("manifest_train.csv", "manifest_val.csv", "manifest_test.csv")
    )
    if manifests_present and files_dir.is_dir() and any(files_dir.glob("p*")):
        print(f"[gcp_finetune_vqa] {mr_dir} already populated — skipping download.")
    else:
        api = HfApi(token=HF_TOKEN)
        all_files = api.list_repo_files(
            repo_id=f"{HF_USER}/cxr-vlm-data", repo_type="dataset"
        )
        tar_files = sorted(
            f for f in all_files
            if f.startswith("MIMIC-CXR_resized/") and f.endswith(".tar")
        )
        print(f"[gcp_finetune_vqa] {len(tar_files)} tar shards on HF")

        snapshot_download(
            repo_id=f"{HF_USER}/cxr-vlm-data",
            repo_type="dataset",
            allow_patterns=[
                "MIMIC-CXR_resized/*.csv",
                "MIMIC-CXR_resized/*.json",
                "MIMIC-CXR_resized/*.txt",
                "MIMIC-CXR_resized/vqa/**",
            ],
            token=HF_TOKEN,
            local_dir=str(DATA_SRC),
        )

        for i, tf in enumerate(tar_files, 1):
            print(f"[gcp_finetune_vqa]  [{i}/{len(tar_files)}] {tf}", flush=True)
            tp = Path(hf_hub_download(
                repo_id=f"{HF_USER}/cxr-vlm-data",
                repo_type="dataset",
                filename=tf,
                token=HF_TOKEN,
                local_dir=str(DATA_SRC),
            ))
            with tarfile.open(tp) as t:
                t.extractall(mr_dir)
            tp.unlink(missing_ok=True)
        print(f"[gcp_finetune_vqa] {mr_dir} ready.")
    DATA_ROOT_RESIZED = mr_dir

else:
    zip_name = f"{DATASET_NAME}.zip"
    marker = DATA_SRC / DATASET_NAME
    if not marker.exists():
        print(f"[gcp_finetune_vqa] downloading {zip_name} ...")
        zpath = hf_hub_download(
            repo_id=f"{HF_USER}/cxr-vlm-data",
            filename=zip_name,
            repo_type="dataset",
            token=HF_TOKEN,
            local_dir=str(DATA_SRC),
        )
        with zipfile.ZipFile(zpath) as zf:
            zf.extractall(DATA_SRC)
        try:
            os.remove(zpath)
        except OSError:
            pass
    else:
        print(f"[gcp_finetune_vqa] {marker} already present — skipping download.")

print(f"[gcp_finetune_vqa] DATA_SRC contents: {sorted(os.listdir(DATA_SRC))}")

# ── 3) Resume-aware: check if TARGET_RUN_ID has Stage-2 progress on HF.
#    If yes → hydrate it locally and SKIP the seed step (train.py will pick
#    up from the latest checkpoint-N). If no → fall through to seed from
#    SOURCE_RUN_ID as a fresh start.
#
#    This handles the resubmit-after-crash case correctly. On the first run
#    the target's HF folder doesn't exist → we seed from source. After the
#    job uploads checkpoint-500 and dies → resubmit pulls checkpoint-500
#    from HF, skipping the seed, and trainer resumes from step 500.
TARGET_DIR  = CKPT_ROOT / TARGET_RUN_ID
TGT_S1_DIR  = TARGET_DIR / "stage1_projection"
TGT_S2_DIR  = TARGET_DIR / "stage2_instruct"

print(f"[gcp_finetune_vqa] checking HF for {TARGET_RUN_ID} (resume state) …")
sys.path.insert(0, str(PROJECT))   # so the import below works
from utils.hf_uploader import hydrate_run_dir_from_hf  # noqa: E402

hydrated = hydrate_run_dir_from_hf(
    repo_id     = HF_RUNS_REPO,
    token       = HF_TOKEN,
    run_id      = TARGET_RUN_ID,
    output_root = str(CKPT_ROOT),
)
# hydrate places: stage1_final_* + stage2_instruct/checkpoint-1/ (mid-resume)
# Decide if we still need to seed from source.
def _has_ckpt(d: Path) -> bool:
    return d.is_dir() and any(d.glob("checkpoint-*"))

stage1_done_local = (TGT_S1_DIR / "stage1_final_projection.pt").is_file()
stage2_mid_local  = _has_ckpt(TGT_S2_DIR)
stage2_done_local = (TGT_S2_DIR / "stage2_final_projection.pt").is_file()

if stage2_done_local:
    print(f"[gcp_finetune_vqa] TARGET already has stage2_final — nothing to do.")
    sys.exit(0)

if stage1_done_local or stage2_mid_local:
    print(f"[gcp_finetune_vqa] RESUME: hydrated from HF "
          f"(stage1_done={stage1_done_local}, stage2_mid={stage2_mid_local}) "
          f"— skipping seed-from-source.")
    # train.py will detect_resume_point → ("stage2", checkpoint-N) and
    # the Trainer's resume_from_checkpoint will reload optimizer state.
else:
    # ── 4) Seed from SOURCE — fresh start of TARGET ──────────────────────
    print(f"[gcp_finetune_vqa] FRESH: no TARGET state on HF, seeding from "
          f"{SOURCE_RUN_ID}/stage2/{SOURCE_CKPT_PICK} …")
    snapshot_download(
        repo_id=HF_RUNS_REPO,
        repo_type="model",
        token=HF_TOKEN,
        allow_patterns=[
            f"{SOURCE_RUN_ID}/configs/**",
            f"{SOURCE_RUN_ID}/run_meta.json",
            f"{SOURCE_RUN_ID}/stage2/{SOURCE_CKPT_PICK}/**",
        ],
        local_dir=str(RUN_PULL_ROOT),
    )
    SRC_DIR     = RUN_PULL_ROOT / SOURCE_RUN_ID
    SRC_S2_DIR  = SRC_DIR / "stage2" / SOURCE_CKPT_PICK
    SRC_PROJ    = SRC_S2_DIR / "checkpoint_projection.pt"
    SRC_LORA    = SRC_S2_DIR / "checkpoint_lora"
    SRC_CHEXPRT = SRC_S2_DIR / "checkpoint_chexpert_classifier.pt"

    assert SRC_PROJ.is_file(), f"source projection missing: {SRC_PROJ}"
    assert (SRC_LORA / "adapter_config.json").is_file(), \
        f"source LoRA dir missing: {SRC_LORA}"

    TGT_S1_DIR.mkdir(parents=True, exist_ok=True)
    shutil.copy2(SRC_PROJ, TGT_S1_DIR / "stage1_final_projection.pt")
    print(f"[gcp_finetune_vqa] seeded projection -> {TGT_S1_DIR / 'stage1_final_projection.pt'}")

    # train.py:1043 gates on the LITERAL file `stage1_final.pt` existing.
    # Content never read; it's a marker only.
    (TGT_S1_DIR / "stage1_final.pt").touch()
    print(f"[gcp_finetune_vqa] seeded sentinel   -> {TGT_S1_DIR / 'stage1_final.pt'}")

    tgt_lora = TGT_S1_DIR / "stage1_final_lora"
    if tgt_lora.exists():
        shutil.rmtree(tgt_lora)
    shutil.copytree(SRC_LORA, tgt_lora)
    print(f"[gcp_finetune_vqa] seeded LoRA       -> {tgt_lora}")

    if SRC_CHEXPRT.is_file():
        shutil.copy2(SRC_CHEXPRT, TGT_S1_DIR / "stage1_final_chexpert_classifier.pt")
        print(f"[gcp_finetune_vqa] seeded chexpert  -> "
              f"{TGT_S1_DIR / 'stage1_final_chexpert_classifier.pt'}")

    # Also pull source's config snapshot for the cfg-build step below to use.
    # (Already pulled by snapshot_download above; nothing more to do.)

# ── 5) Build configs (start from source snapshot to preserve architecture) ───
import torch  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402

# Configs: prefer (a) snapshot we just hydrated into TARGET_DIR/configs/,
# (b) source-run snapshot under RUN_PULL_ROOT, else (c) repo default.
def _find_cfg(name: str):
    for cand in (
        TARGET_DIR / "configs" / name,
        RUN_PULL_ROOT / SOURCE_RUN_ID / "configs" / name,
    ):
        if cand.is_file():
            return cand
    return None

SAVED_TRAIN_CFG = _find_cfg("train_config.yaml")
SAVED_MODEL_CFG = _find_cfg("model_config.yaml")
repo_train_cfg_path = PROJECT / "configs" / "train_config.yaml"
repo_model_cfg_path = PROJECT / "configs" / "model_config.yaml"

if SAVED_TRAIN_CFG is not None:
    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)
    print(f"[gcp_finetune_vqa] train_cfg <- {SAVED_TRAIN_CFG}")
else:
    train_cfg = OmegaConf.load(repo_train_cfg_path)
    print(f"[gcp_finetune_vqa] train_cfg <- repo default")

if SAVED_MODEL_CFG is not None:
    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)
    print(f"[gcp_finetune_vqa] model_cfg <- {SAVED_MODEL_CFG}")
else:
    model_cfg = OmegaConf.load(repo_model_cfg_path)

if REPORT_MODE_OVERRIDE:
    train_cfg.data.report_mode = REPORT_MODE_OVERRIDE
if IMAGE_MODE_OVERRIDE:
    train_cfg.data.image_mode  = IMAGE_MODE_OVERRIDE
print(f"[gcp_finetune_vqa] report_mode={train_cfg.data.report_mode}  image_mode={train_cfg.data.image_mode}")

# Dataset paths (mirror gcp_entrypoint.py)
train_cfg.data.dataset_name           = DATASET_NAME
train_cfg.data.max_images_per_sample  = int(getattr(train_cfg.data, "max_images_per_sample", 2))

out_dir = PROJECT / "data" / "data_files"
out_dir.mkdir(parents=True, exist_ok=True)

if DATASET_NAME == "MIMIC-CXR_resized":
    train_cfg.data.mimic_cxr_resized.root          = str(DATA_ROOT_RESIZED)
    train_cfg.data.mimic_cxr_resized.manifest_dir  = None
    train_cfg.data.mimic_cxr_resized.vqa_dir       = None
    train_cfg.data.mimic_cxr_resized.reports_root  = None
    train_cfg.data.mimic_cxr_resized.auto_build    = True
    train_cfg.data.mimic_cxr_resized.instruct_json = str(
        out_dir / "mimic_cxr_resized_instruct.json")
elif DATASET_NAME == "MIMIC-CXR":
    def _find_mimic_root(root: Path) -> Path:
        for cand in [root / "MIMIC-CXR", root]:
            if (cand / "train").exists() and (cand / "valid").exists() and (cand / "test").exists():
                return cand
        for p in root.rglob("train"):
            if p.is_dir() and (p.parent / "valid").exists() and (p.parent / "test").exists():
                return p.parent
        raise FileNotFoundError(f"MIMIC-CXR train/valid/test not found under {root}")
    cxr_root = _find_mimic_root(DATA_SRC)
    train_cfg.data.mimic_cxr_root   = str(cxr_root)
    train_cfg.data.instruct_json    = str(out_dir / "mimic_cxr_instruct_unified.json")
    train_cfg.data.mimic_auto_build = True
    _cx = sorted(DATA_SRC.rglob("*chexpert*.csv")) or sorted(DATA_SRC.rglob("*chexbert*.csv"))
    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None
    _vqa = list(DATA_SRC.rglob("vqa"))
    train_cfg.data.mimic_vqa_root = str(_vqa[0]) if _vqa else None
else:
    iu_root = DATA_SRC / "IU-Xray"
    train_cfg.data.iu_xray.images_dir    = str(iu_root / "images")
    train_cfg.data.iu_xray.labels_dir    = str(iu_root / "labels")
    train_cfg.data.iu_xray.instruct_json = str(out_dir / "iu_xray_instruct.json")
    train_cfg.data.iu_xray.auto_build    = True

train_cfg.data.train_split       = "train"
train_cfg.data.val_split         = "validate"
train_cfg.data.test_split        = "test"
train_cfg.data.feature_cache_dir = None
train_cfg.training.output_root   = str(CKPT_ROOT)

# ── Task mix: VQA-heavy with rehearsal ───────────────────────────────────────
train_cfg.tasks.findings_generation.enabled    = W_FINDINGS > 0
train_cfg.tasks.findings_generation.weight     = W_FINDINGS
train_cfg.tasks.impression_generation.enabled  = W_IMPRESSION > 0
train_cfg.tasks.impression_generation.weight   = W_IMPRESSION
train_cfg.tasks.vqa.enabled                    = W_VQA > 0
train_cfg.tasks.vqa.weight                     = W_VQA

# ── Stage 1 disabled (we seeded its outputs from source). Stage 2 = short
#    finetune with low LR. Stage1 ITC explicitly off — irrelevant here. ─────
train_cfg.stage1.enabled = False
if "itc" in train_cfg.stage1:
    train_cfg.stage1.itc.enabled = False
train_cfg.stage2.enabled       = True
train_cfg.stage2.num_epochs    = S2_EPOCHS
train_cfg.stage2.learning_rate = S2_LR

# ── GPU auto-profile (mirrors gcp_entrypoint.py) ────────────────────────────
assert torch.cuda.is_available(), "CUDA not available in container"
_props   = torch.cuda.get_device_properties(0)
_cap     = (_props.major, _props.minor)
_vram_gb = _props.total_memory / 1e9
_bf16_ok = torch.cuda.is_bf16_supported()
_fa2_ok  = _cap >= (8, 0)

_flash_attn_installed = False
if _fa2_ok:
    try:
        import flash_attn  # noqa: F401
        _flash_attn_installed = True
    except Exception:
        _flash_attn_installed = False

if _vram_gb >= 70:
    _profile = dict(label="A100/H100 80GB",
                    per_device_train_batch_size=8, per_device_eval_batch_size=8,
                    gradient_accumulation_steps=2, dataloader_num_workers=16,
                    gradient_checkpointing=False)
elif _vram_gb >= 35:
    _profile = dict(label="A100 40GB",
                    per_device_train_batch_size=8, per_device_eval_batch_size=8,
                    gradient_accumulation_steps=2, dataloader_num_workers=12,
                    gradient_checkpointing=False)
elif _vram_gb >= 22:
    _profile = dict(label="3090 / L4 / A10 (24GB)",
                    per_device_train_batch_size=8, per_device_eval_batch_size=8,
                    gradient_accumulation_steps=2, dataloader_num_workers=8,
                    gradient_checkpointing=True)
elif _vram_gb >= 14:
    _profile = dict(label="T4 / V100 (15-16GB)",
                    per_device_train_batch_size=1, per_device_eval_batch_size=1,
                    gradient_accumulation_steps=16, dataloader_num_workers=2,
                    gradient_checkpointing=True)
else:
    _profile = dict(label=f"unknown ({_vram_gb:.0f}GB)",
                    per_device_train_batch_size=1, per_device_eval_batch_size=1,
                    gradient_accumulation_steps=16, dataloader_num_workers=2,
                    gradient_checkpointing=True)

_profile["bf16"] = bool(_bf16_ok)
_profile["fp16"] = not _bf16_ok
_profile["attn_implementation"] = (
    "flash_attention_2" if (_fa2_ok and _flash_attn_installed) else "sdpa"
)
_profile["optim"] = "paged_adamw_8bit" if _cap >= (8, 0) else "adamw_torch"
_profile["bnb_4bit_compute_dtype"] = "bfloat16" if _bf16_ok else "float16"
_profile["torch_dtype"]            = "bfloat16" if _bf16_ok else "float16"

print(f"[gcp_finetune_vqa] GPU: {_props.name}  {_vram_gb:.1f}GB  sm_{_cap[0]}{_cap[1]}  "
      f"bf16={_bf16_ok}  fa2={_fa2_ok}  fa2_wheel={_flash_attn_installed}")
print(f"[gcp_finetune_vqa] → {_profile['label']}")

train_cfg.training.per_device_train_batch_size   = _profile["per_device_train_batch_size"]
train_cfg.training.per_device_eval_batch_size    = _profile["per_device_eval_batch_size"]
train_cfg.training.gradient_accumulation_steps   = _profile["gradient_accumulation_steps"]
train_cfg.training.dataloader_num_workers        = _profile["dataloader_num_workers"]
train_cfg.training.fp16                          = _profile["fp16"]
train_cfg.training.bf16                          = _profile["bf16"]
train_cfg.training.dataloader_pin_memory         = True
train_cfg.training.dataloader_persistent_workers = True
train_cfg.training.optim                         = _profile["optim"]

model_cfg.llm.attn_implementation       = _profile["attn_implementation"]
model_cfg.llm.gradient_checkpointing    = _profile["gradient_checkpointing"]
model_cfg.llm.torch_dtype               = _profile["torch_dtype"]
model_cfg.llm.bnb_4bit_compute_dtype    = _profile["bnb_4bit_compute_dtype"]
model_cfg.llm.bnb_4bit_quant_type       = "nf4"
model_cfg.llm.bnb_4bit_use_double_quant = True
model_cfg.llm.load_in_8bit              = False
model_cfg.llm.load_in_4bit              = True

# CheXpert classifier — enable iff its checkpoint came along with source run.
if (TGT_S1_DIR / "stage1_final_chexpert_classifier.pt").is_file():
    model_cfg.chexpert_classifier.enabled = True
else:
    model_cfg.chexpert_classifier.enabled = False

# HF Hub uploads → TARGET_RUN_ID (so source run_1 stays untouched).
train_cfg.wandb.enabled         = False
train_cfg.hf_hub.enabled        = True
train_cfg.hf_hub.repo_id        = HF_RUNS_REPO
train_cfg.hf_hub.token_env      = "HF_TOKEN"
train_cfg.hf_hub.private        = True
train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / "run_id.txt")

# Pin run_id = target so resolve_run_id picks it up.
(CKPT_ROOT / "run_id.txt").write_text(TARGET_RUN_ID)
print(f"[gcp_finetune_vqa] pinned run_id = {TARGET_RUN_ID}")

OmegaConf.save(train_cfg, repo_train_cfg_path)
OmegaConf.save(model_cfg, repo_model_cfg_path)
print("[gcp_finetune_vqa] configs patched.")

# ── 6) Pre-flight check: VQA samples > 0 ─────────────────────────────────────
# Trigger the builder, then verify the train split actually has vqa samples.
# Bail fast — saves hours of compute if the path-fix didn't work.
import json as _json
from utils.dataset_resolver import resolve_dataset_spec  # noqa: E402

print("[gcp_finetune_vqa] pre-flight: triggering dataset builder + verifying VQA …")
spec = resolve_dataset_spec(train_cfg)
samples = _json.load(open(spec.instruct_json, encoding="utf-8"))
train_counts = {}
for s in samples:
    if s.get("split") == "train":
        train_counts[s["task"]] = train_counts.get(s["task"], 0) + 1
print(f"[gcp_finetune_vqa] train-split task counts: {train_counts}")

if STRICT_VQA_REQUIRED and train_counts.get("vqa", 0) == 0:
    print(f"[gcp_finetune_vqa] !! FATAL: train-split has 0 VQA samples.")
    print(f"   The dataset builder dropped all VQA rows — path-format mismatch.")
    print(f"   Inspect the builder log above for the line:")
    print(f"     [mimic_cxr_resized_builder] vqa added/dropped : N / M")
    print(f"   Fix the builder, push to HF, then re-submit. Set "
          f"STRICT_VQA_REQUIRED=0 to bypass this check (not recommended).")
    sys.exit(2)

# ── 7) Launch training (Stage 2 only — Stage 1 seeded from source) ──────────
cmd = [
    "python", "-u", "-m", "training.train",
    "--model_config", str(repo_model_cfg_path),
    "--train_config", str(repo_train_cfg_path),
    "--mode",         "resume",
    "--run_id",       TARGET_RUN_ID,
]
print(f"[gcp_finetune_vqa] launching: {' '.join(cmd)}", flush=True)
os.chdir(PROJECT)
sys.exit(subprocess.call(cmd))