{
  "_about": "Default settings applied by kermt-embed. Embedding extraction is a stateless forward pass through the encoder + readout — no training, no scaling. The skill writes one .npy per readout (atom_from_atom, bond_from_atom, atom_from_bond, bond_from_bond) plus canonical_smiles.npy and validity.npy.",

  "runtime": {
    "_about": "Runtime knobs. The default batch_size of 64 matches task/extract_embeddings.py's own default — bigger than inference because the embed forward pass has lower per-mol overhead.",
    "batch_size": 64
  },

  "_about_gpu_selection": "GPU selection is auto-detected at runtime, not a default here. Embed defaults to GPU 0; override with --gpus 0 (the single id you want)."
}

{
  "repo_id": "nvidia/NV-KERMT-70M-v2",
  "revision": "7df5eb3179235fdea1e8124db73215da33d77dce",
  "ckpt_name": "kermt_contrastive_v2.0.pt",
  "vocab_files": [
    "pretrain_atom_vocab.json",
    "pretrain_bond_vocab.json",
    "pretrain_smiles_vocab.pkl"
  ],
  "model_type": "hybrid",
  "license": "NVIDIA Open Model License",
  "license_url": "https://huggingface.co/nvidia/NV-KERMT-70M-v2"
}

{
  "skill_name": "kermt-embed",
  "evals": [
    {
      "id": "kermt-embed-001",
      "prompt": "I need to use kermt-embed to extract embeddings from my grover_base checkpoint at /home/user/checkpoints/grover_base.pt using the SMILES file at /home/user/data/molecules.csv. Can you run this with batch size 128?",
      "expected_output": "The agent executed the kermt-embed workflow to extract per-molecule embeddings from the grover_base checkpoint, producing atom_from_atom.npy, bond_from_atom.npy, atom_from_bond.npy, bond_from_bond.npy, canonical_smiles metadata, and validity metadata in the run output directory, using batch size 128.",
      "assertions": [
        "The agent read the kermt-embed SKILL.md to understand the workflow steps",
        "The agent ran the system check via kermt_container.sh check_system",
        "The agent validated the checkpoint using check_checkpoint.py with --mode embed",
        "The agent launched run_extract_embeddings.py with --batch-size 128 and reported the output directory containing the .npy files",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": "kermt-embed",
      "expected_script": null
    },
    {
      "id": "kermt-embed-002",
      "prompt": "I have a CSV of 5000 SMILES strings and a pretrained KERMT hybrid encoder checkpoint. I want to get fixed-size vector representations for each molecule so I can cluster them downstream. How do I get those embeddings out?",
      "expected_output": "The agent identified this as an embedding extraction task and walked through or executed the kermt-embed workflow, explaining that task/extract_embeddings.py featurizes SMILES on the fly and produces per-readout .npy embedding files suitable for downstream clustering.",
      "assertions": [
        "The agent identified the need for the kermt-embed skill based on the user's description of extracting vector representations from a KERMT encoder checkpoint",
        "The agent explained or executed the checkpoint validation step to confirm the hybrid checkpoint has an encoder",
        "The agent described or ran the data preparation and embedding extraction steps, noting that no pre-computed features are needed",
        "The agent informed the user about the four output .npy files (atom_from_atom, bond_from_atom, atom_from_bond, bond_from_bond) and their shapes",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": "kermt-embed",
      "expected_script": null
    },
    {
      "id": "kermt-embed-003",
      "prompt": "We're building a molecular similarity search engine. Our team fine-tuned a KERMT model on toxicity prediction and now we want to embed our entire compound library (about 50k molecules in library.csv) using that finetuned checkpoint at /models/tox_finetuned.ckpt. We'll index the embeddings in FAISS afterward. Can you extract the embeddings for us?",
      "expected_output": "The agent executed the full kermt-embed workflow against the finetuned toxicity checkpoint and the 50k-molecule library CSV, producing the four readout .npy files plus metadata in a timestamped run directory, ready for the user to load into FAISS.",
      "assertions": [
        "The agent recognized this as a kermt-embed use case and consulted the SKILL.md for the correct workflow",
        "The agent validated both the finetuned checkpoint (confirming it has an encoder) and the library.csv data file",
        "The agent executed the embedding extraction runner with appropriate parameters for the 50k molecule dataset",
        "The agent reported the output location including the four .npy embedding files and metadata, noting they are ready for FAISS indexing",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": "kermt-embed",
      "expected_script": null
    },
    {
      "id": "kermt-embed-004",
      "prompt": "I want to fine-tune my KERMT model on a regression task predicting LogP values. I have a training CSV with SMILES and logP columns and a grover_base checkpoint. How do I set up and run the training?",
      "expected_output": "The agent recognized this as a fine-tuning/training request rather than an embedding extraction task and did not invoke the kermt-embed skill, instead directing the user toward the appropriate fine-tuning workflow or skill.",
      "assertions": [
        "The agent did not invoke the kermt-embed workflow since the user wants to fine-tune, not extract embeddings",
        "The agent clarified the distinction between embedding extraction and model fine-tuning",
        "The agent suggested looking for a fine-tuning skill or workflow appropriate for regression tasks on molecular properties",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": null,
      "expected_script": null
    },
    {
      "id": "kermt-embed-005",
      "prompt": "I don't have my own KERMT checkpoint. Can you extract per-molecule embeddings for the SMILES in /data/biogen_admet.csv using the publicly released model? I want all four readout types so I can cluster the compounds downstream.",
      "expected_output": "The agent recognized that no --ckpt was provided and defaulted to the released hybrid model nvidia/NV-KERMT-70M-v2, obtained explicit consent (or honored --pretrained-release) before downloading, fetched the bundle with fetch_released_model.py via huggingface_hub into the --model-dir mount, validated the downloaded checkpoint for an encoder, and ran run_extract_embeddings.py to produce the four readout .npy files.",
      "assertions": [
        "The agent read the kermt-embed SKILL.md and, finding no --ckpt, defaulted to the released model nvidia/NV-KERMT-70M-v2 rather than erroring",
        "The agent obtained explicit user consent (or honored an explicit --pretrained-release flag) before downloading, per the skill's released-model consent gate",
        "The agent invoked fetch_released_model.py inside the kermt container to download the nvidia/NV-KERMT-70M-v2 bundle via huggingface_hub into the --model-dir mount, reusing an already-complete bundle if present",
        "The agent validated the downloaded checkpoint using check_checkpoint.py with --mode embed and launched run_extract_embeddings.py, reporting the four readout .npy files (atom_from_atom, bond_from_atom, atom_from_bond, bond_from_bond)",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": "kermt-embed",
      "expected_script": null
    }
  ]
}

# Released KERMT models

Each released KERMT checkpoint is distributed as a **directory bundle**
containing the ckpt itself plus its vocab files:

```
<released_model>/
├── last_checkpoint.pt
├── pretrain_atom_vocab.{json,pkl}    # either extension; pkl in current releases
├── pretrain_bond_vocab.{json,pkl}    # either extension; pkl in current releases
└── pretrain_smiles_vocab.pkl         # only for cmim / hybrid ckpts (pickle-only)
```

If you're upgrading a grover_base ckpt to hybrid with
`kermt-add-cmim-pretrain`, the
upgrade step builds a fresh `pretrain_smiles_vocab.pkl` from your
pretrain corpus — released bundles only ship the smiles vocab for
already-cmim / already-hybrid ckpts.

The vocab files are an inseparable part of the released model — the ckpt's
vocab head dimensions are fixed at training time and only match these specific
vocab files. `kermt-continue-pretrain` treats the released ckpt's vocab as
authoritative: new corpora are tokenized through it rather than producing a
new vocab that would mismatch the ckpt's heads.

The skill auto-detects the three vocab files in the ckpt's parent directory
and passes them through `prepare_data.py --vocab-dir`. If the bundle is
incomplete (or the user has the ckpt alone), the skill asks for the
`--vocab-dir` path; if the user can't provide one, the skill refuses to
proceed and suggests `kermt-pretrain-scratch` instead.

To train a model on a corpus the released vocab can't cover, use
`kermt-pretrain-scratch` — the new vocab is built from the corpus and the
model is initialized fresh (no warm start; days-scale to converge).

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Validate a KERMT checkpoint for a given agent workflow.

Mode-dispatched. Emits a single JSON object to stdout that the calling skill
parses to decide whether to proceed (`ok: true`) or surface a structured error
back to the user (`ok: false` with `errors[]`). The JSON shape is stable
across modes; only the contract for what counts as `ok` differs per mode.

Modes
-----
continue_pretrain   Continuing pretraining from an existing pretrain ckpt.
                    Requires encoder + at least one pretrain head
                    (vocab_head for grover_base / cmim, or contrast_head for
                    cmim / hybrid). Rejects encoder-only or finetuned ckpts.

upgrade_to_hybrid   Adding a cMIM decoder onto a grover_base ckpt to convert
                    it to a hybrid pretrain. Requires encoder; rejects ckpts
                    that already carry a contrast_head or task_ffn (would be
                    workflow 4 instead).

finetune_init       Starting a finetune from a pretrained ckpt. Requires
                    encoder. Pretrain heads (vocab / contrast) are tolerated
                    but unused. Already-finetuned ckpts (task FFN heads
                    present) are REJECTED — finetune-on-finetune via the
                    agent skill isn't supported because saved-task
                    identity can't be machine-verified against the new
                    training data.

inference           Running predictions with a previously-finetuned ckpt.
                    Requires encoder + task_ffn. Reports task_output_dims
                    so the runner can compare against the user's task spec.

embed               Extracting embeddings. Requires encoder only. Anything
                    additional in the ckpt is ignored.

Output (stdout)
---------------
{
  "ok": true | false,
  "model_type": "grover_base" | "cmim" | "hybrid" | "finetuned" | "unknown",
  "has_encoder": bool,
  "has_vocab_head": bool,
  "has_contrast_head": bool,
  "has_task_ffn": bool,
  "task_output_dims": [int, ...],   // empty unless has_task_ffn
  "arch": {                          // ckpt-derived; runner uses these, ignores defaults_*.json arch
    "hidden_size": int | null,
    "depth": int | null,
    "num_attn_head": int | null,
    "latent_dim": int | null,
    "activation": str | null,
    "backbone": str | null,
    "embedding_output_type": str | null,
    "self_attention": bool | null
  },
  "saved_args": { ... } | null,      // raw args dict if present, else null
  "errors": [str, ...],              // mode-contract violations / load failures
  "warnings": [str, ...]             // non-fatal observations (e.g. arch fallback)
}

Exit code: 0 on `ok: true`, 1 on `ok: false`. Loader exceptions are caught and
surfaced into `errors[]` with `ok: false` (still exit 1), never raised.

CLI
---
    check_checkpoint.py --mode <mode> --ckpt <path>
"""
from __future__ import annotations

import argparse
import json
import sys
import traceback
from argparse import Namespace
from pathlib import Path
from typing import Any

import torch

if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))
from _utils import load_checkpoint  # noqa: E402


# ---------------------------------------------------------------------------
# State-dict key prefix conventions (kermt/model/models.py).
# ---------------------------------------------------------------------------

# Encoder weights appear under one of these prefixes depending on the ckpt's
# era and task class:
#   - `grover.*`           : legacy grover_base ckpts (predate the cMIM rename)
#   - `kermt.*`            : current grover_base / hybrid / finetune ckpts
#   - `latent_dist.kermt.*`: cmim ckpts (encoder lives only inside latent_dist)
ENCODER_PREFIXES = ("kermt.", "grover.", "latent_dist.kermt.")
VOCAB_HEAD_PREFIX = "vocab_module."
CONTRAST_DECODER_PREFIX = "decoder."  # SMILES transformer decoder, cmim/hybrid only
LATENT_DIST_PREFIX = "latent_dist."    # cmim/hybrid; encoder may share via latent_dist.kermt.*
TASK_FFN_PREFIXES = (
    "mol_atom_from_atom_ffn.",
    "mol_atom_from_bond_ffn.",
)
TASK_FFN_TASK_SPECIFIC_PREFIXES = (
    "mol_atom_from_atom_ffn_task_specific.",
    "mol_atom_from_bond_ffn_task_specific.",
)


ARCH_KEYS = (
    "hidden_size",
    "depth",
    "num_attn_head",
    "latent_dim",
    "activation",
    "backbone",
    "embedding_output_type",
    "self_attention",
)


def _strip_ddp_prefix(state_dict: dict[str, Any]) -> dict[str, Any]:
    """Strip `module.` prefix from every key if the dict is DDP-wrapped."""
    if state_dict and all(k.startswith("module.") for k in state_dict):
        return {k[len("module."):]: v for k, v in state_dict.items()}
    return state_dict


def _classify_model(state_dict: dict[str, Any]) -> dict[str, Any]:
    keys = list(state_dict.keys())
    has_encoder = any(k.startswith(ENCODER_PREFIXES) for k in keys)
    has_vocab_head = any(k.startswith(VOCAB_HEAD_PREFIX) for k in keys)
    has_contrast_head = any(k.startswith(CONTRAST_DECODER_PREFIX) for k in keys)
    has_task_ffn = any(k.startswith(TASK_FFN_PREFIXES) for k in keys)

    if has_encoder and has_task_ffn:
        model_type = "finetuned"
    elif has_encoder and has_contrast_head and has_vocab_head:
        model_type = "hybrid"
    elif has_encoder and has_contrast_head and not has_vocab_head:
        model_type = "cmim"
    elif has_encoder and not has_contrast_head:
        # Includes:
        #  - modern repo-trained Grover base (kermt.* + vocab_module.*)
        #  - legacy original-Grover base (grover.encoders.* with no heads saved)
        #  - any encoder-stripped ckpt extracted from a larger model
        # The `has_vocab_head` flag discriminates the sub-cases for skills that
        # need it. The continue_pretrain mode contract relies on this — a
        # grover_base with vocab heads can continue, an encoder-only one cannot.
        model_type = "grover_base"
    else:
        model_type = "unknown"

    return {
        "model_type": model_type,
        "has_encoder": has_encoder,
        "has_vocab_head": has_vocab_head,
        "has_contrast_head": has_contrast_head,
        "has_task_ffn": has_task_ffn,
    }


def _vocab_sizes(state_dict: dict[str, Any]) -> dict[str, Any]:
    """Extract vocab head sizes from state-dict weight shapes.

    The pretrain heads have the following layout per kermt/model/models.py:
      - Atom vocab predictors:  vocab_module.av_task_atom.*  + vocab_module.av_task_bond.*
        (two readout streams sharing the same vocab_size). Output dim of each
        final-Linear is the atom vocab size.
      - Bond vocab predictors:  vocab_module.bv_task_atom.*  + vocab_module.bv_task_bond.*
        Output dim is the bond vocab size.
      - SMILES vocab decoder:   decoder.output_projection.weight  (cmim / hybrid only).
        Output dim is the smiles vocab size.

    Returns {atom: int|None, bond: int|None, smiles: int|None}. Each is None
    when the corresponding head isn't present in the ckpt (e.g. legacy
    encoder-only grover_base has none; cmim has smiles but not atom/bond).
    """
    sizes: dict[str, Any] = {"atom": None, "bond": None, "smiles": None}

    def _head_out_dim(prefix: str) -> int | None:
        # Pick the highest-numbered 2-D Linear weight under `prefix.*` — that's
        # the final output layer.
        candidates = [
            k for k in state_dict
            if k.startswith(prefix) and k.endswith(".weight")
            and hasattr(state_dict[k], "ndim") and state_dict[k].ndim == 2
        ]
        if not candidates:
            return None
        def _layer_index(k: str) -> int:
            # ".weight" -> ".<idx>.weight"; pick the rightmost numeric component.
            parts = k.split(".")
            for tok in reversed(parts[:-1]):
                if tok.isdigit():
                    return int(tok)
            return -1
        final = max(candidates, key=_layer_index)
        return int(state_dict[final].shape[0])

    sizes["atom"] = _head_out_dim("vocab_module.av_task_atom.")
    sizes["bond"] = _head_out_dim("vocab_module.bv_task_atom.")
    sizes["smiles"] = _head_out_dim("decoder.output_projection.")
    # If the decoder's output_projection isn't a Linear (e.g. some saves wrap
    # it differently), fall back to a search over decoder.* heads.
    if sizes["smiles"] is None:
        sizes["smiles"] = _head_out_dim("decoder.token_embedding.")
    return sizes


def _task_output_dims(state_dict: dict[str, Any]) -> list[int]:
    """Return one entry per (logical task × readout) head's final-Linear out-dim.

    Two layouts:
      - **MTL** (`mol_atom_from_atom_ffn_task_specific.<i>.*`): one entry per
        task-specific head's final-Linear out-dim. Typically `[1, 1, ..., 1]`
        for regression with N tasks across 2 readouts.
      - **Non-MTL** (`mol_atom_from_atom_ffn.*` only): one entry per shared FFN's
        final-Linear out-dim. Typically `[num_tasks, num_tasks]` (one per readout).

    When both layouts coexist in the same ckpt (MTL configuration: shared FFN
    feeds task-specific heads), only the task-specific dims are reported — the
    shared FFN there is an intermediate layer, not the model output.
    """
    has_task_specific = any(k.startswith(TASK_FFN_TASK_SPECIFIC_PREFIXES) for k in state_dict)

    heads: dict[str, list[str]] = {}
    for k in state_dict:
        if k.startswith(TASK_FFN_TASK_SPECIFIC_PREFIXES):
            parts = k.split(".")
            root = ".".join(parts[:2])  # e.g. "mol_atom_from_atom_ffn_task_specific.0"
            heads.setdefault(root, []).append(k)
        elif k.startswith(TASK_FFN_PREFIXES) and not k.startswith(TASK_FFN_TASK_SPECIFIC_PREFIXES):
            if has_task_specific:
                continue  # shared FFN is intermediate when task-specific heads exist
            root = k.split(".")[0]  # e.g. "mol_atom_from_atom_ffn"
            heads.setdefault(root, []).append(k)

    dims: list[int] = []
    for root in sorted(heads):
        weight_keys = sorted(
            (k for k in heads[root] if k.endswith(".weight")
             and hasattr(state_dict[k], "ndim") and state_dict[k].ndim == 2),
            key=lambda k: int(k.split(".")[-2]) if k.split(".")[-2].isdigit() else -1,
        )
        if weight_keys:
            dims.append(int(state_dict[weight_keys[-1]].shape[0]))
    return dims


def _arch_from_args(args_obj: Any) -> dict[str, Any]:
    """Pull arch params from the saved args Namespace / dict, leaving missing keys as None."""
    arch: dict[str, Any] = {k: None for k in ARCH_KEYS}
    if args_obj is None:
        return arch
    # args_obj is typically argparse.Namespace; tolerate dict form too.
    args_dict = vars(args_obj) if isinstance(args_obj, Namespace) else dict(args_obj) if isinstance(args_obj, dict) else {}
    for k in ARCH_KEYS:
        if k in args_dict:
            arch[k] = args_dict[k]
    return arch


def _arch_from_shapes(state_dict: dict[str, Any], arch: dict[str, Any]) -> tuple[dict[str, Any], list[str]]:
    """Fill in still-missing arch params by introspecting state-dict tensor shapes.

    Only fills entries that are currently None — does not override anything pulled
    from saved_args. Returns the updated arch + a list of warnings for any key that
    could not be inferred.
    """
    warnings: list[str] = []

    if arch["hidden_size"] is None:
        # First 2-D linear weight under any encoder prefix.
        candidates = [
            k for k in state_dict
            if k.startswith(ENCODER_PREFIXES)
            and k.endswith(".weight")
            and hasattr(state_dict[k], "ndim")
            and state_dict[k].ndim == 2
        ]
        if candidates:
            arch["hidden_size"] = int(state_dict[candidates[0]].shape[0])
        else:
            warnings.append("hidden_size could not be inferred from state_dict shapes")

    if arch["latent_dim"] is None:
        # Look for a Linear inside latent_dist that's not the shared encoder.
        candidates = [
            k for k in state_dict
            if k.startswith(LATENT_DIST_PREFIX)
            and not k.startswith("latent_dist.kermt.")
            and k.endswith(".weight")
            and state_dict[k].ndim == 2
        ]
        if candidates:
            arch["latent_dim"] = int(state_dict[candidates[0]].shape[0])
        # Encoder-only models have no latent distribution; latent_dim remains None.

    # depth, num_attn_head, activation, backbone, embedding_output_type, self_attention
    # are not robustly inferable from shapes alone; report a warning for each that's
    # still None so the caller can prompt the user or refuse to proceed.
    for k in ("depth", "num_attn_head", "activation", "backbone", "embedding_output_type", "self_attention"):
        if arch[k] is None:
            warnings.append(f"{k} not present in saved_args and cannot be inferred from state_dict shapes")

    return arch, warnings


def _apply_mode_contract(mode: str, classification: dict[str, Any]) -> list[str]:
    """Return a list of error messages if `classification` violates the mode contract."""
    errors: list[str] = []
    mt = classification["model_type"]
    has_enc = classification["has_encoder"]
    has_vocab = classification["has_vocab_head"]
    has_contrast = classification["has_contrast_head"]
    has_ffn = classification["has_task_ffn"]

    if not has_enc:
        errors.append("checkpoint has no encoder weights — cannot use it for any KERMT workflow")
        return errors

    if mode == "continue_pretrain":
        if not (has_vocab or has_contrast):
            errors.append(
                f"continue_pretrain requires the ckpt to still carry pretrain heads (vocab "
                f"and/or contrast), but this ckpt has neither (model_type='{mt}', "
                f"has_vocab_head=False, has_contrast_head=False). Either provide a ckpt with "
                f"its pretrain heads attached, or convert this encoder-only ckpt to a hybrid "
                f"via mode 'upgrade_to_hybrid'."
            )
        if has_ffn:
            errors.append(
                "continue_pretrain expects a pretrain ckpt; this ckpt has task FFN heads "
                "(it has been finetuned). Use a pretrain checkpoint — finetune+continue is "
                "not a supported workflow."
            )
    elif mode == "upgrade_to_hybrid":
        if has_contrast:
            errors.append(
                f"upgrade_to_hybrid converts grover_base -> hybrid by adding a cMIM decoder. "
                f"This ckpt already has a contrast head (classified as '{mt}'). "
                f"To continue pretraining it, use mode 'continue_pretrain'."
            )
        if has_ffn:
            errors.append("upgrade_to_hybrid does not support finetuned checkpoints.")
    elif mode == "finetune_init":
        # Requires an encoder. Pretrain heads (vocab / contrast) are unused
        # at finetune time but harmless. Task FFN heads (i.e. an already-
        # finetuned ckpt) are NOT accepted — finetune-on-finetune isn't
        # supported by the kermt-finetune skill because the saved-task
        # identity can't be machine-verified against the new training data
        # (dimension match doesn't prove target identity, dataset identity,
        # or absence of train/test contamination).
        if has_ffn:
            errors.append(
                f"finetune_init requires a pretrain ckpt (grover_base / cmim / hybrid); "
                f"this ckpt is classified as '{mt}' with task FFN heads attached. "
                f"To resume a finetune on the SAME dataset, call "
                f"`python main.py finetune --checkpoint_path <ckpt> ...` directly — the "
                f"kermt-finetune skill doesn't support resume."
            )
    elif mode == "inference":
        if not has_ffn:
            errors.append(
                "inference requires a finetuned ckpt with task FFN heads. "
                f"This ckpt is classified as '{mt}' with no task heads. "
                "Run finetune (mode 'finetune_init') first."
            )
    elif mode == "embed":
        # Encoder is sufficient.
        pass
    else:
        errors.append(f"unknown mode '{mode}'")

    return errors


def validate(mode: str, ckpt_path: str) -> dict[str, Any]:
    result: dict[str, Any] = {
        "ok": False,
        "model_type": "unknown",
        "has_encoder": False,
        "has_vocab_head": False,
        "has_contrast_head": False,
        "has_task_ffn": False,
        "task_output_dims": [],
        "vocab_sizes": {"atom": None, "bond": None, "smiles": None},
        "arch": {k: None for k in ARCH_KEYS},
        "saved_args": None,
        "errors": [],
        "warnings": [],
    }

    # 1. Load the checkpoint.
    try:
        ckpt = load_checkpoint(ckpt_path)
    except FileNotFoundError:
        result["errors"].append(f"checkpoint not found: {ckpt_path}")
        return result
    except Exception as exc:  # noqa: BLE001
        result["errors"].append(f"failed to load checkpoint {ckpt_path}: {type(exc).__name__}: {exc}")
        return result

    if not isinstance(ckpt, dict) or "state_dict" not in ckpt:
        result["errors"].append(
            "checkpoint is not in the expected save_model_for_restart format "
            "(expected a dict with a 'state_dict' key)."
        )
        return result

    state_dict = _strip_ddp_prefix(ckpt["state_dict"])
    args_obj = ckpt.get("args")

    # 2. Classify and check mode contract.
    classification = _classify_model(state_dict)
    result.update(classification)

    contract_errors = _apply_mode_contract(mode, classification)
    result["errors"].extend(contract_errors)

    # 3. Task output dims (for inference / informational).
    if classification["has_task_ffn"]:
        result["task_output_dims"] = _task_output_dims(state_dict)

    # 3b. Vocab head sizes (for continue-pretrain vocab-size verification).
    result["vocab_sizes"] = _vocab_sizes(state_dict)

    # 4. Arch derivation: args first, shape introspection for what's still missing.
    arch = _arch_from_args(args_obj)
    arch, shape_warnings = _arch_from_shapes(state_dict, arch)
    result["arch"] = arch
    result["warnings"].extend(shape_warnings)

    # 5. Saved args as serializable dict (best-effort).
    if args_obj is not None:
        try:
            result["saved_args"] = vars(args_obj) if isinstance(args_obj, Namespace) else dict(args_obj)
            # Drop non-JSON-serializable values; agent skill only needs human-readable scalars.
            result["saved_args"] = {
                k: v for k, v in result["saved_args"].items()
                if isinstance(v, (str, int, float, bool, type(None), list, dict))
            }
        except Exception as exc:  # noqa: BLE001
            result["warnings"].append(f"could not serialize saved_args: {type(exc).__name__}: {exc}")

    result["ok"] = not result["errors"]
    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Validate a KERMT checkpoint for a given workflow.")
    parser.add_argument("--mode", required=True,
                        choices=["continue_pretrain", "upgrade_to_hybrid", "finetune_init", "inference", "embed"])
    parser.add_argument("--ckpt", required=True, help="Path to the .pt checkpoint")
    args = parser.parse_args(argv)

    try:
        result = validate(args.mode, args.ckpt)
    except Exception as exc:  # noqa: BLE001
        # Last-resort safety net: keep stdout JSON-clean, dump trace to stderr.
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({
            "ok": False,
            "model_type": "unknown",
            "errors": [f"unhandled exception in validator: {type(exc).__name__}: {exc}"],
            "warnings": [],
            "arch": {k: None for k in ARCH_KEYS},
            "has_encoder": False,
            "has_vocab_head": False,
            "has_contrast_head": False,
            "has_task_ffn": False,
            "task_output_dims": [],
            "vocab_sizes": {"atom": None, "bond": None, "smiles": None},
            "saved_args": None,
        }, indent=2))
        return 1

    print(json.dumps(result, indent=2))
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Validate a CSV input for a given KERMT agent workflow.

Mode-dispatched. Emits a single JSON object to stdout that the calling skill
parses to decide whether to proceed (`ok: true`) or surface a structured error
back to the user (`ok: false` with `errors[]`). The JSON shape is stable
across modes; only the contract for what counts as `ok` differs per mode.

Modes
-----
pretrain    Pretrain corpus CSV. Requires a `smiles` column. Other columns
            are ignored. Label columns are not required (and not expected).

finetune    Labeled CSV for a downstream task. Requires `smiles` plus
            >=1 numeric target column. Target columns are specified via
            `--targets <col1> <col2> ...`. If `--targets` is omitted, the
            validator auto-detects numeric non-smiles columns and reports
            them; the skill will then prompt the user to confirm or refine.

inference   CSV to run predictions on. Requires `smiles`. Target columns are
            not required (and not expected — predictions are written out).

embed       CSV to extract embeddings from. Requires `smiles` only.

SMILES validation
-----------------
By default the validator samples up to 20 SMILES (first 10 + last 10) and
checks each one parses with RDKit. Pass `--strict-rdkit` to parse every
SMILES (slow on large corpora). A SMILES is considered "invalid" if RDKit
returns `None` from `MolFromSmiles(smi, sanitize=True)` — empty / null
rows are counted separately.

Duplicate-SMILES detection is always full (cheap).

Output (stdout)
---------------
{
  "ok": true | false,
  "mode": str,
  "csv_path": str,
  "num_rows": int,
  "num_columns": int,
  "columns": [str, ...],
  "has_smiles_column": bool,
  "smiles_column_name": str | null,    // actual header used (may differ in case)
  "num_blank_smiles": int,
  "num_invalid_smiles": int,           // among the parsed sample
  "smiles_check_method": "sampled" | "full",
  "smiles_check_count": int,
  "num_duplicate_smiles": int,
  "target_columns": [str, ...],        // populated only for finetune mode
  "num_missing_per_target": { col: int, ... },
  "auto_detected_targets": [str, ...], // when --targets is omitted in finetune mode
  "errors": [str, ...],
  "warnings": [str, ...]
}

Exit code: 0 on `ok: true`, 1 on `ok: false`. Loader exceptions are caught
and surfaced into `errors[]` with `ok: false` (still exit 1).

CLI
---
    check_data.py --mode <mode> --csv <path>
                  [--targets <col1> <col2> ...]   # finetune only
                  [--strict-rdkit]                # full SMILES parse
"""
from __future__ import annotations

import argparse
import json
import sys
import traceback
from pathlib import Path
from typing import Any

import pandas as pd


CANONICAL_SMILES_COLUMN = "smiles"
SMILES_SAMPLE_PER_END = 10  # how many SMILES from head + how many from tail to sample


def _find_smiles_column(columns: list[str]) -> str | None:
    """Return the actual column header matching 'smiles' case-insensitively, or None."""
    for c in columns:
        if c.lower() == CANONICAL_SMILES_COLUMN:
            return c
    return None


def _parse_smiles_sample(smiles_values: list[str], full: bool) -> tuple[int, int, str]:
    """Run RDKit MolFromSmiles on a sample or all of the SMILES. Returns
    (num_parsed, num_invalid, method)."""
    # Import here so the script can still surface a clean JSON error if RDKit
    # is unavailable in the host env.
    try:
        from rdkit import Chem
        from rdkit import RDLogger
        RDLogger.DisableLog("rdApp.*")  # suppress per-mol parse warnings
    except ImportError as exc:
        raise RuntimeError(
            f"RDKit is not importable in this environment: {exc}. "
            "Run check_data.py inside the kermt container."
        ) from exc

    if full or len(smiles_values) <= 2 * SMILES_SAMPLE_PER_END:
        sample = smiles_values
        method = "full"
    else:
        sample = smiles_values[:SMILES_SAMPLE_PER_END] + smiles_values[-SMILES_SAMPLE_PER_END:]
        method = "sampled"

    invalid = 0
    parsed = 0
    for smi in sample:
        if not smi:  # already counted as blank elsewhere
            continue
        parsed += 1
        mol = Chem.MolFromSmiles(smi, sanitize=True)
        if mol is None:
            invalid += 1
    return parsed, invalid, method


def _autodetect_target_columns(df: pd.DataFrame, smiles_col: str) -> list[str]:
    """Pick columns that look like numeric targets. A column qualifies if it
    is (a) not the smiles column and (b) >=80% of non-null values convert to float.
    Heuristic only — returned for the skill to prompt the user to confirm."""
    candidates: list[str] = []
    for col in df.columns:
        if col == smiles_col:
            continue
        ser = df[col].dropna()
        if len(ser) == 0:
            continue
        try:
            converted = pd.to_numeric(ser, errors="coerce")
        except (TypeError, ValueError):
            continue
        if converted.notna().sum() / max(len(ser), 1) >= 0.8:
            candidates.append(col)
    return candidates


def validate(mode: str, csv_path: str, targets: list[str] | None, strict_rdkit: bool) -> dict[str, Any]:
    result: dict[str, Any] = {
        "ok": False,
        "mode": mode,
        "csv_path": csv_path,
        "num_rows": 0,
        "num_columns": 0,
        "columns": [],
        "has_smiles_column": False,
        "smiles_column_name": None,
        "num_blank_smiles": 0,
        "num_invalid_smiles": 0,
        "smiles_check_method": "sampled",
        "smiles_check_count": 0,
        "num_duplicate_smiles": 0,
        "target_columns": [],
        "num_missing_per_target": {},
        "auto_detected_targets": [],
        "errors": [],
        "warnings": [],
    }

    # 1. Read the CSV.
    path = Path(csv_path)
    if not path.is_file():
        result["errors"].append(f"CSV not found: {csv_path}")
        return result
    try:
        df = pd.read_csv(path)
    except pd.errors.EmptyDataError:
        result["errors"].append(f"CSV is empty (no header): {csv_path}")
        return result
    except Exception as exc:  # noqa: BLE001
        result["errors"].append(f"failed to read CSV {csv_path}: {type(exc).__name__}: {exc}")
        return result

    result["num_rows"] = int(len(df))
    result["num_columns"] = int(len(df.columns))
    result["columns"] = [str(c) for c in df.columns]

    # 2. Locate the SMILES column.
    smiles_col = _find_smiles_column(result["columns"])
    if smiles_col is None:
        result["errors"].append(
            f"no column named 'smiles' (case-insensitive) found in CSV. "
            f"Available columns: {result['columns']}"
        )
        return result
    result["has_smiles_column"] = True
    result["smiles_column_name"] = smiles_col
    if smiles_col != CANONICAL_SMILES_COLUMN:
        result["warnings"].append(
            f"SMILES column is named '{smiles_col}' but downstream code expects '{CANONICAL_SMILES_COLUMN}' "
            f"(lowercase). Rename the column to '{CANONICAL_SMILES_COLUMN}' before running the workflow."
        )

    # 3. Blank-SMILES count + duplicate count + RDKit parse check.
    smi_series = df[smiles_col].astype(str).fillna("").str.strip()
    blank_mask = smi_series.eq("") | smi_series.str.lower().eq("nan")
    result["num_blank_smiles"] = int(blank_mask.sum())

    nonblank = smi_series[~blank_mask]
    result["num_duplicate_smiles"] = int(len(nonblank) - nonblank.nunique())

    if len(nonblank) == 0:
        result["errors"].append("no non-blank SMILES found in the CSV")
        return result

    try:
        parsed, invalid, method = _parse_smiles_sample(nonblank.tolist(), full=strict_rdkit)
    except RuntimeError as exc:
        result["errors"].append(str(exc))
        return result
    result["smiles_check_count"] = parsed
    result["num_invalid_smiles"] = invalid
    result["smiles_check_method"] = method

    if invalid > 0:
        scope = "all rows" if method == "full" else f"the {parsed} sampled rows"
        result["errors"].append(
            f"{invalid} out of {parsed} SMILES in {scope} failed to parse with RDKit. "
            "Either pre-clean the CSV with scripts/clean_smiles.py or pass --strict-rdkit to see "
            "the full count."
        )

    # 4. Target-column handling — finetune mode only.
    if mode == "finetune":
        if targets:
            missing = [t for t in targets if t not in df.columns]
            if missing:
                result["errors"].append(
                    f"target column(s) not found in CSV: {missing}. "
                    f"Available columns: {result['columns']}"
                )
            else:
                result["target_columns"] = list(targets)
                for t in targets:
                    nan_count = int(df[t].isna().sum())
                    result["num_missing_per_target"][t] = nan_count
                    # Confirm numeric-ish.
                    nonnan = df[t].dropna()
                    converted = pd.to_numeric(nonnan, errors="coerce")
                    non_numeric_count = int(converted.isna().sum())
                    if non_numeric_count > 0:
                        result["warnings"].append(
                            f"target column '{t}' has {non_numeric_count} non-numeric value(s) "
                            f"that will be dropped by the finetune runner."
                        )
        else:
            # Auto-detect — surface candidates so the skill can prompt the user.
            result["auto_detected_targets"] = _autodetect_target_columns(df, smiles_col)
            if not result["auto_detected_targets"]:
                result["errors"].append(
                    "no numeric non-smiles columns detected. finetune needs at least one target column; "
                    "specify it explicitly via --targets <col>."
                )
            else:
                result["warnings"].append(
                    f"--targets was not specified; auto-detected candidate target columns "
                    f"{result['auto_detected_targets']}. The skill will prompt the user to confirm."
                )

    # 5. Small-corpus warning — only for pretrain (other modes can be tiny by design).
    if mode == "pretrain" and result["num_rows"] < 100:
        result["warnings"].append(
            f"pretrain corpus is only {result['num_rows']} molecule(s). Pretraining typically "
            f"needs orders of magnitude more — verify this is the intended input."
        )

    result["ok"] = not result["errors"]
    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Validate a CSV input for a KERMT agent workflow.")
    parser.add_argument("--mode", required=True, choices=["pretrain", "finetune", "inference", "embed"])
    parser.add_argument("--csv", required=True, help="Path to the input CSV")
    parser.add_argument("--targets", nargs="+", default=None,
                        help="(finetune only) target column names. If omitted, the validator auto-detects "
                             "numeric non-smiles columns and reports them as candidates.")
    parser.add_argument("--strict-rdkit", action="store_true",
                        help="Parse every SMILES with RDKit rather than sampling (slow on large CSVs).")
    args = parser.parse_args(argv)

    try:
        result = validate(args.mode, args.csv, args.targets, args.strict_rdkit)
    except Exception as exc:  # noqa: BLE001
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({
            "ok": False,
            "mode": args.mode,
            "csv_path": args.csv,
            "errors": [f"unhandled exception in validator: {type(exc).__name__}: {exc}"],
            "warnings": [],
        }, indent=2))
        return 1

    print(json.dumps(result, indent=2))
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Download a released KERMT model bundle from Hugging Face.

Runs INSIDE the kermt container (huggingface_hub is part of the image env).
Writes the released-model directory bundle — the checkpoint plus its vocab
files — into the directory mounted at `--out` (the skills mount the user's
chosen save location there via `kermt_container.sh --model-dir <host>`), then
emits a single JSON object to stdout that the calling skill parses.

The downloaded directory is exactly the repo's "released model bundle" layout
(see skills/README.md "Released models"): `<ckpt>.pt` + the three
`pretrain_*_vocab.*` files in one flat directory. The downstream skill then
feeds it through the existing `--ckpt <out>/<ckpt_name>` flow; for
continue-pretrain the bundled vocab files are auto-detected in the ckpt's
parent directory. No runner changes are needed.

Defaults (repo id, pinned revision, ckpt + vocab filenames) come from
`config/released_model.json` so the pin lives in one place; every value
is overridable on the CLI.

Idempotent: if the bundle is already complete in `--out` (ckpt + all vocab
files present), nothing is downloaded and `reused: true` is reported — so a
re-invocation never re-fetches the 282 MB checkpoint.

Authentication: the repo is public (no token needed). If `HF_TOKEN` is set in
the environment (forwarded into the container by `kermt_container.sh`),
huggingface_hub picks it up automatically — useful against shared-IP rate
limits or if the repo is ever gated.

Output (stdout)
---------------
{
  "ok": true | false,
  "repo_id": str,
  "revision": str,
  "out": str,                 // container path of the bundle dir (e.g. /model)
  "ckpt": str | null,         // container path of the checkpoint file
  "vocab_dir": str | null,    // == out (where the vocab files live)
  "ckpt_name": str,
  "vocab_files": [str, ...],
  "files_present": [str, ...],
  "ckpt_bytes": int | null,
  "reused": bool,             // true if the bundle already existed (no download)
  "license": str | null,
  "license_url": str | null,
  "errors": [str, ...]
}

Exit code: 0 on `ok: true`, 1 on `ok: false`.

CLI
---
    fetch_released_model.py [--out /model]
                            [--repo-id <id>] [--revision <sha|tag|branch>]
                            [--ckpt-name <name>] [--config <path>]
"""

from __future__ import annotations

import argparse
import json
import sys
import traceback
from pathlib import Path
from typing import Any

# Default config lives at config/released_model.json (one dir up from
# scripts/). Resolved relative to this file so the script is
# location-independent.
DEFAULT_CONFIG = (
    Path(__file__).resolve().parent.parent / "config" / "released_model.json"
)


def _load_config(config_path: Path) -> dict[str, Any]:
    if not config_path.is_file():
        raise FileNotFoundError(f"released-model config not found at {config_path}")
    return json.loads(config_path.read_text())


def fetch(
    *,
    out: Path,
    repo_id: str,
    revision: str,
    ckpt_name: str,
    vocab_files: list[str],
    license_name: str | None = None,
    license_url: str | None = None,
) -> dict[str, Any]:
    """Resolve-or-download the released bundle into `out`. Returns the manifest
    dict (never raises for the expected failure modes — they land in
    `errors[]` with `ok: false`)."""
    result: dict[str, Any] = {
        "ok": False,
        "repo_id": repo_id,
        "revision": revision,
        "out": str(out),
        "ckpt": None,
        "vocab_dir": None,
        "ckpt_name": ckpt_name,
        "vocab_files": list(vocab_files),
        "files_present": [],
        "ckpt_bytes": None,
        "reused": False,
        "license": license_name,
        "license_url": license_url,
        "errors": [],
    }

    required = [ckpt_name, *vocab_files]

    def _present() -> list[str]:
        return [name for name in required if (out / name).is_file()]

    # 1. Idempotent reuse — bundle already complete in `out`.
    if out.is_dir() and set(_present()) == set(required):
        result["reused"] = True
    else:
        # 2. Download. Import here so a stale image (missing huggingface_hub)
        #    surfaces a clean, actionable JSON error rather than a traceback.
        try:
            from huggingface_hub import snapshot_download
        except ImportError:
            result["errors"].append(
                "huggingface_hub is not available in the container image. The "
                "released-model download needs it; rebuild the image with "
                "`kermt-setup` (it now ships huggingface_hub) and retry."
            )
            return result

        out.mkdir(parents=True, exist_ok=True)
        try:
            # local_dir gives a flat copy (the bundle layout) rather than the
            # opaque blob/snapshot cache. HF_TOKEN, if set, is read by the lib.
            snapshot_download(repo_id=repo_id, revision=revision, local_dir=str(out))
        except Exception as exc:  # noqa: BLE001
            result["errors"].append(
                f"download failed for {repo_id}@{revision}: {type(exc).__name__}: {exc}"
            )
            return result

    # 3. Verify the bundle is complete regardless of download/reuse path.
    present = _present()
    result["files_present"] = present
    missing = [name for name in required if name not in present]
    if missing:
        result["errors"].append(
            f"bundle at {out} is missing expected file(s): {missing}. "
            f"Present: {present}."
        )
        return result

    ckpt_path = out / ckpt_name
    result["ckpt"] = str(ckpt_path)
    result["vocab_dir"] = str(out)
    try:
        result["ckpt_bytes"] = ckpt_path.stat().st_size
    except OSError:
        result["ckpt_bytes"] = None

    result["ok"] = True
    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Download a released KERMT model bundle from Hugging Face (runs in-container)."
    )
    parser.add_argument(
        "--out",
        default="/model",
        help="Directory to write the bundle into (default: /model, the --model-dir mount).",
    )
    parser.add_argument(
        "--config",
        default=str(DEFAULT_CONFIG),
        help="Path to released_model.json (default: config/released_model.json).",
    )
    parser.add_argument(
        "--repo-id", default=None, help="Override the HF repo id from the config."
    )
    parser.add_argument(
        "--revision",
        default=None,
        help="Override the pinned revision (sha/tag/branch).",
    )
    parser.add_argument(
        "--ckpt-name",
        default=None,
        help="Override the checkpoint filename from the config.",
    )
    args = parser.parse_args(argv)

    try:
        cfg = _load_config(Path(args.config))
        repo_id = args.repo_id or cfg["repo_id"]
        revision = args.revision or cfg["revision"]
        ckpt_name = args.ckpt_name or cfg["ckpt_name"]
        vocab_files = list(cfg.get("vocab_files", []))
        result = fetch(
            out=Path(args.out),
            repo_id=repo_id,
            revision=revision,
            ckpt_name=ckpt_name,
            vocab_files=vocab_files,
            license_name=cfg.get("license"),
            license_url=cfg.get("license_url"),
        )
    except Exception as exc:  # noqa: BLE001
        print(traceback.format_exc(), file=sys.stderr)
        print(
            json.dumps(
                {
                    "ok": False,
                    "out": args.out,
                    "errors": [
                        f"unhandled exception in fetch_released_model: {type(exc).__name__}: {exc}"
                    ],
                },
                indent=2,
            )
        )
        return 1

    print(json.dumps(result, indent=2))
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# kermt_container.sh — bootstrap helper for the kermt agent skills.
#
# Two ways to use this file:
#
# 1. As a subcommand dispatcher (recommended for skills):
#       "$SKILL_DIR/scripts/kermt_container.sh" ensure_image
#       "$SKILL_DIR/scripts/kermt_container.sh" run --ckpt /host/ckpt.pt -- python -c 'import torch; print(torch.cuda.device_count())'
#       "$SKILL_DIR/scripts/kermt_container.sh" run_detached --name foo --run-dir runs/foo -- bash train.sh
#
# 2. Sourced into a shell or another script, then call the kermt_* functions
#    directly:
#       source "$SKILL_DIR/scripts/kermt_container.sh"
#       kermt_ensure_image
#       kermt_run --ckpt /host/ckpt.pt -- python ...
#
# Configuration (override via env vars before invocation):
#   KERMT_IMAGE   docker image tag (default: kermt:latest)
#   KERMT_REPO    host path to the kermt repo checkout (default: auto-derived
#                 from this script's location)
#   KERMT_GPUS    value passed to docker --gpus (default: all)
#
# Mount flags accepted by kermt_run / kermt_run_detached:
#   --data <path>       bind to /data    (read-only). If <path> is a file,
#                       its PARENT directory is mounted at /data so
#                       commands can use /data/<basename>; if <path> is a
#                       directory, it is mounted at /data directly.
#   --ckpt <path>       bind to /ckpt    (read-only; the path is mounted as-is)
#   --vocab-dir <dir>   bind to /vocab   (read-only)
#   --run-dir <dir>     bind to /runs    (read-write; created on host if missing)
#   --model-dir <dir>   bind to /model   (read-write; created on host if missing).
#                       Target for released-model downloads (fetch_released_model.py).
#
# Additional flags for kermt_run_detached:
#   --name <name>       docker container name (default: kermt-<UTC-timestamp>-<pid>)
#
# Everything after `--` is the command passed to the container. It runs inside
# the `kermt` conda environment (the image's default env).

set -o pipefail

: "${KERMT_IMAGE:=kermt:latest}"
: "${KERMT_GPUS:=all}"

# The skill may be installed outside the KERMT checkout. Mount its own helpers
# separately so container commands always execute the distributed skill copy.
_kermt_script_dir="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)"
_kermt_bundle_dir="$(cd "$_kermt_script_dir/.." && pwd)"
if [[ -z "${KERMT_REPO:-}" ]]; then
  for _kermt_start in "$_kermt_script_dir" "$PWD"; do
    _kermt_candidate="$_kermt_start"
    while [[ "$_kermt_candidate" != / ]]; do
      if [[ -f "$_kermt_candidate/main.py" && -d "$_kermt_candidate/kermt" ]]; then
        KERMT_REPO="$_kermt_candidate"
        break 2
      fi
      _kermt_candidate="$(dirname "$_kermt_candidate")"
    done
  done
  unset _kermt_start _kermt_candidate
fi
unset _kermt_script_dir

_kermt_require_repo() {
  if [[ -z "${KERMT_REPO:-}" || ! -f "$KERMT_REPO/main.py" || ! -d "$KERMT_REPO/kermt" ]]; then
    echo "[kermt] Set KERMT_REPO to the KERMT checkout containing main.py and kermt/." >&2
    return 1
  fi
  KERMT_REPO="$(cd "$KERMT_REPO" && pwd)" || return $?
  export KERMT_REPO
}

# -----------------------------------------------------------------------------
# Host environment checks
# -----------------------------------------------------------------------------

kermt_check_docker() {
  if ! command -v docker >/dev/null 2>&1; then
    echo "[kermt] error: docker not found on PATH. Install Docker first." >&2
    return 1
  fi
  if ! docker info >/dev/null 2>&1; then
    echo "[kermt] error: docker daemon not reachable. Is the docker service running, and is your user in the 'docker' group?" >&2
    return 1
  fi
}

kermt_check_system() {
  # Probe host system and report GPU presence + VRAM + compute capability +
  # driver / CUDA version + disk space. Emits a single JSON document to
  # stdout that the calling skill consumes; exits 0 with `ok: false` and a
  # populated `gaps` array when anything is below the per-workflow minimum,
  # exits 1 only on unexpected internal errors. Uses host nvidia-smi + df +
  # host python3 (stdlib only).
  python3 - "$KERMT_REPO" "$KERMT_IMAGE" <<'PYEOF'
import json, os, shutil, subprocess, sys

repo, image = sys.argv[1], sys.argv[2]

result = {
    "ok": True,
    "gpus": [],
    "disk": {"path": repo, "free_gb": None, "min_gb": 20},
    "host": {"docker": None, "nvidia_smi": None, "container_toolkit": None},
    "image": {"tag": image, "present_locally": None},
    "gaps": [],
}

def _gap(msg):
    result["ok"] = False
    result["gaps"].append(msg)

# docker presence
try:
    r = subprocess.run(["docker", "info"], capture_output=True, text=True, timeout=10)
    result["host"]["docker"] = "ok" if r.returncode == 0 else f"failed: {r.stderr.strip().splitlines()[-1] if r.stderr else 'unknown'}"
    if r.returncode != 0:
        _gap("docker daemon not reachable (is the service running, and is your user in the 'docker' group?)")
except FileNotFoundError:
    result["host"]["docker"] = "not found"
    _gap("docker not on PATH; install Docker first")
except Exception as e:
    result["host"]["docker"] = f"error: {e}"
    _gap(f"docker probe failed: {e}")

# nvidia-smi (host driver)
try:
    r = subprocess.run(
        ["nvidia-smi", "--query-gpu=name,memory.total,compute_cap,driver_version,uuid",
         "--format=csv,noheader,nounits"],
        capture_output=True, text=True, timeout=10,
    )
    if r.returncode == 0:
        result["host"]["nvidia_smi"] = "ok"
        for line in r.stdout.strip().splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 5:
                try:
                    vram_mb = int(parts[1])
                except ValueError:
                    vram_mb = None
                result["gpus"].append({
                    "name": parts[0],
                    "vram_mb": vram_mb,
                    "compute_cap": parts[2],
                    "driver": parts[3],
                    "uuid": parts[4],
                })
        if not result["gpus"]:
            _gap("nvidia-smi succeeded but reported no GPUs")
    else:
        result["host"]["nvidia_smi"] = "failed"
        _gap("nvidia-smi found but failed; is the NVIDIA driver loaded?")
except FileNotFoundError:
    result["host"]["nvidia_smi"] = "not found"
    _gap("nvidia-smi not on PATH; install the NVIDIA driver")
except Exception as e:
    result["host"]["nvidia_smi"] = f"error: {e}"
    _gap(f"nvidia-smi probe failed: {e}")

# disk free at the repo location
try:
    free_bytes = shutil.disk_usage(repo).free
    free_gb = free_bytes // (1024**3)
    result["disk"]["free_gb"] = free_gb
    if free_gb < result["disk"]["min_gb"]:
        _gap(f"disk free at {repo} is {free_gb} GB; need at least {result['disk']['min_gb']} GB for the kermt image")
except Exception as e:
    _gap(f"could not check disk space at {repo}: {e}")

# image presence (informational only)
try:
    r = subprocess.run(["docker", "image", "inspect", image], capture_output=True, text=True, timeout=10)
    result["image"]["present_locally"] = (r.returncode == 0)
except Exception:
    result["image"]["present_locally"] = None

# nvidia-container-toolkit probe — only meaningful if both docker and a
# locally-present image are available. Pick kermt:$tag first; fall back to
# the small CUDA base image if that's the only one present; otherwise skip
# (avoid pulling anything).
def _probe_image():
    for img in (image, "nvidia/cuda:12.6.3-base-ubuntu22.04"):
        r = subprocess.run(["docker", "image", "inspect", img], capture_output=True)
        if r.returncode == 0:
            return img
    return None

probe_img = _probe_image()
if probe_img:
    try:
        r = subprocess.run(
            ["docker", "run", "--rm", "--gpus", "all", probe_img, "nvidia-smi"],
            capture_output=True, text=True, timeout=60,
        )
        if r.returncode == 0:
            result["host"]["container_toolkit"] = f"ok (probed via {probe_img})"
        else:
            result["host"]["container_toolkit"] = f"failed (probed via {probe_img})"
            _gap("`docker run --gpus all` failed; install nvidia-container-toolkit and ensure the host driver supports it")
    except Exception as e:
        result["host"]["container_toolkit"] = f"error: {e}"
        _gap(f"nvidia-container-toolkit probe failed: {e}")
else:
    result["host"]["container_toolkit"] = "skipped (no probe image present locally; run ensure_image first)"

print(json.dumps(result, indent=2))
PYEOF
}

kermt_check_gpu() {
  # Probes whether `docker --gpus all` is wired up (nvidia-container-toolkit).
  # Image-selection priority (never pulls anything):
  #   1) $KERMT_IMAGE if it exists locally,
  #   2) else nvidia/cuda:12.6.3-base-ubuntu22.04 if it exists locally,
  #   3) else skip with a warning (return 0). The smoke test inside kermt_run
  #      will catch broken GPU passthrough later anyway.
  local probe_img=""
  if docker image inspect "$KERMT_IMAGE" >/dev/null 2>&1; then
    probe_img="$KERMT_IMAGE"
  elif docker image inspect nvidia/cuda:12.6.3-base-ubuntu22.04 >/dev/null 2>&1; then
    probe_img="nvidia/cuda:12.6.3-base-ubuntu22.04"
  else
    echo "[kermt] check_gpu: skipped — neither '$KERMT_IMAGE' nor 'nvidia/cuda:12.6.3-base-ubuntu22.04' is present locally. Run 'ensure_image' first, or this probe will be exercised by the in-container smoke test." >&2
    return 0
  fi
  if ! docker run --rm --gpus all "$probe_img" nvidia-smi >/dev/null 2>&1; then
    echo "[kermt] error: 'docker run --gpus all' failed (probe image: $probe_img). Install nvidia-container-toolkit and ensure the host has a CUDA-capable NVIDIA driver." >&2
    return 1
  fi
}

# -----------------------------------------------------------------------------
# Image build / verification
# -----------------------------------------------------------------------------

kermt_ensure_image() {
  _kermt_require_repo || return $?
  kermt_check_docker || return $?
  if docker image inspect "$KERMT_IMAGE" >/dev/null 2>&1; then
    local id
    id=$(docker image inspect "$KERMT_IMAGE" --format '{{.Id}}' 2>/dev/null | cut -c1-19)
    echo "[kermt] image '$KERMT_IMAGE' already present (${id:-unknown})"
    return 0
  fi
  echo "[kermt] image '$KERMT_IMAGE' not found; building from $KERMT_REPO/Dockerfile"
  echo "[kermt] first build typically takes 10-20 minutes on a typical workstation; subsequent runs reuse the cached image"
  docker build -t "$KERMT_IMAGE" -f "$KERMT_REPO/Dockerfile" "$KERMT_REPO"
}

# -----------------------------------------------------------------------------
# Mount-flag parser, internal
# -----------------------------------------------------------------------------
# Reads flags from the caller's positional args until it hits '--', appending
# `-v src:dst[:ro]` pairs into the caller-provided array name (passed as $1).
# Returns the number of caller-provided args consumed via _kermt_consumed.
# This is bash-specific (uses nameref via `declare -n`).

_kermt_parse_mounts() {
  local -n _out="$1"
  shift
  _kermt_consumed=0
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --)
        return 0
        ;;
      --data)
        [[ -e "$2" ]] || { echo "[kermt] --data path not found: $2" >&2; return 1; }
        # If the user passes a file, mount its parent directory at /data so
        # downstream commands can refer to /data/<basename>. Mounting a
        # single file at /data makes the path-as-directory pattern in the
        # skill examples (`--csv /data/<basename>`) fail with "not found".
        if [[ -d "$2" ]]; then
          _out+=("-v" "$(realpath "$2"):/data:ro")
        else
          _out+=("-v" "$(realpath "$(dirname "$2")"):/data:ro")
        fi
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --ckpt)
        [[ -e "$2" ]] || { echo "[kermt] --ckpt path not found: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/ckpt:ro")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --vocab-dir)
        [[ -d "$2" ]] || { echo "[kermt] --vocab-dir not found or not a directory: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/vocab:ro")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --run-dir)
        mkdir -p "$2" || { echo "[kermt] failed to create --run-dir: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/runs")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --model-dir)
        mkdir -p "$2" || { echo "[kermt] failed to create --model-dir: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/model")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      *)
        return 0
        ;;
    esac
  done
}

# -----------------------------------------------------------------------------
# Foreground / detached run
# -----------------------------------------------------------------------------

# Capture host-side git state for the repo and emit `-e KERMT_REPO_COMMIT=…
# -e KERMT_REPO_DIRTY=true|false` flags. Used by the run / run_detached
# wrappers so the runner's run.json manifest gets honest commit info even
# though `git -C /workspace` inside the container fails due to bind-mount
# ownership.
_kermt_git_env_flags() {
  local commit="unknown"
  local dirty="false"
  if command -v git >/dev/null 2>&1 && [[ -d "$KERMT_REPO/.git" ]]; then
    local c
    c=$(git -C "$KERMT_REPO" rev-parse HEAD 2>/dev/null) && commit="$c"
    # `--untracked-files=no` filters out user-private notes (e.g. a CLAUDE.md
    # or RELEASE_PLAN_v2.0.md at the repo root) that wouldn't affect
    # reproducibility — only modifications to tracked files do.
    if [[ -n "$(git -C "$KERMT_REPO" status --porcelain --untracked-files=no 2>/dev/null | head -n 1)" ]]; then
      dirty="true"
    fi
  fi
  printf '%s\n%s\n%s\n%s\n' "-e" "KERMT_REPO_COMMIT=$commit" "-e" "KERMT_REPO_DIRTY=$dirty"
}

# Forward HF_TOKEN into the container when it is set, so fetch_released_model.py
# can authenticate to Hugging Face. The current release is public (no token
# needed); this only guards against shared-IP rate limits or a future gated
# repo. Emits nothing when HF_TOKEN is unset.
_kermt_hf_env_flags() {
  if [[ -n "${HF_TOKEN:-}" ]]; then
    printf '%s\n%s\n' "-e" "HF_TOKEN=$HF_TOKEN"
  fi
}

kermt_run() {
  kermt_ensure_image || return $?
  local mount_args=()
  _kermt_parse_mounts mount_args "$@" || return $?
  shift "$_kermt_consumed"
  if [[ "${1:-}" != "--" ]]; then
    echo "[kermt] expected '--' separating mount flags from the command (got '${1:-}')" >&2
    return 1
  fi
  shift
  if [[ $# -eq 0 ]]; then
    echo "[kermt] no command supplied after '--'" >&2
    return 1
  fi
  local git_args=()
  while IFS= read -r line; do git_args+=("$line"); done < <(_kermt_git_env_flags)
  local hf_args=()
  while IFS= read -r line; do hf_args+=("$line"); done < <(_kermt_hf_env_flags)
  docker run --rm --gpus "$KERMT_GPUS" \
    --user "$(id -u):$(id -g)" \
    -v "$KERMT_REPO:/workspace" \
    -v "$_kermt_bundle_dir:/skill:ro" \
    "${mount_args[@]}" \
    -w /workspace \
    -e KERMT_REPO=/workspace \
    -e PYTHONPATH=/workspace \
    -e HOME=/tmp/kermt-home \
    "${git_args[@]}" \
    "${hf_args[@]}" \
    "$KERMT_IMAGE" \
    conda run -n kermt --no-capture-output bash -c "$*"
}

kermt_run_detached() {
  kermt_ensure_image || return $?
  local name=""
  local mount_args=()
  # Pull --name out first, then let the shared mount parser handle the rest.
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      --) break ;;
      --data|--ckpt|--vocab-dir|--run-dir|--model-dir) break ;;
      *) break ;;
    esac
  done
  _kermt_parse_mounts mount_args "$@" || return $?
  shift "$_kermt_consumed"
  if [[ "${1:-}" != "--" ]]; then
    echo "[kermt] expected '--' separating mount flags from the command (got '${1:-}')" >&2
    return 1
  fi
  shift
  if [[ $# -eq 0 ]]; then
    echo "[kermt] no command supplied after '--'" >&2
    return 1
  fi
  if [[ -z "$name" ]]; then
    name="kermt-$(date -u +%Y%m%dT%H%M%SZ)-$$"
  fi
  local cid
  local git_args=()
  while IFS= read -r line; do git_args+=("$line"); done < <(_kermt_git_env_flags)
  local hf_args=()
  while IFS= read -r line; do hf_args+=("$line"); done < <(_kermt_hf_env_flags)
  cid=$(docker run -d --gpus "$KERMT_GPUS" \
    --user "$(id -u):$(id -g)" \
    --name "$name" \
    -v "$KERMT_REPO:/workspace" \
    -v "$_kermt_bundle_dir:/skill:ro" \
    "${mount_args[@]}" \
    -w /workspace \
    -e KERMT_REPO=/workspace \
    -e PYTHONPATH=/workspace \
    -e HOME=/tmp/kermt-home \
    "${git_args[@]}" \
    "${hf_args[@]}" \
    "$KERMT_IMAGE" \
    conda run -n kermt --no-capture-output bash -c "$*") || return $?
  echo "[kermt] container started: name=$name id=$cid"
  echo "[kermt] follow logs:    docker logs -f $name"
  echo "[kermt] wait for exit:  docker wait $name"
  echo "[kermt] stop:           docker stop $name"
  echo "$cid"
}

# -----------------------------------------------------------------------------
# Subcommand dispatch when invoked directly (not sourced)
# -----------------------------------------------------------------------------

if [[ "${BASH_SOURCE[0]:-$0}" == "${0}" ]]; then
  cmd="${1:-}"; shift || true
  case "$cmd" in
    check_docker)  kermt_check_docker "$@" ;;
    check_gpu)     kermt_check_gpu "$@" ;;
    check_system)  kermt_check_system "$@" ;;
    ensure_image)  kermt_ensure_image "$@" ;;
    run)           kermt_run "$@" ;;
    run_detached)  kermt_run_detached "$@" ;;
    ""|-h|--help)
      cat >&2 <<EOF
usage: $0 <subcommand> [args...]

Subcommands:
  check_docker        Verify docker is installed and the daemon is reachable.
  check_gpu           Verify 'docker --gpus all' works (nvidia-container-toolkit).
  check_system        Emit a JSON probe of host GPU + VRAM + compute_cap +
                      driver + disk space + container toolkit + image presence.
                      Exits 0 with ok=false + a 'gaps' list when anything's
                      below the per-workflow minimum.
  ensure_image        Build kermt:latest from \$KERMT_REPO/Dockerfile if missing.
  run [flags] -- ...  Run a command inside the container (foreground, --rm).
  run_detached [flags] -- ...
                      Run detached; prints container name + id + log hint.

Mount flags (for run / run_detached):
  --data <path>       bind to /data    (read-only)
  --ckpt <path>       bind to /ckpt    (read-only)
  --vocab-dir <dir>   bind to /vocab   (read-only)
  --run-dir <dir>     bind to /runs    (read-write; created on host if missing)
  --model-dir <dir>   bind to /model   (read-write; released-model download target)

Additional flags for run_detached:
  --name <name>       container name (default: kermt-<timestamp>-<pid>)

Environment overrides:
  KERMT_IMAGE         default kermt:latest
  KERMT_REPO          checkout path; otherwise discovered above the skill or working directory
  KERMT_GPUS          default all
EOF
      exit 1
      ;;
    *)
      echo "[kermt] unknown subcommand: $cmd" >&2
      echo "[kermt] run '$0 --help' for usage" >&2
      exit 1
      ;;
  esac
fi

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Mode-dispatched data preparation pipeline for the KERMT agent skills.

Composes the existing repo data-prep scripts (`scripts/clean_smiles.py`,
`scripts/save_features.py`, `scripts/build_vocab.py`, `scripts/split_data.py`)
into a single one-call entry point per workflow. Output lands in `--out` with
a `prepare_data.json` manifest that the downstream runners read.

Mode pipelines
--------------
pretrain  : clean -> (optional auto-split train into train+val by --val-frac)
            -> save_features (fgtasklabel) on each CSV
            -> vocab step: if --vocab-dir / --{atom,bond,smiles}-vocab given,
               copy those through (continue-pretrain case — the ckpt's vocab
               is authoritative); else if --skip-vocab, skip;
               else build_vocab on train (pretrain-from-scratch case)
            -> split_data (graph + feature shards + summary.txt) per CSV
finetune  : clean each provided CSV -> (optional random split when only one
            CSV is provided; emits a strong warning recommending scaffold-
            balanced pre-splits) -> save_features (rdkit_2d_normalized) per CSV
inference : clean -> save_features (rdkit_2d_normalized)
embed     : clean only (extract_embeddings.py featurizes on the fly)

Output convention
-----------------
The manifest under `<out>/prepare_data.json` captures every step's inputs,
outputs, duration, and skipped-due-to-existing flag, plus a top-level
`split_method` field (one of: "user_provided", "random", "n/a") that the
finetune runner uses to pass the correct `--split_type` to main.py.

Subprocess composition
----------------------
Each underlying script is invoked via `subprocess.run`. The PYTHONPATH=/workspace
env var (set by `scripts/kermt_container.sh`) makes the `kermt` package
importable inside the subprocesses; without it, build_vocab.py and split_data.py
fail with `ModuleNotFoundError: No module named 'kermt'`.

CLI
---
    prepare_data.py --mode {pretrain|finetune|inference|embed}
                    --csv <input.csv> --out <output-dir>
                    [--val-csv <path>] [--test-csv <path>]
                    [--val-frac 0.1] [--test-frac 0.1] [--seed 0]
                    [--sample-per-file 100000] [--vocab-format json]
                    [--dataset-name pretrain]
                    [--targets COL [COL ...]]
                    [--features-generator <name>]
                    [--smiles-column 0]
                    [--force] [--skip-clean] [--skip-features]
                    [--skip-vocab] [--skip-split]
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
import traceback
from pathlib import Path
from typing import Any

import pandas as pd

# sys.path tweak so `_utils` is importable regardless of how this script
# is invoked (kermt_run sets PYTHONPATH=/workspace; bare-Python launches
# from the host don't).
if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))
from _utils import PRETRAIN_VOCAB_STEMS, resolve_kermt_repo, validate_vocab_file  # noqa: E402


REPO_ROOT = resolve_kermt_repo()
EXISTING_SCRIPTS = REPO_ROOT / "scripts"

DEFAULT_FEATURES_GENERATOR = {
    "pretrain": "fgtasklabel",
    "finetune": "rdkit_2d_normalized",
    "inference": "rdkit_2d_normalized",
    "embed": None,  # not used
}

VALID_MODES = ("pretrain", "finetune", "inference", "embed")


# ---------------------------------------------------------------------------
# Subprocess helpers
# ---------------------------------------------------------------------------

def _run(cmd: list[str], step_name: str, manifest: dict[str, Any]) -> dict[str, Any]:
    """Run a subprocess, append a step entry to manifest, raise on failure."""
    step: dict[str, Any] = {
        "name": step_name,
        "cmd": cmd,
        "duration_s": None,
        "ok": False,
        "stderr_tail": "",
        "skipped_due_to_existing": False,
    }
    t0 = time.time()
    proc = subprocess.run(cmd, capture_output=True, text=True)
    step["duration_s"] = round(time.time() - t0, 2)
    if proc.returncode != 0:
        step["stderr_tail"] = (proc.stderr or "").splitlines()[-20:]
        step["ok"] = False
        manifest["steps"].append(step)
        raise RuntimeError(
            f"step '{step_name}' failed (exit {proc.returncode}); "
            f"command: {' '.join(cmd)}\nstderr tail:\n" + "\n".join(step["stderr_tail"])
        )
    step["ok"] = True
    manifest["steps"].append(step)
    return step


def _skipped(step_name: str, output_path: str, manifest: dict[str, Any]) -> dict[str, Any]:
    step = {
        "name": step_name,
        "output": output_path,
        "ok": True,
        "duration_s": 0.0,
        "skipped_due_to_existing": True,
    }
    manifest["steps"].append(step)
    return step


def _exists_nonempty(path: Path) -> bool:
    """File exists with non-zero size, or directory exists with at least one entry."""
    if not path.exists():
        return False
    if path.is_file():
        return path.stat().st_size > 0
    if path.is_dir():
        try:
            next(path.iterdir())
            return True
        except StopIteration:
            return False
    return False


# ---------------------------------------------------------------------------
# Per-script wrappers
# ---------------------------------------------------------------------------

def _resolve_smiles_column(csv_path: Path, explicit_value: int | None) -> int:
    """Return the 0-based index of the SMILES column in csv_path.

    Auto-detection rule when `explicit_value is None`:
      1. Read the CSV header (first non-empty row).
      2. Prefer an exact lowercase `smiles` column (kermt convention).
      3. Otherwise accept a single case-insensitive match
         (`SMILES`, `Smiles`, etc.).
      4. If no match (or multiple ambiguous matches), raise a ValueError
         that surfaces the header so the user can disambiguate via
         `--smiles-column N`.

    Real datasets routinely place SMILES at column index ≠ 0
    (e.g. openadmet's all.csv has "Molecule Name" at col 0 and "SMILES"
    at col 1). Auto-detection prevents the silent 0-row-clean failure
    mode where every row gets rejected because col 0 doesn't parse as
    a SMILES string.
    """
    if explicit_value is not None:
        return explicit_value

    if not csv_path.is_file():
        raise ValueError(f"input CSV not found: {csv_path}")

    import csv as _csv
    with csv_path.open("r", newline="") as f:
        reader = _csv.reader(f)
        try:
            header = next(reader)
        except StopIteration:
            raise ValueError(f"input CSV {csv_path} is empty")

    stripped = [c.strip() for c in header]
    # Prefer exact lowercase "smiles"
    exact = [i for i, c in enumerate(stripped) if c == "smiles"]
    if exact:
        return exact[0]
    # Then case-insensitive
    ci = [i for i, c in enumerate(stripped) if c.lower() == "smiles"]
    if len(ci) == 1:
        return ci[0]
    if len(ci) > 1:
        raise ValueError(
            f"input CSV {csv_path} has multiple SMILES-named columns: "
            f"{[header[i] for i in ci]} at indices {ci}. "
            "Pass --smiles-column N (0-based) to disambiguate."
        )
    raise ValueError(
        f"could not auto-detect a SMILES column in {csv_path}. "
        f"Header columns: {header}. "
        "Pass --smiles-column N (0-based) to specify which column holds SMILES."
    )


def _clean_smiles(
    input_csv: Path, output_csv: Path, smiles_column: int, manifest: dict[str, Any], force: bool
) -> Path:
    if not force and _exists_nonempty(output_csv):
        _skipped(f"clean_smiles({input_csv.name})", str(output_csv), manifest)
        return output_csv
    output_csv.parent.mkdir(parents=True, exist_ok=True)
    if force and output_csv.exists():
        # clean_smiles.py prompts interactively (input()) when the output file
        # already exists — that's an EOFError in a non-TTY subprocess. Pre-delete.
        output_csv.unlink()
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "clean_smiles.py"),
        "--input", str(input_csv),
        "--output", str(output_csv),
        "--smiles_column", str(smiles_column),
    ]
    _run(cmd, f"clean_smiles({input_csv.name})", manifest)
    return output_csv


def _reduce_to_smiles_column(
    csv_path: Path, smiles_column: int, manifest: dict[str, Any]
) -> Path:
    """Rewrite an inference CSV to keep only the SMILES column (at index 0).

    Downstream `kermt.util.utils.get_data` -> `MoleculeDatapoint.__init__`
    floats every column after SMILES, which crashes on non-numeric passthrough
    columns (e.g. a 'split' label of 'train'/'val'/'test', or a 'Molecule Name'
    string). Inference does not need target columns, so drop them here.

    Note on skip semantics: this step is idempotent — running it on an
    already-single-column file is a no-op. We record that with
    `skipped_due_to_idempotent: True`, NOT `skipped_due_to_existing: True`.
    The two fields have different meanings: `_existing` means "I found a
    cached output file from a prior run and reused it" (overridden by
    `--force`); `_idempotent` means "the input is already in the desired
    state, so re-executing changes nothing" (safe to skip even under
    `--force`).
    """
    step_name = f"reduce_to_smiles_only({csv_path.name})"
    start = time.time()
    df = pd.read_csv(csv_path)
    if df.shape[1] == 1:
        manifest["steps"].append({
            "name": step_name,
            "output": str(csv_path),
            "ok": True,
            "duration_s": time.time() - start,
            "skipped_due_to_idempotent": True,
            "note": "already single-column",
        })
        return csv_path
    effective_col = smiles_column if 0 <= smiles_column < df.shape[1] else 0
    df.iloc[:, [effective_col]].to_csv(csv_path, index=False)
    manifest["steps"].append({
        "name": step_name,
        "output": str(csv_path),
        "ok": True,
        "duration_s": time.time() - start,
        "input_cols": int(df.shape[1]),
        "kept_col": effective_col,
        "kept_col_name": str(df.columns[effective_col]),
    })
    return csv_path


def _save_features(
    csv_path: Path, npz_path: Path, generator: str, manifest: dict[str, Any], force: bool
) -> Path:
    if not force and _exists_nonempty(npz_path):
        _skipped(f"save_features({csv_path.name}, {generator})", str(npz_path), manifest)
        return npz_path
    npz_path.parent.mkdir(parents=True, exist_ok=True)
    if force and npz_path.exists():
        npz_path.unlink()  # --restart still loads partial state if file exists; pre-delete to be safe
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "save_features.py"),
        "--data_path", str(csv_path),
        "--save_path", str(npz_path),
        "--features_generator", generator,
        "--restart",
    ]
    _run(cmd, f"save_features({csv_path.name}, {generator})", manifest)
    return npz_path


def _resolve_vocab_inputs(args: argparse.Namespace) -> dict[str, Path | None] | None:
    """Returns {atom, bond, smiles}->Path|None when the user supplied vocab
    inputs (via --vocab-dir or --atom-vocab/--bond-vocab/--smiles-vocab),
    else None (signal to fall through to build_vocab).

    Conventional filenames inside --vocab-dir:
      pretrain_atom_vocab.{json,pkl}
      pretrain_bond_vocab.{json,pkl}
      pretrain_smiles_vocab.pkl
    """
    if args.vocab_dir:
        d = Path(args.vocab_dir).resolve()
        if not d.is_dir():
            raise FileNotFoundError(f"--vocab-dir not found or not a directory: {d}")
        def _find(stem: str, exts: tuple[str, ...]) -> Path | None:
            for ext in exts:
                p = d / f"{stem}.{ext}"
                if p.is_file():
                    return p
            return None
        atom = _find(PRETRAIN_VOCAB_STEMS["atom"], ("json", "pkl"))
        bond = _find(PRETRAIN_VOCAB_STEMS["bond"], ("json", "pkl"))
        smiles = _find(PRETRAIN_VOCAB_STEMS["smiles"], ("pkl",))
        if atom is None and bond is None and smiles is None:
            stems = [PRETRAIN_VOCAB_STEMS[k] for k in ("atom", "bond", "smiles")]
            raise FileNotFoundError(
                f"--vocab-dir {d} contained no {{ {', '.join(stems) }}}.{{json,pkl}} "
                f"files. Expected at least {PRETRAIN_VOCAB_STEMS['atom']} + "
                f"{PRETRAIN_VOCAB_STEMS['bond']}."
            )
        return {"atom": atom, "bond": bond, "smiles": smiles}

    if args.atom_vocab or args.bond_vocab or args.smiles_vocab:
        return {
            "atom":   Path(args.atom_vocab).resolve()   if args.atom_vocab   else None,
            "bond":   Path(args.bond_vocab).resolve()   if args.bond_vocab   else None,
            "smiles": Path(args.smiles_vocab).resolve() if args.smiles_vocab else None,
        }

    return None


def _copy_provided_vocab(
    src: dict[str, Path | None], dst_dir: Path, dataset_name: str, manifest: dict[str, Any],
    force: bool,
) -> dict[str, Path]:
    """When the user supplies vocab files (use ckpt's vocab as-is),
    copy them into `<dst_dir>/<dataset_name>_<which>_vocab.<ext>` so the
    downstream pretrain command sees the conventional filenames.

    `src` is `{atom: Path|None, bond: Path|None, smiles: Path|None}`. The atom
    and bond entries must be both present or both absent (paired). smiles is
    optional (cmim/hybrid only).

    Returns the same dict of (resolved) destination paths.
    """
    import shutil
    if (src["atom"] is None) != (src["bond"] is None):
        raise ValueError(
            "vocab pass-through requires atom and bond vocab paths to be paired; "
            "got atom=" + str(src["atom"]) + ", bond=" + str(src["bond"])
        )
    out: dict[str, Path] = {}
    dst_dir.mkdir(parents=True, exist_ok=True)
    for which, path in src.items():
        if path is None:
            continue
        # Validate the source file IS a loadable KERMT vocab before copying.
        # Catches the "user pointed --smiles-vocab at a random pickle" case
        # early, with a clear error, instead of letting it surface as a cryptic
        # SMILESVocab.load_vocab failure at pretrain_ddp.py launch time.
        validate_vocab_file(path, kind=which)
        ext = path.suffix.lstrip(".")
        if which == "smiles":
            ext = "pkl"  # smiles vocab is always pickle
        dst = dst_dir / f"{dataset_name}_{which}_vocab.{ext}"
        if not force and _exists_nonempty(dst):
            _skipped(f"copy_vocab({which})", str(dst), manifest)
            out[which] = dst
            continue
        if force and dst.exists():
            dst.unlink()
        shutil.copy2(path, dst)
        manifest["steps"].append({
            "name": f"copy_vocab({which})",
            "src": str(path), "dst": str(dst), "ok": True,
            "duration_s": 0.0, "skipped_due_to_existing": False,
        })
        out[which] = dst
    return out


def _build_vocab(
    csv_path: Path, vocab_dir: Path, dataset_name: str, vocab_format: str,
    manifest: dict[str, Any], force: bool,
) -> dict[str, Path]:
    """Builds atom + bond (in --vocab-format) and smiles (always pickle) vocabs.
    Returns a dict of {atom, bond, smiles} -> Path."""
    suffix = "json" if vocab_format == "json" else "pkl"
    expected = {
        "atom": vocab_dir / f"{dataset_name}_atom_vocab.{suffix}",
        "bond": vocab_dir / f"{dataset_name}_bond_vocab.{suffix}",
        "smiles": vocab_dir / f"{dataset_name}_smiles_vocab.pkl",
    }
    if not force and all(_exists_nonempty(p) for p in expected.values()):
        _skipped(f"build_vocab({csv_path.name})", str(vocab_dir), manifest)
        return expected
    vocab_dir.mkdir(parents=True, exist_ok=True)
    if force:
        for p in expected.values():
            if p.exists():
                p.unlink()
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "build_vocab.py"),
        "--data_path", str(csv_path),
        "--vocab_save_folder", str(vocab_dir),
        "--dataset_name", dataset_name,
        "--vocab_format", vocab_format,
    ]
    _run(cmd, f"build_vocab({csv_path.name})", manifest)
    return expected


def _split_data(
    csv_path: Path, features_path: Path | None, sample_per_file: int, output_dir: Path,
    manifest: dict[str, Any], force: bool,
) -> Path:
    """Run split_data.py to produce shard dirs (graph/ + optionally feature/ + summary.txt)."""
    summary = output_dir / "summary.txt"
    if not force and _exists_nonempty(summary):
        _skipped(f"split_data({csv_path.name})", str(output_dir), manifest)
        return output_dir
    if force and output_dir.exists():
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "split_data.py"),
        "--data_path", str(csv_path),
        "--sample_per_file", str(sample_per_file),
        "--output_path", str(output_dir),
    ]
    if features_path is not None:
        cmd += ["--features_path", str(features_path)]
    _run(cmd, f"split_data({csv_path.name})", manifest)
    return output_dir


# ---------------------------------------------------------------------------
# Random splitter (used only when the user supplies a single CSV)
# ---------------------------------------------------------------------------

def _random_split_csv(
    src_csv: Path, dst_csvs: dict[str, Path], fractions: dict[str, float], seed: int,
    manifest: dict[str, Any], force: bool,
) -> None:
    """Shuffle src_csv and partition rows into dst_csvs by fractions.
    `dst_csvs` and `fractions` are dicts keyed by the split name (e.g. 'train', 'val').
    Sum of fractions must be 1.0 (within float tolerance). Writes each dst_csv with the
    same header as the input."""
    step = {
        "name": f"random_split({src_csv.name})",
        "seed": seed,
        "fractions": fractions,
        "ok": False,
        "duration_s": None,
        "skipped_due_to_existing": False,
        "row_counts": {},
    }
    if not force and all(_exists_nonempty(p) for p in dst_csvs.values()):
        step["skipped_due_to_existing"] = True
        step["ok"] = True
        manifest["steps"].append(step)
        return

    if abs(sum(fractions.values()) - 1.0) > 1e-6:
        raise ValueError(f"split fractions must sum to 1.0 (got {sum(fractions.values())})")

    t0 = time.time()
    df = pd.read_csv(src_csv).sample(frac=1.0, random_state=seed).reset_index(drop=True)
    n = len(df)
    sizes: dict[str, int] = {}
    remaining = n
    split_names = list(fractions.keys())
    for name in split_names[:-1]:
        sizes[name] = int(round(fractions[name] * n))
        remaining -= sizes[name]
    sizes[split_names[-1]] = remaining

    start = 0
    for name in split_names:
        dst = dst_csvs[name]
        dst.parent.mkdir(parents=True, exist_ok=True)
        df.iloc[start:start + sizes[name]].to_csv(dst, index=False)
        step["row_counts"][name] = sizes[name]
        start += sizes[name]

    step["duration_s"] = round(time.time() - t0, 2)
    step["ok"] = True
    manifest["steps"].append(step)


def _emit_random_split_warning(
    src_csv: Path, fractions: dict[str, float], seed: int, manifest: dict[str, Any]
) -> None:
    row_counts = manifest["steps"][-1].get("row_counts", {})
    n = sum(row_counts.values()) if row_counts else "?"
    lines = [
        f"WARNING: Auto-splitting {n} rows from {src_csv.name} into:",
    ]
    for name, frac in fractions.items():
        cnt = row_counts.get(name, "?")
        lines.append(f"  {name}: {cnt} rows ({frac * 100:.1f}%)")
    lines += [
        f"using random split with seed {seed}.",
        "",
        "This is a RANDOM split. For rigorous ADMET evaluation, scaffold-balanced",
        "(or other structure-aware) splits are strongly preferred — molecules with",
        "similar scaffolds can leak across splits and inflate apparent generalization.",
        "",
        "To use your own pre-computed splits instead, pass:",
        "    --train-csv <train.csv> --val-csv <val.csv> --test-csv <test.csv>",
        "",
        "To customize fractions:",
        "    --val-frac 0.15 --test-frac 0.15",
    ]
    warning = "\n".join(lines)
    print(warning, file=sys.stderr)
    manifest["warnings"].append(warning)


# ---------------------------------------------------------------------------
# Mode pipelines
# ---------------------------------------------------------------------------

def _prepare_embed(args, out: Path, manifest: dict[str, Any]) -> None:
    manifest["split_method"] = "n/a"
    if args.skip_clean:
        clean = Path(args.csv)
        manifest["steps"].append({"name": "clean_smiles", "skipped_by_flag": True, "ok": True})
    else:
        clean = _clean_smiles(Path(args.csv), out / "clean.csv", args.smiles_column, manifest, args.force)
    manifest["outputs"]["clean_csv"] = str(clean)


def _prepare_inference(args, out: Path, manifest: dict[str, Any]) -> None:
    manifest["split_method"] = "n/a"
    clean = _clean_smiles(Path(args.csv), out / "clean.csv", args.smiles_column, manifest, args.force)
    # Reduce to SMILES-only: downstream get_data/MoleculeDatapoint floats every
    # non-SMILES column, which crashes on non-numeric passthrough columns
    # (e.g. a 'split' label). Inference does not need target columns.
    _reduce_to_smiles_column(clean, args.smiles_column, manifest)
    manifest["outputs"]["clean_csv"] = str(clean)
    if args.skip_features:
        manifest["steps"].append({"name": "save_features", "skipped_by_flag": True, "ok": True})
        return
    generator = args.features_generator or DEFAULT_FEATURES_GENERATOR["inference"]
    npz = _save_features(clean, out / "clean.npz", generator, manifest, args.force)
    manifest["outputs"]["clean_npz"] = str(npz)


def _prepare_finetune(args, out: Path, manifest: dict[str, Any]) -> None:
    src_train = Path(args.csv)
    has_val = args.val_csv is not None
    has_test = args.test_csv is not None
    split_type = args.split_type

    if has_val and has_test:
        # User supplied explicit val + test CSVs: trust them, just clean + featurize.
        # split_type is irrelevant when val/test are given separately.
        manifest["split_method"] = "user_provided"
        clean_train = _clean_smiles(src_train, out / "clean_train.csv", args.smiles_column, manifest, args.force)
        clean_val = _clean_smiles(Path(args.val_csv), out / "clean_val.csv", args.smiles_column, manifest, args.force)
        clean_test = _clean_smiles(Path(args.test_csv), out / "clean_test.csv", args.smiles_column, manifest, args.force)
        manifest["outputs"]["clean_train_csv"] = str(clean_train)
        manifest["outputs"]["clean_val_csv"] = str(clean_val)
        manifest["outputs"]["clean_test_csv"] = str(clean_test)
        per_split = (("train", clean_train), ("val", clean_val), ("test", clean_test))
    elif has_val or has_test:
        raise ValueError(
            "for finetune mode, either provide BOTH --val-csv and --test-csv (user-provided splits) "
            "or NEITHER (run with --split-type {random|scaffold_balanced|index_predetermined}). "
            "Got one but not both."
        )
    elif split_type == "random":
        # Random auto-split — done here in prep so train.py gets ready-made CSVs.
        manifest["split_method"] = "random"
        manifest["split_seed"] = args.seed
        train_frac = max(0.0, 1.0 - args.val_frac - args.test_frac)
        manifest["split_fractions"] = {"train": train_frac, "val": args.val_frac, "test": args.test_frac}
        clean_full = _clean_smiles(src_train, out / "_clean_full.csv", args.smiles_column, manifest, args.force)
        dst = {
            "train": out / "clean_train.csv",
            "val": out / "clean_val.csv",
            "test": out / "clean_test.csv",
        }
        _random_split_csv(clean_full, dst, manifest["split_fractions"], args.seed, manifest, args.force)
        clean_train, clean_val, clean_test = dst["train"], dst["val"], dst["test"]
        manifest["outputs"]["clean_train_csv"] = str(clean_train)
        manifest["outputs"]["clean_val_csv"] = str(clean_val)
        manifest["outputs"]["clean_test_csv"] = str(clean_test)
        _emit_random_split_warning(src_train, manifest["split_fractions"], args.seed, manifest)
        per_split = (("train", clean_train), ("val", clean_val), ("test", clean_test))
    else:
        # Scaffold-balanced or index-predetermined: prep cleans + featurizes the full
        # CSV and defers actual splitting to task/train.py, which calls split_data
        # with the user-supplied seed and split_sizes.
        manifest["split_method"] = "deferred_to_runner"
        manifest["split_type"] = split_type
        manifest["split_seed"] = args.seed
        manifest["split_fractions"] = {
            "train": max(0.0, 1.0 - args.val_frac - args.test_frac),
            "val": args.val_frac,
            "test": args.test_frac,
        }
        clean_full = _clean_smiles(src_train, out / "clean_full.csv", args.smiles_column, manifest, args.force)
        manifest["outputs"]["clean_full_csv"] = str(clean_full)
        per_split = (("full", clean_full),)

    if args.skip_features:
        manifest["steps"].append({"name": "save_features", "skipped_by_flag": True, "ok": True})
        return
    generator = args.features_generator or DEFAULT_FEATURES_GENERATOR["finetune"]
    for split_name, csv in per_split:
        npz = _save_features(csv, csv.with_suffix(".npz"), generator, manifest, args.force)
        manifest["outputs"][f"clean_{split_name}_npz"] = str(npz)


def _prepare_pretrain(args, out: Path, manifest: dict[str, Any]) -> None:
    src_train = Path(args.csv)
    if args.val_csv is not None:
        manifest["split_method"] = "user_provided"
        clean_train = _clean_smiles(src_train, out / "clean_train.csv", args.smiles_column, manifest, args.force)
        clean_val = _clean_smiles(Path(args.val_csv), out / "clean_val.csv", args.smiles_column, manifest, args.force)
    else:
        manifest["split_method"] = "random"
        manifest["split_seed"] = args.seed
        train_frac = max(0.0, 1.0 - args.val_frac)
        manifest["split_fractions"] = {"train": train_frac, "val": args.val_frac}
        clean_full = _clean_smiles(src_train, out / "_clean_full.csv", args.smiles_column, manifest, args.force)
        dst = {"train": out / "clean_train.csv", "val": out / "clean_val.csv"}
        _random_split_csv(clean_full, dst, manifest["split_fractions"], args.seed, manifest, args.force)
        clean_train, clean_val = dst["train"], dst["val"]

    manifest["outputs"]["clean_train_csv"] = str(clean_train)
    manifest["outputs"]["clean_val_csv"] = str(clean_val)

    generator = args.features_generator or DEFAULT_FEATURES_GENERATOR["pretrain"]
    if args.skip_features:
        manifest["steps"].append({"name": "save_features", "skipped_by_flag": True, "ok": True})
        train_npz: Path | None = None
        val_npz: Path | None = None
    else:
        train_npz = _save_features(clean_train, out / "clean_train.npz", generator, manifest, args.force)
        val_npz = _save_features(clean_val, out / "clean_val.npz", generator, manifest, args.force)
        manifest["outputs"]["clean_train_npz"] = str(train_npz)
        manifest["outputs"]["clean_val_npz"] = str(val_npz)

    if args.skip_vocab:
        manifest["steps"].append({"name": "build_vocab", "skipped_by_flag": True, "ok": True})
        manifest["vocab_source"] = "skipped"
    else:
        # Resolve user-provided vocab paths from --vocab-dir or explicit flags.
        provided = _resolve_vocab_inputs(args)
        if provided:
            # Use the user-supplied (ckpt's) vocab as-is. Copy into the
            # conventional filenames the downstream pretrain command expects.
            vocabs = _copy_provided_vocab(provided, out, args.dataset_name, manifest, args.force)
            manifest["vocab_source"] = "user_provided"
        else:
            # Fall back to the existing build-from-corpus behavior. Used by
            # pretrain-from-scratch and by any continue case where the user
            # explicitly wants a fresh vocab (rare, usually wrong).
            vocabs = _build_vocab(clean_train, out, args.dataset_name, args.vocab_format, manifest, args.force)
            manifest["vocab_source"] = "built_fresh"
        if "atom" in vocabs:
            manifest["outputs"]["atom_vocab"] = str(vocabs["atom"])
        if "bond" in vocabs:
            manifest["outputs"]["bond_vocab"] = str(vocabs["bond"])
        if "smiles" in vocabs:
            manifest["outputs"]["smiles_vocab"] = str(vocabs["smiles"])

    if args.skip_split:
        manifest["steps"].append({"name": "split_data", "skipped_by_flag": True, "ok": True})
    else:
        train_dir = _split_data(clean_train, train_npz, args.sample_per_file, out / "train", manifest, args.force)
        val_dir = _split_data(clean_val, val_npz, args.sample_per_file, out / "val", manifest, args.force)
        manifest["outputs"]["train_dir"] = str(train_dir)
        manifest["outputs"]["val_dir"] = str(val_dir)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def prepare(args: argparse.Namespace) -> dict[str, Any]:
    out = Path(args.out).resolve()
    out.mkdir(parents=True, exist_ok=True)
    manifest: dict[str, Any] = {
        "mode": args.mode,
        "input_csv": str(Path(args.csv).resolve()),
        "val_csv": str(Path(args.val_csv).resolve()) if args.val_csv else None,
        "test_csv": str(Path(args.test_csv).resolve()) if args.test_csv else None,
        "output_dir": str(out),
        "split_method": None,
        "steps": [],
        "outputs": {},
        "errors": [],
        "warnings": [],
    }
    try:
        if args.mode == "pretrain":
            _prepare_pretrain(args, out, manifest)
        elif args.mode == "finetune":
            _prepare_finetune(args, out, manifest)
        elif args.mode == "inference":
            _prepare_inference(args, out, manifest)
        elif args.mode == "embed":
            _prepare_embed(args, out, manifest)
        manifest["ok"] = True
    except Exception as exc:  # noqa: BLE001
        manifest["ok"] = False
        manifest["errors"].append(f"{type(exc).__name__}: {exc}")
    # Always write the manifest so partial-failure state is visible to the agent.
    (out / "prepare_data.json").write_text(json.dumps(manifest, indent=2))
    return manifest


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Mode-dispatched data prep for the KERMT agent skills.")
    p.add_argument("--mode", required=True, choices=VALID_MODES)
    p.add_argument("--csv", required=True, help="Primary input CSV (train CSV for pretrain/finetune)")
    p.add_argument("--out", required=True, help="Output directory")
    p.add_argument("--val-csv", default=None, help="Optional separate val CSV (pretrain/finetune)")
    p.add_argument("--test-csv", default=None, help="Optional separate test CSV (finetune only)")
    p.add_argument("--val-frac", type=float, default=0.1, help="Auto-split val fraction (default 0.1)")
    p.add_argument("--test-frac", type=float, default=0.1, help="Auto-split test fraction (finetune only, default 0.1)")
    p.add_argument("--seed", type=int, default=0, help="Random split seed (default 0)")
    p.add_argument("--split-type", choices=["random", "scaffold_balanced", "index_predetermined"],
                   default="random",
                   help="(finetune only, when --val-csv/--test-csv are not given) how to split. "
                        "'random' splits in prep using --val-frac/--test-frac/--seed. "
                        "'scaffold_balanced' and 'index_predetermined' defer the actual split to the "
                        "runner (task/train.py invokes split_data with the appropriate algorithm "
                        "using the user-supplied seed); prep only cleans + featurizes the full CSV.")
    p.add_argument("--sample-per-file", type=int, default=100_000,
                   help="split_data shard size (pretrain only, default 100000)")
    p.add_argument("--vocab-format", choices=["json", "pkl"], default="json",
                   help="atom/bond vocab format (default json); smiles vocab is always pkl")
    # Vocab pass-through (pretrain mode): when continuing from a released ckpt,
    # pass its bundled vocab files in so we don't rebuild a mismatched vocab.
    p.add_argument("--vocab-dir", default=None,
                   help="(pretrain) directory containing pretrain_{atom,bond}_vocab.{json,pkl} "
                        "(+ pretrain_smiles_vocab.pkl for cmim/hybrid). When given, prepare_data "
                        "skips build_vocab and copies these files into the output dir under the "
                        "expected filenames. Used by kermt-continue-pretrain to bind the released "
                        "ckpt's vocab to the new corpus (the ckpt's vocab is authoritative).")
    p.add_argument("--atom-vocab", default=None,
                   help="(pretrain) explicit atom vocab path; pairs with --bond-vocab. Overrides "
                        "--vocab-dir's pretrain_atom_vocab.* discovery if both are given.")
    p.add_argument("--bond-vocab", default=None,
                   help="(pretrain) explicit bond vocab path; pairs with --atom-vocab.")
    p.add_argument("--smiles-vocab", default=None,
                   help="(pretrain, cmim/hybrid) explicit smiles vocab .pkl path. Optional for "
                        "vocab-only pretrain.")
    p.add_argument("--dataset-name", default="pretrain",
                   help="vocab filename prefix (default 'pretrain' so downstream pretrain commands "
                        "can reference pretrain_{atom,bond}_vocab.{json|pkl}, pretrain_smiles_vocab.pkl)")
    p.add_argument("--targets", nargs="+", default=None,
                   help="(finetune only) target column names; forwarded to the finetune runner via the manifest")
    p.add_argument("--features-generator", default=None,
                   help="Override the per-mode default (pretrain: fgtasklabel; finetune/inference: rdkit_2d_normalized)")
    p.add_argument("--smiles-column", type=int, default=None,
                   help="0-based column index of SMILES in the input CSV. "
                        "When omitted, auto-detected by header name "
                        "(prefers lowercase `smiles`; accepts case-insensitive "
                        "`SMILES`/`Smiles`). Pass explicitly to override.")
    p.add_argument("--force", action="store_true",
                   help="Re-run every step even if its outputs already exist")
    p.add_argument("--skip-clean", action="store_true", help="(embed mode) skip the cleaning step")
    p.add_argument("--skip-features", action="store_true", help="Skip feature generation")
    p.add_argument("--skip-vocab", action="store_true", help="(pretrain) skip vocab build")
    p.add_argument("--skip-split", action="store_true", help="(pretrain) skip shard split")
    args = p.parse_args(argv)

    # Forward --targets through the manifest so the finetune runner can see them.
    if args.mode == "finetune" and args.targets:
        pass  # captured in manifest below

    # Resolve the SMILES column index (auto-detect from header when the user
    # didn't pass --smiles-column). This is the only point where args.csv is
    # touched before downstream _clean_smiles calls fan it out.
    try:
        resolved_smiles_col = _resolve_smiles_column(Path(args.csv), args.smiles_column)
    except ValueError as exc:
        err_manifest = {
            "ok": False,
            "mode": args.mode,
            "errors": [f"smiles-column resolution failed: {exc}"],
        }
        Path(args.out).mkdir(parents=True, exist_ok=True)
        (Path(args.out) / "prepare_data.json").write_text(json.dumps(err_manifest, indent=2))
        print(json.dumps(err_manifest, indent=2))
        return 1
    if args.smiles_column is None:
        print(f"[prepare_data] auto-detected --smiles-column {resolved_smiles_col} "
              f"from {Path(args.csv).name} header", file=sys.stderr)
    args.smiles_column = resolved_smiles_col

    try:
        manifest = prepare(args)
    except Exception as exc:  # noqa: BLE001
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({"ok": False, "errors": [f"unhandled: {type(exc).__name__}: {exc}"]}, indent=2))
        return 1
    if args.targets:
        manifest["targets"] = list(args.targets)
    # Record the resolved SMILES column so the manifest is self-describing.
    manifest["smiles_column"] = args.smiles_column
    (Path(args.out) / "prepare_data.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2))
    return 0 if manifest.get("ok") else 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Workstation embedding-extraction runner — wraps task/extract_embeddings.py
and emits a reproducible `run.json` manifest alongside the outputs.

Accepts any encoder-bearing ckpt (grover_base / cmim / hybrid / finetuned).
Validates via check_checkpoint.py --mode embed (encoder-only sufficient).
Reads the prepare_data manifest (mode=embed: clean CSV only — featurization
happens on the fly inside extract_embeddings.py).

Output layout:
    <out>/out/atom_from_atom.npy
    <out>/out/bond_from_atom.npy
    <out>/out/atom_from_bond.npy
    <out>/out/bond_from_bond.npy
    <out>/out/canonical_smiles.npy
    <out>/out/validity.npy

Blocking-by-default. Embedding extraction is minutes-scale.

CLI
---
    run_extract_embeddings.py
        --ckpt <encoder-bearing.pt>   # required
        --prepare-manifest <path>     # prepare_data.json (mode=embed)
        --out <run-dir>               # output dir
        [--ckpt-validator-out <path>]
        [--gpus 0]                    # single GPU id (default 0)
        [--batch-size N]              # override defaults_embed.runtime.batch_size
        [--dry-run]
"""
from __future__ import annotations

import argparse
import datetime
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any

if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))
from _utils import (  # noqa: E402
    resolve_kermt_repo, assert_prepare_manifest_basics, docker_image_digest, format_cmd_replay,
    git_commit_with_env_override, load_json, merge_default_into_applied,
    resolve_single_gpu, run_checkpoint_validator, runner_environment,
)


