feat: enrich training deploy via mimir sdg

scripts/deploy_training_data.sh

@@ -5,9 +5,22 @@ cd "$(dirname "$0")/.."
 
 export COPYFILE_DISABLE=1
 
+if [[ -f .env ]]; then
+  set -a
+  # shellcheck disable=SC1091
+  source .env
+  set +a
+fi
+
 CUSTOM_LOCAL_LLM_ROOT="${CUSTOM_LOCAL_LLM_ROOT:-/Volumes/Zoe/custom-local-llm}"
 TRAINING_DATA_DIR="${TRAINING_DATA_DIR:-$CUSTOM_LOCAL_LLM_ROOT/training-data}"
 TRAINING_DATA_MANIFEST="${TRAINING_DATA_MANIFEST:-$CUSTOM_LOCAL_LLM_ROOT/manifest.llm.json}"
+ENABLE_MIMIR_SDG="${ENABLE_MIMIR_SDG:-0}"
+MIMIR_SDG_HOST="${MIMIR_SDG_HOST:-100.80.52.47}"
+MIMIR_SDG_USER="${MIMIR_SDG_USER:-aaron-pressey}"
+MIMIR_SDG_KEY="${MIMIR_SDG_KEY:-$HOME/.ssh/silma_orson_ed25519}"
+MIMIR_SDG_MODEL_ALIAS="${MIMIR_SDG_MODEL_ALIAS:-nvidia-text}"
+MIMIR_SDG_NUM_RECORDS="${MIMIR_SDG_NUM_RECORDS:-1000}"
 
 if [[ ! -d "$CUSTOM_LOCAL_LLM_ROOT" ]]; then
   echo "missing custom local LLM root: $CUSTOM_LOCAL_LLM_ROOT" >&2
@@ -22,10 +35,28 @@ if [[ ! -s "$TRAINING_DATA_MANIFEST" ]]; then
 fi
 
 cargo run -- dictionary generate
-python3 scripts/generate_training_data.py \
+bridge_args=(
+  scripts/generate_training_data.py
   --dictionary-dir dictionary/ \
   --output-dir "$TRAINING_DATA_DIR" \
   --manifest "$TRAINING_DATA_MANIFEST"
+)
+
+if [[ "$ENABLE_MIMIR_SDG" = "1" ]]; then
+  if [[ ! -r "$MIMIR_SDG_KEY" ]]; then
+    echo "missing Mimir SDG SSH key: $MIMIR_SDG_KEY" >&2
+    exit 1
+  fi
+  bridge_args+=(
+    --sdg-host "$MIMIR_SDG_HOST"
+    --sdg-user "$MIMIR_SDG_USER"
+    --sdg-key "$MIMIR_SDG_KEY"
+    --sdg-model-alias "$MIMIR_SDG_MODEL_ALIAS"
+    --num-records "$MIMIR_SDG_NUM_RECORDS"
+  )
+fi
+
+python3 "${bridge_args[@]}"
 
 python3 - <<'PY'
 import os

scripts/generate_training_data.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import shlex
 import subprocess
 import sys
 import tempfile
@@ -18,6 +19,8 @@ DEFAULT_MIN_RECORDS_PER_ACTION = 20
 CUSTOM_LOCAL_LLM_ROOT = Path("/Volumes/Zoe/custom-local-llm")
 DEFAULT_OUTPUT_DIR = CUSTOM_LOCAL_LLM_ROOT / "training-data"
 DEFAULT_MANIFEST = CUSTOM_LOCAL_LLM_ROOT / "manifest.llm.json"
+MIMIR_DATA_DESIGNER_BIN = "/home/aaron-pressey/.venvs/data-designer/bin/data-designer"
+MIMIR_DATA_DESIGNER_ENV = "/home/aaron-pressey/.config/home-grown-llm-data/data-designer.env"
 
 
 def main() -> int:
@@ -41,6 +44,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--sdg-host")
     parser.add_argument("--sdg-user")
     parser.add_argument("--sdg-key")
+    parser.add_argument("--sdg-model-alias", default="nvidia-text")
     parser.add_argument("--num-records", type=int, default=1000)
     parser.add_argument(
         "--timestamp",
@@ -71,6 +75,7 @@ def run_bridge(args: argparse.Namespace) -> dict[str, Any]:
         "output_dir": str(output_root / dataset_key),
         "manifest": str(manifest_path),
         "sdg_requested": has_sdg_args(args),
+        "sdg_model_alias": args.sdg_model_alias,
     }
 
     if args.dry_run:
@@ -93,6 +98,7 @@ def run_bridge(args: argparse.Namespace) -> dict[str, Any]:
                 generated_at=generated_at,
                 seed_records=train_records + validation_records,
                 action_ids=sorted(artifacts["actions"].keys()),
