final commit

training/config.yaml

@@ -80,15 +80,16 @@ generation_do_sample: true
 
 # ---- Evaluation ----
 eval_interval: 50                     # Evaluate every N iterations
-eval_episodes: 3                      # Episodes per task during eval
-eval_max_steps: 60
+eval_episodes: 2                      # Episodes per task during eval
+eval_max_steps: 30                     # 30 steps × 2 episodes per task = 60 total
 
 # ---- Checkpointing ----
 checkpoint_interval: 5                # Save checkpoint every 5 iters (every ~15 min)
 save_total_limit: 5                   # Max checkpoints on Hub
 
-# ---- Metrics ----
-push_interval: 10                     # Push metrics to Hub every N iterations
+# ---- Metrics & Logging ----
+# Metrics, logs, eval results, and plots are pushed to hub_model_repo/<run_id>/
+# alongside checkpoints. No separate dataset repo needed.
 
 # ---- Output ----
 output_dir: "/workspace/antiatropos_checkpoints"  # Base dir — run_id subfolder auto-created
@@ -100,7 +101,8 @@ plot_dpi: 150                         # Resolution
 
 # ---- Hugging Face Hub ----
 hub_model_repo: ""                    # e.g. "pranavkk/antiatropos-qlora-qwen3.5-4b"
-hub_metrics_dataset: ""               # e.g. "pranavkk/antiatropos-training-metrics"
+# All metrics, logs, eval results, and plots are pushed to this repo under <run_id>/
+# along with checkpoints.
 push_to_hub: true                     # Push adapter + metrics after training
 
 # ---- Reproducibility ----

training/eval.py

@@ -5,7 +5,8 @@ Runs episodes with:
   1. The fine-tuned model (current LoRA adapter)
   2. The heuristic baseline
 
-Compares average rewards across tasks. Pushes results to Hub metrics dataset.
+Compares average rewards across tasks. Results are saved locally and pushed
+via push_run_files_to_hub (in train.py) under hub_model_repo/<run_id>/eval_results.json.
 """
 
 from __future__ import annotations
@@ -128,70 +129,12 @@ def evaluate(
     # Save eval results
     import os
     os.makedirs(output_dir, exist_ok=True)
-    with open(f"{output_dir}/eval_results.json", "w") as f:
+    eval_path = f"{output_dir}/eval_results.json"
+    with open(eval_path, "w") as f:
         json.dump(summary, f, indent=2)
+    print(f"  [eval] Saved results → {eval_path}")
 
     return summary
 
 
-def push_eval_results(
-    results: Dict[str, Any],
-    hub_dataset: str,
-    run_id: str,
-    iteration: int,
-) -> None:
-    """Push eval results as a row to the HF metrics dataset."""
-    if not hub_dataset:
-        return
-
-    row = {
-        "run_id": run_id,
-        "step": iteration,
-        "type": "eval",
-        **{f"eval_{k}": v for k, v in results.items() if not isinstance(v, dict)},
-    }
-    # Flatten per-task results
-    for task_id, task_results in results.get("per_task", {}).items():
-        for metric, value in task_results.items():
-            row[f"eval_{task_id}_{metric}"] = value
-
-    _append_to_dataset(row, hub_dataset)
-
 
-def _append_to_dataset(row: Dict[str, Any], hub_dataset: str) -> None:
-    """Append a row to a JSONL file on Hub (creates if not exists)."""
-    try:
-        from huggingface_hub import HfApi
-        api = HfApi()
-
-        # Download existing data or start fresh
-        import tempfile, os
-        tmp_dir = tempfile.mkdtemp()
-        jsonl_path = os.path.join(tmp_dir, "metrics.jsonl")
-
-        try:
-            api.hf_hub_download(
-                repo_id=hub_dataset,
-                filename="metrics.jsonl",
-                repo_type="dataset",
-                local_dir=tmp_dir,
-            )
-        except Exception:
-            pass  # File doesn't exist yet — that's fine
-
-        # Append row
-        with open(jsonl_path, "a") as f:
-            f.write(json.dumps(row) + "\n")
-
-        # Upload back
-        api.upload_file(
-            path_or_fileobj=jsonl_path,
-            path_in_repo="metrics.jsonl",
-            repo_id=hub_dataset,
-            repo_type="dataset",
-            commit_message=f"AntiAtropos metrics — {row.get('run_id', 'unknown')} step {row.get('step', '?')}",
-        )
-        print(f"[eval] Metrics pushed to {hub_dataset}")
-
-    except Exception as e:
-        print(f"[eval] Failed to push metrics: {e}")

training/launch_train.py

@@ -2,7 +2,7 @@
 """
 launch_train.py — Launch full AntiAtropos training on Hugging Face Jobs.
 