REPO_ROOT = resolve_kermt_repo()
SKILL_ROOT = Path(__file__).resolve().parent.parent
DEFAULTS_PATH = SKILL_ROOT / "config" / "defaults_embed.json"
CHECK_CHECKPOINT_PATH = SKILL_ROOT / "scripts" / "check_checkpoint.py"
EXTRACT_EMBEDDINGS_PATH = REPO_ROOT / "task" / "extract_embeddings.py"


def _verify_prepare_manifest(manifest: dict[str, Any]) -> None:
    assert_prepare_manifest_basics(manifest, "embed")
    out = manifest.get("outputs", {})
    if "clean_csv" not in out:
        raise ValueError(
            "prepare_data manifest is missing required output 'clean_csv'. "
            "Was prepare_data run successfully?"
        )


def _apply_defaults(args: argparse.Namespace, defaults: dict[str, Any]) -> dict[str, dict[str, Any]]:
    """Merge defaults_embed.json with CLI overrides."""
    applied: dict[str, dict[str, Any]] = {}
    runtime = defaults.get("runtime", {})

    merge_default_into_applied(applied, args, "batch_size", runtime)

    return applied


def _build_argv(
    *, gpu: int, ckpt: Path, manifest: dict[str, Any], out_dir: Path,
    applied: dict[str, dict[str, Any]],
) -> list[str]:
    """Constructs the task/extract_embeddings.py argv. Note: extract_embeddings
    uses its own CLI (--checkpoint, --input_file, --output_path) — NOT main.py.
    --format defaults to npy inside extract_embeddings.py, so we don't pass it."""
    outputs_dir = out_dir / "out"
    outputs_dir.mkdir(parents=True, exist_ok=True)

    argv: list[str] = [sys.executable, "-u", str(EXTRACT_EMBEDDINGS_PATH)]
    argv += ["--checkpoint", str(ckpt)]
    argv += ["--input_file", manifest["outputs"]["clean_csv"]]
    argv += ["--output_path", str(outputs_dir)]
    argv += ["--device", "cuda"]
    # clean_smiles.py preserves the input CSV's column layout, so the cleaned
    # CSV has SMILES at whichever column the input had. Forward the
    # auto-detected smiles_column from prepare_data.json so extract_embeddings
    # doesn't fall back to its default of column 0 (which would index a
    # non-SMILES column for inputs like openadmet/all.csv where SMILES is at
    # column 1). Default to 0 if the manifest is from an older prepare_data
    # version that didn't record the field.
    argv += ["--smiles_column", str(manifest.get("smiles_column", 0))]
    if "batch_size" in applied:
        argv += ["--batch_size", str(applied["batch_size"]["value"])]

    return argv


