feat: comprehensive W&B logging + retry=0 hardening

Changes:
- evaluation.py: Add pre_rank (rank at 1 trace) and max_rank (worst rank)
- mtl_trainer.py: Log per-byte ranks to W&B (ranks/byte_X/{pre,final,min,max,@500,@1000})
- trainer.py: Add GradientNormLogger to single-byte trainer, add pre_rank/max_rank
- schemas.py: Change JobCreate max_retries default from 2 to 0
- database.py: Change DB max_retries default from 2 to 0
- tq.py: Change CLI max_retries defaults from 2 to 0
- agent.py: Change worker max_retries defaults from 3/2 to 0

All changes are additive - backward compatible with existing training runs.

orchestrator/cli/tq.py

@@ -314,7 +314,7 @@ def add(config_file: str):
         "name": config.get("name", os.path.basename(config_file)),
         "config": config.get("config", config),
         "priority": config.get("priority", 0),
-        "max_retries": config.get("max_retries", 2),
+        "max_retries": config.get("max_retries", 0),
         "tags": config.get("tags", []),
     }
 
@@ -344,7 +344,7 @@ def batch(jobs_file: str):
                 "name": j.get("name", f"job-{i}"),
                 "config": j.get("config", j),
                 "priority": j.get("priority", 0),
-                "max_retries": j.get("max_retries", 2),
+                "max_retries": j.get("max_retries", 0),
                 "tags": j.get("tags", []),
             }
             for i, j in enumerate(jobs_list)

orchestrator/server/database.py

@@ -61,7 +61,7 @@ def init_db() -> None:
                 completed_at TEXT,
                 result TEXT,
                 retry_count INTEGER NOT NULL DEFAULT 0,
-                max_retries INTEGER NOT NULL DEFAULT 2,
+                max_retries INTEGER NOT NULL DEFAULT 0,
                 tags TEXT NOT NULL DEFAULT '[]'
             );
 

orchestrator/server/schemas.py

@@ -16,7 +16,7 @@ class JobCreate(BaseModel):
     name: str = Field(..., description="Human-readable job name")
     config: Dict[str, Any] = Field(..., description="Training configuration")
     priority: int = Field(default=0, description="Job priority (higher = first)")
-    max_retries: int = Field(default=2, description="Max training retries")
+    max_retries: int = Field(default=0, description="Max training retries (0 = no retries)")
     tags: List[str] = Field(default_factory=list, description="Tags for filtering")
 
 

orchestrator/worker/agent.py

@@ -829,7 +829,7 @@ class WorkerAgent:
                 "--desync", str(config.get("desync", 0)),
                 "--variant", str(variant or model_type or "hps"),
                 "--seed", str(config.get("seed", 42)),
-                "--max-retries", str(config.get("max_retries", 3)),
+                "--max-retries", str(config.get("max_retries", 0)),
                 "--data-dir", self.data_dir,
                 "--output-dir", f"/root/jobs/{job_id[:8]}",
             ]
@@ -911,7 +911,7 @@ class WorkerAgent:
                 "--byte", str(config.get("target_byte", 0)),
                 "--desync", str(config.get("desync", 0)),
                 "--seed", str(config.get("seed", 42)),
-                "--max-retries", str(config.get("max_retries", 2)),
+                "--max-retries", str(config.get("max_retries", 0)),
                 "--data-dir", self.data_dir,
                 "--output-dir", f"/root/jobs/{job_id[:8]}",
             ]

src/evaluation.py

