feat: add USE_MODEL_REGISTRY env var and run.sh for CI

Introduces an explicit USE_MODEL_REGISTRY flag to decouple model source
from Opik tracing. When unset (default), distilgpt2 loads from HuggingFace
so CI can run without CometML registry credentials while still sending
real traces to Opik. Both the script and notebook are updated to match.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: LeoRoccoBreedt

guides/tracing_finetuned_models/run.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+set -e
+
+export OPIK_PROJECT_NAME="tracing-finetuned-models"
+
+# Set USE_MODEL_REGISTRY=true and export COMET_API_KEY + COMET_WORKSPACE to
+# download the fine-tuned model from the CometML registry instead of HuggingFace.
+# export USE_MODEL_REGISTRY=true
+
+uv sync
+
+# Runs inference tracing only. train_and_register.py requires GPU and is excluded from CI.
+uv run python use_registered_model.py

guides/tracing_finetuned_models/tracing_finetuned_models.ipynb

@@ -19,19 +19,19 @@
     "```\n",
     "Section 1                       Section 2\n",
     "train_and_register              fetch_and_trace\n",
-    "       \u2502                               \u2502\n",
-    "       \u2502  fine-tunes distilgpt2        \u2502  fetches registered model\n",
-    "       \u2502  logs checkpoints             \u2502  runs Opik-traced inference\n",
-    "       \u25bc                               \u25bc\n",
-    " CometML Model Registry \u2500\u2500\u2500\u2500\u2500\u2500\u25ba Opik trace\n",
+    "       │                               │\n",
+    "       │  fine-tunes distilgpt2        │  fetches registered model\n",
+    "       │  logs checkpoints             │  runs Opik-traced inference\n",
+    "       ▼                               ▼\n",
+    " CometML Model Registry ──────► Opik trace\n",
     "   sft-distilgpt2 v1.0.0          metadata.model_registry_url\n",
     "```\n",
     "\n",
     "**Already have a registered model?** Skip straight to [Section 2](#section-2).\n",
     "\n",
     "---\n",
     "\n",
-    "**Recommended runtime:** GPU (T4 or better). In Colab: *Runtime \u2192 Change runtime type \u2192 T4 GPU*."
+    "**Recommended runtime:** GPU (T4 or better). In Colab: *Runtime → Change runtime type → T4 GPU*."
    ],
    "outputs": [],
    "execution_count": null
@@ -61,7 +61,7 @@
    "id": "cell-restart-note",
    "metadata": {},
    "source": [
-    "> **Note:** If you just installed the packages for the first time, restart the Colab runtime now (*Runtime \u2192 Restart session*) then run all cells again from the top."
+    "> **Note:** If you just installed the packages for the first time, restart the Colab runtime now (*Runtime → Restart session*) then run all cells again from the top."
    ],
    "outputs": [],
    "execution_count": null
@@ -72,7 +72,7 @@
    "id": "cell-credentials",
    "metadata": {},
    "outputs": [],
