qcoder 0.1.0a0__py3-none-any.whl

qcoder/tools/train_baseline_models.py

@@ -0,0 +1,619 @@
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import hashlib
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+try:
+    import numpy as np
+except Exception:  # pragma: no cover - fallback path
+    import jax.numpy as np  # type: ignore
+
+
+DEFAULT_INPUT = Path("data/training/joined_observations_ok.jsonl")
+DEFAULT_OUTPUT_DIR = Path("data/artifacts/models/qr12/schema_0.4.0")
+DEFAULT_RUNTIME_OUT = DEFAULT_OUTPUT_DIR / "runtime_ridge.json"
+DEFAULT_BUILD_OUT = DEFAULT_OUTPUT_DIR / "build_ridge.json"
+DEFAULT_FIDELITY_OUT = DEFAULT_OUTPUT_DIR / "fidelity_ridge.json"
+DEFAULT_MEMORY_OUT = DEFAULT_OUTPUT_DIR / "memory_ridge.json"
+
+
+def sha256_file(path: str | Path) -> str:
+    h = hashlib.sha256()
+    with Path(path).open("rb") as f:
+        for chunk in iter(lambda: f.read(65536), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _now_utc() -> str:
+    return dt.datetime.now(dt.timezone.utc).replace(microsecond=0).isoformat()
+
+
+def _load_jsonl(path: str | Path) -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    with Path(path).open("r", encoding="utf-8") as f:
+        for line in f:
+            s = line.strip()
+            if not s:
+                continue
+            rows.append(json.loads(s))
+    return rows
+
+
+def _precision_is_double(precision: Any) -> float:
+    p = str(precision or "").strip().lower()
+    return 1.0 if p in {"double", "fp64"} else 0.0
+
+
+def _safe_log2_threshold(x: Any) -> float:
+    try:
+        v = float(x)
+    except Exception:
+        return 0.0
+    return math.log2(v) if v > 0 else 0.0
+
+
+def _safe_log2_shots_runner(x: Any) -> float:
+    s = _safe_int(x)
+    if s is None or s <= 0:
+        return 0.0
+    return math.log2(float(s))
+
+
+def _content_hash_is_test(content_hash: str, test_fraction: float) -> bool:
+    if not content_hash:
+        return False
+    digest = hashlib.sha256(content_hash.encode("utf-8")).hexdigest()
+    n = int(digest[:8], 16) / float(0xFFFFFFFF)
+    return n < test_fraction
+
+
+def _safe_int(value: Any) -> int | None:
+    try:
+        return int(value)
+    except Exception:
+        return None
+
+
+def _is_total_runtime_row(r: dict[str, Any]) -> bool:
+    if r.get("run_kind") != "forward_runner_execution":
+        return False
+    if r.get("runner_wall_s") is None:
+        return False
+    shots_runner = _safe_int(r.get("shots_runner"))
+    if shots_runner is None:
+        return False
+    shots_hist = _safe_int(r.get("shots_hist"))
+    if shots_hist is not None and shots_hist > 1:
+        return shots_runner == shots_hist
+    return shots_runner > 1
+
+
+def _derive_execution_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """
+    Derive execution-time rows from forward_runner_execution rows.
+    total_runtime ≈ build_time + execution_time; execution_time ≈ slope * shots.
+    Match state row (shots_runner==1) and hist rows (shots_runner>1) by
+    content_hash, backend_id, precision, threshold; then
+    execution_time_s = max(0, runner_wall_s_hist - runner_wall_s_state).
+    Returns one row per (group, hist_row) with execution_time_s and hist row's features.
+    """
+    forward = [r for r in rows if r.get("run_kind") == "forward_runner_execution"]
+    key_fn = lambda r: (
+        str(r.get("content_hash") or ""),
+        str(r.get("backend_id") or r.get("backend") or ""),
+        str(r.get("precision") or "").strip() or "double",
+        _safe_int(r.get("threshold")),
+    )
+    groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {}
+    for r in forward:
+        k = key_fn(r)
+        if k[3] is None:
+            continue
+        groups.setdefault(k, []).append(r)
+
+    out: list[dict[str, Any]] = []
+    for (_ch, _bid, _prec, thr), group in groups.items():
+        state_row = next((r for r in group if _safe_int(r.get("shots_runner")) == 1), None)
+        if state_row is None:
+            continue
+        runner_wall_state = state_row.get("runner_wall_s")
+        if runner_wall_state is None:
+            continue
+        try:
+            build_s = float(runner_wall_state)
+        except (TypeError, ValueError):
+            continue
+        for hist_row in group:
+            shots = _safe_int(hist_row.get("shots_runner"))
+            if shots is None or shots <= 1:
+                continue
+            runner_wall_hist = hist_row.get("runner_wall_s")
+            if runner_wall_hist is None:
+                continue
+            try:
+                total_s = float(runner_wall_hist)
+            except (TypeError, ValueError):
+                continue
+            execution_time_s = max(0.0, total_s - build_s)
+            out.append({
+                "content_hash": _ch,
+                "backend_id": _bid,
+                "precision": _prec,
+                "threshold": thr,
+                "shots_runner": shots,
+                "execution_time_s": execution_time_s,
+                "feature_names": list(hist_row.get("feature_names") or []),
+                "features": list(hist_row.get("features") or []),
+            })
+    return out
+
+
+def _is_build_runtime_row(r: dict[str, Any]) -> bool:
+    if r.get("run_kind") != "forward_runner_execution":
+        return False
+    if r.get("runner_wall_s") is None:
+        return False
+    shots_runner = _safe_int(r.get("shots_runner"))
+    return shots_runner == 1
+
+
+def _is_memory_row(r: dict[str, Any]) -> bool:
+    if not _is_total_runtime_row(r):
+        return False
+    return r.get("peak_rss_mb") is not None
+
+
+def _build_xy(
+    rows: list[dict[str, Any]],
+    *,
+    for_runtime: bool,
+    eps: float,
+    runtime_selector: Any = None,
+    include_log2_shots_runner: bool = False,
+) -> tuple[np.ndarray, np.ndarray, list[str], np.ndarray]:
+    selected: list[dict[str, Any]] = []
+    for r in rows:
+        if for_runtime:
+            if runtime_selector is None:
+                keep = r.get("run_kind") == "forward_runner_execution" and r.get("runner_wall_s") is not None
+            else:
+                keep = bool(runtime_selector(r))
+            if not keep:
+                continue
+        else:
+            if r.get("run_kind") != "mirror_threshold_attempt":
+                continue
+            if r.get("fidelity") is None:
+                continue
+        selected.append(r)
+
+    if not selected:
+        raise ValueError("no rows selected for model training")
+
+    first = selected[0]
+    circuit_feature_names = list(first.get("feature_names") or [])
+    if not circuit_feature_names:
+        raise ValueError("joined rows missing feature_names")
+
+    if include_log2_shots_runner:
+        feature_names = ["bias", "log2_threshold", "precision_is_double", "log2_shots_runner"] + circuit_feature_names
+    else:
+        feature_names = ["bias", "log2_threshold", "precision_is_double"] + circuit_feature_names
+
+    x_rows: list[list[float]] = []
+    y_vals: list[float] = []
+    y_eval_vals: list[float] = []
+
+    for r in selected:
+        feats = list(r.get("features") or [])
+        if len(feats) != len(circuit_feature_names):
+            raise ValueError("feature length mismatch in joined rows")
+
+        x = [1.0, _safe_log2_threshold(r.get("threshold")), _precision_is_double(r.get("precision"))]
+        if include_log2_shots_runner:
+            x.append(_safe_log2_shots_runner(r.get("shots_runner")))
+        x.extend(float(v) for v in feats)
+        x_rows.append(x)
+
+        if for_runtime:
+            y_sec = float(r["runner_wall_s"])
+            y_vals.append(math.log1p(max(0.0, y_sec)))
+            y_eval_vals.append(max(0.0, y_sec))
+        else:
+            fid = float(r["fidelity"])
+            fid = min(1.0 - eps, max(eps, fid))
+            y_vals.append(math.log(fid / (1.0 - fid)))
+            y_eval_vals.append(float(r["fidelity"]))
+
+    return np.asarray(x_rows, dtype=float), np.asarray(y_vals, dtype=float), feature_names, np.asarray(y_eval_vals, dtype=float)
+
+
+def _build_xy_execution(
+    execution_rows: list[dict[str, Any]],
+) -> tuple[np.ndarray, np.ndarray, list[str], np.ndarray]:
+    """
+    Build X and y from derived execution rows. y = log1p(execution_time_s).
+    total_runtime ≈ build_time + execution_time; we model execution_time.
+    """
+    if not execution_rows:
+        raise ValueError("no execution rows for runtime model training")
+    circuit_feature_names = list((execution_rows[0].get("feature_names") or []))
+    if not circuit_feature_names:
+        raise ValueError("execution rows missing feature_names")
+    feature_names = ["bias", "log2_threshold", "precision_is_double", "log2_shots_runner"] + circuit_feature_names
+    x_rows: list[list[float]] = []
+    y_vals: list[float] = []
+    y_eval_vals: list[float] = []
+    for r in execution_rows:
+        feats = list(r.get("features") or [])
+        if len(feats) != len(circuit_feature_names):
+            raise ValueError("feature length mismatch in execution rows")
+        x = [1.0, _safe_log2_threshold(r.get("threshold")), _precision_is_double(r.get("precision"))]
+        x.append(_safe_log2_shots_runner(r.get("shots_runner")))
+        x.extend(float(v) for v in feats)
+        x_rows.append(x)
+        exec_s = max(0.0, float(r["execution_time_s"]))
+        y_vals.append(math.log1p(exec_s))
+        y_eval_vals.append(exec_s)
+    return (
+        np.asarray(x_rows, dtype=float),
+        np.asarray(y_vals, dtype=float),
+        feature_names,
+        np.asarray(y_eval_vals, dtype=float),
+    )
+
+
+def _build_xy_memory(
+    rows: list[dict[str, Any]],
+    *,
+    selector: Any,
+) -> tuple[np.ndarray, np.ndarray, list[str], np.ndarray]:
+    selected: list[dict[str, Any]] = []
+    for r in rows:
+        if not bool(selector(r)):
+            continue
+        selected.append(r)
+
+    if not selected:
+        raise ValueError("no rows selected for memory model training")
+
+    first = selected[0]
+    feature_names = list(first.get("feature_names") or [])
+    if not feature_names:
+        raise ValueError("joined rows missing feature_names")
+
+    x_rows: list[list[float]] = []
+    y_vals: list[float] = []
+    y_eval_vals: list[float] = []
+    for r in selected:
+        feats = list(r.get("features") or [])
+        if len(feats) != len(feature_names):
+            raise ValueError("feature length mismatch in joined rows")
+        x = [1.0, _safe_log2_threshold(r.get("threshold")), _precision_is_double(r.get("precision"))]
+        x.extend(float(v) for v in feats)
+        x_rows.append(x)
+
+        y_mb = max(0.0, float(r["peak_rss_mb"]))
+        y_vals.append(math.log1p(y_mb))
+        y_eval_vals.append(y_mb)
+
+    return np.asarray(x_rows, dtype=float), np.asarray(y_vals, dtype=float), feature_names, np.asarray(y_eval_vals, dtype=float)
+
+
+def _split_indices(rows: list[dict[str, Any]], test_fraction: float) -> tuple[list[int], list[int]]:
+    train_idx: list[int] = []
+    test_idx: list[int] = []
+    for i, r in enumerate(rows):
+        ch = str(r.get("content_hash") or "")
+        if _content_hash_is_test(ch, test_fraction):
+            test_idx.append(i)
+        else:
+            train_idx.append(i)
+    if not train_idx:
+        train_idx = list(range(len(rows)))
+        test_idx = []
+    return train_idx, test_idx
+
+
+def _standardize_from_train(x: np.ndarray, train_idx: list[int]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    x_mean = np.mean(x[train_idx], axis=0)
+    x_std = np.std(x[train_idx], axis=0)
+    x_mean[0] = 0.0
+    x_std[0] = 1.0
+    x_std = np.where(x_std == 0.0, 1.0, x_std)
+    x_scaled = (x - x_mean) / x_std
+    return x_scaled, x_mean, x_std
+
+
+def _fit_ridge_closed_form(x_train: np.ndarray, y_train: np.ndarray, alpha: float) -> np.ndarray:
+    n_features = int(x_train.shape[1])
+    xtx = x_train.T @ x_train
+    reg = np.eye(n_features, dtype=float) * float(alpha)
+    reg[0, 0] = 0.0  # do not regularize bias
+    xty = x_train.T @ y_train
+    try:
+        coef = np.linalg.solve(xtx + reg, xty)
+    except Exception:
+        coef = np.linalg.lstsq(xtx + reg, xty, rcond=None)[0]
+    return coef
+
+
+def _runtime_metrics(y_log_true: np.ndarray, y_log_pred: np.ndarray, y_sec_true: np.ndarray) -> dict[str, float | None]:
+    if y_log_true.size == 0:
+        return {"rmse_log": None, "median_multiplicative_error_seconds": None}
+    rmse_log = float(np.sqrt(np.mean((y_log_pred - y_log_true) ** 2)))
+    y_sec_pred = np.expm1(y_log_pred)
+    eps = 1e-9
+    ratios = np.maximum((y_sec_pred + eps) / (y_sec_true + eps), (y_sec_true + eps) / (y_sec_pred + eps))
+    med_mul = float(np.median(ratios))
+    return {"rmse_log": rmse_log, "median_multiplicative_error_seconds": med_mul}
+
+
+def _sigmoid(x: np.ndarray) -> np.ndarray:
+    return 1.0 / (1.0 + np.exp(-x))
+
+
+def _fidelity_metrics(y_fid_true: np.ndarray, y_logit_pred: np.ndarray) -> dict[str, float | None]:
+    if y_fid_true.size == 0:
+        return {"mae_fidelity": None}
+    y_fid_pred = _sigmoid(y_logit_pred)
+    mae = float(np.mean(np.abs(y_fid_pred - y_fid_true)))
+    return {"mae_fidelity": mae}
+
+
+def _memory_metrics(y_log_true: np.ndarray, y_log_pred: np.ndarray, y_mb_true: np.ndarray) -> dict[str, float | None]:
+    if y_log_true.size == 0:
+        return {"rmse_log": None, "median_multiplicative_error_mb": None}
+    rmse_log = float(np.sqrt(np.mean((y_log_pred - y_log_true) ** 2)))
+    y_mb_pred = np.expm1(y_log_pred)
+    eps = 1e-9
+    ratios = np.maximum((y_mb_pred + eps) / (y_mb_true + eps), (y_mb_true + eps) / (y_mb_pred + eps))
+    med_mul = float(np.median(ratios))
+    return {"rmse_log": rmse_log, "median_multiplicative_error_mb": med_mul}
+
+
+def _build_artifact(
+    *,
+    schema_version: str,
+    feature_names: list[str],
+    coef: np.ndarray,
+    x_mean: np.ndarray,
+    x_std: np.ndarray,
+    alpha: float,
+    joined_input_path: str,
+    joined_input_sha256: str,
+    rows_total: int,
+    rows_train: int,
+    rows_test: int,
+    model_kind: str,
+    metrics: dict[str, float | None],
+) -> dict[str, Any]:
+    return {
+        "schema_version": schema_version,
+        "feature_names": feature_names,
+        "coef": [float(v) for v in coef.tolist()],
+        "x_mean": [float(v) for v in x_mean.tolist()],
+        "x_std": [float(v) for v in x_std.tolist()],
+        "alpha": float(alpha),
+        "training_metadata": {
+            "model_kind": model_kind,
+            "created_utc": _now_utc(),
+            "joined_input_path": str(joined_input_path),
+            "joined_input_sha256": joined_input_sha256,
+            "rows_used": int(rows_total),
+            "rows_train": int(rows_train),
+            "rows_test": int(rows_test),
+            "metrics": metrics,
+        },
+    }
+
+
+def train_from_joined(
+    input_path: str | Path = DEFAULT_INPUT,
+    *,
+    runtime_out: str | Path = DEFAULT_RUNTIME_OUT,
+    build_out: str | Path = DEFAULT_BUILD_OUT,
+    fidelity_out: str | Path = DEFAULT_FIDELITY_OUT,
+    memory_out: str | Path = DEFAULT_MEMORY_OUT,
+    alpha: float = 1.0,
+    eps: float = 1e-6,
+    test_fraction: float = 0.2,
+) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any] | None]:
+    rows = _load_jsonl(input_path)
+    if not rows:
+        raise ValueError("no joined rows found")
+
+    schema_version = str(rows[0].get("schema_version") or "")
+    joined_sha = sha256_file(input_path)
+    full_feature_names = ["bias", "log2_threshold", "precision_is_double"] + list(rows[0].get("feature_names") or [])
+
+    # Execution-time model: total_runtime ≈ build_time + execution_time; execution_time ≈ slope * shots.
+    # Derive execution_time_s = max(0, runner_wall_s_hist - runner_wall_s_state) per (content_hash, backend_id, precision, threshold).
+    execution_rows = _derive_execution_rows(rows)
+    x_rt, y_rt, full_feature_names_runtime, y_rt_sec = _build_xy_execution(execution_rows)
+    rt_train_idx, rt_test_idx = _split_indices(execution_rows, test_fraction)
+    x_rt_scaled, x_rt_mean, x_rt_std = _standardize_from_train(x_rt, rt_train_idx)
+    rt_coef = _fit_ridge_closed_form(x_rt_scaled[rt_train_idx], y_rt[rt_train_idx], alpha=float(alpha))
+    y_rt_pred_test = x_rt_scaled[rt_test_idx] @ rt_coef if rt_test_idx else np.asarray([], dtype=float)
+    rt_metrics = _runtime_metrics(
+        y_rt[rt_test_idx] if rt_test_idx else np.asarray([], dtype=float),
+        y_rt_pred_test,
+        y_rt_sec[rt_test_idx] if rt_test_idx else np.asarray([], dtype=float),
+    )
+    runtime_artifact = _build_artifact(
+        schema_version=schema_version,
+        feature_names=full_feature_names_runtime,
+        coef=rt_coef,
+        x_mean=x_rt_mean,
+        x_std=x_rt_std,
+        alpha=float(alpha),
+        joined_input_path=str(input_path),
+        joined_input_sha256=joined_sha,
+        rows_total=len(execution_rows),
+        rows_train=len(rt_train_idx),
+        rows_test=len(rt_test_idx),
+        model_kind="execution_runtime_ridge",
+        metrics=rt_metrics,
+    )
+
+    # Build-time model (1-shot probe forward rows)
+    x_bd, y_bd, _, y_bd_sec = _build_xy(rows, for_runtime=True, eps=eps, runtime_selector=_is_build_runtime_row)
+    bd_rows = [
+        r
+        for r in rows
+        if _is_build_runtime_row(r)
+    ]
+    bd_train_idx, bd_test_idx = _split_indices(bd_rows, test_fraction)
+    x_bd_scaled, x_bd_mean, x_bd_std = _standardize_from_train(x_bd, bd_train_idx)
+    bd_coef = _fit_ridge_closed_form(x_bd_scaled[bd_train_idx], y_bd[bd_train_idx], alpha=float(alpha))
+    y_bd_pred_test = x_bd_scaled[bd_test_idx] @ bd_coef if bd_test_idx else np.asarray([], dtype=float)
+    bd_metrics = _runtime_metrics(
+        y_bd[bd_test_idx] if bd_test_idx else np.asarray([], dtype=float),
+        y_bd_pred_test,
+        y_bd_sec[bd_test_idx] if bd_test_idx else np.asarray([], dtype=float),
+    )
+    build_artifact = _build_artifact(
+        schema_version=schema_version,
+        feature_names=full_feature_names,
+        coef=bd_coef,
+        x_mean=x_bd_mean,
+        x_std=x_bd_std,
+        alpha=float(alpha),
+        joined_input_path=str(input_path),
+        joined_input_sha256=joined_sha,
+        rows_total=len(bd_rows),
+        rows_train=len(bd_train_idx),
+        rows_test=len(bd_test_idx),
+        model_kind="build_ridge_log1p",
+        metrics=bd_metrics,
+    )
+
+    # Fidelity model
+    x_fd, y_fd, _, y_fd_true = _build_xy(rows, for_runtime=False, eps=eps)
+    fd_rows = [
+        r
+        for r in rows
+        if r.get("run_kind") == "mirror_threshold_attempt" and r.get("fidelity") is not None
+    ]
+    fd_train_idx, fd_test_idx = _split_indices(fd_rows, test_fraction)
+    x_fd_scaled, x_fd_mean, x_fd_std = _standardize_from_train(x_fd, fd_train_idx)
+    fd_coef = _fit_ridge_closed_form(x_fd_scaled[fd_train_idx], y_fd[fd_train_idx], alpha=float(alpha))
+    y_fd_pred_test = x_fd_scaled[fd_test_idx] @ fd_coef if fd_test_idx else np.asarray([], dtype=float)
+    fd_metrics = _fidelity_metrics(
+        y_fd_true[fd_test_idx] if fd_test_idx else np.asarray([], dtype=float),
+        y_fd_pred_test,
+    )
+    fidelity_artifact = _build_artifact(
+        schema_version=schema_version,
+        feature_names=full_feature_names,
+        coef=fd_coef,
+        x_mean=x_fd_mean,
+        x_std=x_fd_std,
+        alpha=float(alpha),
+        joined_input_path=str(input_path),
+        joined_input_sha256=joined_sha,
+        rows_total=len(fd_rows),
+        rows_train=len(fd_train_idx),
+        rows_test=len(fd_test_idx),
+        model_kind="fidelity_ridge_logit",
+        metrics=fd_metrics,
+    )
+
+    # Memory model (optional; uses total-runtime row selection)
+    memory_artifact: dict[str, Any] | None = None
+    mem_rows = [r for r in rows if _is_memory_row(r)]
+    if mem_rows:
+        x_mb, y_mb, _, y_mb_true = _build_xy_memory(rows, selector=_is_memory_row)
+        mb_train_idx, mb_test_idx = _split_indices(mem_rows, test_fraction)
+        x_mb_scaled, x_mb_mean, x_mb_std = _standardize_from_train(x_mb, mb_train_idx)
+        mb_coef = _fit_ridge_closed_form(x_mb_scaled[mb_train_idx], y_mb[mb_train_idx], alpha=float(alpha))
+        y_mb_pred_test = x_mb_scaled[mb_test_idx] @ mb_coef if mb_test_idx else np.asarray([], dtype=float)
+        mb_metrics = _memory_metrics(
+            y_mb[mb_test_idx] if mb_test_idx else np.asarray([], dtype=float),
+            y_mb_pred_test,
+            y_mb_true[mb_test_idx] if mb_test_idx else np.asarray([], dtype=float),
+        )
+        memory_artifact = _build_artifact(
+            schema_version=schema_version,
+            feature_names=full_feature_names,
+            coef=mb_coef,
+            x_mean=x_mb_mean,
+            x_std=x_mb_std,
+            alpha=float(alpha),
+            joined_input_path=str(input_path),
+            joined_input_sha256=joined_sha,
+            rows_total=len(mem_rows),
+            rows_train=len(mb_train_idx),
+            rows_test=len(mb_test_idx),
+            model_kind="memory_ridge_log1p_mb",
+            metrics=mb_metrics,
+        )
+
+    runtime_out = Path(runtime_out)
+    build_out = Path(build_out)
+    fidelity_out = Path(fidelity_out)
+    memory_out = Path(memory_out)
+    runtime_out.parent.mkdir(parents=True, exist_ok=True)
+    build_out.parent.mkdir(parents=True, exist_ok=True)
+    fidelity_out.parent.mkdir(parents=True, exist_ok=True)
+    memory_out.parent.mkdir(parents=True, exist_ok=True)
+    runtime_out.write_text(json.dumps(runtime_artifact, indent=2, sort_keys=True), encoding="utf-8")
+    build_out.write_text(json.dumps(build_artifact, indent=2, sort_keys=True), encoding="utf-8")
+    fidelity_out.write_text(json.dumps(fidelity_artifact, indent=2, sort_keys=True), encoding="utf-8")
+    if memory_artifact is not None:
+        memory_out.write_text(json.dumps(memory_artifact, indent=2, sort_keys=True), encoding="utf-8")
+    return runtime_artifact, build_artifact, fidelity_artifact, memory_artifact
+
+
+def main(argv: list[str] | None = None) -> int:
+    ap = argparse.ArgumentParser(description="Train baseline runtime/fidelity ridge regressors from joined observations.")
+    ap.add_argument("--input", default=str(DEFAULT_INPUT), help="Joined observation JSONL path.")
+    ap.add_argument("--runtime-out", default=str(DEFAULT_RUNTIME_OUT), help="Runtime artifact JSON path.")
+    ap.add_argument("--build-out", default=str(DEFAULT_BUILD_OUT), help="Build-time artifact JSON path.")
+    ap.add_argument("--fidelity-out", default=str(DEFAULT_FIDELITY_OUT), help="Fidelity artifact JSON path.")
+    ap.add_argument("--memory-out", default=str(DEFAULT_MEMORY_OUT), help="Memory artifact JSON path (optional if labels exist).")
+    ap.add_argument("--alpha", type=float, default=1.0, help="Ridge L2 strength.")
+    ap.add_argument("--test-fraction", type=float, default=0.2, help="Deterministic content-hash holdout fraction.")
+    ap.add_argument("--eps", type=float, default=1e-6, help="Clipping epsilon for fidelity logit transform.")
+    args = ap.parse_args(argv)
+
+    rt, bd, fd, mb = train_from_joined(
+        input_path=args.input,
+        runtime_out=args.runtime_out,
+        build_out=args.build_out,
+        fidelity_out=args.fidelity_out,
+        memory_out=args.memory_out,
+        alpha=args.alpha,
+        eps=args.eps,
+        test_fraction=args.test_fraction,
+    )
+    print(
+        f"[ok] wrote runtime model: {args.runtime_out} (rows={rt['training_metadata']['rows_used']})",
+        flush=True,
+    )
+    print(
+        f"[ok] wrote build-time model: {args.build_out} (rows={bd['training_metadata']['rows_used']})",
+        flush=True,
+    )
+    print(
+        f"[ok] wrote fidelity model: {args.fidelity_out} (rows={fd['training_metadata']['rows_used']})",
+        flush=True,
+    )
+    if mb is not None:
+        print(
+            f"[ok] wrote memory model: {args.memory_out} (rows={mb['training_metadata']['rows_used']})",
+            flush=True,
+        )
+    else:
+        print("[info] skipped memory model: no eligible rows with peak_rss_mb labels", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())