+                model_alias=args.sdg_model_alias,
             )
         except Exception as exc:  # noqa: BLE001 - keep local seed output and report remote failure.
             sdg_error = str(exc)
@@ -108,6 +114,7 @@ def run_bridge(args: argparse.Namespace) -> dict[str, Any]:
         artifacts=artifacts,
         expanded_path=expanded_path,
         sdg_requested=has_sdg_args(args),
+        sdg_model_alias=args.sdg_model_alias,
         sdg_error=sdg_error,
     )
     summary["expanded_path"] = str(expanded_path) if expanded_path else None
@@ -339,19 +346,23 @@ def run_mimir_expansion(
     generated_at: int,
     seed_records: list[dict[str, Any]],
     action_ids: list[str],
+    model_alias: str,
 ) -> Path:
     require_sdg_args(args)
     remote_base = "/mnt/storage/data-designer/managed-assets"
+    remote_artifact_path = f"{remote_base}/module_controller_artifacts_{generated_at}"
     remote_seed = f"{remote_base}/module_controller_seed_{generated_at}.jsonl"
-    remote_config = f"{remote_base}/module_controller_{generated_at}.yaml"
+    remote_config = f"{remote_base}/module_controller_{generated_at}.py"
+    remote_expanded = f"{remote_artifact_path}/{dataset_key}.jsonl"
     expanded_path = dataset_dir / "expanded.jsonl"
+    raw_expanded_path = dataset_dir / "expanded.raw.jsonl"
 
     with tempfile.TemporaryDirectory() as tmp:
         tmp_path = Path(tmp)
         local_seed = tmp_path / "seed.jsonl"
-        local_config = tmp_path / "module_controller.yaml"
+        local_config = tmp_path / "module_controller.py"
         write_jsonl(local_seed, seed_records)
-        local_config.write_text(data_designer_config(dataset_key, generated_at, remote_seed, action_ids))
+        local_config.write_text(data_designer_config(dataset_key, generated_at, remote_seed, action_ids, model_alias))
 
         target = f"{args.sdg_user}@{args.sdg_host}"
         scp_base = ["scp", "-i", str(Path(args.sdg_key).expanduser())]