# ---------------------------------------------------------------------------
# Main flow
# ---------------------------------------------------------------------------

def run(args: argparse.Namespace) -> dict[str, Any]:
    out_dir = Path(args.out).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "logs").mkdir(parents=True, exist_ok=True)
    (out_dir / "out").mkdir(parents=True, exist_ok=True)

    defaults = load_json(DEFAULTS_PATH, name="defaults_embed.json")
    prep_manifest_path = Path(args.prepare_manifest).resolve()
    manifest = load_json(prep_manifest_path, name="prepare_data.json")
    _verify_prepare_manifest(manifest)

    ckpt = Path(args.ckpt).resolve()
    if args.ckpt_validator_out:
        validator_out = load_json(Path(args.ckpt_validator_out), name="ckpt validator output")
    else:
        validator_out = run_checkpoint_validator(ckpt, mode="embed", script_path=CHECK_CHECKPOINT_PATH)
    if not validator_out.get("ok"):
        raise ValueError(
            f"check_checkpoint.py rejected the input ckpt: {validator_out.get('errors')}"
        )

    model_type = validator_out.get("model_type")
    arch = validator_out.get("arch", {})

    gpu = resolve_single_gpu(args.gpus, workflow="embed")
    applied = _apply_defaults(args, defaults)

    argv = _build_argv(gpu=gpu, ckpt=ckpt, manifest=manifest, out_dir=out_dir, applied=applied)

    commit, dirty = git_commit_with_env_override(REPO_ROOT)
    image_tag = os.environ.get("KERMT_IMAGE", "kermt:latest")
    image_digest = docker_image_digest(image_tag)
    cmd_replay = format_cmd_replay(argv, env={"CUDA_VISIBLE_DEVICES": gpu})
    run_manifest: dict[str, Any] = {
        "workflow": "embed",
        "started_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "container": {"image_tag": image_tag, "image_digest": image_digest},
        "repo": {"commit": commit, "dirty": dirty},
        "inputs": {
            "ckpt": str(ckpt),
            "prepare_data_manifest": str(prep_manifest_path),
            "ckpt_validator_out": (
                str(Path(args.ckpt_validator_out).resolve()) if args.ckpt_validator_out else None
            ),
        },
        "model_type": model_type,
        "gpu": gpu,
        "args_applied": applied,
        "arch": arch,
        "output_dir": str(out_dir / "out"),
        "logs_dir": str(out_dir / "logs"),
        "argv": argv,
        "cmd_replay": cmd_replay,
        "ok_to_replay": (not dirty) and (commit != "unknown"),
        "dry_run": bool(args.dry_run),
    }
    (out_dir / "run.json").write_text(json.dumps(run_manifest, indent=2))

    if args.dry_run:
        run_manifest["status"] = "dry_run"
        return run_manifest

    env = runner_environment(REPO_ROOT)
    env["CUDA_VISIBLE_DEVICES"] = str(gpu)
    log_file = out_dir / "logs" / "embed.log"
    with log_file.open("w") as logf:
        proc = subprocess.run(argv, env=env, stdout=logf, stderr=subprocess.STDOUT)
    run_manifest["exit_code"] = proc.returncode
    run_manifest["status"] = "ok" if proc.returncode == 0 else "failed"
    (out_dir / "run.json").write_text(json.dumps(run_manifest, indent=2))
    return run_manifest


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(
        description="Workstation embedding-extraction runner — wraps task/extract_embeddings.py."
    )
    p.add_argument("--ckpt", required=True, help="Path to encoder-bearing ckpt (any model_type).")
    p.add_argument("--prepare-manifest", required=True,
                   help="Path to a prepare_data.json produced with --mode embed.")
    p.add_argument("--out", required=True, help="Output run directory.")
    p.add_argument("--ckpt-validator-out", default=None,
                   help="Optional cached check_checkpoint.py JSON.")
    p.add_argument("--gpus", default=None, help="Single GPU id (default 0). Multi-GPU rejected.")
    p.add_argument("--dry-run", action="store_true")

    p.add_argument("--batch-size", type=int, default=None)

    args = p.parse_args(argv)

    try:
        manifest = run(args)
    except (FileNotFoundError, ValueError, RuntimeError) as exc:
        print(json.dumps({"ok": False, "errors": [f"{type(exc).__name__}: {exc}"]}, indent=2))
        return 1
    except Exception as exc:  # noqa: BLE001
        import traceback
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({"ok": False, "errors": [f"unhandled: {type(exc).__name__}: {exc}"]},
                         indent=2))
        return 1

    print(json.dumps({"ok": True, "manifest": manifest}, indent=2))
    return 0 if manifest.get("status") != "failed" else 1