-   "source": "import os\nimport getpass\n\n# CometML and Opik share the same workspace and API key on the Comet platform.\n# Locally: export COMET_API_KEY and COMET_WORKSPACE before running.\nAPI_KEY   = os.environ.get('COMET_API_KEY', '') or getpass.getpass('Enter your Comet API key: ')\nWORKSPACE = os.environ.get('COMET_WORKSPACE', '') or getpass.getpass('Enter your Comet workspace: ')\n\n# Propagate to both SDKs\nos.environ['COMET_API_KEY']   = API_KEY\nos.environ['COMET_WORKSPACE'] = WORKSPACE\nos.environ['OPIK_API_KEY']    = API_KEY\nos.environ['OPIK_WORKSPACE']  = WORKSPACE\n\nCOMET_API_KEY   = API_KEY\nCOMET_WORKSPACE = WORKSPACE\nOPIK_API_KEY    = API_KEY\nOPIK_WORKSPACE  = WORKSPACE\n\n# Update these to match your use case.\nREGISTRY_NAME = 'sft-distilgpt2'\nMODEL_VERSION = '1.0.0'\nOPIK_PROJECT  = 'tracing-finetuned-models'\n\nprint(f'Workspace : {WORKSPACE or \"(not set)\"}')"
+   "source": "import os\nimport getpass\n\n# CometML and Opik share the same workspace and API key on the Comet platform.\n# Locally: export COMET_API_KEY and COMET_WORKSPACE before running.\nAPI_KEY   = os.environ.get('COMET_API_KEY', '') or getpass.getpass('Enter your Comet API key: ')\nWORKSPACE = os.environ.get('COMET_WORKSPACE', '') or getpass.getpass('Enter your Comet workspace: ')\n\n# Propagate to both SDKs\nos.environ['COMET_API_KEY']   = API_KEY\nos.environ['COMET_WORKSPACE'] = WORKSPACE\nos.environ['OPIK_API_KEY']    = API_KEY\nos.environ['OPIK_WORKSPACE']  = WORKSPACE\n\nCOMET_API_KEY   = API_KEY\nCOMET_WORKSPACE = WORKSPACE\nOPIK_API_KEY    = API_KEY\nOPIK_WORKSPACE  = WORKSPACE\n\n# Set to True to download the fine-tuned model from the CometML registry.\n# False (default) loads distilgpt2 directly from HuggingFace — Opik tracing still runs.\nUSE_MODEL_REGISTRY = os.environ.get('USE_MODEL_REGISTRY', 'false').lower() in ('1', 'true', 'yes')\n\n# Update these to match your registered model when USE_MODEL_REGISTRY=True.\nREGISTRY_NAME = 'sft-distilgpt2'\nMODEL_VERSION = '1.0.0'\nOPIK_PROJECT  = 'tracing-finetuned-models'\n\nprint(f'Workspace          : {WORKSPACE or \"(not set)\"}')\nprint(f'USE_MODEL_REGISTRY : {USE_MODEL_REGISTRY}')"
   },
   {
    "cell_type": "markdown",
@@ -110,7 +110,7 @@
    "source": [
     "# comet_ml MUST be imported first.\n",
     "os.environ['COMET_API_KEY'] = COMET_API_KEY or ''\n",
-    "import comet_ml  # noqa: F401 \u2014 triggers auto-integration with Trainer"
+    "import comet_ml  # noqa: F401 — triggers auto-integration with Trainer"
    ]
   },
   {
@@ -196,7 +196,7 @@
     "            metadata={'epoch': epoch},\n",
     "        )\n",
     "        shutil.rmtree(checkpoint_path)\n",
-    "        print(f'  \u2192 Logged checkpoint_epoch_{epoch} to CometML')"
+    "        print(f'  → Logged checkpoint_epoch_{epoch} to CometML')"
    ]
   },
   {
@@ -274,7 +274,7 @@
     "    print(f'  REGISTRY_NAME = \"{REGISTRY_NAME}\"')\n",
     "    print(f'  MODEL_VERSION = \"{MODEL_VERSION}\"')\n",
     "else:\n",
-    "    print('No active CometML experiment \u2014 check that COMET_API_KEY is set.')"
+    "    print('No active CometML experiment — check that COMET_API_KEY is set.')"
    ]
   },
   {
@@ -299,7 +299,7 @@
    "id": "cell-download",
    "metadata": {},
    "outputs": [],
-   "source": "import comet_ml\n\nMODEL_LOCAL_DIR = './downloaded_model'\n\nif COMET_API_KEY and COMET_WORKSPACE:\n    api = comet_ml.API(api_key=COMET_API_KEY)\n    print(f'Downloading {REGISTRY_NAME} v{MODEL_VERSION} from CometML Model Registry...')\n    registered_model = api.get_model(workspace=COMET_WORKSPACE, model_name=REGISTRY_NAME)\n    registered_model.download(version=MODEL_VERSION, output_folder=MODEL_LOCAL_DIR, expand=True)\n    model_source = MODEL_LOCAL_DIR\n    print(f'Downloaded to {MODEL_LOCAL_DIR}')\nelse:\n    print('[No credentials] Loading distilgpt2 from HuggingFace as a stand-in.')\n    model_source = 'distilgpt2'"
+   "source": "import comet_ml\n\nMODEL_LOCAL_DIR = './downloaded_model'\n\nif USE_MODEL_REGISTRY:\n    api = comet_ml.API(api_key=COMET_API_KEY)\n    print(f'Downloading {REGISTRY_NAME} v{MODEL_VERSION} from CometML Model Registry...')\n    registered_model = api.get_model(workspace=COMET_WORKSPACE, model_name=REGISTRY_NAME)\n    registered_model.download(version=MODEL_VERSION, output_folder=MODEL_LOCAL_DIR, expand=True)\n    model_source = MODEL_LOCAL_DIR\n    print(f'Downloaded to {MODEL_LOCAL_DIR}')\nelse:\n    print('USE_MODEL_REGISTRY not set — loading distilgpt2 from HuggingFace.')\n    print('Set USE_MODEL_REGISTRY=true (+ COMET_API_KEY / COMET_WORKSPACE) to use a registered model.')\n    model_source = 'distilgpt2'"
   },
   {
    "cell_type": "code",
@@ -338,32 +338,7 @@
    "id": "cell-generate",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "@opik.track(project_name=OPIK_PROJECT)\n",
-    "def generate(prompt: str, max_new_tokens: int = 100) -> str:\n",
-    "    \"\"\"\n",
-    "    Run inference and attach the registered model version to the Opik trace.\n",
-    "    The model_registry_url field lets you navigate from any trace directly\n",
-    "    to the CometML experiment and checkpoint that produced it.\n",
-    "    \"\"\"\n",
-    "    opik.update_current_trace(\n",
-    "        metadata={\n",
-    "            'model_registry_name': REGISTRY_NAME,\n",
-    "            'model_version':       MODEL_VERSION,\n",
-    "            'model_registry_url':  _registry_url(),\n",
-    "        }\n",
-    "    )\n",
-    "\n",
-    "    inputs = inference_tokenizer(prompt, return_tensors='pt')\n",
-    "    with torch.no_grad():\n",
-    "        output_ids = inference_model.generate(\n",
-    "            **inputs,\n",
-    "            max_new_tokens=max_new_tokens,\n",
-    "            pad_token_id=inference_tokenizer.eos_token_id,\n",
-    "            do_sample=False,\n",
-    "        )\n",
-    "    return inference_tokenizer.decode(output_ids[0], skip_special_tokens=True)"
-   ]
+   "source": "@opik.track(project_name=OPIK_PROJECT)\ndef generate(prompt: str, max_new_tokens: int = 100) -> str:\n    \"\"\"Run inference and optionally attach registry metadata to the Opik trace.\"\"\"\n    if USE_MODEL_REGISTRY:\n        opik.update_current_trace(\n            metadata={\n                'model_registry_name': REGISTRY_NAME,\n                'model_version':       MODEL_VERSION,\n                'model_registry_url':  _registry_url(),\n            }\n        )\n\n    inputs = inference_tokenizer(prompt, return_tensors='pt')\n    with torch.no_grad():\n        output_ids = inference_model.generate(\n            **inputs,\n            max_new_tokens=max_new_tokens,\n            pad_token_id=inference_tokenizer.eos_token_id,\n            do_sample=False,\n        )\n    return inference_tokenizer.decode(output_ids[0], skip_special_tokens=True)"
   },
   {
    "cell_type": "code",

guides/tracing_finetuned_models/use_registered_model.py

@@ -6,17 +6,21 @@
 that produced it — so you can always trace a prediction back to the exact
 checkpoint.
 
-Dry-run (no credentials): loads distilgpt2 from HuggingFace directly and
-prints traces to the console instead of sending to Opik.
+Set USE_MODEL_REGISTRY=true to download from the CometML registry (requires
+COMET_API_KEY, COMET_WORKSPACE, and a registered model). Without it the script
+loads distilgpt2 directly from HuggingFace. Opik tracing runs in either case
+as long as OPIK_API_KEY + OPIK_WORKSPACE are set.
 
 Run:
     pip install comet_ml opik transformers torch
-    export COMET_API_KEY="<your-api-key>"
-    export COMET_WORKSPACE="<your-workspace>"
-    export COMET_REGISTRY_NAME="sft-distilgpt2"
-    export COMET_MODEL_VERSION="1.0.0"
     export OPIK_API_KEY="<your-api-key>"
     export OPIK_WORKSPACE="<your-workspace>"
+    # Optional — use a registered model instead of HuggingFace:
+    # export USE_MODEL_REGISTRY=true
+    # export COMET_API_KEY="<your-api-key>"
+    # export COMET_WORKSPACE="<your-workspace>"
+    # export COMET_REGISTRY_NAME="sft-distilgpt2"
+    # export COMET_MODEL_VERSION="1.0.0"
     python use_registered_model.py
 """
 
@@ -27,17 +31,18 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # ── Credentials ───────────────────────────────────────────────────────────────
-COMET_API_KEY    = os.environ.get("COMET_API_KEY")
-COMET_WORKSPACE  = os.environ.get("COMET_WORKSPACE")
-OPIK_API_KEY     = os.environ.get("OPIK_API_KEY")
-OPIK_WORKSPACE   = os.environ.get("OPIK_WORKSPACE")
+OPIK_API_KEY    = os.environ.get("OPIK_API_KEY")
+OPIK_WORKSPACE  = os.environ.get("OPIK_WORKSPACE")
+COMET_API_KEY   = os.environ.get("COMET_API_KEY")
+COMET_WORKSPACE = os.environ.get("COMET_WORKSPACE")
 
-REGISTRY_NAME    = os.environ.get("COMET_REGISTRY_NAME", "sft-distilgpt2")
-MODEL_VERSION    = os.environ.get("COMET_MODEL_VERSION", "1.0.0")
-MODEL_LOCAL_DIR  = "./downloaded_model"
-OPIK_PROJECT     = os.environ.get("OPIK_PROJECT_NAME", "tracing-finetuned-models")
+REGISTRY_NAME   = os.environ.get("COMET_REGISTRY_NAME", "sft-distilgpt2")
+MODEL_VERSION   = os.environ.get("COMET_MODEL_VERSION", "1.0.0")
+MODEL_LOCAL_DIR = "./downloaded_model"
+OPIK_PROJECT    = os.environ.get("OPIK_PROJECT_NAME", "tracing-finetuned-models")
 
-DRY_RUN = not (COMET_API_KEY and COMET_WORKSPACE and OPIK_API_KEY and OPIK_WORKSPACE)
+USE_REGISTRY = os.environ.get("USE_MODEL_REGISTRY", "").lower() in ("1", "true", "yes")
+DRY_RUN      = not (OPIK_API_KEY and OPIK_WORKSPACE)
 
 # Module-level model handles — populated in main() before any inference call.
 tokenizer = None
@@ -73,18 +78,15 @@ def load_model(model_dir: str) -> None:
 # ── Step 3: Opik-traced inference ─────────────────────────────────────────────
 @opik.track(project_name=OPIK_PROJECT)
 def generate(prompt: str, max_new_tokens: int = 100) -> str:
-    """
-    Run inference and attach a link to the registered model version on the trace.
-    The registry URL in metadata lets you navigate from any Opik trace directly
-    to the CometML experiment and model version that produced it.
-    """
-    opik.update_current_trace(
-        metadata={
-            "model_registry_name": REGISTRY_NAME,
-            "model_version":       MODEL_VERSION,
-            "model_registry_url":  registry_url(),
-        }
-    )
+    """Run inference; attach registry metadata to the trace when a registry model is loaded."""
+    if USE_REGISTRY:
+        opik.update_current_trace(
+            metadata={
+                "model_registry_name": REGISTRY_NAME,
+                "model_version":       MODEL_VERSION,
+                "model_registry_url":  registry_url(),
+            }
+        )
 
     inputs = tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
@@ -99,12 +101,12 @@ def generate(prompt: str, max_new_tokens: int = 100) -> str:
 
 # ── Main ──────────────────────────────────────────────────────────────────────
 def main() -> None:
-    if DRY_RUN:
-        print("[DRY RUN] Missing credentials — loading distilgpt2 from HuggingFace directly.")
-        print("Set COMET_API_KEY, COMET_WORKSPACE, OPIK_API_KEY, OPIK_WORKSPACE to use a registered model.\n")
-        model_dir = "distilgpt2"
-    else:
+    if USE_REGISTRY:
         model_dir = download_registered_model()
+    else:
+        print("USE_MODEL_REGISTRY not set — loading distilgpt2 from HuggingFace.")
+        print("Set USE_MODEL_REGISTRY=true (+ COMET_API_KEY / COMET_WORKSPACE) to use a registered model.\n")
+        model_dir = "distilgpt2"
 
     load_model(model_dir)
 
@@ -120,12 +122,10 @@ def main() -> None:
         print(f"Response: {response}\n")
 
     if DRY_RUN:
-        print("[DRY RUN] Traces printed locally. Set credentials to send to Opik.")
-        print(f"Each trace would include: model_registry_url = {registry_url()}")
+        print("[DRY RUN] OPIK_API_KEY / OPIK_WORKSPACE not set — traces were not sent to Opik.")
     else:
         opik.flush_tracker()
         print(f"Traces sent to Opik project '{OPIK_PROJECT}'.")
-        print(f"Each trace links to: {registry_url()}")
 
 
 if __name__ == "__main__":