Artifex/control_plane/model_studio/models.py
2026-08-17 02:04:43 +07:00

443 lines
21 KiB
Python

from __future__ import annotations
from django.db import models
from control_plane.common import TimestampedModel
class TrainingProjectStatus(models.TextChoices):
IMPORTING = "IMPORTING"
ARCHAEOLOGY = "ARCHAEOLOGY"
NEEDS_REPAIR = "NEEDS_REPAIR"
BASELINING = "BASELINING"
READY = "READY"
OVERNIGHT_RUNNING = "OVERNIGHT_RUNNING"
PAUSED = "PAUSED"
FAILED = "FAILED"
FINISHED = "FINISHED"
EVOLVING = "EVOLVING"
class DatasetValidationStatus(models.TextChoices):
VALID = "VALID"
WARNING = "WARNING"
BLOCKED = "BLOCKED"
UNKNOWN = "UNKNOWN"
class CheckpointType(models.TextChoices):
BASE = "BASE"
IMPORTED = "IMPORTED"
CHAMPION = "CHAMPION"
CHALLENGER = "CHALLENGER"
INTERMEDIATE = "INTERMEDIATE"
class CheckpointValidityStatus(models.TextChoices):
UNKNOWN = "UNKNOWN"
VALID = "VALID"
INVALID = "INVALID"
PARTIAL = "PARTIAL"
CORRUPT = "CORRUPT"
class ExperimentStatus(models.TextChoices):
PROPOSED = "PROPOSED"
QUEUED = "QUEUED"
RUNNING = "RUNNING"
EVALUATING = "EVALUATING"
VALIDATING = "VALIDATING"
PROMOTED = "PROMOTED"
REJECTED = "REJECTED"
INCONCLUSIVE = "INCONCLUSIVE"
FAILED = "FAILED"
CANCELLED = "CANCELLED"
SUPERSEDED = "SUPERSEDED"
class TrainingRunStatus(models.TextChoices):
QUEUED = "QUEUED"
STARTING = "STARTING"
RUNNING = "RUNNING"
CHECKPOINTING = "CHECKPOINTING"
SUCCEEDED = "SUCCEEDED"
FAILED = "FAILED"
OOM = "OOM"
TIMEOUT = "TIMEOUT"
CANCELLED = "CANCELLED"
INTERRUPTED = "INTERRUPTED"
class EvaluationGroup(models.TextChoices):
PRIMARY = "PRIMARY"
CAPABILITY_SUBSET = "CAPABILITY_SUBSET"
HOLDOUT = "HOLDOUT"
REGRESSION = "REGRESSION"
ADVERSARIAL = "ADVERSARIAL"
FORMAT = "FORMAT"
PERFORMANCE = "PERFORMANCE"
class EvaluationRunStatus(models.TextChoices):
QUEUED = "QUEUED"
RUNNING = "RUNNING"
SUCCEEDED = "SUCCEEDED"
FAILED = "FAILED"
INVALID = "INVALID"
class ProgramStatus(models.TextChoices):
CREATED = "CREATED"
ARCHAEOLOGY = "ARCHAEOLOGY"
BASELINING = "BASELINING"
PLANNING = "PLANNING"
RUNNING_EXPERIMENT = "RUNNING_EXPERIMENT"
EVALUATING = "EVALUATING"
ADAPTING = "ADAPTING"
FINALIZING = "FINALIZING"
COMPLETED = "COMPLETED"
PARTIAL = "PARTIAL"
FAILED = "FAILED"
PAUSED = "PAUSED"
class PromotionDecision(models.TextChoices):
PROMOTE = "PROMOTE"
REJECT = "REJECT"
INCONCLUSIVE = "INCONCLUSIVE"
REQUIRE_REPLICATION = "REQUIRE_REPLICATION"
class Conclusion(models.TextChoices):
SUPPORTED = "SUPPORTED"
WEAKLY_SUPPORTED = "WEAKLY_SUPPORTED"
REFUTED = "REFUTED"
AMBIGUOUS = "AMBIGUOUS"
EXECUTION_FAILED = "EXECUTION_FAILED"
class TrainingProject(TimestampedModel):
project = models.ForeignKey("projects.Project", on_delete=models.PROTECT, related_name="training_projects")
studio_type = models.CharField(max_length=32, default="MODEL")
name = models.CharField(max_length=200)
slug = models.SlugField(max_length=120, unique=True)
description = models.TextField(blank=True)
goal = models.TextField()
capability_target = models.TextField(blank=True)
model_family = models.CharField(max_length=160, blank=True)
model_size = models.CharField(max_length=80, blank=True)
base_model = models.TextField(blank=True)
repository_path = models.TextField(blank=True)
repository_reference = models.TextField(blank=True)
working_directory = models.TextField(blank=True)
status = models.CharField(max_length=32, choices=TrainingProjectStatus.choices, default=TrainingProjectStatus.IMPORTING)
default_profile = models.CharField(max_length=120, default="guard")
training_backend = models.CharField(max_length=120, default="guard_subprocess")
current_champion = models.ForeignKey("ModelCheckpoint", on_delete=models.SET_NULL, null=True, blank=True, related_name="champion_for_projects")
baseline_evaluation = models.ForeignKey("EvaluationRun", on_delete=models.SET_NULL, null=True, blank=True, related_name="baseline_for_projects")
metadata = models.JSONField(default=dict, blank=True)
class Dataset(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="datasets")
name = models.CharField(max_length=200)
description = models.TextField(blank=True)
metadata = models.JSONField(default=dict, blank=True)
class Meta:
constraints = [models.UniqueConstraint(fields=["training_project", "name"], name="unique_model_studio_dataset")]
class DatasetVersion(TimestampedModel):
dataset = models.ForeignKey(Dataset, on_delete=models.CASCADE, related_name="versions")
version = models.CharField(max_length=120)
parent_version = models.ForeignKey("self", on_delete=models.SET_NULL, null=True, blank=True, related_name="children")
manifest_reference = models.TextField()
content_hash = models.CharField(max_length=128)
record_count = models.PositiveIntegerField(null=True, blank=True)
split_metadata = models.JSONField(default=dict, blank=True)
source_metadata = models.JSONField(default=dict, blank=True)
generation_metadata = models.JSONField(default=dict, blank=True)
tags = models.JSONField(default=list, blank=True)
validation_status = models.CharField(max_length=16, choices=DatasetValidationStatus.choices, default=DatasetValidationStatus.UNKNOWN)
contamination_status = models.CharField(max_length=16, choices=DatasetValidationStatus.choices, default=DatasetValidationStatus.UNKNOWN)
immutable = models.BooleanField(default=False)
class Meta:
constraints = [models.UniqueConstraint(fields=["dataset", "version"], name="unique_model_studio_dataset_version")]
def save(self, *args, **kwargs):
if self.pk and self.immutable:
original = type(self).objects.get(pk=self.pk)
immutable_fields = ["version", "parent_version_id", "manifest_reference", "content_hash", "record_count", "split_metadata", "source_metadata", "generation_metadata", "tags"]
if any(getattr(self, field) != getattr(original, field) for field in immutable_fields):
raise ValueError("DatasetVersion is immutable after completed training use.")
super().save(*args, **kwargs)
class DatasetCurationProposalStatus(models.TextChoices):
PROPOSED = "PROPOSED"
VALIDATED = "VALIDATED"
REJECTED = "REJECTED"
MATERIALIZED = "MATERIALIZED"
class DatasetCurationProposal(TimestampedModel):
source_versions = models.ManyToManyField(DatasetVersion, related_name="curation_proposals")
title = models.CharField(max_length=255)
hypothesis = models.TextField()
evidence = models.JSONField(default=dict)
proposed_operations = models.JSONField(default=list)
expected_capability_effect = models.TextField(blank=True)
expected_risks = models.JSONField(default=list, blank=True)
validation_plan = models.JSONField(default=dict)
contamination_plan = models.JSONField(default=dict)
status = models.CharField(max_length=32, choices=DatasetCurationProposalStatus.choices, default=DatasetCurationProposalStatus.PROPOSED)
created_by_agent = models.CharField(max_length=120, default="DATASET_CURATOR")
model_evidence = models.JSONField(default=dict, blank=True)
materialized_version = models.ForeignKey(DatasetVersion, on_delete=models.SET_NULL, null=True, blank=True, related_name="materialized_by_proposals")
class TrainingRecipe(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="recipes")
name = models.CharField(max_length=200)
configuration = models.JSONField(default=dict)
recipe_hash = models.CharField(max_length=128, unique=True)
immutable = models.BooleanField(default=False)
def save(self, *args, **kwargs):
if self.pk and self.immutable:
original = type(self).objects.get(pk=self.pk)
if self.configuration != original.configuration or self.recipe_hash != original.recipe_hash:
raise ValueError("TrainingRecipe is immutable after a TrainingRun starts.")
super().save(*args, **kwargs)
class TrainingExperiment(TimestampedModel):
experiment_id = models.CharField(max_length=120, unique=True)
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="experiments")
title = models.CharField(max_length=255)
hypothesis = models.TextField()
reasoning = models.TextField()
intervention = models.JSONField(default=dict)
controls = models.JSONField(default=dict)
expected_result = models.TextField()
primary_success_metric = models.CharField(max_length=160)
success_threshold = models.JSONField(default=dict)
regression_constraints = models.JSONField(default=dict)
rejection_condition = models.TextField()
ambiguity_policy = models.TextField()
estimated_runtime_seconds = models.PositiveIntegerField(default=0)
maximum_runtime_seconds = models.PositiveIntegerField(default=0)
compute_budget = models.JSONField(default=dict)
parent_experiment = models.ForeignKey("self", on_delete=models.SET_NULL, null=True, blank=True, related_name="derived_experiments")
derived_from_failure_cluster = models.ForeignKey("FailureCluster", on_delete=models.SET_NULL, null=True, blank=True, related_name="experiments")
status = models.CharField(max_length=32, choices=ExperimentStatus.choices, default=ExperimentStatus.PROPOSED)
priority = models.FloatField(default=0)
expected_information_gain = models.FloatField(default=0)
expected_improvement = models.FloatField(default=0)
estimated_compute_cost = models.FloatField(default=0)
experiment_value_score = models.FloatField(default=0)
fingerprint = models.CharField(max_length=128)
created_by_agent = models.CharField(max_length=120, default="MODEL_DIRECTOR")
approved_by_model_director = models.BooleanField(default=False)
result_summary = models.TextField(blank=True)
conclusion = models.CharField(max_length=32, choices=Conclusion.choices, blank=True)
class Meta:
indexes = [models.Index(fields=["training_project", "fingerprint"])]
class ExperimentDependency(TimestampedModel):
experiment = models.ForeignKey(TrainingExperiment, on_delete=models.CASCADE, related_name="dependencies")
depends_on = models.ForeignKey(TrainingExperiment, on_delete=models.CASCADE, related_name="dependents")
required_conclusions = models.JSONField(default=list, blank=True)
rationale = models.TextField(blank=True)
class Meta:
constraints = [models.UniqueConstraint(fields=["experiment", "depends_on"], name="unique_experiment_dependency")]
class ModelCheckpoint(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="checkpoints")
name = models.CharField(max_length=255)
checkpoint_type = models.CharField(max_length=32, choices=CheckpointType.choices)
reference = models.TextField()
content_hash = models.CharField(max_length=128, blank=True)
base_checkpoint = models.ForeignKey("self", on_delete=models.SET_NULL, null=True, blank=True, related_name="derived_checkpoints")
training_run = models.ForeignKey("TrainingRun", on_delete=models.SET_NULL, null=True, blank=True, related_name="output_checkpoints")
recipe = models.ForeignKey(TrainingRecipe, on_delete=models.SET_NULL, null=True, blank=True, related_name="checkpoints")
dataset_versions = models.ManyToManyField(DatasetVersion, blank=True, related_name="checkpoints")
parameter_count = models.BigIntegerField(null=True, blank=True)
dtype = models.CharField(max_length=80, blank=True)
adapter_type = models.CharField(max_length=80, blank=True)
quantization = models.CharField(max_length=80, blank=True)
validity_status = models.CharField(max_length=16, choices=CheckpointValidityStatus.choices, default=CheckpointValidityStatus.UNKNOWN)
load_verified = models.BooleanField(default=False)
evaluation_status = models.CharField(max_length=32, default="NOT_EVALUATED")
metadata = models.JSONField(default=dict, blank=True)
class Meta:
constraints = [models.UniqueConstraint(fields=["training_project", "reference"], name="unique_model_checkpoint_reference")]
class TrainingRun(TimestampedModel):
experiment = models.ForeignKey(TrainingExperiment, on_delete=models.PROTECT, related_name="runs")
recipe = models.ForeignKey(TrainingRecipe, on_delete=models.PROTECT, related_name="runs")
input_checkpoint = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="input_runs")
output_checkpoint = models.ForeignKey(ModelCheckpoint, on_delete=models.SET_NULL, null=True, blank=True, related_name="producing_run")
status = models.CharField(max_length=32, choices=TrainingRunStatus.choices, default=TrainingRunStatus.QUEUED)
command = models.JSONField(default=list, blank=True)
working_directory = models.TextField(blank=True)
environment_snapshot = models.JSONField(default=dict, blank=True)
host = models.CharField(max_length=255, blank=True)
gpu_device = models.CharField(max_length=255, blank=True)
allocated_memory_mb = models.PositiveIntegerField(null=True, blank=True)
started_at = models.DateTimeField(null=True, blank=True)
completed_at = models.DateTimeField(null=True, blank=True)
wall_seconds = models.FloatField(default=0)
exit_code = models.IntegerField(null=True, blank=True)
stdout_reference = models.TextField(blank=True)
stderr_reference = models.TextField(blank=True)
training_log_reference = models.TextField(blank=True)
peak_memory_mb = models.PositiveIntegerField(null=True, blank=True)
gpu_utilization = models.FloatField(null=True, blank=True)
failure_category = models.CharField(max_length=80, blank=True)
failure_details = models.TextField(blank=True)
resume_source = models.TextField(blank=True)
retry_count = models.PositiveIntegerField(default=0)
pid = models.IntegerField(null=True, blank=True)
heartbeat_at = models.DateTimeField(null=True, blank=True)
class EvaluationSuite(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="evaluation_suites")
name = models.CharField(max_length=200)
description = models.TextField(blank=True)
metadata = models.JSONField(default=dict, blank=True)
class EvaluationSuiteVersion(TimestampedModel):
suite = models.ForeignKey(EvaluationSuite, on_delete=models.CASCADE, related_name="versions")
version = models.CharField(max_length=120)
reference = models.TextField()
content_hash = models.CharField(max_length=128)
command_template = models.JSONField(default=list, blank=True)
groups = models.JSONField(default=list, blank=True)
immutable = models.BooleanField(default=False)
integrity_status = models.CharField(max_length=16, choices=DatasetValidationStatus.choices, default=DatasetValidationStatus.UNKNOWN)
integrity_evidence = models.JSONField(default=dict, blank=True)
class Meta:
constraints = [models.UniqueConstraint(fields=["suite", "version"], name="unique_evaluation_suite_version")]
class EvaluationRun(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="evaluation_runs")
checkpoint = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="evaluation_runs")
suite_version = models.ForeignKey(EvaluationSuiteVersion, on_delete=models.PROTECT, related_name="runs")
status = models.CharField(max_length=16, choices=EvaluationRunStatus.choices, default=EvaluationRunStatus.QUEUED)
command = models.JSONField(default=list, blank=True)
output_reference = models.TextField(blank=True)
started_at = models.DateTimeField(null=True, blank=True)
completed_at = models.DateTimeField(null=True, blank=True)
wall_seconds = models.FloatField(default=0)
parser_version = models.CharField(max_length=120, blank=True)
integrity_evidence = models.JSONField(default=dict, blank=True)
summary = models.JSONField(default=dict, blank=True)
failure_details = models.TextField(blank=True)
class BenchmarkResult(TimestampedModel):
evaluation_run = models.ForeignKey(EvaluationRun, on_delete=models.CASCADE, related_name="results")
group = models.CharField(max_length=32, choices=EvaluationGroup.choices, default=EvaluationGroup.PRIMARY)
metric = models.CharField(max_length=160)
value = models.FloatField()
unit = models.CharField(max_length=80, blank=True)
subset = models.CharField(max_length=160, blank=True)
sample_count = models.PositiveIntegerField(null=True, blank=True)
passed = models.BooleanField(null=True, blank=True)
provenance = models.JSONField(default=dict, blank=True)
class RegressionBankItem(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="regression_bank_items")
reference = models.TextField()
content_hash = models.CharField(max_length=128, blank=True)
category = models.CharField(max_length=160, blank=True)
kind = models.CharField(max_length=80, default="IMPORTED")
provenance = models.JSONField(default=dict, blank=True)
evaluation_only = models.BooleanField(default=True)
class FailureCluster(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="failure_clusters")
name = models.CharField(max_length=200)
description = models.TextField(blank=True)
failure_type = models.CharField(max_length=160, blank=True)
severity = models.CharField(max_length=32, default="UNKNOWN")
sample_count = models.PositiveIntegerField(default=0)
representative_examples = models.JSONField(default=list, blank=True)
affected_benchmarks = models.JSONField(default=list, blank=True)
suspected_causes = models.JSONField(default=list, blank=True)
confidence = models.FloatField(default=0)
training_data_coverage = models.JSONField(default=dict, blank=True)
priority = models.FloatField(default=0)
class ModelPromotionPolicy(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="promotion_policies")
name = models.CharField(max_length=160)
version = models.CharField(max_length=80)
criteria = models.JSONField(default=dict)
active = models.BooleanField(default=True)
class ModelPromotionDecision(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="promotion_decisions")
from_champion = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="promotion_sources")
candidate = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="promotion_candidates")
experiment = models.ForeignKey(TrainingExperiment, on_delete=models.PROTECT, related_name="promotion_decisions")
decision = models.CharField(max_length=32, choices=PromotionDecision.choices)
policy = models.ForeignKey(ModelPromotionPolicy, on_delete=models.PROTECT, related_name="decisions")
evaluation_evidence = models.JSONField(default=dict)
judge_result = models.JSONField(default=dict)
reason = models.TextField()
judge_actor = models.CharField(max_length=120, default="MODEL_JUDGE")
class OvernightTrainingProgram(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="overnight_programs")
starting_champion = models.ForeignKey(ModelCheckpoint, on_delete=models.PROTECT, related_name="starting_programs")
ending_champion = models.ForeignKey(ModelCheckpoint, on_delete=models.SET_NULL, null=True, blank=True, related_name="ending_programs")
status = models.CharField(max_length=32, choices=ProgramStatus.choices, default=ProgramStatus.CREATED)
start_time = models.DateTimeField(null=True, blank=True)
deadline = models.DateTimeField()
maximum_wall_seconds = models.PositiveIntegerField()
maximum_training_runs = models.PositiveIntegerField(default=8)
maximum_failed_runs = models.PositiveIntegerField(default=3)
maximum_single_run_seconds = models.PositiveIntegerField()
evaluation_reserve_seconds = models.PositiveIntegerField()
allowed_experiment_types = models.JSONField(default=list, blank=True)
promotion_policy = models.ForeignKey(ModelPromotionPolicy, on_delete=models.PROTECT, related_name="programs")
started_at = models.DateTimeField(null=True, blank=True)
completed_at = models.DateTimeField(null=True, blank=True)
termination_reason = models.TextField(blank=True)
pause_after_current_run = models.BooleanField(default=False)
telemetry = models.JSONField(default=dict, blank=True)
class OvernightResearchReport(TimestampedModel):
program = models.OneToOneField(OvernightTrainingProgram, on_delete=models.CASCADE, related_name="report")
markdown = models.TextField()
payload = models.JSONField(default=dict)
artifact_reference = models.TextField(blank=True)
class ModelStudioArtifact(TimestampedModel):
training_project = models.ForeignKey(TrainingProject, on_delete=models.CASCADE, related_name="artifacts")
artifact_type = models.CharField(max_length=120)
name = models.CharField(max_length=255)
content = models.JSONField(default=dict)
readable = models.TextField(blank=True)
source_reference = models.TextField(blank=True)