if __name__ == "__main__":
    sys.exit(main())

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Shared utilities for the agent scripts.

Kept intentionally small — only logic that appears (or would otherwise be
duplicated) in two or more `scripts/*.py` modules. Each script
maintains its own primary CLI + main flow.
"""
from __future__ import annotations

import argparse
from collections import Counter
import json
import os
import pickle
import re
import shlex
import subprocess
import sys
from pathlib import Path
from typing import Any


# Conventional pretrain vocab filename stems. Used by prepare_data.py +
# upgrade_to_hybrid.py + the README "Released models" bundling docs +
# the test helpers. Centralized here so a future rename only touches one
# spot.
PRETRAIN_VOCAB_STEMS = {
    "atom":   "pretrain_atom_vocab",
    "bond":   "pretrain_bond_vocab",
    "smiles": "pretrain_smiles_vocab",
}


def resolve_kermt_repo() -> Path:
    """Find the runtime checkout independently of the installed skill location.

    An explicit KERMT_REPO takes precedence. In a repository checkout, walking
    up from this helper or the working directory also supports local use.
    """
    explicit = os.environ.get("KERMT_REPO")
    if explicit:
        candidates = [Path(explicit).expanduser().resolve()]
    else:
        candidates = []
        for start in (Path(__file__).resolve().parent, Path.cwd()):
            candidates.extend((start, *start.parents))
    for candidate in candidates:
        if (candidate / "main.py").is_file() and (candidate / "kermt").is_dir():
            return candidate
    raise FileNotFoundError(
        "KERMT checkout not found. Set KERMT_REPO to the checkout containing "
        "main.py and kermt/; the installed skill directory is separate."
    )


def load_json(path: Path, *, name: str) -> dict[str, Any]:
    """Load a JSON file with consistent error messages.

    `name` is a human-readable label for the document (e.g. "prepare_data.json")
    so the error tells the user which schema we expected at that path.
    """
    if not path.is_file():
        raise FileNotFoundError(f"{name} not found at {path}")
    try:
        return json.loads(path.read_text())
    except json.JSONDecodeError as exc:
        raise ValueError(f"{name} at {path} is not valid JSON: {exc}") from exc


def count_vocab_entries(vocab_path: Path) -> int:
    """Return the number of entries in a KERMT vocab file.

    Handles three layouts:
      - JSON with `{stoi: {token: idx}, ...}` (MolVocab.save_vocab default)
      - JSON as a raw `{token: idx}` dict (legacy / hand-edited)
      - Legacy MolVocab / SMILESVocab pickles, read as inert vocabulary state.

    The pickle reader accepts only the known vocabulary containers and their
    Counter/regex metadata. It cannot import arbitrary classes or run reducers
    supplied by the artifact, and it never falls back to an unrestricted loader.
    """
    if vocab_path.suffix == ".json":
        data = json.loads(vocab_path.read_text())
        if isinstance(data, dict) and "stoi" in data:
            return len(data["stoi"])
        if isinstance(data, dict):
            return len(data)
        raise ValueError(f"unsupported JSON vocab shape at {vocab_path}: {type(data).__name__}")

    with vocab_path.open("rb") as f:
        data = _VocabUnpickler(f).load()
    if isinstance(data, _VocabState) and isinstance(data.stoi, dict):
        return len(data.stoi)
    if isinstance(data, (dict, list, tuple)):
        return len(data)
    raise ValueError(f"could not count entries in {vocab_path}")


class _VocabState:
    """Data-only stand-in: counting tokens does not require tokenizer methods."""


class _VocabUnpickler(pickle.Unpickler):
    def find_class(self, module: str, name: str) -> Any:
        if module in {"kermt.data.torchvocab", "grover.data.torchvocab"} and name in {
            "TorchVocab", "MolVocab", "SMILESVocab",
        }:
            return _VocabState
        if (module, name) == ("collections", "Counter"):
            return Counter
        if (module, name) == ("re", "_compile"):
            return re.compile
        raise pickle.UnpicklingError(f"unsupported vocabulary object: {module}.{name}")


def load_checkpoint(path: Path | str) -> dict[str, Any]:
    """Read KERMT tensors and known metadata with PyTorch's restricted loader.

    Saved arguments use argparse.Namespace; finetuned checkpoints also contain
    numeric NumPy scaler arrays. Explicit globals cover those formats, including
    NumPy 1/2 module names, without accepting artifact-selected imports.
    """
    import numpy as np
    import torch

    multiarray = np._core.multiarray if hasattr(np, "_core") else np.core.multiarray
    allowed = [argparse.Namespace, np.ndarray, np.dtype]
    for module in ("numpy.core.multiarray", "numpy._core.multiarray"):
        allowed.extend([
            (multiarray._reconstruct, f"{module}._reconstruct"),
            (multiarray.scalar, f"{module}.scalar"),
        ])
    allowed.extend(type(np.dtype(name)) for name in (
        "bool", "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64",
        "float16", "float32", "float64",
    ))
    with torch.serialization.safe_globals(allowed):
        return torch.load(path, map_location="cpu", weights_only=True)


def runner_environment(repo: Path, *, wandb: bool = False) -> dict[str, str]:
    """Forward named runtime settings, keeping unrelated credentials out of jobs.

    W&B credentials/settings are included only for an explicitly enabled W&B
    run. Hugging Face authentication belongs to the separate download helper.
    """
    names = (
        "PATH", "HOME", "TMPDIR", "TEMP", "TMP", "LANG", "LC_ALL", "LC_CTYPE", "TZ",
        "LD_LIBRARY_PATH", "LIBRARY_PATH", "CUDA_HOME", "CUDA_PATH", "PYTHONPATH",
        "PYTHONDONTWRITEBYTECODE", "PYTHONUNBUFFERED", "PYTHONWARNINGS",
        "CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER", "CUDA_LAUNCH_BLOCKING", "NVIDIA_VISIBLE_DEVICES",
        "NVIDIA_DRIVER_CAPABILITIES", "CUBLAS_WORKSPACE_CONFIG", "OMP_NUM_THREADS",
        "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS",
        "PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_ALLOC_CONF", "PYTORCH_NO_CUDA_MEMORY_CACHING",
        "TORCH_CPP_LOG_LEVEL", "TORCH_DISTRIBUTED_DEBUG", "NCCL_DEBUG", "NCCL_SOCKET_IFNAME",
        "NCCL_IB_DISABLE", "NCCL_P2P_DISABLE", "NCCL_SHM_DISABLE", "GLOO_SOCKET_IFNAME",
        "MASTER_ADDR", "MASTER_PORT",
        "KERMT_REPO", "KERMT_REPO_COMMIT", "KERMT_REPO_DIRTY",
        "SSL_CERT_FILE", "REQUESTS_CA_BUNDLE",
    )
    if wandb:
        names += (
            "WANDB_API_KEY", "WANDB_BASE_URL", "WANDB_MODE", "WANDB_DIR", "WANDB_ENTITY",
            "WANDB_PROJECT", "WANDB_RUN_ID", "WANDB_RESUME", "WANDB_CACHE_DIR",
            "WANDB_CONFIG_DIR", "WANDB_DATA_DIR", "WANDB_DISABLED",
        )
    env = {name: value for name in names if (value := os.environ.get(name)) is not None}
    env["PYTHONPATH"] = os.pathsep.join(filter(None, (str(repo), env.get("PYTHONPATH"))))
    return env


def validate_vocab_file(vocab_path: Path, *, kind: str) -> None:
    """Verify a user-provided vocab file is loadable BEFORE copying it into a
    run directory. Raises ValueError on failure with a clear, user-facing message.

    `kind` is one of {"atom", "bond", "smiles"} — used only in the error message
    so the user knows which file is wrong.
    """
    if not vocab_path.is_file():
        raise FileNotFoundError(f"{kind} vocab file not found: {vocab_path}")
    try:
        n = count_vocab_entries(vocab_path)
    except Exception as exc:  # noqa: BLE001
        raise ValueError(
            f"{kind} vocab file {vocab_path} is not loadable as a KERMT vocab "
            f"({type(exc).__name__}: {exc}). Expected a MolVocab JSON or pickle "
            f"(or a SMILESVocab pickle for the smiles vocab)."
        ) from exc
    if n <= 0:
        raise ValueError(f"{kind} vocab file {vocab_path} contains zero entries")


# ---------------------------------------------------------------------------
# Runner-shared helpers (run.json manifest fields)
# ---------------------------------------------------------------------------

def git_commit_with_env_override(repo: Path) -> tuple[str, bool]:
    """Returns (commit_sha, dirty_tree). Honors `KERMT_REPO_COMMIT` /
    `KERMT_REPO_DIRTY` env vars first — set by `scripts/kermt_container.sh`
    from the host before launching docker (necessary because `git -C /workspace`
    inside the container fails due to bind-mount ownership). Falls back to the
    in-container git probe when the env vars aren't set."""
    env_commit = os.environ.get("KERMT_REPO_COMMIT")
    if env_commit:
        env_dirty = os.environ.get("KERMT_REPO_DIRTY", "false").strip().lower() == "true"
        return env_commit, env_dirty
    try:
        sha = subprocess.run(
            ["git", "-C", str(repo), "rev-parse", "HEAD"],
            capture_output=True, text=True, check=True,
        ).stdout.strip()
        diff = subprocess.run(
            ["git", "-C", str(repo), "status", "--porcelain"],
            capture_output=True, text=True, check=True,
        )
        return sha, bool(diff.stdout.strip())
    except Exception:
        return "unknown", False