-Pushes model checkpoints, metrics dataset, and plots to HF Hub.
+Pushes model checkpoints, metrics, logs, and plots to HF Hub model repo.
 The local server is co-located for zero-latency environment interaction.
 Supports automatic resume from latest Hub checkpoint.
 
@@ -10,56 +10,38 @@ Prerequisites:
   1. pip install "huggingface_hub>=0.25.0"
   2. huggingface-cli login   (or set HF_TOKEN env var)
   3. HF Pro/Team account (required for GPU jobs)
-  4. The Hub model and dataset repos are auto-created if they don't exist.
-     Alternatively create them manually:
+  4. The Hub model repo is auto-created if it doesn't exist.
+     Alternatively create it manually:
        hf repo create <hub-model-repo> --type model
-       hf repo create <hub-metrics-dataset> --type dataset
 
 Lifecycle:
-  ┌─────────────────────────────────────────────┐
-  │  HF Job (A10G, ~4h)                        │
-  │  ┌──────────────────────────────────────┐   │
-  │  │  uvicorn :8000  ←──→  train.py       │   │
-  │  │  (simulator)          (GPU model)     │   │
-  │  └──────────┬───────────────────────────┘   │
-  │             │ push adapter + plots           │
-  │             ▼                                │
-  │      HF Hub Model Repo                       │
-  │      (checkpoint-25, checkpoint-50, ...)     │
-  │             │ push metrics.jsonl              │
-  │             ▼                                │
-  │      HF Hub Metrics Dataset                  │
-  └─────────────────────────────────────────────┘
+  All run artifacts (checkpoints, metrics, logs, eval results, plots)
+  are pushed to <hub-model-repo>/<run_id>/ on the Hub.
 
 Usage:
-  # Quick test (∼10 min):
+  # Quick test (~10 min):
   python training/launch_train.py \
     --hub-model-repo Keshav051/antiatropos-qlora \
-    --hub-metrics-dataset Keshav051/antiatropos-training-metrics \
     --num-iterations 20 --num-episodes 4
 
-  # Full training (a10g-large ≈ $0.34/hr, ∼2h):
+  # Full training (a10g-large ~ $0.34/hr, ~2h):
   python training/launch_train.py \
-    --hub-model-repo Keshav051/antiatropos-qlora \
-    --hub-metrics-dataset Keshav051/antiatropos-training-metrics
+    --hub-model-repo Keshav051/antiatropos-qlora
 
   # Custom flavor / longer timeout:
   python training/launch_train.py \
     --hub-model-repo Keshav051/antiatropos-qlora \
-    --hub-metrics-dataset Keshav051/antiatropos-training-metrics \
     --flavor a10g-xlarge --timeout 12h \
     --num-iterations 2000 --num-episodes 24
 
   # Resume from latest Hub checkpoint:
   python training/launch_train.py \
     --hub-model-repo Keshav051/antiatropos-qlora \
-    --hub-metrics-dataset Keshav051/antiatropos-training-metrics \
     --run-id exp_002
 
   # Dry run (prints job command without launching):
   python training/launch_train.py \
     --hub-model-repo Keshav051/antiatropos-qlora \
