atlas-iac/testing/tests/test_hermes_execution_pool_placement.py

"""Placement and rollout-cost contracts for the distributed worker pool.

The independent review found that mediator-N pinned itself hard to worker-N while
sharing a ReadWriteOnce claim with it, so a drain or preemption that moved only
the worker deadlocked the ordinal on Multi-Attach until an operator deleted a
Pod. It also found that the provider-CLI install was redone on every Pod start
inside the 10m Flux health window for the whole ``hermes`` app.
"""

from __future__ import annotations

import json

import yaml

from testing.tests.test_hermes_execution_pool_manifests import (  # noqa: E402
    FLUX,
    HERMES,
    documents,
    mediators,
    node,
    placeable,
    worker,
)


SCRIPTS = HERMES / "scripts"


def test_the_shared_workspace_claim_cannot_deadlock_on_multi_attach():
    claims = {
        item["metadata"]["name"]: item["spec"]["accessModes"]
        for item in worker()["spec"]["volumeClaimTemplates"]
    }
    assert claims["workspace"] == ["ReadWriteMany"]
    assert claims["tools"] == ["ReadWriteOnce"]
    assert claims["provider-access"] == ["ReadWriteOnce"]

    # Every claim the mediator co-mounts with the worker must be multi-attachable.
    for ordinal, deployment in enumerate(mediators()):
        shared = {
            volume["persistentVolumeClaim"]["claimName"]
            for volume in deployment["spec"]["template"]["spec"]["volumes"]
            if "persistentVolumeClaim" in volume
        }
        assert f"workspace-hermes-execution-worker-{ordinal}" in shared
        for name in shared:
            if name.startswith("workspace-hermes-execution-worker-"):
                assert claims["workspace"] == ["ReadWriteMany"]
            else:
                assert name == f"hermes-execution-mediator-state-{ordinal}"


def test_mediator_placement_is_a_preference_and_shares_the_worker_priority():
    stateful = worker()["spec"]["template"]["spec"]
    worker_terms = stateful["affinity"]["nodeAffinity"][
        "requiredDuringSchedulingIgnoredDuringExecution"
    ]["nodeSelectorTerms"]
    for ordinal, deployment in enumerate(mediators()):
        pod = deployment["spec"]["template"]["spec"]
        assert pod["priorityClassName"] == stateful["priorityClassName"] == "scavenger"
        assert deployment["spec"]["strategy"]["type"] == "Recreate"
        affinity = pod["affinity"]
        # No hard podAffinity: the mediator must never be unschedulable purely
        # because its worker is Pending or placed on a different node.
        assert "requiredDuringSchedulingIgnoredDuringExecution" not in affinity["podAffinity"]
        preferred = affinity["podAffinity"][
            "preferredDuringSchedulingIgnoredDuringExecution"
        ]
        assert preferred[0]["podAffinityTerm"]["labelSelector"]["matchLabels"] == {
            "statefulset.kubernetes.io/pod-name": f"hermes-execution-worker-{ordinal}"
        }
        assert affinity["nodeAffinity"][
            "requiredDuringSchedulingIgnoredDuringExecution"
        ]["nodeSelectorTerms"] == worker_terms


def test_a_preemption_that_moves_the_worker_leaves_both_pods_schedulable():
    """The exact drain/pressure sequence that used to need a manual Pod delete."""
    nodes = [node(f"titan-{index:02d}") for index in (5, 6, 7, 8)]
    stateful = worker()["spec"]["template"]["spec"]
    mediator = mediators()[1]["spec"]["template"]["spec"]

    # Steady state: worker-1 and mediator-1 are colocated on titan-05.
    assert "titan-05" in placeable(stateful, nodes)
    assert "titan-05" in placeable(mediator, nodes)

    # titan-05 comes under pressure and evicts the scavenger-priority worker,
    # which reschedules elsewhere while the mediator has not moved yet.
    survivors = placeable(stateful, nodes, occupied=("titan-05",))
    assert survivors and "titan-05" not in survivors

    # The mediator still serves from titan-05 -- the shared claim is RWX, so the
    # rescheduled worker does not block on Multi-Attach.
    assert "titan-05" in placeable(mediator, nodes)

    # And when titan-05 itself is drained the mediator follows independently,
    # because nothing pins it to a Pod that may still be Pending.
    drained = [
        node(item["name"], cordoned=item["name"] == "titan-05") for item in nodes
    ]
    rescheduled = placeable(mediator, drained)
    assert rescheduled and "titan-05" not in rescheduled


