105 lines
5.5 KiB
Python
105 lines
5.5 KiB
Python
"""Persist actor-scoped Model Lab execution identities."""
|
|
|
|
import sqlalchemy as sa
|
|
|
|
from alembic import op
|
|
|
|
revision = "20260924_0038"
|
|
down_revision = "20260923_0037"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
op.create_table(
|
|
"lab_execution_runs",
|
|
sa.Column("id", sa.Uuid(), nullable=False),
|
|
sa.Column("actor_subject", sa.String(length=255), nullable=False),
|
|
sa.Column("idempotency_key", sa.String(length=128), nullable=False),
|
|
sa.Column("request_fingerprint", sa.String(length=64), nullable=False),
|
|
sa.Column("status", sa.String(length=24), server_default="accepted", nullable=False),
|
|
sa.Column("capability", sa.String(length=128), server_default="rag.embedding@1", nullable=False),
|
|
sa.Column("target_snapshot", sa.JSON(), nullable=False),
|
|
sa.Column("input_sha256", sa.String(length=64), nullable=False),
|
|
sa.Column("input_count", sa.Integer(), nullable=False),
|
|
sa.Column("input_encoded_bytes", sa.Integer(), nullable=False),
|
|
sa.Column("payload_reference", sa.String(length=128), nullable=True),
|
|
sa.Column("input_payload_expires_at", sa.DateTime(timezone=True), nullable=False),
|
|
sa.Column("gateway_request_id", sa.Uuid(), nullable=False),
|
|
sa.Column("timeout_seconds", sa.Integer(), server_default="60", nullable=False),
|
|
sa.Column("queue_timeout_seconds", sa.Integer(), server_default="10", nullable=False),
|
|
sa.Column("retention_seconds", sa.Integer(), server_default="300", nullable=False),
|
|
sa.Column("output_metadata", sa.JSON(), nullable=True),
|
|
sa.Column("worker_observations", sa.JSON(), nullable=True),
|
|
sa.Column("failure_code", sa.String(length=64), nullable=True),
|
|
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("dispatch_started_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("heartbeat_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("cancellation_requested_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("cancellation_acknowledged_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.CheckConstraint(
|
|
"status IN ('accepted', 'preparing', 'queued', 'running', 'cancel_requested', "
|
|
"'completed', 'failed', 'cancelled', 'interrupted')",
|
|
name="ck_lab_execution_run_status",
|
|
),
|
|
sa.CheckConstraint("input_count BETWEEN 1 AND 8", name="ck_lab_execution_run_input_count"),
|
|
sa.CheckConstraint("timeout_seconds BETWEEN 5 AND 300", name="ck_lab_execution_run_timeout"),
|
|
sa.CheckConstraint("queue_timeout_seconds BETWEEN 1 AND 60", name="ck_lab_execution_run_queue_timeout"),
|
|
sa.CheckConstraint("retention_seconds BETWEEN 60 AND 3600", name="ck_lab_execution_run_retention"),
|
|
sa.ForeignKeyConstraint(
|
|
["gateway_request_id"], ["gateway_requests.id"], ondelete="RESTRICT"
|
|
),
|
|
sa.PrimaryKeyConstraint("id"),
|
|
sa.UniqueConstraint("actor_subject", "idempotency_key", name="uq_lab_run_actor_idempotency"),
|
|
sa.UniqueConstraint("gateway_request_id", name="uq_lab_execution_run_gateway_request"),
|
|
)
|
|
op.create_index("ix_lab_execution_runs_actor_subject", "lab_execution_runs", ["actor_subject"])
|
|
op.create_index("ix_lab_execution_runs_status", "lab_execution_runs", ["status"])
|
|
op.create_index("ix_lab_execution_runs_gateway_request_id", "lab_execution_runs", ["gateway_request_id"])
|
|
op.create_index("ix_lab_execution_runs_actor_status_created", "lab_execution_runs", ["actor_subject", "status", "created_at"])
|
|
op.add_column(
|
|
"serving_jobs",
|
|
sa.Column("lab_execution_run_id", sa.Uuid(), nullable=True),
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_serving_jobs_lab_execution_run_id_lab_execution_runs",
|
|
"serving_jobs",
|
|
"lab_execution_runs",
|
|
["lab_execution_run_id"],
|
|
["id"],
|
|
ondelete="SET NULL",
|
|
)
|
|
op.create_index(
|
|
"ix_serving_jobs_lab_execution_run_id", "serving_jobs", ["lab_execution_run_id"]
|
|
)
|
|
if op.get_bind().dialect.name == "postgresql":
|
|
op.execute(
|
|
"REVOKE DELETE, TRUNCATE ON lab_execution_runs FROM modelforge_runtime"
|
|
)
|
|
op.execute(
|
|
"GRANT SELECT, INSERT, UPDATE ON lab_execution_runs TO modelforge_runtime"
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
bind = op.get_bind()
|
|
if bind.scalar(sa.select(sa.func.count()).select_from(sa.table("lab_execution_runs", sa.column("id")))):
|
|
raise RuntimeError(
|
|
"LAB run history exists; restore the pre-migration backup instead of deleting it."
|
|
)
|
|
op.drop_index("ix_serving_jobs_lab_execution_run_id", table_name="serving_jobs")
|
|
op.drop_constraint(
|
|
"fk_serving_jobs_lab_execution_run_id_lab_execution_runs",
|
|
"serving_jobs",
|
|
type_="foreignkey",
|
|
)
|
|
op.drop_column("serving_jobs", "lab_execution_run_id")
|
|
op.drop_index("ix_lab_execution_runs_actor_status_created", table_name="lab_execution_runs")
|
|
op.drop_index("ix_lab_execution_runs_gateway_request_id", table_name="lab_execution_runs")
|
|
op.drop_index("ix_lab_execution_runs_status", table_name="lab_execution_runs")
|
|
op.drop_index("ix_lab_execution_runs_actor_subject", table_name="lab_execution_runs")
|
|
op.drop_table("lab_execution_runs")
|