1# /// script
2# dependencies = [
3#   "openai>=1.0",
4#   "azure-identity",
5#   "azure-ai-projects",
6# ]
7# ///
8"""
9monitor_training.py — Monitor a fine-tuning job until completion.
10 
11Polls the job status and streams training events (reward, loss, errors)
12in real time. Exits when the job reaches a terminal state.
13 
14Usage:
15  python monitor_training.py --job-id ftjob-abc123
16  python monitor_training.py --base-url https://<resource>.services.ai.azure.com/api/projects/<project>/openai/v1/ --api-key KEY --job-id ftjob-abc123
17  python monitor_training.py --job-id ftjob-abc123 --poll-interval 30
18"""
19 
20import argparse
21import os
22import sys
23 
24try:
25    sys.stdout.reconfigure(encoding="utf-8")
26    sys.stderr.reconfigure(encoding="utf-8")
27except (AttributeError, OSError):
28    pass  # Stream not reconfigurable (older Python or non-tty); default encoding is fine
29import time
30 
31sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
32from common import HelpOnErrorParser, get_clients
33 
34TERMINAL_STATUSES = {"succeeded", "failed", "cancelled"}
35 
36 
37def monitor_job(client, job_id, poll_interval=15):
38    """Poll a fine-tuning job until it reaches a terminal state."""
39    # Cap memory for long-running jobs (RFT can run hours/days, accumulating thousands of events)
40    seen_events = set()
41    MAX_SEEN_EVENTS = 5000
42 
43    print(f"Monitoring job: {job_id}")
44    print(f"Polling every {poll_interval}s. Ctrl+C to stop.\n")
45 
46    while True:
47        try:
48            job = client.fine_tuning.jobs.retrieve(job_id)
49        except Exception as e:
50            print(f"⚠️ Error retrieving job: {e}")
51            time.sleep(poll_interval)
52            continue
53 
54        status = (job.status or "").lower()
55 
56        # Fetch and display new events
57        try:
58            events = list(client.fine_tuning.jobs.list_events(job_id, limit=20))
59            for event in reversed(events):
60                event_key = (event.created_at, event.message)
61                if event_key not in seen_events:
62                    if len(seen_events) >= MAX_SEEN_EVENTS:
63                        # Keep only the most recent half — a fully-flushed dedup window
64                        # would risk re-printing old events on transient API hiccups, but
65                        # without trimming this set grows unbounded for long RFT runs.
66                        seen_events = set(list(seen_events)[-(MAX_SEEN_EVENTS // 2):])
67                    seen_events.add(event_key)
68                    ts = time.strftime("%H:%M:%S", time.localtime(event.created_at))
69                    level = event.level or "info"
70 
71                    # Highlight step events
72                    if "Step" in event.message and "reward" in event.message:
73                        print(f"  📈 [{ts}] {event.message}")
74                    elif "Step" in event.message and "loss" in event.message:
75                        print(f"  📉 [{ts}] {event.message}")
76                    elif "error" in event.message.lower() or level == "error":
77                        print(f"  ❌ [{ts}] {event.message}")
78                    elif "started" in event.message.lower() or "completed" in event.message.lower():
79                        print(f"  🔔 [{ts}] {event.message}")
80                    else:
81                        print(f"  ℹ️ [{ts}] {event.message}")
82        except Exception:
83            pass  # Events API may not be available for all job states
84 
85        # Check terminal state
86        if status in TERMINAL_STATUSES:
87            print(f"\n{'='*50}")
88            if status == "succeeded":
89                model = job.fine_tuned_model or "unknown"
90                print(f"  ✅ Job succeeded!")
91                print(f"  Fine-tuned model: {model}")
92                if job.trained_tokens:
93                    print(f"  Trained tokens: {job.trained_tokens:,}")
94            elif status == "failed":
95                print(f"  ❌ Job failed.")
96                if hasattr(job, "error") and job.error:
97                    print(f"  Error: {job.error}")
98            elif status == "cancelled":
99                print(f"  ⚠️ Job was cancelled.")
100            print(f"{'='*50}")
101            return status
102 
103        time.sleep(poll_interval)
104 
105 
106def build_parser():
107    parser = HelpOnErrorParser(
108        description="Monitor a fine-tuning job until completion",
109        epilog=(
110            "Example:\n"
111            "  python monitor_training.py --job-id ftjob-abc123\n"
112            "  python monitor_training.py --base-url https://<resource>.services.ai.azure.com/api/projects/<project>/openai/v1/ --api-key KEY --job-id ftjob-abc123"
113        ),
114        formatter_class=argparse.RawTextHelpFormatter,
115    )
116    parser.add_argument("--base-url", default=os.environ.get("OPENAI_BASE_URL"), help="Project /v1/ endpoint URL")
117    parser.add_argument("--endpoint", default=os.environ.get("AZURE_OPENAI_ENDPOINT"),
118                        help="Azure OpenAI endpoint (fallback)")
119    parser.add_argument("--api-key", default=os.environ.get("AZURE_OPENAI_API_KEY"), help="API key")
120    parser.add_argument("--project-endpoint", default=os.environ.get("AZURE_AI_PROJECT_ENDPOINT"),
121                        help="Azure AI project endpoint (alternative to --base-url)")
122    parser.add_argument("--job-id", required=True, help="Fine-tuning job ID (e.g., ftjob-abc123)")
123    parser.add_argument("--poll-interval", type=int, default=15, help="Seconds between status checks (default: 15)")
124    return parser
125 
126 
127if __name__ == "__main__":
128    parser = build_parser()
129    if len(sys.argv) == 1:
130        parser.print_help()
131        sys.exit(0)
132 
133    args = parser.parse_args()
134    client, method = get_clients(base_url=args.base_url, azure_endpoint=args.endpoint, project_endpoint=args.project_endpoint, api_key=args.api_key)
135    status = monitor_job(client, args.job_id, args.poll_interval)
136    sys.exit(0 if status == "succeeded" else 1)
137