-    --hub-metrics-dataset Keshav051/antiatropos-training-metrics \
     --dry-run
 """
 
@@ -76,12 +58,12 @@ TRAINING_DIR = Path(__file__).resolve().parent
 
 DOCKER_IMAGE = "pytorch/pytorch:2.10.0-cuda12.6-cudnn9-devel"
 
-DEFAULT_NUM_ITERATIONS = 500
-DEFAULT_NUM_EPISODES = 4
+DEFAULT_NUM_ITERATIONS = 15
+DEFAULT_NUM_EPISODES = 6
 DEFAULT_MAX_STEPS = 20
 DEFAULT_EVAL_INTERVAL = 50
-DEFAULT_CHECKPOINT_INTERVAL = 25
-DEFAULT_PLOT_INTERVAL = 25
+DEFAULT_CHECKPOINT_INTERVAL = 5
+DEFAULT_PLOT_INTERVAL = 10
 
 
 def build_job_command() -> str:
@@ -123,7 +105,6 @@ def build_job_command() -> str:
         "echo '[bootstrap] Launching training (local server, Hub persistence)...'\n"
         "export PYTORCH_ALLOC_CONF='expandable_segments:True'  # required by Qwen3.5 to avoid OOM fragmentation\n"
         "ANTIATROPOS_HUB_MODEL_REPO=$HUB_MODEL_REPO "
-        "ANTIATROPOS_HUB_METRICS_DATASET=$HUB_METRICS_DATASET "
         "ANTIATROPOS_ENV_URL=http://localhost:8000 "
         "python training/train.py "
         "--run-id $RUN_ID "
@@ -146,48 +127,38 @@ def build_job_command() -> str:
 
 def ensure_hub_repos(
     hub_model_repo: str,
-    hub_metrics_dataset: str,
     hf_token: Optional[str],
 ) -> None:
-    """Check if Hub repos exist; create them automatically if not."""
+    """Check if the Hub model repo exists; create it automatically if not."""
     if not hf_token:
         print("  [hub] No HF_TOKEN available, skipping repo check")
         return
 
+    if not hub_model_repo:
+        return
+
     try:
         from huggingface_hub import HfApi
 
         api = HfApi()
 
-        for repo_id, repo_type in [
-            (hub_model_repo, "model"),
-            (hub_metrics_dataset, "dataset"),
-        ]:
-            base_url = (
-                "https://huggingface.co"
-                if repo_type == "model"
-                else "https://huggingface.co/datasets"
+        try:
+            info = api.repo_info(repo_id=hub_model_repo, repo_type="model")
+            print(f"  [hub] Repo OK: https://huggingface.co/{hub_model_repo}")
+        except Exception:
+            print(f"  [hub] Creating repo: {hub_model_repo} (model)...")
+            api.create_repo(
+                repo_id=hub_model_repo,
+                repo_type="model",
+                private=True,
+                exist_ok=True,
             )
-            try:
-                info = api.repo_info(repo_id=repo_id, repo_type=repo_type)
-                print(f"  [hub] Repo OK: {base_url}/{repo_id}")
-            except Exception:
-                print(f"  [hub] Creating repo: {repo_id} ({repo_type})...")
-                api.create_repo(
-                    repo_id=repo_id,
-                    repo_type=repo_type,
-                    private=True,
-                    exist_ok=True,
-                )
-                print(f"  [hub] Created: {base_url}/{repo_id}")
+            print(f"  [hub] Created: https://huggingface.co/{hub_model_repo}")
     except Exception as e:
-        print(f"\n  [hub] WARNING: Could not verify/create Hub repos: {e}")
-        print("  [hub] Create them manually:")
+        print(f"\n  [hub] WARNING: Could not verify/create Hub repo: {e}")
+        print("  [hub] Create it manually:")
         print(f"    hf repo create {hub_model_repo} --type model")
-        print(f"    hf repo create {hub_metrics_dataset} --type dataset")
-        print(f"    Then visit:")
-        print(f"      https://huggingface.co/{hub_model_repo}")
-        print(f"      https://huggingface.co/datasets/{hub_metrics_dataset}")
+        print(f"    Then visit: https://huggingface.co/{hub_model_repo}")
 
 
 def main() -> None:
@@ -212,14 +183,8 @@ def main() -> None:
     parser.add_argument(
         "--hub-model-repo",
         required=True,
-        help="HF Hub model repo for checkpoints, final adapter, and plots "
-        "(e.g. Keshav051/antiatropos-qlora)",
-    )
-    parser.add_argument(
-        "--hub-metrics-dataset",
-        required=True,
-        help="HF Hub dataset repo for training metrics.jsonl "
-        "(e.g. Keshav051/antiatropos-training-metrics)",
+        help="HF Hub model repo for checkpoints, metrics, logs, and plots "
+        "(e.g. Keshav051/antiatropos-qlora). All run artifacts go under <run_id>/.",
     )
     parser.add_argument(
         "--run-id",
@@ -300,7 +265,6 @@ def main() -> None:
     print(f"  Timeout:             {args.timeout}")
     print(f"  Code repo:           {args.repo}")
     print(f"  Hub model repo:      {args.hub_model_repo}")
-    print(f"  Hub metrics dataset: {args.hub_metrics_dataset}")
     print(f"  Run ID:              {run_id}")
     print(f"  Loss type:           {args.loss_type}")
     if args.loss_type == "grpo":
@@ -361,7 +325,7 @@ def main() -> None:
     # ---- Ensure Hub repos exist ----
     if not args.no_create_repos and hf_token:
         ensure_hub_repos(
-            args.hub_model_repo, args.hub_metrics_dataset, hf_token
+            args.hub_model_repo, hf_token
         )
 
     # ---- Launch via run_job ----
@@ -385,7 +349,6 @@ def main() -> None:
             "REPO": args.repo,
             "RUN_ID": run_id,
             "HUB_MODEL_REPO": args.hub_model_repo,
-            "HUB_METRICS_DATASET": args.hub_metrics_dataset,
             "NUM_ITERATIONS": str(args.num_iterations),
             "NUM_EPISODES": str(args.num_episodes),
             "MAX_STEPS": str(args.max_steps),

training/plotting.py

@@ -520,6 +520,7 @@ def push_plots_to_hub(
     plot_paths: List[str],
     hub_repo: str,
     iteration: int,
+    run_id: str = "",
 ) -> None:
     if not hub_repo or not plot_paths:
         return
@@ -528,14 +529,15 @@ def push_plots_to_hub(
         api = HfApi()
         for path in plot_paths:
             filename = Path(path).name
+            prefix = f"{run_id}/" if run_id else ""
             api.upload_file(
                 path_or_fileobj=path,
-                path_in_repo=f"plots/iter_{iteration}/{filename}",
+                path_in_repo=f"{prefix}plots/{filename}",
                 repo_id=hub_repo,
                 repo_type="model",
-                commit_message=f"Training plots - iteration {iteration}",
+                commit_message=f"[{run_id}] Training plots - iteration {iteration}",
             )
-        print(f"[plotting] Pushed {len(plot_paths)} plots to {hub_repo}")
+        print(f"[plotting] Pushed {len(plot_paths)} plots to {hub_repo}/{prefix}plots/")
     except Exception as e:
         print(f"[plotting] Push failed: {e}")
 

training/train.py

@@ -66,7 +66,7 @@ from openenv_loop import (
     rollout_episode,
     rollout_heuristic_episode,
 )
-from eval import evaluate, push_eval_results
+from eval import evaluate
 from plotting import (
     generate_all_plots,
     push_plots_to_hub,
@@ -516,63 +516,22 @@ def grpo_loss_fn(
 
 
 # ────────────────────────────────────────────────────────────
-# Metrics Push
+# Run Files Push (to hub_model_repo/<run_id>/)
 # ────────────────────────────────────────────────────────────
 
-def push_train_metrics(
-    metrics: Dict[str, Any],
-    hub_dataset: str,
-) -> None:
-    """Push training metrics row to the Hub dataset."""
-    if not hub_dataset:
-        return
-
-    try:
-        from huggingface_hub import HfApi
-        api = HfApi()
-        import tempfile
-
-        tmp_dir = tempfile.mkdtemp()
-        jsonl_path = os.path.join(tmp_dir, "metrics.jsonl")
-
-        try:
-            api.hf_hub_download(
-                repo_id=hub_dataset,
-                filename="metrics.jsonl",
-                repo_type="dataset",
-                local_dir=tmp_dir,
-            )
-        except Exception:
-            pass
-
-        with open(jsonl_path, "a") as f:
-            f.write(json.dumps(metrics) + "\n")
-
-        api.upload_file(
-            path_or_fileobj=jsonl_path,
-            path_in_repo="metrics.jsonl",
-            repo_id=hub_dataset,
-            repo_type="dataset",
-            commit_message=f"train metrics — {metrics.get('run_id')} iter {metrics.get('iteration')}",
-        )
-    except Exception as e:
-        print(f"[train] Metrics push failed: {e}")
-
 
 def push_run_files_to_hub(
     run_id: str,
     output_dir: Path,
-    hub_dataset: str,
+    hub_model_repo: str,
     iteration: int,
 ) -> None:
-    """Upload step_metrics.jsonl, iter_metrics.jsonl, and training.log to HF Hub.
+    """Upload step_metrics.jsonl, iter_metrics.jsonl, training.log, and eval results.
 
-    Files are uploaded under <run_id>/ prefix in the metrics dataset so each
-    run's data is isolated and you can diff runs on the Hub viewer.
-
-    Called every push_interval iterations and at the end of training.
+    Files are uploaded under <run_id>/ in the model repo alongside checkpoints.
+    Called every checkpoint_interval iterations and at the end of training.
     """
