Artifex/control_plane/model_studio/management/commands/modelstudio_overnight.py
2026-08-17 01:04:02 +07:00

38 lines
2.4 KiB
Python

import json
from django.core.management.base import BaseCommand, CommandError
from control_plane.model_studio.models import ModelPromotionPolicy, TrainingProject
from control_plane.model_studio.services import ModelStudioService
class Command(BaseCommand):
help = "Create a bounded Model Studio overnight program. --dry-run never allocates training compute."
def add_arguments(self, parser):
parser.add_argument("--project", required=True, help="TrainingProject slug")
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--wall-seconds", type=int, default=8 * 3600)
parser.add_argument("--max-runs", type=int, default=8)
def handle(self, *args, **options):
project = TrainingProject.objects.filter(slug=options["project"]).first()
if project is None:
raise CommandError("TrainingProject not found.")
policy = ModelPromotionPolicy.objects.filter(training_project=project, active=True).order_by("-created_at").first()
if policy is None:
raise CommandError("Create a ModelPromotionPolicy before starting an overnight program.")
service = ModelStudioService()
if options["dry_run"]:
rows = list(project.experiments.filter(status__in=["PROPOSED", "QUEUED"]).order_by("-experiment_value_score").values("experiment_id", "title", "experiment_value_score", "estimated_runtime_seconds"))
blockers = []
if project.status != "READY":
blockers.append(f"training project status is {project.status}")
if not project.metadata.get("spark_working_directory_verified"):
blockers.append("Spark Guard working directory is not verified")
if project.baseline_evaluation_id is None:
blockers.append("fresh Champion baseline is missing")
self.stdout.write(json.dumps({"dry_run": True, "champion": str(project.current_champion_id or ""), "experiments": rows, "status": project.status, "blockers": blockers, "estimated_window_seconds": options["wall_seconds"], "maximum_runs": options["max_runs"]}, indent=2, default=str))
return
program = service.create_program(project, policy, wall_seconds=options["wall_seconds"], max_runs=options["max_runs"])
self.stdout.write(self.style.SUCCESS(f"Created overnight program {program.id}; execution is intentionally queued for scoped worker supervision."))