Upload run_self_healing_job.py

run_self_healing_job.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""
+Production Job Runner: Self-Healing Training
+============================================
+Complete self-healing training job with Trackio monitoring.
+
+Pre-flight:
+  ✓ Dataset: trl-lib/Capybara (messages format, SFT-compatible)
+  ✓ Model: Qwen2.5-0.5B
+  ✓ Trackio monitoring
+  ✓ Timeout: 2h for 0.5B model on a10g-large
+"""
+import os, sys, json, time, math, gc
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from self_healing import SelfHealingTrainer, HealingConfig
+
+def main():
+    print("\n" + "=" * 70)
+    print("  SELF-HEALING TRAINING JOB")
+    print("=" * 70 + "\n")
+    
+    model_id = "Qwen/Qwen2.5-0.5B"
+    print(f"[1/4] Loading {model_id}")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, torch_dtype=torch.bfloat16, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
+    
+    print("[2/4] Loading trl-lib/Capybara")
+    dataset = load_dataset("trl-lib/Capybara", split="train[:5000]")
+    print(f"  {len(dataset)} samples")
+    
+    from trl import SFTConfig, SFTTrainer
+    training_args = SFTConfig(
+        output_dir="./output",
+        per_device_train_batch_size=4,
+        gradient_accumulation_steps=4,
+        learning_rate=2e-5,
+        max_steps=500,
+        logging_steps=10,
+        logging_strategy="steps",
+        logging_first_step=True,
+        save_steps=200, save_total_limit=3,
+        bf16=True, gradient_checkpointing=True,
+        warmup_steps=50, lr_scheduler_type="cosine",
+        report_to="trackio",
+        run_name="selfheal-qwen0.5b-sft",
+        project="self-healing-system",
+        trackio_space_id=os.environ.get("TRACKIO_SPACE_ID", ""),
+        push_to_hub=False, disable_tqdm=True,
+    )
+    
+    trainer = SFTTrainer(
+        model=model, args=training_args,
+        train_dataset=dataset, tokenizer=tokenizer)
+    
+    print("[3/4] Configuring self-healing")
+    healing_config = HealingConfig(
+        nan_patience=3, loss_spike_factor=5.0, divergence_patience=50,
+        grad_explosion_threshold=100.0, zclip_enabled=True,
+        zclip_z_threshold=3.0, max_recovery_attempts=5,
+        max_lr_reductions=3, max_batch_reductions=2,
+        postmortem_path="./postmortem.json",
+    )
+    sh_trainer = SelfHealingTrainer(trainer, healing_config)
+    
+    print("[4/4] Dry-run validation")
+    try:
+        sh_trainer.dry_run(num_steps=2)
+        print("  ✓ Dry-run passed!\n")
+    except Exception as e:
+        print(f"  ✗ Dry-run FAILED: {e}")
+        sys.exit(1)
+    
+    print("Training with autonomous self-healing...\n")
+    start = time.time()
+    result = sh_trainer.train()
+    elapsed = time.time() - start
+    
+    print("\n" + "=" * 70)
+    print("  TRAINING COMPLETE")
+    print("=" * 70)
+    report = sh_trainer.get_report()
+    print(f"  Converged: {sh_trainer.converged} | Attempts: {sh_trainer.attempt}")
+    print(f"  Recoveries: {report['total_recoveries']} | ZClip: {report['zclip_total_clips']}")
+    print(f"  NaN: {report['nan_count']} | LR cuts: {report['lr_reductions']}")
+    print(f"  Elapsed: {elapsed:.0f}s ({elapsed/60:.1f}min)")
+    
+    if report["recovery_history"]:
+        print(f"\n  Recovery log ({len(report['recovery_history'])}):")
+        for i, rec in enumerate(report["recovery_history"]):
+            print(f"    [{i+1}] {rec['failure']}: {rec['actions']}")
+    
+    if os.path.exists(healing_config.postmortem_path):
+        with open(healing_config.postmortem_path) as f:
+            pm = json.load(f)
+        print(f"\n  Postmortem: {pm.get('exit_reason')} at step {pm.get('last_step')}")
+    
+    space = os.environ.get("TRACKIO_SPACE_ID", "mlintern/selfheal-demo")
+    print(f"\n  Dashboard: https://huggingface.co/spaces/{space}")
+
+if __name__ == "__main__":
+    main()