def docker_image_digest(tag: str) -> str | None:
    """Return the docker image's content-addressable Id (sha256:…) for the given
    tag, or None if docker isn't available / the image isn't local."""
    try:
        r = subprocess.run(
            ["docker", "image", "inspect", tag, "--format", "{{.Id}}"],
            capture_output=True, text=True,
        )
        if r.returncode == 0:
            return r.stdout.strip()
    except FileNotFoundError:
        pass
    return None


def format_cmd_replay(argv: list[str], *, env: dict[str, str] | None = None) -> str:
    """Render a copy-pasteable env-prefix + command for the cmd_replay manifest
    field. `env` is the set of environment variables to prefix (typically
    {CUDA_VISIBLE_DEVICES, WORLD_SIZE})."""
    env = env or {}
    env_prefix = [f"{k}={shlex.quote(str(v))}" for k, v in env.items()]
    quoted = " ".join(shlex.quote(a) for a in argv)
    return " ".join(env_prefix + [quoted])


def resolve_single_gpu(override: str | None, *, workflow: str) -> int:
    """Returns a single GPU id (int). The finetune/inference/embed workflows are
    single-GPU only; `--gpus '0,1'` or multi-id CUDA_VISIBLE_DEVICES is rejected
    with a workflow-specific error. (The pretrain runner has its own multi-GPU
    `_detect_gpus` helper — see run_pretrain_local.py.)"""
    if override is None:
        env_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
        if env_visible:
            ids = [g for g in env_visible.split(",") if g]
            if len(ids) > 1:
                raise ValueError(
                    f"CUDA_VISIBLE_DEVICES='{env_visible}' selects multiple GPUs but "
                    f"the {workflow} workflow is single-GPU only. Restrict to one id."
                )
            return int(ids[0])
        return 0
    parts = [p.strip() for p in override.split(",") if p.strip()]
    if len(parts) != 1:
        raise ValueError(
            f"--gpus '{override}' selects {len(parts)} GPUs; the {workflow} workflow is single-GPU only."
        )
    return int(parts[0])


