434 lines
15 KiB
Python
434 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from pydantic import SecretStr
|
|
from sqlalchemy import create_engine, select
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
from modelforge_api.api.routes.agent import (
|
|
get_agent_service,
|
|
get_decommission_service,
|
|
)
|
|
from modelforge_api.api.routes.agent import (
|
|
router as agent_router,
|
|
)
|
|
from modelforge_api.domain.agent_protocol import (
|
|
AgentMetadata,
|
|
HeartbeatRequest,
|
|
NodeManagementUpdate,
|
|
)
|
|
from modelforge_api.domain.node_decommission import NodeDecommissionExecute
|
|
from modelforge_api.domain.node_lifecycle import (
|
|
NodeLifecycleChange,
|
|
NodeLifecyclePreviewRequest,
|
|
)
|
|
from modelforge_api.main import agent_protocol_error, node_decommission_error
|
|
from modelforge_api.persistence.models import (
|
|
Accelerator,
|
|
AuditEvent,
|
|
Base,
|
|
ComputeNode,
|
|
GatewayRequest,
|
|
GpuLease,
|
|
NodeCredential,
|
|
)
|
|
from modelforge_api.services.node_agent import (
|
|
AgentAuthorizationError,
|
|
AgentConflictError,
|
|
NodeAgentService,
|
|
secret_hash,
|
|
)
|
|
from modelforge_api.services.node_decommission import (
|
|
NodeDecommissionError,
|
|
NodeDecommissionService,
|
|
)
|
|
from modelforge_api.settings import Settings, get_settings
|
|
|
|
|
|
@pytest.fixture
|
|
def session() -> Session:
|
|
engine = create_engine(
|
|
"sqlite+pysqlite:///:memory:",
|
|
connect_args={"check_same_thread": False},
|
|
poolclass=StaticPool,
|
|
)
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as value:
|
|
yield value
|
|
|
|
|
|
def _node(session: Session, *, enabled: bool = True, liveness: str = "offline") -> ComputeNode:
|
|
node = ComputeNode(
|
|
key=f"node-{uuid.uuid4()}",
|
|
hostname="node-lifecycle-test",
|
|
display_name="Lifecycle test node",
|
|
enabled=enabled,
|
|
status="active",
|
|
liveness_state=liveness,
|
|
inventory={"source": "test"},
|
|
observation_source="remote_agent",
|
|
agent_protocol_version=1,
|
|
agent_capabilities=["hardware.inventory"],
|
|
agent_version="1.3.0-rc.1",
|
|
last_heartbeat_at=datetime.now(UTC),
|
|
)
|
|
session.add(node)
|
|
session.commit()
|
|
return node
|
|
|
|
|
|
def _request(preview, *, target_enabled: bool | None = None) -> NodeLifecycleChange:
|
|
return NodeLifecycleChange(
|
|
target_enabled=preview.target_enabled if target_enabled is None else target_enabled,
|
|
expected_generation=preview.node_generation,
|
|
dependency_digest=preview.dependency_digest,
|
|
operator="test-operator",
|
|
reason="Planned node maintenance with workloads drained.",
|
|
confirmation=preview.persisted_identity,
|
|
)
|
|
|
|
|
|
def test_disable_requires_preview_confirmation_and_audits_transition(session: Session) -> None:
|
|
node = _node(session, liveness="online")
|
|
service = NodeDecommissionService(session)
|
|
|
|
preview = service.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=False)
|
|
)
|
|
assert preview.can_apply is True
|
|
assert preview.current_enabled is True
|
|
assert preview.target_enabled is False
|
|
assert preview.blockers == []
|
|
|
|
result = service.apply_enabled_change(node.id, _request(preview))
|
|
|
|
session.refresh(node)
|
|
assert result.enabled is False
|
|
assert result.liveness == "disabled"
|
|
assert node.enabled is False
|
|
assert node.liveness_state == "disabled"
|
|
event = session.scalar(
|
|
select(AuditEvent).where(
|
|
AuditEvent.action == "NODE_DISABLED",
|
|
AuditEvent.resource_id == str(node.id),
|
|
)
|
|
)
|
|
assert event is not None
|
|
assert event.details["reason"] == "Planned node maintenance with workloads drained."
|
|
|
|
|
|
def test_disable_fails_closed_for_active_gateway_request_and_gpu_lease(session: Session) -> None:
|
|
node = _node(session)
|
|
accelerator = Accelerator(
|
|
compute_node_id=node.id,
|
|
device_index=0,
|
|
device_uuid="GPU-node-lifecycle",
|
|
name="Test GPU",
|
|
)
|
|
session.add(accelerator)
|
|
session.flush()
|
|
gateway = GatewayRequest(
|
|
request_id=uuid.uuid4(),
|
|
capability_key="assistant.general",
|
|
capability_version=1,
|
|
compute_node_id=node.id,
|
|
status="running",
|
|
priority="production",
|
|
input_sha256="a" * 64,
|
|
input_count=1,
|
|
)
|
|
lease = GpuLease(
|
|
accelerator_id=accelerator.id,
|
|
deployment_id=uuid.uuid4(),
|
|
priority="production",
|
|
reserved_vram_mb=1024,
|
|
state="active",
|
|
)
|
|
session.add_all([gateway, lease])
|
|
session.commit()
|
|
service = NodeDecommissionService(session)
|
|
|
|
preview = service.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=False)
|
|
)
|
|
assert preview.can_apply is False
|
|
assert {item.code for item in preview.blockers} >= {
|
|
"active_gateway_request",
|
|
"active_gpu_lease",
|
|
}
|
|
with pytest.raises(NodeDecommissionError) as raised:
|
|
service.apply_enabled_change(node.id, _request(preview))
|
|
assert raised.value.code == "node_lifecycle_blocked"
|
|
session.refresh(node)
|
|
assert node.enabled is True
|
|
|
|
|
|
def test_disable_rejects_preview_when_work_arrives_after_impact_review(session: Session) -> None:
|
|
node = _node(session)
|
|
service = NodeDecommissionService(session)
|
|
preview = service.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=False)
|
|
)
|
|
request = _request(preview)
|
|
session.add(
|
|
GatewayRequest(
|
|
request_id=uuid.uuid4(),
|
|
capability_key="assistant.general",
|
|
capability_version=1,
|
|
compute_node_id=node.id,
|
|
status="queued",
|
|
priority="production",
|
|
input_sha256="b" * 64,
|
|
input_count=1,
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
with pytest.raises(NodeDecommissionError) as raised:
|
|
service.apply_enabled_change(node.id, request)
|
|
assert raised.value.code == "node_lifecycle_preview_stale"
|
|
session.refresh(node)
|
|
assert node.enabled is True
|
|
assert node.liveness_state == "offline"
|
|
|
|
|
|
def test_node_generation_change_invalidates_lifecycle_preview(session: Session) -> None:
|
|
node = _node(session)
|
|
service = NodeDecommissionService(session)
|
|
preview = service.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=False)
|
|
)
|
|
request = _request(preview)
|
|
node.display_name = "Renamed after preview"
|
|
session.commit()
|
|
|
|
with pytest.raises(NodeDecommissionError) as raised:
|
|
service.apply_enabled_change(node.id, request)
|
|
assert raised.value.code == "node_generation_conflict"
|
|
session.refresh(node)
|
|
assert node.enabled is True
|
|
|
|
|
|
def test_enable_requires_active_remote_credential_and_preserves_unknown_readiness(
|
|
session: Session,
|
|
) -> None:
|
|
node = _node(session, enabled=False, liveness="disabled")
|
|
session.add(
|
|
NodeCredential(
|
|
compute_node_id=node.id,
|
|
secret_hash="c" * 64,
|
|
scope="node.publish",
|
|
)
|
|
)
|
|
session.commit()
|
|
service = NodeDecommissionService(session)
|
|
|
|
preview = service.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=True)
|
|
)
|
|
assert preview.can_apply is True
|
|
result = service.apply_enabled_change(node.id, _request(preview))
|
|
|
|
session.refresh(node)
|
|
assert result.enabled is True
|
|
assert result.liveness == "offline"
|
|
assert result.agent_evidence_state == "awaiting_fresh_authenticated_heartbeat"
|
|
assert node.enabled is True
|
|
assert node.liveness_state == "offline"
|
|
assert node.agent_protocol_version is None
|
|
assert node.agent_capabilities == []
|
|
assert node.agent_version is None
|
|
assert node.last_heartbeat_at is None
|
|
assert "awaiting fresh authenticated heartbeat" in node.last_connection_error
|
|
|
|
|
|
def test_enable_without_active_credential_and_enable_of_tombstone_are_blocked(
|
|
session: Session,
|
|
) -> None:
|
|
node = _node(session, enabled=False, liveness="disabled")
|
|
service = NodeDecommissionService(session)
|
|
preview = service.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=True)
|
|
)
|
|
assert preview.can_apply is False
|
|
assert "active_node_credential_missing" in {item.code for item in preview.blockers}
|
|
with pytest.raises(NodeDecommissionError) as raised:
|
|
service.apply_enabled_change(node.id, _request(preview))
|
|
assert raised.value.code == "node_lifecycle_blocked"
|
|
with pytest.raises(AgentConflictError, match="impact preview and typed confirmation"):
|
|
NodeAgentService(session, Settings(_env_file=None)).update_node(
|
|
node.id, NodeManagementUpdate(enabled=True)
|
|
)
|
|
|
|
|
|
def test_direct_node_patch_rejects_enabled_but_keeps_metadata_updates(session: Session) -> None:
|
|
node = _node(session)
|
|
service = NodeAgentService(session, Settings(_env_file=None))
|
|
|
|
updated = service.update_node(node.id, NodeManagementUpdate(display_name="Renamed safely"))
|
|
assert updated.display_name == "Renamed safely"
|
|
with pytest.raises(AgentConflictError, match="impact preview and typed confirmation"):
|
|
service.update_node(node.id, NodeManagementUpdate(enabled=False))
|
|
session.refresh(node)
|
|
assert node.enabled is True
|
|
|
|
|
|
def test_agent_heartbeat_and_post_body_auth_refuse_node_disabled_after_prebody_auth(
|
|
session: Session,
|
|
) -> None:
|
|
node = _node(session)
|
|
credential_id = uuid.uuid4()
|
|
raw_credential = f"mfnode_{credential_id}_race-test-secret"
|
|
session.add(
|
|
NodeCredential(
|
|
id=credential_id,
|
|
compute_node_id=node.id,
|
|
secret_hash=secret_hash(raw_credential),
|
|
scope="node.publish",
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
# Keep the pre-body identity cached like a request whose lifecycle changes while its body is
|
|
# being read. A later admin transaction disables the same row before request dispatch.
|
|
agent_session = Session(session.get_bind(), expire_on_commit=False)
|
|
agent_service = NodeAgentService(agent_session, Settings(_env_file=None))
|
|
(_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
|
|
f"Bearer {raw_credential}"
|
|
)
|
|
agent_session.commit()
|
|
assert stale_node.enabled is True
|
|
|
|
lifecycle_session = Session(session.get_bind())
|
|
lifecycle = NodeDecommissionService(lifecycle_session)
|
|
preview = lifecycle.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=False)
|
|
)
|
|
lifecycle.apply_enabled_change(node.id, _request(preview))
|
|
lifecycle_session.close()
|
|
|
|
with pytest.raises(AgentAuthorizationError, match="compute node is disabled"):
|
|
agent_service.reuse_authentication(evidence, f"Bearer {raw_credential}")
|
|
|
|
heartbeat = HeartbeatRequest(
|
|
identity_key=node.key,
|
|
metadata=AgentMetadata(
|
|
agent_version="1.3.0-rc.1",
|
|
protocol_version=1,
|
|
supported_capabilities=["hardware.inventory"],
|
|
started_at=datetime.now(UTC),
|
|
),
|
|
observed_at=datetime.now(UTC),
|
|
)
|
|
with pytest.raises(AgentAuthorizationError, match="compute node is disabled"):
|
|
agent_service.heartbeat(stale_node, heartbeat)
|
|
|
|
session.refresh(node)
|
|
assert node.enabled is False
|
|
assert node.liveness_state == "disabled"
|
|
agent_session.close()
|
|
|
|
|
|
def test_decommissioned_identity_cannot_be_reenabled(session: Session) -> None:
|
|
node = _node(session, enabled=False, liveness="disabled")
|
|
service = NodeDecommissionService(session)
|
|
retirement = service.preview(node.id)
|
|
service.execute(
|
|
node.id,
|
|
NodeDecommissionExecute(
|
|
expected_generation=retirement.node_generation,
|
|
preview_digest=retirement.dependency_digest,
|
|
idempotency_key=f"decom-{uuid.uuid4()}",
|
|
operator="test-operator",
|
|
reason="This test node has left service permanently.",
|
|
confirmation=retirement.persisted_identity,
|
|
),
|
|
)
|
|
|
|
preview = service.preview_enabled_change(
|
|
node.id, NodeLifecyclePreviewRequest(target_enabled=True)
|
|
)
|
|
assert preview.can_apply is False
|
|
assert "node_decommissioned" in {item.code for item in preview.blockers}
|
|
with pytest.raises(NodeDecommissionError) as raised:
|
|
service.apply_enabled_change(node.id, _request(preview))
|
|
assert raised.value.code == "node_lifecycle_blocked"
|
|
|
|
|
|
def test_admin_api_requires_preview_for_enabled_patch_but_preserves_metadata_patch(
|
|
session: Session,
|
|
) -> None:
|
|
node = _node(session)
|
|
settings = Settings(_env_file=None, operator_api_key=SecretStr("admin-key"))
|
|
agent = NodeAgentService(session, settings)
|
|
lifecycle = NodeDecommissionService(session)
|
|
test_app = FastAPI()
|
|
test_app.include_router(agent_router)
|
|
test_app.add_exception_handler(AgentConflictError, agent_protocol_error)
|
|
test_app.add_exception_handler(NodeDecommissionError, node_decommission_error)
|
|
test_app.dependency_overrides[get_agent_service] = lambda: agent
|
|
test_app.dependency_overrides[get_decommission_service] = lambda: lifecycle
|
|
test_app.dependency_overrides[get_settings] = lambda: settings
|
|
|
|
try:
|
|
with TestClient(test_app) as client:
|
|
path = f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle/preview"
|
|
assert client.post(path, json={"target_enabled": False}).status_code == 401
|
|
headers = {"X-ModelForge-Admin-Token": "admin-key"}
|
|
preview_response = client.post(path, headers=headers, json={"target_enabled": False})
|
|
assert preview_response.status_code == 200
|
|
preview = preview_response.json()
|
|
assert preview["can_apply"] is True
|
|
|
|
bypass = client.patch(
|
|
f"/api/v1/admin/hardware/nodes/{node.id}",
|
|
headers=headers,
|
|
json={"enabled": False},
|
|
)
|
|
assert bypass.status_code == 409
|
|
assert bypass.json()["error"]["code"] == "agent_conflict"
|
|
metadata = client.patch(
|
|
f"/api/v1/admin/hardware/nodes/{node.id}",
|
|
headers=headers,
|
|
json={"display_name": "Updated metadata"},
|
|
)
|
|
assert metadata.status_code == 200
|
|
assert node.display_name == "Updated metadata"
|
|
|
|
stale_change = {
|
|
"target_enabled": False,
|
|
"expected_generation": preview["node_generation"],
|
|
"dependency_digest": preview["dependency_digest"],
|
|
"operator": "test-operator",
|
|
"reason": "Planned node maintenance after draining workloads.",
|
|
"confirmation": preview["persisted_identity"],
|
|
}
|
|
stale = client.post(
|
|
f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle",
|
|
headers=headers,
|
|
json=stale_change,
|
|
)
|
|
assert stale.status_code == 409
|
|
assert stale.json()["error"]["code"] == "node_generation_conflict"
|
|
|
|
refreshed_preview = client.post(
|
|
path, headers=headers, json={"target_enabled": False}
|
|
).json()
|
|
change = {**stale_change, "expected_generation": refreshed_preview["node_generation"], "dependency_digest": refreshed_preview["dependency_digest"]}
|
|
applied = client.post(
|
|
f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle",
|
|
headers=headers,
|
|
json=change,
|
|
)
|
|
assert applied.status_code == 200
|
|
assert applied.json()["enabled"] is False
|
|
assert applied.json()["liveness"] == "disabled"
|
|
finally:
|
|
test_app.dependency_overrides.clear()
|