@@ -361,59 +372,130 @@ def run_mimir_expansion(
         subprocess.run(
             [
                 *ssh_base,
-                "/home/aaron-pressey/.venvs/data-designer/bin/data-designer",
-                "create",
-                remote_config,
-                "--num-records",
-                str(args.num_records),
-                "--dataset-name",
-                dataset_key,
+                remote_bash_command(
+                    remote_data_designer_create_command(
+                        remote_config=remote_config,
+                        num_records=args.num_records,
+                        dataset_key=dataset_key,
+                        artifact_path=remote_artifact_path,
+                    )
+                ),
             ],
             check=True,
         )
-        subprocess.run([*scp_base, f"{target}:{remote_base}/{dataset_key}.jsonl", str(expanded_path)], check=True)
+        subprocess.run([*scp_base, f"{target}:{remote_expanded}", str(raw_expanded_path)], check=True)
+        normalize_expanded_jsonl(raw_expanded_path, expanded_path)
     return expanded_path
 
 
-def data_designer_config(dataset_key: str, generated_at: int, remote_seed: str, action_ids: list[str]) -> str:
-    action_ids_json = json.dumps(action_ids)
-    return f"""# auto-generated by generate_training_data.py - do not edit by hand
+def remote_data_designer_create_command(
+    remote_config: str,
+    num_records: int,
+    dataset_key: str,
+    artifact_path: str,
+) -> str:
+    command = [
+        MIMIR_DATA_DESIGNER_BIN,
+        "create",
+        remote_config,
+        "--num-records",
+        str(num_records),
+        "--dataset-name",
+        dataset_key,
+        "--artifact-path",
+        artifact_path,
+        "--output-format",
+        "jsonl",
+        "--no-tui",
+    ]
+    quoted_command = " ".join(shlex.quote(part) for part in command)
+    quoted_env = shlex.quote(MIMIR_DATA_DESIGNER_ENV)
+    return f"set -a; [ ! -f {quoted_env} ] || . {quoted_env}; set +a; {quoted_command}"
+
+
+def remote_bash_command(command: str) -> str:
+    return f"bash -lc {shlex.quote(command)}"
+
+
+def data_designer_config(
+    dataset_key: str,
+    generated_at: int,
+    remote_seed: str,
+    action_ids: list[str],
+    model_alias: str = "nvidia-text",
+) -> str:
+    action_ids_repr = repr(action_ids)
+    action_ids_json = json.dumps(action_ids, separators=(",", ":"))
+    return f'''# auto-generated by generate_training_data.py - do not edit by hand
 # generated_at: {generated_at}
 
-model_config_path: /mnt/storage/data-designer/model_configs.yaml
-model_providers_path: /mnt/storage/data-designer/model_providers.yaml
+from data_designer.config import DataDesignerConfigBuilder, LocalFileSeedSource
 
-dataset:
-  name: {dataset_key}
-  schema_profile: instruction
-  seed_file: {remote_seed}
 
-columns:
-  - name: input
-    type: seed_passthrough
+ACTION_IDS = {action_ids_repr}
 
-  - name: output
-    type: llm_text
-    model_alias: nvidia-text
-    prompt: |
-      You are generating training data for a local action dispatcher.
-      Given the natural language request below, produce valid JSON with
-      exactly two fields: "action_id" (string) and "payload" (object).
-      The action_id must be one of: {action_ids_json}.
-      Vary the phrasing of the input naturally but keep the output schema strict.
 
-      Request: {{{{input}}}}
-    output_schema:
-      type: object
-      required: [action_id, payload]
-      additionalProperties: false
-      properties:
-        action_id:
-          type: string
-          enum: {action_ids_json}
-        payload:
-          type: object
-"""
+def load_config_builder():
+    builder = DataDesignerConfigBuilder(
+        model_configs="/mnt/storage/data-designer/model_configs.yaml",
+    )
+    builder.with_seed_dataset(LocalFileSeedSource(path="{remote_seed}"))
+    builder.add_column(
+        name="expanded_record",
+        column_type="llm-structured",
+        model_alias="{model_alias}",
+        prompt="""You are generating additional instruction-tuning rows for a local action dispatcher.
+
+Given this seed natural-language request and target JSON, create one new varied request
+that maps to the same action_id. Keep the output schema strict.
+
+Allowed action_id values: {action_ids_json}
+
+Seed request:
+{{{{ input }}}}
+
+Seed output JSON:
+{{{{ output }}}}
+
+Return a complete instruction dataset row. The output field must be a JSON string with
+exactly action_id and payload. The metadata.action_id must match that output action_id.""",
+        output_format={{
+            "type": "object",
+            "required": ["instruction", "input", "output", "metadata"],
+            "additionalProperties": False,
+            "properties": {{
+                "instruction": {{"type": "string", "minLength": 1}},
+                "input": {{"type": "string", "minLength": 1}},
+                "output": {{"type": "string", "minLength": 1}},
+                "metadata": {{
+                    "type": "object",
+                    "required": ["action_id", "source", "split", "generated_at"],
+                    "additionalProperties": True,
+                    "properties": {{
+                        "action_id": {{"type": "string", "enum": ACTION_IDS}},
+                        "source": {{"type": "string"}},
+                        "split": {{"type": "string"}},
+                        "generated_at": {{"type": "integer"}},
+                    }},
+                }},
+            }},
+        }},
+    )
+    return builder
+'''
+
+
+def normalize_expanded_jsonl(raw_path: Path, output_path: Path) -> None:
+    normalized: list[dict[str, Any]] = []
+    for row in read_jsonl(raw_path):
+        candidate = row.get("expanded_record") if isinstance(row, dict) else None
+        if isinstance(candidate, dict):
+            normalized.append(candidate)
+        elif isinstance(row, dict) and {"instruction", "input", "output", "metadata"}.issubset(row):
+            normalized.append(row)
+    if not normalized:
+        raise ValueError(f"no expanded instruction rows found in {raw_path}")
+    write_jsonl(output_path, normalized)
 
 
 def has_sdg_args(args: argparse.Namespace) -> bool:
@@ -434,6 +516,7 @@ def patch_manifest(
     artifacts: dict[str, Any],
     expanded_path: Path | None,
     sdg_requested: bool,
+    sdg_model_alias: str | None,
     sdg_error: str | None = None,
 ) -> None:
     manifest = read_json(manifest_path)
@@ -458,7 +541,8 @@ def patch_manifest(
         "generator": "mp-ai-module-controller/scripts/generate_training_data.py",
         "sdg_engine": "data-designer" if sdg_requested else None,
         "sdg_host": "mimir (100.80.52.47)" if sdg_requested else None,
-        "sdg_model": "nvidia/nemotron-3-nano-30b-a3b" if sdg_requested else None,
+        "sdg_model": sdg_model_alias if sdg_requested else None,
+        "sdg_model_alias": sdg_model_alias if sdg_requested else None,
         "expanded_path": str(expanded_path) if expanded_path else None,
         "sdg_error": sdg_error,
         "codebook_version": codebook["dictionary_version"],