def assert_prepare_manifest_basics(manifest: dict[str, Any], expected_mode: str) -> None:
    """Standard pre-check for a prepare_data.json before a runner consumes it:
    verify `mode` matches and `ok` is True. Raises ValueError with a consistent
    error message on either mismatch.

    Each runner is responsible for its own required-outputs check after this
    (those vary per-mode — e.g. pretrain wants train_dir/val_dir/atom_vocab/
    bond_vocab; finetune has the split-method branch; inference/embed want
    clean_csv)."""
    if manifest.get("mode") != expected_mode:
        raise ValueError(
            f"prepare_data manifest is mode='{manifest.get('mode')}', expected '{expected_mode}'. "
            f"Run `prepare_data.py --mode {expected_mode}` to produce a valid manifest."
        )
    if not manifest.get("ok"):
        raise ValueError(
            f"prepare_data manifest reports ok=False: {manifest.get('errors')}"
        )


def merge_default_into_applied(
    applied: dict[str, dict[str, Any]],
    args: argparse.Namespace,
    name: str,
    defaults_group: dict[str, Any],
) -> None:
    """Standard CLI-override / default-config merge for one hyperparameter.

    Mutates `applied` in place:
      - If the user passed `--<name>` on the CLI (so `getattr(args, name)` is
        not None), records `{"value": cli_val, "source": "user"}`.
      - Else if `name` is present in `defaults_group`, records
        `{"value": defaults_group[name], "source": "default-config"}`.
      - Else `applied[name]` is left absent — the runner's argv-builder skips
        the flag, and the downstream argparse default takes effect.

    `name` is the snake_case argparse dest (same form used as the dict key);
    argparse automatically converts CLI `--<name-with-hyphens>` to that dest,
    so `getattr(args, name, None)` is the correct CLI lookup."""
    cli_val = getattr(args, name, None)
    if cli_val is not None:
        applied[name] = {"value": cli_val, "source": "user"}
    elif name in defaults_group:
        applied[name] = {"value": defaults_group[name], "source": "default-config"}


def run_checkpoint_validator(ckpt: Path, *, mode: str, script_path: Path) -> dict[str, Any]:
    """Invoke `check_checkpoint.py --mode <mode> --ckpt <path>` as a subprocess
    and return the parsed JSON. Raises RuntimeError on non-JSON output (e.g. the
    validator crashed before printing). `script_path` is the absolute path to
    `scripts/check_checkpoint.py` — passed in so this helper has no
    dependency on the caller's layout."""
    r = subprocess.run(
        [sys.executable, str(script_path), "--mode", mode, "--ckpt", str(ckpt)],
        capture_output=True, text=True,
    )
    try:
        return json.loads(r.stdout)
    except json.JSONDecodeError as exc:
        raise RuntimeError(
            f"check_checkpoint.py emitted non-JSON output (exit {r.returncode}). "
            f"stdout (first 200 chars): {r.stdout[:200]}\n"
            f"stderr (first 200 chars): {r.stderr[:200]}"
        ) from exc