feat(kermt): sync open-model skills + scripts to KERMT main

open-models-skills/kermt/.skillsource.json

@@ -2,17 +2,19 @@
   "name": "kermt",
   "subpath": "agent",
   "type": "dir",
-  "resolved_sha": "f423e9db359ec63815dd4f3ea6fb8fcc003cc44f",
+  "resolved_sha": "85f255fd00b270997f5eead1755d6c606bbd8edd",
   "generated_by": "build_repo/autogenerate.py",
   "files": [
     "README.md",
     "config/defaults_embed.json",
     "config/defaults_finetune.json",
     "config/defaults_inference.json",
     "config/defaults_pretrain.json",
+    "config/released_model.json",
     "scripts/_utils.py",
     "scripts/check_checkpoint.py",
     "scripts/check_data.py",
+    "scripts/fetch_released_model.py",
     "scripts/kermt_container.sh",
     "scripts/prepare_data.py",
     "scripts/run_extract_embeddings.py",
@@ -32,6 +34,8 @@
     "tests/conftest.py",
     "tests/test_check_checkpoint.py",
     "tests/test_check_data.py",
+    "tests/test_e2e_released_download.py",
+    "tests/test_fetch_released_model.py",
     "tests/test_prepare_data.py",
     "tests/test_run_extract_embeddings.py",
     "tests/test_run_finetune_local.py",

open-models-skills/kermt/config/released_model.json

@@ -0,0 +1,13 @@
+{
+  "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"
+}

open-models-skills/kermt/scripts/fetch_released_model.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0 OR CC-BY-4.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 agent/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
+`agent/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 agent/config/released_model.json (one dir up from
+# agent/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: agent/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())

open-models-skills/kermt/scripts/kermt_container.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
+# SPDX-License-Identifier: Apache-2.0 OR CC-BY-4.0
 
 # kermt_container.sh — bootstrap helper for the kermt agent skills.
 #
@@ -31,6 +31,8 @@
 #   --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>)
@@ -273,6 +275,11 @@ _kermt_parse_mounts() {
         _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
         ;;
@@ -305,6 +312,16 @@ _kermt_git_env_flags() {
   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=()
@@ -321,6 +338,8 @@ kermt_run() {
   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" \
@@ -329,6 +348,7 @@ kermt_run() {
     -e PYTHONPATH=/workspace \
     -e HOME=/tmp/kermt-home \
     "${git_args[@]}" \
+    "${hf_args[@]}" \
     "$KERMT_IMAGE" \
     conda run -n kermt --no-capture-output bash -c "$*"
 }
@@ -342,7 +362,7 @@ kermt_run_detached() {
     case "$1" in
       --name) name="$2"; shift 2 ;;
       --) break ;;
-      --data|--ckpt|--vocab-dir|--run-dir) break ;;
+      --data|--ckpt|--vocab-dir|--run-dir|--model-dir) break ;;
       *) break ;;
     esac
   done
@@ -363,6 +383,8 @@ kermt_run_detached() {
   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" \
@@ -372,6 +394,7 @@ kermt_run_detached() {
     -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"
@@ -415,6 +438,7 @@ Mount flags (for run / run_detached):
   --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>)