def test_three_workers_still_spread_across_distinct_eligible_nodes():
    stateful = worker()["spec"]["template"]["spec"]
    nodes = [node(f"titan-{index:02d}") for index in (5, 6, 7)] + [
        node("titan-13"), node("titan-22"),
    ]
    occupied: list[str] = []
    for _replica in range(3):
        allowed = placeable(stateful, nodes, occupied=tuple(occupied))
        assert allowed, f"replica {len(occupied)} is unschedulable"
        occupied.append(allowed[0])
    assert sorted(occupied) == ["titan-05", "titan-06", "titan-07"]


def test_provider_clis_install_once_onto_a_durable_verified_volume():
    stateful = worker()
    pod = stateful["spec"]["template"]["spec"]
    claims = {item["metadata"]["name"] for item in stateful["spec"]["volumeClaimTemplates"]}
    assert "tools" in claims
    ephemeral = {
        volume["name"] for volume in pod["volumes"] if "emptyDir" in volume
    }
    assert "tools" not in ephemeral, "the CLI cache must survive a Pod restart"

    install = next(
        item for item in pod["initContainers"] if item["name"] == "install-provider-clis"
    )
    script = install["args"][0]
    # The marker alone is not trusted: a pruned or partial cache reinstalls.
    assert 'if [ ! -f "${marker}" ] || [ ! -x "${tools}/bin/codex" ]' in script
    assert 'rm -f "${marker}"' in script
    # And the install itself is bounded, so a stalled registry fails fast.
    assert "timeout 900 npm install" in script
    assert "--fetch-timeout=120000" in script
    assert script.count("npm install") == 1

    worker_mount = next(
        item for item in pod["containers"][0]["volumeMounts"]
        if item["name"] == "tools"
    )
    assert worker_mount["readOnly"] is True


def test_flux_health_for_hermes_does_not_wait_on_the_best_effort_pool():
    kustomization = documents(FLUX / "hermes/kustomization.yaml")[0]
    gated = {
        (item["kind"], item["name"])
        for item in kustomization["spec"]["healthChecks"]
    }
    assert ("StatefulSet", "hermes-execution-worker") not in gated
    assert not any(name.startswith("hermes-execution-") for _kind, name in gated)
    # The rest of the app's health contract is untouched, so its dependents keep
    # the same guarantees they had before the pool existed.
    assert ("Deployment", "hermes-agent") in gated
    assert ("Deployment", "hermes-switchyard") in gated
    assert kustomization["spec"]["timeout"] == "10m"

    for name in ("hermes-chat", "hermes-observer-bindings"):
        dependent = documents(FLUX / name / "kustomization.yaml")[0]
        assert {item["name"] for item in dependent["spec"]["dependsOn"]} >= {"hermes"}


def test_the_pool_still_reports_its_own_readiness():
    pod = worker()["spec"]["template"]["spec"]
    assert pod["containers"][0]["readinessProbe"]["exec"]["command"][-1].startswith(
        "test -w /workspace"
    )
    for deployment in mediators():
        container = deployment["spec"]["template"]["spec"]["containers"][0]
        assert container["readinessProbe"]["httpGet"] == {
            "path": "/ready", "port": "mediator"
        }
        assert container["startupProbe"]["failureThreshold"] == 60


def test_the_rendered_pool_configmap_carries_every_mounted_module():
    kustomization = yaml.safe_load((HERMES / "kustomization.yaml").read_text())
    generator = next(
        item for item in kustomization["configMapGenerator"]
        if item["name"] == "hermes-execution-pool"
    )
    keys = {entry.split("=", 1)[0] for entry in generator["files"]}
    assert {
        "execution_pool_store.py",
        "execution_pool_maintenance.py",
        "execution_pool_coordinator.py",
        "execution_pool_server.py",
    } <= keys
    modules = sorted(
        path.name for path in SCRIPTS.glob("execution_pool_*.py")
    )
    assert set(modules) <= keys, "a pool module is missing from its own mount"
    assert json.dumps(sorted(keys))