Files
ModelForge/backend/alembic/versions/20260826_0014_m10_advanced_scheduling.py
T

233 lines
9.4 KiB
Python

"""M10 advanced residency and GPU scheduling.
Revision ID: 20260826_0014
Revises: 20260826_0013
"""
import uuid
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
revision: str = "20260826_0014"
down_revision: str | None = "20260826_0013"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
op.add_column(
"residency_allocations",
sa.Column("generation", sa.Integer(), server_default="1", nullable=False),
)
op.add_column(
"residency_allocations", sa.Column("transition_reason", sa.String(length=64), nullable=True)
)
op.add_column(
"serving_gpu_leases",
sa.Column(
"lease_type", sa.String(length=32), server_default="request_execution", nullable=False
),
)
op.add_column(
"serving_gpu_leases",
sa.Column("materialized_vram_bytes", sa.BigInteger(), server_default="0", nullable=False),
)
op.add_column("serving_gpu_leases", sa.Column("serving_job_id", sa.Uuid(), nullable=True))
op.add_column(
"serving_gpu_leases",
sa.Column("generation", sa.Integer(), server_default="1", nullable=False),
)
op.create_foreign_key(
"fk_serving_gpu_leases_job",
"serving_gpu_leases",
"serving_jobs",
["serving_job_id"],
["id"],
ondelete="SET NULL",
)
op.create_index(
"ix_serving_gpu_leases_serving_job_id", "serving_gpu_leases", ["serving_job_id"]
)
op.create_table(
"scheduler_policy_revisions",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("revision", sa.String(length=64), nullable=False),
sa.Column("configuration", sa.JSON(), nullable=False),
sa.Column("active", sa.Boolean(), nullable=False),
sa.Column(
"created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint("revision"),
)
op.create_index(
"ix_scheduler_policy_revisions_active", "scheduler_policy_revisions", ["active"]
)
policy_table = sa.table(
"scheduler_policy_revisions",
sa.column("id", sa.Uuid()),
sa.column("revision", sa.String()),
sa.column("configuration", sa.JSON()),
sa.column("active", sa.Boolean()),
)
op.bulk_insert(
policy_table,
[
{
"id": uuid.UUID("8e4d6a85-8e50-4f78-a182-000000000010"),
"revision": "m10-v1",
"configuration": {
"reserve_minimum_bytes": 1073741824,
"reserve_percentage": 0.05,
"runtime_margin_bytes": 268435456,
"deployment_margin_minimum_bytes": 134217728,
"deployment_margin_percentage": 0.1,
"request_execution_floor_bytes": 67108864,
"pressure_stable_seconds": 30,
"eviction_cooldown_seconds": 60,
"global_queue_limit": 128,
"placement_history_limit": 500,
"lab_paused": False,
},
"active": True,
}
],
)
op.create_table(
"scheduler_accelerator_states",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("accelerator_id", sa.Uuid(), nullable=False),
sa.Column("pressure_state", sa.String(length=16), nullable=False),
sa.Column("recovery_candidate", sa.String(length=16), nullable=True),
sa.Column("recovery_candidate_since", sa.DateTime(timezone=True), nullable=True),
sa.Column(
"pressure_changed_at",
sa.DateTime(timezone=True),
server_default=sa.func.now(),
nullable=False,
),
sa.Column(
"last_observed_at",
sa.DateTime(timezone=True),
server_default=sa.func.now(),
nullable=False,
),
sa.ForeignKeyConstraint(["accelerator_id"], ["accelerators.id"], ondelete="CASCADE"),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint("accelerator_id"),
)
op.create_index(
"ix_scheduler_accelerator_states_accelerator_id",
"scheduler_accelerator_states",
["accelerator_id"],
unique=True,
)
op.create_table(
"placement_plans",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("capability_deployment_id", sa.Uuid(), nullable=False),
sa.Column("request_id", sa.Uuid(), nullable=True),
sa.Column("policy_revision", sa.String(length=64), nullable=False),
sa.Column("verdict", sa.String(length=32), nullable=False),
sa.Column("reason_codes", sa.JSON(), nullable=False),
sa.Column("evidence", sa.JSON(), nullable=False),
sa.Column("evidence_fingerprint", sa.String(length=64), nullable=False),
sa.Column("dry_run", sa.Boolean(), nullable=False),
sa.Column(
"created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
),
sa.ForeignKeyConstraint(
["capability_deployment_id"], ["capability_deployments.id"], ondelete="RESTRICT"
),
sa.PrimaryKeyConstraint("id"),
)
for column in (
"capability_deployment_id",
"request_id",
"verdict",
"evidence_fingerprint",
"created_at",
):
op.create_index(f"ix_placement_plans_{column}", "placement_plans", [column])
op.create_table(
"co_residency_evidence",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("left_deployment_id", sa.Uuid(), nullable=False),
sa.Column("right_deployment_id", sa.Uuid(), nullable=False),
sa.Column("status", sa.String(length=32), nullable=False),
sa.Column("expected_combined_bytes", sa.BigInteger(), nullable=False),
sa.Column("measured_combined_bytes", sa.BigInteger(), nullable=True),
sa.Column("evidence", sa.JSON(), nullable=False),
sa.Column("measured_at", sa.DateTime(timezone=True), nullable=True),
sa.Column(
"created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
),
sa.ForeignKeyConstraint(
["left_deployment_id"], ["capability_deployments.id"], ondelete="RESTRICT"
),
sa.ForeignKeyConstraint(
["right_deployment_id"], ["capability_deployments.id"], ondelete="RESTRICT"
),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint(
"left_deployment_id", "right_deployment_id", name="uq_co_residency_pair"
),
)
for column in ("left_deployment_id", "right_deployment_id", "status"):
op.create_index(f"ix_co_residency_evidence_{column}", "co_residency_evidence", [column])
op.create_table(
"scheduler_evictions",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("capability_deployment_id", sa.Uuid(), nullable=False),
sa.Column("requested_deployment_id", sa.Uuid(), nullable=True),
sa.Column("reason_code", sa.String(length=64), nullable=False),
sa.Column("expected_reclaimed_bytes", sa.BigInteger(), nullable=False),
sa.Column("actual_reclaimed_bytes", sa.BigInteger(), nullable=True),
sa.Column("state", sa.String(length=32), nullable=False),
sa.Column(
"created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
),
sa.ForeignKeyConstraint(
["capability_deployment_id"], ["capability_deployments.id"], ondelete="RESTRICT"
),
sa.ForeignKeyConstraint(
["requested_deployment_id"], ["capability_deployments.id"], ondelete="SET NULL"
),
sa.PrimaryKeyConstraint("id"),
)
for column in ("capability_deployment_id", "requested_deployment_id", "reason_code"):
op.create_index(f"ix_scheduler_evictions_{column}", "scheduler_evictions", [column])
def downgrade() -> None:
op.drop_index(
"ix_scheduler_accelerator_states_accelerator_id", table_name="scheduler_accelerator_states"
)
op.drop_table("scheduler_accelerator_states")
for column in ("capability_deployment_id", "requested_deployment_id", "reason_code"):
op.drop_index(f"ix_scheduler_evictions_{column}", table_name="scheduler_evictions")
op.drop_table("scheduler_evictions")
for column in ("left_deployment_id", "right_deployment_id", "status"):
op.drop_index(f"ix_co_residency_evidence_{column}", table_name="co_residency_evidence")
op.drop_table("co_residency_evidence")
for column in (
"capability_deployment_id",
"request_id",
"verdict",
"evidence_fingerprint",
"created_at",
):
op.drop_index(f"ix_placement_plans_{column}", table_name="placement_plans")
op.drop_table("placement_plans")
op.drop_index("ix_scheduler_policy_revisions_active", table_name="scheduler_policy_revisions")
op.drop_table("scheduler_policy_revisions")
op.drop_index("ix_serving_gpu_leases_serving_job_id", table_name="serving_gpu_leases")
op.drop_constraint("fk_serving_gpu_leases_job", "serving_gpu_leases", type_="foreignkey")
for column in ("generation", "serving_job_id", "materialized_vram_bytes", "lease_type"):
op.drop_column("serving_gpu_leases", column)
op.drop_column("residency_allocations", "transition_reason")
op.drop_column("residency_allocations", "generation")