614 lines
22 KiB
Python
614 lines
22 KiB
Python
"""M16 concurrency and race tests.
|
|
|
|
Races are the failure class unit tests miss most reliably, because a single-threaded test can
|
|
satisfy every assertion while the same code loses an update under load. These tests run real
|
|
threads against a shared database and assert on the *outcome* — exactly one winner, no duplicate
|
|
authoritative object, no lost update — rather than on timing.
|
|
|
|
SQLite serialises writers, so a lost update here would be a logic error rather than an isolation
|
|
one. The PostgreSQL-specific behaviour (deadlock retry, serialisation failure) is characterised in
|
|
the live chaos evidence; what these tests pin down is that the *code* claims a single winner.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from pydantic import SecretStr
|
|
from sqlalchemy import create_engine, event, func, select, text
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.exc import IntegrityError, OperationalError
|
|
from sqlalchemy.orm import Session
|
|
|
|
from modelforge_api.domain.agent_protocol import (
|
|
AGENT_PROTOCOL_CAPABILITIES,
|
|
AgentMetadata,
|
|
EnrollmentRequest,
|
|
EnrollmentTokenCreate,
|
|
)
|
|
from modelforge_api.persistence.models import (
|
|
AuditChainHead,
|
|
AuditEvent,
|
|
Base,
|
|
ComputeNode,
|
|
NodeCredential,
|
|
NodeEnrollment,
|
|
)
|
|
from modelforge_api.services.audit import AuditWriter, audit_chain_violations
|
|
from modelforge_api.services.invariants import InvariantStatus, check_invariants
|
|
from modelforge_api.services.node_agent import (
|
|
AgentAuthenticationError,
|
|
AgentConflictError,
|
|
NodeAgentService,
|
|
)
|
|
from modelforge_api.settings import Settings
|
|
|
|
|
|
@pytest.fixture
|
|
def engine(tmp_path: Path) -> Engine:
|
|
"""A file-backed database so each thread contends on a real connection.
|
|
|
|
An in-memory SQLite database behind a StaticPool shares one connection between threads, which
|
|
is an API misuse rather than a concurrency test: the contention being measured would be the
|
|
driver's, not the platform's.
|
|
"""
|
|
|
|
value = create_engine(f"sqlite+pysqlite:///{tmp_path / 'm16.sqlite'}")
|
|
|
|
@event.listens_for(value, "connect")
|
|
def _enable_busy_timeout(connection: Any, _record: Any) -> None:
|
|
cursor = connection.cursor()
|
|
cursor.execute("PRAGMA busy_timeout = 15000")
|
|
cursor.close()
|
|
|
|
Base.metadata.create_all(value)
|
|
return value
|
|
|
|
|
|
def settings() -> Settings:
|
|
return Settings(
|
|
_env_file=None,
|
|
operator_api_key=SecretStr("m16-operator-key"),
|
|
node_stale_after_seconds=10,
|
|
node_offline_after_seconds=20,
|
|
)
|
|
|
|
|
|
def metadata() -> AgentMetadata:
|
|
return AgentMetadata(
|
|
agent_version="0.1.0",
|
|
protocol_version=1,
|
|
supported_capabilities=AGENT_PROTOCOL_CAPABILITIES,
|
|
started_at=datetime.now(UTC),
|
|
)
|
|
|
|
|
|
def run_concurrently(
|
|
worker: Callable[[int], Any], count: int, *, max_workers: int | None = None
|
|
) -> list[Any]:
|
|
"""Start every worker at the same moment so they actually contend."""
|
|
|
|
barrier = threading.Barrier(count)
|
|
|
|
def wrapped(index: int) -> Any:
|
|
barrier.wait(timeout=30)
|
|
return worker(index)
|
|
|
|
with ThreadPoolExecutor(max_workers=max_workers or count) as pool:
|
|
return list(pool.map(wrapped, range(count)))
|
|
|
|
|
|
# --------------------------------------------------------------------- enrolment storm
|
|
|
|
|
|
@pytest.mark.parametrize("attempts", [20, 60])
|
|
def test_a_single_use_enrolment_token_survives_a_concurrent_storm(
|
|
engine: Engine, attempts: int
|
|
) -> None:
|
|
"""M15 closed this race for two threads; M16 proves it holds at storm scale."""
|
|
|
|
with Session(engine) as session:
|
|
created = NodeAgentService(session, settings()).create_enrollment(
|
|
EnrollmentTokenCreate(display_name="Storm target")
|
|
)
|
|
token = created.enrollment_token
|
|
|
|
def attempt(index: int) -> str:
|
|
with Session(engine) as session:
|
|
service = NodeAgentService(session, settings())
|
|
try:
|
|
service.enroll(
|
|
EnrollmentRequest(
|
|
enrollment_token=token,
|
|
identity_key=f"storm-node-{index}",
|
|
identity_source="persisted_uuid",
|
|
hostname=f"storm-{index}",
|
|
display_name=f"storm-{index}",
|
|
metadata=metadata(),
|
|
)
|
|
)
|
|
return "ENROLLED"
|
|
except AgentAuthenticationError:
|
|
return "REFUSED"
|
|
except (IntegrityError, OperationalError):
|
|
# A database-level refusal is still a refusal; what matters is that it is not a
|
|
# second successful identity.
|
|
session.rollback()
|
|
return "REFUSED"
|
|
|
|
outcomes = run_concurrently(attempt, attempts)
|
|
assert outcomes.count("ENROLLED") == 1, outcomes
|
|
assert outcomes.count("REFUSED") == attempts - 1
|
|
|
|
with Session(engine) as session:
|
|
assert session.scalar(select(func.count()).select_from(ComputeNode)) == 1
|
|
active = session.scalar(
|
|
select(func.count()).select_from(NodeCredential).where(NodeCredential.revoked_at.is_(None))
|
|
)
|
|
assert active == 1
|
|
enrollment = session.scalar(select(NodeEnrollment))
|
|
assert enrollment is not None
|
|
assert enrollment.used_at is not None
|
|
assert enrollment.enrolled_node_id is not None
|
|
report = check_invariants(session)
|
|
assert report.violated == 0, [
|
|
item.key for item in report.results if item.status is InvariantStatus.VIOLATED
|
|
]
|
|
|
|
|
|
def test_managed_postgresql_claim_and_revocation_linearize_on_the_row_lock() -> None:
|
|
"""Managed gate for the PostgreSQL row-lock semantics SQLite cannot reproduce.
|
|
|
|
Set MODELFORGE_TEST_POSTGRES_URL to an isolated, disposable PostgreSQL database. The test
|
|
creates and removes one UUID-named schema, pauses enrolment after its conditional claim has
|
|
acquired the row lock, then proves a concurrent revocation loses after the claim commits.
|
|
"""
|
|
|
|
database_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL")
|
|
if not database_url:
|
|
pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL")
|
|
if not database_url.startswith(("postgresql://", "postgresql+psycopg://")):
|
|
pytest.fail("MODELFORGE_TEST_POSTGRES_URL must use PostgreSQL")
|
|
if database_url.startswith("postgresql://"):
|
|
database_url = database_url.replace("postgresql://", "postgresql+psycopg://", 1)
|
|
|
|
schema = f"test_enrollment_linearization_{uuid.uuid4().hex}"
|
|
admin_engine = create_engine(database_url, pool_pre_ping=True)
|
|
with admin_engine.begin() as connection:
|
|
connection.execute(text(f'CREATE SCHEMA "{schema}"'))
|
|
engine = create_engine(
|
|
database_url,
|
|
pool_pre_ping=True,
|
|
connect_args={"options": f"-csearch_path={schema}"},
|
|
)
|
|
claim_written = threading.Event()
|
|
allow_claim_commit = threading.Event()
|
|
revoke_update_started = threading.Event()
|
|
|
|
@event.listens_for(engine, "before_cursor_execute")
|
|
def _observe_revoke_update(
|
|
_connection: Any,
|
|
_cursor: Any,
|
|
statement: str,
|
|
_parameters: Any,
|
|
_context: Any,
|
|
_executemany: bool,
|
|
) -> None:
|
|
normalized = statement.lower()
|
|
if normalized.startswith("update node_enrollments set revoked_at"):
|
|
revoke_update_started.set()
|
|
|
|
try:
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
created = NodeAgentService(session, settings()).create_enrollment(
|
|
EnrollmentTokenCreate(display_name="PostgreSQL row-lock target")
|
|
)
|
|
enrollment_id = uuid.UUID(created.id)
|
|
|
|
def claim() -> str:
|
|
with Session(engine) as session:
|
|
service = NodeAgentService(session, settings())
|
|
original_claim = service._claim_enrollment
|
|
|
|
def claim_then_pause(
|
|
*, enrollment_id: uuid.UUID, token_hash: str, claim_now: datetime
|
|
) -> bool:
|
|
won = original_claim(
|
|
enrollment_id=enrollment_id,
|
|
token_hash=token_hash,
|
|
claim_now=claim_now,
|
|
)
|
|
if won:
|
|
claim_written.set()
|
|
if not allow_claim_commit.wait(timeout=30):
|
|
raise RuntimeError("timed out while holding enrollment claim")
|
|
return won
|
|
|
|
service._claim_enrollment = claim_then_pause # type: ignore[method-assign]
|
|
response = service.enroll(
|
|
EnrollmentRequest(
|
|
enrollment_token=created.enrollment_token,
|
|
identity_key="postgres-race-node",
|
|
identity_source="persisted_uuid",
|
|
hostname="postgres-race-node",
|
|
display_name="postgres-race-node",
|
|
metadata=metadata(),
|
|
)
|
|
)
|
|
return response.credential_id
|
|
|
|
def revoke() -> str:
|
|
with Session(engine) as session:
|
|
try:
|
|
NodeAgentService(session, settings()).revoke_enrollment(enrollment_id)
|
|
except AgentConflictError:
|
|
return "CONFLICT"
|
|
return "REVOKED"
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
claim_future = pool.submit(claim)
|
|
revoke_future = None
|
|
try:
|
|
assert claim_written.wait(timeout=30), "claim never acquired its row lock"
|
|
revoke_future = pool.submit(revoke)
|
|
assert revoke_update_started.wait(timeout=30), "revoke UPDATE never started"
|
|
assert not revoke_future.done(), "revoke did not wait for the claim row lock"
|
|
finally:
|
|
allow_claim_commit.set()
|
|
credential_id = claim_future.result(timeout=30)
|
|
assert revoke_future is not None
|
|
assert revoke_future.result(timeout=30) == "CONFLICT"
|
|
|
|
with Session(engine) as session:
|
|
enrollment = session.get(NodeEnrollment, enrollment_id)
|
|
credential = session.get(NodeCredential, uuid.UUID(credential_id))
|
|
assert enrollment is not None and enrollment.used_at is not None
|
|
assert enrollment.revoked_at is None
|
|
assert credential is not None and credential.revoked_at is None
|
|
assert (
|
|
session.scalar(
|
|
select(func.count())
|
|
.select_from(NodeCredential)
|
|
.where(NodeCredential.revoked_at.is_(None))
|
|
)
|
|
== 1
|
|
)
|
|
assert (
|
|
session.scalar(
|
|
select(func.count())
|
|
.select_from(AuditEvent)
|
|
.where(AuditEvent.action == "NODE_ENROLLMENT_TOKEN_REVOKED")
|
|
)
|
|
== 0
|
|
)
|
|
finally:
|
|
allow_claim_commit.set()
|
|
engine.dispose()
|
|
with admin_engine.begin() as connection:
|
|
connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
|
admin_engine.dispose()
|
|
|
|
|
|
def test_a_storm_of_distinct_tokens_creates_exactly_one_node_each(engine: Engine) -> None:
|
|
"""The guard must reject a reused token without rejecting legitimate parallel enrolments."""
|
|
|
|
count = 12
|
|
with Session(engine) as session:
|
|
service = NodeAgentService(session, settings())
|
|
tokens = [
|
|
service.create_enrollment(
|
|
EnrollmentTokenCreate(display_name=f"Node {index}")
|
|
).enrollment_token
|
|
for index in range(count)
|
|
]
|
|
|
|
def attempt(index: int) -> str:
|
|
with Session(engine) as session:
|
|
try:
|
|
NodeAgentService(session, settings()).enroll(
|
|
EnrollmentRequest(
|
|
enrollment_token=tokens[index],
|
|
identity_key=f"parallel-node-{index}",
|
|
identity_source="persisted_uuid",
|
|
hostname=f"parallel-{index}",
|
|
display_name=f"parallel-{index}",
|
|
metadata=metadata(),
|
|
)
|
|
)
|
|
return "ENROLLED"
|
|
except (
|
|
AgentAuthenticationError,
|
|
AgentConflictError,
|
|
IntegrityError,
|
|
OperationalError,
|
|
):
|
|
session.rollback()
|
|
return "REFUSED"
|
|
|
|
outcomes = run_concurrently(attempt, count)
|
|
assert outcomes.count("ENROLLED") == count, outcomes
|
|
|
|
with Session(engine) as session:
|
|
assert session.scalar(select(func.count()).select_from(ComputeNode)) == count
|
|
assert (
|
|
session.scalar(
|
|
select(func.count())
|
|
.select_from(NodeCredential)
|
|
.where(NodeCredential.revoked_at.is_(None))
|
|
)
|
|
== count
|
|
)
|
|
assert check_invariants(session).violated == 0
|
|
|
|
|
|
def test_re_enrolment_of_one_identity_under_load_keeps_one_active_credential(
|
|
engine: Engine,
|
|
) -> None:
|
|
"""Repeated re-enrolment of the same hardware must never leave two usable credentials."""
|
|
|
|
identity = "gpu_node-hardware"
|
|
rounds = 8
|
|
with Session(engine) as session:
|
|
service = NodeAgentService(session, settings())
|
|
tokens = [
|
|
service.create_enrollment(
|
|
EnrollmentTokenCreate(display_name="GPU Node")
|
|
).enrollment_token
|
|
for _ in range(rounds)
|
|
]
|
|
|
|
def attempt(index: int) -> str:
|
|
with Session(engine) as session:
|
|
try:
|
|
NodeAgentService(session, settings()).enroll(
|
|
EnrollmentRequest(
|
|
enrollment_token=tokens[index],
|
|
identity_key=identity,
|
|
identity_source="persisted_uuid",
|
|
hostname="gpu_node",
|
|
display_name="GPU Node",
|
|
metadata=metadata(),
|
|
)
|
|
)
|
|
return "ENROLLED"
|
|
except (
|
|
AgentAuthenticationError,
|
|
AgentConflictError,
|
|
IntegrityError,
|
|
OperationalError,
|
|
):
|
|
session.rollback()
|
|
return "REFUSED"
|
|
|
|
outcomes = run_concurrently(attempt, rounds)
|
|
assert "ENROLLED" in outcomes
|
|
|
|
with Session(engine) as session:
|
|
nodes = list(session.scalars(select(ComputeNode)))
|
|
assert len(nodes) == 1
|
|
assert nodes[0].key == identity
|
|
active = list(
|
|
session.scalars(select(NodeCredential).where(NodeCredential.revoked_at.is_(None)))
|
|
)
|
|
assert len(active) == 1
|
|
report = check_invariants(session)
|
|
assert report.violated == 0
|
|
|
|
|
|
# --------------------------------------------------------------------- idempotency storm
|
|
|
|
|
|
def test_the_same_idempotency_key_in_parallel_produces_one_logical_operation(
|
|
engine: Engine,
|
|
) -> None:
|
|
"""A retried request under load must converge on one row, not many."""
|
|
|
|
from modelforge_api.persistence.models import ServingJob
|
|
|
|
key = uuid.uuid4().hex
|
|
deployment = uuid.uuid4()
|
|
node = uuid.uuid4()
|
|
attempts = 24
|
|
|
|
def attempt(_index: int) -> str:
|
|
with Session(engine) as session:
|
|
existing = session.scalar(
|
|
select(ServingJob).where(ServingJob.idempotency_key == key)
|
|
)
|
|
if existing is not None:
|
|
return "REUSED"
|
|
session.add(
|
|
ServingJob(
|
|
capability_deployment_id=deployment,
|
|
compute_node_id=node,
|
|
operation="invoke",
|
|
status="queued",
|
|
priority="production",
|
|
idempotency_key=key,
|
|
)
|
|
)
|
|
try:
|
|
session.commit()
|
|
return "CREATED"
|
|
except (IntegrityError, OperationalError):
|
|
session.rollback()
|
|
return "REUSED"
|
|
|
|
outcomes = run_concurrently(attempt, attempts)
|
|
assert outcomes.count("CREATED") >= 1
|
|
|
|
with Session(engine) as session:
|
|
rows = session.scalar(
|
|
select(func.count()).select_from(ServingJob).where(ServingJob.idempotency_key == key)
|
|
)
|
|
assert rows == 1, f"{outcomes.count('CREATED')} creators produced {rows} rows"
|
|
|
|
|
|
# --------------------------------------------------------------------- audit-chain serialization
|
|
|
|
|
|
def test_concurrent_audit_appends_form_one_strict_chain(engine: Engine) -> None:
|
|
attempts = 24
|
|
|
|
def append(index: int) -> int:
|
|
with Session(engine) as session:
|
|
event = AuditWriter(session, "operator", f"writer-{index}").write(
|
|
"CONCURRENT_APPEND",
|
|
"audit-test",
|
|
str(index),
|
|
{"attempt": index},
|
|
)
|
|
session.commit()
|
|
return event.sequence
|
|
|
|
sequences = run_concurrently(append, attempts)
|
|
assert sorted(sequences) == list(range(1, attempts + 1))
|
|
|
|
with Session(engine) as session:
|
|
events = list(session.scalars(select(AuditEvent).order_by(AuditEvent.sequence)))
|
|
assert len(events) == attempts
|
|
assert audit_chain_violations(events, session.get(AuditChainHead, 1)) == []
|
|
result = next(
|
|
item
|
|
for item in check_invariants(session).results
|
|
if item.key == "audit_chain_intact"
|
|
)
|
|
assert result.status is InvariantStatus.HOLDS
|
|
|
|
|
|
# --------------------------------------------------------------------- credential races
|
|
|
|
|
|
def test_concurrent_revocation_and_authentication_never_accepts_a_revoked_credential(
|
|
engine: Engine,
|
|
) -> None:
|
|
"""Whichever order the race resolves in, a revoked credential must never authenticate."""
|
|
|
|
with Session(engine) as session:
|
|
service = NodeAgentService(session, settings())
|
|
created = service.create_enrollment(EnrollmentTokenCreate(display_name="Race target"))
|
|
response = service.enroll(
|
|
EnrollmentRequest(
|
|
enrollment_token=created.enrollment_token,
|
|
identity_key="race-node",
|
|
identity_source="persisted_uuid",
|
|
hostname="race",
|
|
display_name="race",
|
|
metadata=metadata(),
|
|
)
|
|
)
|
|
credential = response.node_credential
|
|
node_id = session.scalar(select(ComputeNode.id))
|
|
|
|
results: list[str] = []
|
|
lock = threading.Lock()
|
|
|
|
def authenticate(_index: int) -> None:
|
|
with Session(engine) as session:
|
|
try:
|
|
NodeAgentService(session, settings()).authenticate(f"Bearer {credential}")
|
|
outcome = "ACCEPTED"
|
|
except AgentAuthenticationError:
|
|
outcome = "REFUSED"
|
|
except (IntegrityError, OperationalError):
|
|
session.rollback()
|
|
outcome = "REFUSED"
|
|
with lock:
|
|
results.append(outcome)
|
|
|
|
def revoke(_index: int) -> None:
|
|
with Session(engine) as session:
|
|
try:
|
|
NodeAgentService(session, settings()).revoke_credential(node_id)
|
|
except (IntegrityError, OperationalError):
|
|
session.rollback()
|
|
|
|
def worker(index: int) -> None:
|
|
if index == 0:
|
|
revoke(index)
|
|
else:
|
|
authenticate(index)
|
|
|
|
run_concurrently(worker, 16)
|
|
|
|
with Session(engine) as session:
|
|
stored = session.scalar(select(NodeCredential))
|
|
assert stored is not None
|
|
assert stored.revoked_at is not None
|
|
with pytest.raises(AgentAuthenticationError):
|
|
NodeAgentService(session, settings()).authenticate(f"Bearer {credential}")
|
|
report = check_invariants(session)
|
|
assert report.violated == 0
|
|
|
|
|
|
def test_authentication_after_revocation_is_refused_for_every_attempt(engine: Engine) -> None:
|
|
"""No window exists in which a revoked credential is intermittently accepted."""
|
|
|
|
with Session(engine) as session:
|
|
service = NodeAgentService(session, settings())
|
|
created = service.create_enrollment(EnrollmentTokenCreate(display_name="Revoked target"))
|
|
response = service.enroll(
|
|
EnrollmentRequest(
|
|
enrollment_token=created.enrollment_token,
|
|
identity_key="revoked-node",
|
|
identity_source="persisted_uuid",
|
|
hostname="revoked",
|
|
display_name="revoked",
|
|
metadata=metadata(),
|
|
)
|
|
)
|
|
credential = response.node_credential
|
|
service.revoke_credential(session.scalar(select(ComputeNode.id)))
|
|
|
|
def attempt(_index: int) -> str:
|
|
with Session(engine) as session:
|
|
try:
|
|
NodeAgentService(session, settings()).authenticate(f"Bearer {credential}")
|
|
return "ACCEPTED"
|
|
except AgentAuthenticationError:
|
|
return "REFUSED"
|
|
|
|
outcomes = run_concurrently(attempt, 24)
|
|
assert set(outcomes) == {"REFUSED"}, outcomes
|
|
|
|
|
|
# --------------------------------------------------------------------- invariants under load
|
|
|
|
|
|
def test_invariants_stay_readable_while_the_database_is_being_written(engine: Engine) -> None:
|
|
"""The safety check must not need a quiet system to answer."""
|
|
|
|
stop = threading.Event()
|
|
|
|
def writer() -> None:
|
|
index = 0
|
|
while not stop.is_set() and index < 200:
|
|
with Session(engine) as session:
|
|
session.add(ComputeNode(key=f"writer-{index}", hostname=f"writer-{index}"))
|
|
try:
|
|
session.commit()
|
|
except (IntegrityError, OperationalError):
|
|
session.rollback()
|
|
index += 1
|
|
|
|
thread = threading.Thread(target=writer)
|
|
thread.start()
|
|
try:
|
|
for _ in range(15):
|
|
with Session(engine) as session:
|
|
report = check_invariants(session)
|
|
assert report.checked == 16
|
|
assert report.violated == 0
|
|
finally:
|
|
stop.set()
|
|
thread.join(timeout=30)
|