-    if not hub_dataset:
+    if not hub_model_repo:
         return
 
     files_to_push = [
@@ -582,6 +541,14 @@ def push_run_files_to_hub(
         ("run_info.json",       f"{run_id}/run_info.json"),
     ]
 
+    # Also push eval results if they exist
+    eval_path = output_dir / "eval" / "eval_results.json"
+    if eval_path.exists():
+        files_to_push.append(("eval/eval_results.json", f"{run_id}/eval_results.json"))
+    final_eval_path = output_dir / "final_eval" / "eval_results.json"
+    if final_eval_path.exists():
+        files_to_push.append(("final_eval/eval_results.json", f"{run_id}/final_eval_results.json"))
+
     try:
         from huggingface_hub import HfApi
         api = HfApi()
@@ -594,15 +561,15 @@ def push_run_files_to_hub(
                 api.upload_file(
                     path_or_fileobj=str(local_path),
                     path_in_repo=hub_path,
-                    repo_id=hub_dataset,
-                    repo_type="dataset",
+                    repo_id=hub_model_repo,
+                    repo_type="model",
                     commit_message=f"[{run_id}] iter {iteration}: {local_name}",
                 )
                 pushed.append(local_name)
             except Exception as e:
                 print(f"  [push] Failed to push {local_name}: {e}")
         if pushed:
-            print(f"  [push] → HF dataset {hub_dataset}/{run_id}/: {', '.join(pushed)}",
+            print(f"  [push] \u2192 HF model {hub_model_repo}/{run_id}/: {', '.join(pushed)}",
                   flush=True)
     except Exception as e:
         print(f"[train] Hub file push failed: {e}")
@@ -802,7 +769,6 @@ def train(cfg: Dict[str, Any]) -> None:
     print(f"[train] Run manifest:  {run_info_path}")
 
     hub_model_repo = cfg.get("hub_model_repo", "")
-    hub_metrics_dataset = cfg.get("hub_metrics_dataset", "")
     push_to_hub_flag = cfg.get("push_to_hub", True)
 
     # ---- Verify environment ----
@@ -863,7 +829,6 @@ def train(cfg: Dict[str, Any]) -> None:
     max_grad_norm = cfg.get("max_grad_norm", 1.0)
     checkpoint_interval = cfg.get("checkpoint_interval", 10)  # default: every 10 iters
     eval_interval = cfg.get("eval_interval", 50)
-    push_interval = cfg.get("push_interval", 10)
     plot_interval = cfg.get("plot_interval", 25)
 
     # ---- Training loop ----
@@ -878,7 +843,6 @@ def train(cfg: Dict[str, Any]) -> None:
     print(f"  Max steps:     {max_steps}")
     print(f"  Learning rate: {lr}")
     print(f"  Hub model:     {hub_model_repo or '(not configured)'}")
-    print(f"  Hub metrics:   {hub_metrics_dataset or '(not configured)'}")
     print(f"  Output dir:    {output_dir}")
     print(f"{'='*70}\n")
 
@@ -1017,12 +981,7 @@ def train(cfg: Dict[str, Any]) -> None:
         if len(recent_episodes_data) > 200:  # Keep last ~200 episodes
             recent_episodes_data = recent_episodes_data[-200:]
 
-        # ---- Push metrics + run files to Hub ----
-        if (iteration + 1) % push_interval == 0 and hub_metrics_dataset:
-            # Push step_metrics.jsonl, iter_metrics.jsonl, training.log, run_info.json
-            push_run_files_to_hub(run_id, output_dir, hub_metrics_dataset, iteration + 1)
-
-        # ---- Checkpoint ----
+        # ---- Checkpoint + push run files ----
         if (iteration + 1) % checkpoint_interval == 0:
             # Pad iteration number so ls sorts correctly: checkpoint-0010, checkpoint-0050, ...
             ckpt_name = f"checkpoint-{iteration + 1:04d}"
@@ -1041,7 +1000,7 @@ def train(cfg: Dict[str, Any]) -> None:
             (ckpt_dir / "checkpoint_meta.json").write_text(
                 _json.dumps(ckpt_meta, indent=2)
             )
-            print(f"  [ckpt] Saved → {ckpt_dir}  "
+            print(f"  [ckpt] Saved \u2192 {ckpt_dir}  "
                   f"(reward={avg_reward:.4f}  loss={loss.item():.4f})", flush=True)
             if push_to_hub_flag and hub_model_repo:
                 push_to_hub(
@@ -1050,6 +1009,8 @@ def train(cfg: Dict[str, Any]) -> None:
                     commit_message=f"[{run_id}] {ckpt_name}",
                     path_in_repo=f"{run_id}/{ckpt_name}",
                 )
+                # Push run files (metrics, logs) alongside checkpoint
+                push_run_files_to_hub(run_id, output_dir, hub_model_repo, iteration + 1)
 
         # ---- Evaluation ----
         if (iteration + 1) % eval_interval == 0:
@@ -1071,10 +1032,6 @@ def train(cfg: Dict[str, Any]) -> None:
                     eval_row[f"eval_{tid}_{mk}"] = mv
             eval_metrics_history.append(eval_row)
 
-            if hub_metrics_dataset:
-                push_eval_results(
-                    eval_results, hub_metrics_dataset, run_id, iteration
-                )
             # Re-enable training mode
             model.train()
 
@@ -1089,7 +1046,7 @@ def train(cfg: Dict[str, Any]) -> None:
                     cfg=cfg,
                 )
                 if push_to_hub_flag and hub_model_repo:
-                    push_plots_to_hub(plot_paths, hub_model_repo, iteration)
+                    push_plots_to_hub(plot_paths, hub_model_repo, iteration, run_id=run_id)
             except Exception as e:
                 print(f"  [iter {iteration}] Plotting failed: {e}")
 
@@ -1123,12 +1080,6 @@ def train(cfg: Dict[str, Any]) -> None:
         client, model, tokenizer, cfg,
         output_dir=str(output_dir / "final_eval"),
     )
-
-    if hub_metrics_dataset:
-        push_eval_results(
-            final_eval, hub_metrics_dataset, run_id, num_iterations
-        )
-
     # Final plots (full training history)
     try:
         final_eval_row = {
@@ -1152,19 +1103,17 @@ def train(cfg: Dict[str, Any]) -> None:
             cfg=cfg,
         )
         if push_to_hub_flag and hub_model_repo:
-            push_plots_to_hub(plot_paths, hub_model_repo, num_iterations)
+            push_plots_to_hub(plot_paths, hub_model_repo, num_iterations, run_id=run_id)
     except Exception as e:
         print(f"[train] Final plotting failed: {e}")
 
     print(f"\n[train] All done. Final adapter: {final_dir}")
     if hub_model_repo:
         print(f"[train] Hub repo: https://huggingface.co/{hub_model_repo}")
-    if hub_metrics_dataset:
-        print(f"[train] Metrics:  https://huggingface.co/datasets/{hub_metrics_dataset}/tree/main/{run_id}")
-
-    # ── Final push of all run files (captures last iterations even if push_interval missed them)
-    if hub_metrics_dataset:
-        push_run_files_to_hub(run_id, output_dir, hub_metrics_dataset, num_iterations)
+    
+    # \u2500\u2500 Final push of all run files
+    if hub_model_repo:
+        push_run_files_to_hub(run_id, output_dir, hub_model_repo, num_iterations)
 
     # ── Flush and close the TeeLogger ────────────────────────────────────────
     # Restore original stdout/stderr so any code after train() works normally.
@@ -1265,7 +1214,6 @@ def main():
     if args.no_push:
         cfg["push_to_hub"] = False
         cfg["hub_model_repo"] = ""
-        cfg["hub_metrics_dataset"] = ""
 
     # Allow HF_TOKEN from env
     hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")