@@ -178,7 +178,8 @@ def evaluate_model(
 
     Returns:
         Dictionary with evaluation results:
-            'final_rank', 'ranks', 'min_rank', 'rank_at_500', 'rank_at_1000'.
+            'final_rank', 'ranks', 'pre_rank', 'min_rank', 'max_rank',
+            'rank_at_500', 'rank_at_1000'.
     """
     if cached_predictions is not None:
         raw_predictions = cached_predictions
@@ -238,19 +239,25 @@ def evaluate_model(
     )
 
     min_rank = int(np.min(ranks_array[:, 1])) if len(ranks_array) > 0 else 256
+    max_rank = int(np.max(ranks_array[:, 1])) if len(ranks_array) > 0 else 256
+    pre_rank = int(ranks_array[0, 1]) if len(ranks_array) > 0 else 256
     rank_at_500 = _get_rank_at_n(ranks_array, 500)
     rank_at_1000 = _get_rank_at_n(ranks_array, 1000)
 
     result = {
         "final_rank": final_rank,
         "ranks": ranks_array,
+        "pre_rank": pre_rank,
         "min_rank": min_rank,
+        "max_rank": max_rank,
         "rank_at_500": rank_at_500,
         "rank_at_1000": rank_at_1000,
     }
 
     logger.info(
-        "Byte %d: final_rank=%d, min_rank=%d, rank@500=%d, rank@1000=%d",
-        target_byte, final_rank, min_rank, rank_at_500, rank_at_1000,
+        "Byte %d: pre_rank=%d, final_rank=%d, min_rank=%d, max_rank=%d, "
+        "rank@500=%d, rank@1000=%d",
+        target_byte, pre_rank, final_rank, min_rank, max_rank,
+        rank_at_500, rank_at_1000,
     )
     return result

src/training/mtl_trainer.py

@@ -967,9 +967,12 @@ class MTLTrainer:
                     "all_final_ranks": eval_results["all_final_ranks"],
                     "per_byte_results": {
                         str(k): {
+                            "pre_rank": v.get("pre_rank", 256),
                             "final_rank": v["final_rank"],
                             "min_rank": v["min_rank"],
+                            "max_rank": v.get("max_rank", 256),
                             "rank_at_500": v.get("rank_at_500", -1),
+                            "rank_at_1000": v.get("rank_at_1000", -1),
                         }
                         for k, v in eval_results["byte_results"].items()
                     },
@@ -1005,13 +1008,22 @@ class MTLTrainer:
                 if self.wandb_project:
                     try:
                         import wandb
-                        wandb.log({
+                        rank_metrics = {
                             "max_final_rank": eval_results["max_final_rank"],
                             "mean_final_rank": eval_results["mean_final_rank"],
                             "num_rank0": eval_results["num_rank0"],
                             "final_train_loss": result["final_train_loss"],
                             "final_val_loss": result["final_val_loss"],
-                        })
+                        }
+                        # Log per-byte rank metrics
+                        for byte_idx, byte_res in eval_results["byte_results"].items():
+                            rank_metrics[f"ranks/byte_{byte_idx}_pre"] = byte_res.get("pre_rank", 256)
+                            rank_metrics[f"ranks/byte_{byte_idx}_final"] = byte_res["final_rank"]
+                            rank_metrics[f"ranks/byte_{byte_idx}_min"] = byte_res["min_rank"]
+                            rank_metrics[f"ranks/byte_{byte_idx}_max"] = byte_res.get("max_rank", 256)
+                            rank_metrics[f"ranks/byte_{byte_idx}_at500"] = byte_res.get("rank_at_500", -1)
+                            rank_metrics[f"ranks/byte_{byte_idx}_at1000"] = byte_res.get("rank_at_1000", -1)
+                        wandb.log(rank_metrics)
                         wandb.finish()
                     except Exception:
                         pass

src/training/trainer.py

@@ -173,6 +173,21 @@ class Trainer:
                     reinit=True,
                 )
                 callbacks.append(WandbMetricsLogger())
+
+                # Gradient norm logging to W&B
+                from ..gradient_logger import GradientNormLogger
+                grad_logger = GradientNormLogger(
+                    val_data=(
+                        data["atk_traces_reshaped"],
+                        data["atk_labels"],
+                    ),
+                    log_every_n_epochs=1,
+                    batch_size=128,
+                )
+                callbacks.append(grad_logger)
+                logger.info(
+                    "Gradient norm logging ENABLED for single-byte trainer."
+                )
             except ImportError:
                 logger.warning("wandb not installed; skipping logging.")
 
@@ -216,8 +231,10 @@ class Trainer:
             "final_train_accuracy": float(history.history["accuracy"][-1]),
             "final_val_loss": float(history.history["val_loss"][-1]),
             "final_val_accuracy": float(history.history["val_accuracy"][-1]),
+            "post_train_pre_rank": eval_result.get("pre_rank", 256),
             "post_train_final_rank": eval_result["final_rank"],
             "post_train_min_rank": eval_result["min_rank"],
+            "post_train_max_rank": eval_result.get("max_rank", 256),
             "post_train_rank_at_500": eval_result["rank_at_500"],
             "post_train_rank_at_1000": eval_result["rank_at_1000"],
         }
@@ -227,10 +244,12 @@ class Trainer:
             try:
                 import wandb
                 wandb.log({
-                    "final_rank": eval_result["final_rank"],
-                    "min_rank": eval_result["min_rank"],
-                    "rank_at_500": eval_result["rank_at_500"],
-                    "rank_at_1000": eval_result["rank_at_1000"],
+                    "ranks/pre_rank": eval_result.get("pre_rank", 256),
+                    "ranks/final_rank": eval_result["final_rank"],
+                    "ranks/min_rank": eval_result["min_rank"],
+                    "ranks/max_rank": eval_result.get("max_rank", 256),
+                    "ranks/rank_at_500": eval_result["rank_at_500"],
+                    "ranks/rank_at_1000": eval_result["rank_at_1000"],
                 })
                 wandb.finish()
             except Exception: