137 lines
5.7 KiB
Python
137 lines
5.7 KiB
Python
# /// script
|
||
# dependencies = [
|
||
# "openai>=1.0",
|
||
# "azure-identity",
|
||
# "azure-ai-projects",
|
||
# ]
|
||
# ///
|
||
"""
|
||
monitor_training.py — Monitor a fine-tuning job until completion.
|
||
|
||
Polls the job status and streams training events (reward, loss, errors)
|
||
in real time. Exits when the job reaches a terminal state.
|
||
|
||
Usage:
|
||
python monitor_training.py --job-id ftjob-abc123
|
||
python monitor_training.py --base-url https://<resource>.services.ai.azure.com/api/projects/<project>/openai/v1/ --api-key KEY --job-id ftjob-abc123
|
||
python monitor_training.py --job-id ftjob-abc123 --poll-interval 30
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
import sys
|
||
|
||
try:
|
||
sys.stdout.reconfigure(encoding="utf-8")
|
||
sys.stderr.reconfigure(encoding="utf-8")
|
||
except (AttributeError, OSError):
|
||
pass # Stream not reconfigurable (older Python or non-tty); default encoding is fine
|
||
import time
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
from common import HelpOnErrorParser, get_clients
|
||
|
||
TERMINAL_STATUSES = {"succeeded", "failed", "cancelled"}
|
||
|
||
|
||
def monitor_job(client, job_id, poll_interval=15):
|
||
"""Poll a fine-tuning job until it reaches a terminal state."""
|
||
# Cap memory for long-running jobs (RFT can run hours/days, accumulating thousands of events)
|
||
seen_events = set()
|
||
MAX_SEEN_EVENTS = 5000
|
||
|
||
print(f"Monitoring job: {job_id}")
|
||
print(f"Polling every {poll_interval}s. Ctrl+C to stop.\n")
|
||
|
||
while True:
|
||
try:
|
||
job = client.fine_tuning.jobs.retrieve(job_id)
|
||
except Exception as e:
|
||
print(f"⚠️ Error retrieving job: {e}")
|
||
time.sleep(poll_interval)
|
||
continue
|
||
|
||
status = (job.status or "").lower()
|
||
|
||
# Fetch and display new events
|
||
try:
|
||
events = list(client.fine_tuning.jobs.list_events(job_id, limit=20))
|
||
for event in reversed(events):
|
||
event_key = (event.created_at, event.message)
|
||
if event_key not in seen_events:
|
||
if len(seen_events) >= MAX_SEEN_EVENTS:
|
||
# Keep only the most recent half — a fully-flushed dedup window
|
||
# would risk re-printing old events on transient API hiccups, but
|
||
# without trimming this set grows unbounded for long RFT runs.
|
||
seen_events = set(list(seen_events)[-(MAX_SEEN_EVENTS // 2):])
|
||
seen_events.add(event_key)
|
||
ts = time.strftime("%H:%M:%S", time.localtime(event.created_at))
|
||
level = event.level or "info"
|
||
|
||
# Highlight step events
|
||
if "Step" in event.message and "reward" in event.message:
|
||
print(f" 📈 [{ts}] {event.message}")
|
||
elif "Step" in event.message and "loss" in event.message:
|
||
print(f" 📉 [{ts}] {event.message}")
|
||
elif "error" in event.message.lower() or level == "error":
|
||
print(f" ❌ [{ts}] {event.message}")
|
||
elif "started" in event.message.lower() or "completed" in event.message.lower():
|
||
print(f" 🔔 [{ts}] {event.message}")
|
||
else:
|
||
print(f" ℹ️ [{ts}] {event.message}")
|
||
except Exception:
|
||
pass # Events API may not be available for all job states
|
||
|
||
# Check terminal state
|
||
if status in TERMINAL_STATUSES:
|
||
print(f"\n{'='*50}")
|
||
if status == "succeeded":
|
||
model = job.fine_tuned_model or "unknown"
|
||
print(f" ✅ Job succeeded!")
|
||
print(f" Fine-tuned model: {model}")
|
||
if job.trained_tokens:
|
||
print(f" Trained tokens: {job.trained_tokens:,}")
|
||
elif status == "failed":
|
||
print(f" ❌ Job failed.")
|
||
if hasattr(job, "error") and job.error:
|
||
print(f" Error: {job.error}")
|
||
elif status == "cancelled":
|
||
print(f" ⚠️ Job was cancelled.")
|
||
print(f"{'='*50}")
|
||
return status
|
||
|
||
time.sleep(poll_interval)
|
||
|
||
|
||
def build_parser():
|
||
parser = HelpOnErrorParser(
|
||
description="Monitor a fine-tuning job until completion",
|
||
epilog=(
|
||
"Example:\n"
|
||
" python monitor_training.py --job-id ftjob-abc123\n"
|
||
" python monitor_training.py --base-url https://<resource>.services.ai.azure.com/api/projects/<project>/openai/v1/ --api-key KEY --job-id ftjob-abc123"
|
||
),
|
||
formatter_class=argparse.RawTextHelpFormatter,
|
||
)
|
||
parser.add_argument("--base-url", default=os.environ.get("OPENAI_BASE_URL"), help="Project /v1/ endpoint URL")
|
||
parser.add_argument("--endpoint", default=os.environ.get("AZURE_OPENAI_ENDPOINT"),
|
||
help="Azure OpenAI endpoint (fallback)")
|
||
parser.add_argument("--api-key", default=os.environ.get("AZURE_OPENAI_API_KEY"), help="API key")
|
||
parser.add_argument("--project-endpoint", default=os.environ.get("AZURE_AI_PROJECT_ENDPOINT"),
|
||
help="Azure AI project endpoint (alternative to --base-url)")
|
||
parser.add_argument("--job-id", required=True, help="Fine-tuning job ID (e.g., ftjob-abc123)")
|
||
parser.add_argument("--poll-interval", type=int, default=15, help="Seconds between status checks (default: 15)")
|
||
return parser
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = build_parser()
|
||
if len(sys.argv) == 1:
|
||
parser.print_help()
|
||
sys.exit(0)
|
||
|
||
args = parser.parse_args()
|
||
client, method = get_clients(base_url=args.base_url, azure_endpoint=args.endpoint, project_endpoint=args.project_endpoint, api_key=args.api_key)
|
||
status = monitor_job(client, args.job_id, args.poll_interval)
|
||
sys.exit(0 if status == "succeeded" else 1)
|