UmbraTrans/umbratrans_web/tests/test_assignments.py

"""The controller assigns work, and a job is done when the controller can prove it.

Two rules, and the second is the one with teeth:

* **The controller assigns.** A node no longer claims work; it asks what it has been
  given. ``jobs.lease_owner`` could name only one holder, so a job fanned out to several
  machines was inexpressible, and "what was assigned and what is not complete" had no
  answer.

* **A job is complete when the controller has the artifacts to show it.** Not when a node
  says so. A node reporting success with nothing behind it is a failure, recorded as one,
  with the reason naming what was missing.

The second rule is easy to state and easy to implement wrongly: the tempting version
marks the job complete inside the same transaction that accepted the node's report, and
then verifies afterwards "as a check". That version cannot fail, because the only thing it
could have compared against is the report it just trusted.
"""

from __future__ import annotations

import os
from uuid import uuid4

import pytest

from umbratrans_web.compute import COMPLETE, FAILED, QUEUED, RUNNING, ComputeRepository
from umbratrans_web.tests.test_compute import _device, _node, _online

INTEGRATION_ENV = "UMBRATRANS_TEST_DSN"


def _integration_enabled() -> bool:
    return bool(os.environ.get(INTEGRATION_ENV)) and os.environ.get(
        "UMBRATRANS_TEST_DATABASE_IS_DISPOSABLE", ""
    ) in {"1", "true", "yes", "on"}


integration = pytest.mark.skipif(
    not _integration_enabled(),
    reason=f"needs {INTEGRATION_ENV} and an acknowledged disposable database",
)


@pytest.fixture
def compute():
    """A repository against the scratch database, emptied on teardown.

    Duplicated from `test_compute.py` rather than moved to `conftest.py`: those fixtures
    are module-local on purpose, and hoisting them would make every file's fixtures depend
    on one file's imports. The cost is that a change to the teardown list has to be made
    twice, which is the smaller of the two mistakes.
    """

    from umbratrans_web.tests.test_compute import _require_scratch_database

    _require_scratch_database()
    repository = ComputeRepository(
        os.environ[INTEGRATION_ENV], min_size=1, max_size=6, passwords=None
    )
    try:
        yield repository
        with repository._connection() as connection, connection.transaction():  # noqa: SLF001
            connection.execute(
                "TRUNCATE public.job_events, public.jobs, public.job_assignments, "
                "public.job_artifacts, public.artifacts, public.node_heartbeats, "
                "public.node_models, public.node_credentials, public.compute_devices, "
                "public.compute_nodes, public.model_registry CASCADE"
            )
    finally:
        repository.close()


@pytest.fixture
def organization(compute):
    """One organization, seeded directly.

    Created in the repository rather than through the service because these tests are
    about dispatch, and the service path would need an authenticated administrator to
    exist first.
    """

    identifier = uuid4()
    with compute._connection() as connection, connection.transaction():  # noqa: SLF001
        connection.execute(
            "INSERT INTO public.organizations (id, name, slug) VALUES (%s, %s, %s)",
            (identifier, "Assignment Test", f"assign-{identifier.hex[:8]}"),
        )
    return identifier


@pytest.fixture
def world(compute, organization):
    """Two online machines, a device each, and a project the jobs belong to.

    The project is not incidental. `artifacts.project_id` is NOT NULL while
    `jobs.project_id` is nullable, so a job with no project cannot produce an artifact and
    therefore cannot be completed - every test here would be measuring a constraint
    violation instead of the thing under test. Which is itself worth knowing: work that
    produces output has to belong to something.
    """

    from umbratrans_web.tenancy import tenant_search_path

    project = uuid4()
    with compute._connection() as connection, compute._tx(connection):
        connection.execute(
            "SELECT set_config('search_path', %s, true)", (tenant_search_path(organization),)
        )
        connection.execute(
            "INSERT INTO projects (id, name, slug) VALUES (%s, %s, %s)",
            (project, "Assigned Book", f"assigned-{project.hex[:8]}"),
        )

    first = _online(
        compute,
        _node(compute, organization, "assign-a"),
        [_device(_node(compute, organization, "device-a"))],
    )
    second = _online(
        compute,
        _node(compute, organization, "assign-b"),
        [_device(_node(compute, organization, "device-b"))],
    )
    return {
        "compute": compute,
        "organization": organization,
        "project": project,
        "a": first,
        "b": second,
    }


def _job(world, **overrides):
    """A job that belongs to the project's organization, so its artifacts can exist."""

    kwargs = {"kind": "speak", "requested_by": "test", "project_id": world["project"]}
    kwargs.update(overrides)
    return world["compute"].submit_job(world["organization"], **kwargs)


def _artifact(**overrides) -> dict:
    """A well-formed artifact. Everything verification needs is present by default."""

    spec = {
        "kind": "audio",
        "content_hash": "a" * 64,
        "storage_key": "audio/chapter-001.mp3",
        "byte_size": 1024,
        "media_type": "audio/mpeg",
        "role": "produced",
    }
    spec.update(overrides)
    return spec


# -- the controller assigns ---------------------------------------------------


@integration
def test_a_job_goes_to_the_node_the_controller_named(world) -> None:
    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)

    created = compute.assign_job(job.id, node_ids=[a.id])

    assert len(created) == 1
    assert created[0].node_id == a.id
    assert created[0].state == "assigned"


@integration
def test_a_job_is_not_running_until_a_node_starts_it(world) -> None:
    """Assignment is not progress.

    A console that says ``running`` for a job nobody has picked up yet is claiming work
    that has not happened, and the difference matters when an operator is deciding whether
    the queue is moving.
    """

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)

    compute.assign_job(job.id, node_ids=[a.id])

    assert compute.get_job(job.id).status == QUEUED, "an unstarted job claims to be running"


@integration
def test_a_node_sees_only_what_was_assigned_to_it(world) -> None:
    """A node's whole view of the queue, and it is only its own work."""

    compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"]
    mine = _job(world)
    theirs = _job(world)
    compute.assign_job(mine.id, node_ids=[a.id])
    compute.assign_job(theirs.id, node_ids=[b.id])

    seen_by_a = compute.assignments_for_node(a.id)

    assert [one.job_id for one in seen_by_a] == [mine.id], (
        f"node A was shown work that is not its own: {[str(one.job_id) for one in seen_by_a]}"
    )


@integration
def test_polling_does_not_hand_back_work_already_in_hand(world) -> None:
    """An agent polls in a loop. Without this it would start the same job every time."""

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)

    assert compute.assignments_for_node(a.id) == [], (
        "a running assignment is offered again to the node already running it"
    )


@integration
def test_work_comes_in_priority_order(world) -> None:
    """The controller expresses urgency through priority, so the node must honour it.

    Ordered by assignment time instead, a node picks up whatever arrived first and the
    priority column becomes decoration.
    """

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    low = _job(world, priority=900)
    high = _job(world, priority=1)
    compute.assign_job(low.id, node_ids=[a.id])
    compute.assign_job(high.id, node_ids=[a.id])

    seen = compute.assignments_for_node(a.id)

    assert [one.job_id for one in seen] == [high.id, low.id], (
        "the node was given the low-priority work first"
    )


@integration
def test_a_drained_node_is_not_given_work_by_assignment(world) -> None:
    """The claim path checks status; the assignment path must too, or draining is theatre.

    Draining a machine and then assigning to it by name would be the worst of both: the
    operator's control appears to work, and a node explicitly switched off keeps receiving
    work.
    """


    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    compute.set_node_status(a.id, "draining")

    created = compute.assign_job(job.id, node_ids=[a.id])

    assert not any(one.node_id == a.id for one in created), (
        "a drained node was assigned work"
    )
    assert compute.assignments_for_job(job.id) == []


@integration
def test_one_node_cannot_start_another_nodes_assignment(world) -> None:
    """Guessing an id must not put a machine's name on work it never did."""

    from umbratrans_web.errors import ConflictError

    compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]

    with pytest.raises(ConflictError) as error:
        compute.start_assignment(assignment.id, b.id)

    assert error.value.code == "assignment_not_yours", (
        f"the refusal does not say the assignment is another node's: {error.value}"
    )
    assert compute.assignments_for_job(job.id)[0].state == "assigned"


# -- a job is complete when the controller has the artifacts ------------------


@integration
def test_completing_with_artifacts_completes_the_job(world) -> None:
    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)

    outcome = compute.complete_assignment(
        assignment.id, a.id, artifacts=[_artifact()], result={"chapters": 1}
    )

    assert outcome.verified is True
    assert outcome.job_status == COMPLETE
    assert compute.get_job(job.id).status == COMPLETE
    assert outcome.assignment.artifacts_verified_at is not None


@integration
def test_a_node_reporting_success_with_nothing_behind_it_does_not_complete_the_job(world) -> None:
    """The rule with teeth.

    The node says it finished. The controller looks for artifacts, finds none, and records
    a failure - because "a job is complete when the controller has the artifacts to show
    it" means the node's report is a claim to be checked, not a fact to be recorded.
    """

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)

    outcome = compute.complete_assignment(assignment.id, a.id, artifacts=[], result={"chapters": 1})

    assert outcome.verified is False
    assert compute.get_job(job.id).status != COMPLETE, (
        "the job was completed on the node's word alone, with no artifact behind it"
    )
    assert outcome.missing, "the failure does not say what was missing"


@integration
def test_an_artifact_without_a_hash_is_not_stored_and_does_not_count(world) -> None:
    """Otherwise verification would be checking the node's report against itself.

    An artifact with neither a hash nor a storage key is a claim. Accepting it would mean
    the verification step compares a node's report with the same node's report, which is
    not a check at all - and it would pass.
    """

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)

    outcome = compute.complete_assignment(
        assignment.id, a.id, artifacts=[{"kind": "audio", "byte_size": 10}]
    )

    assert outcome.verified is False, "a hashless artifact was accepted as proof"
    assert outcome.artifacts_recorded == []


@integration
def test_verification_is_recorded_separately_from_the_node_saying_so(world) -> None:
    """`succeeded` and `verified` are different facts and must not look alike.

    If they collapse, then a node's report is the definition of done, and the column that
    records independent confirmation has nothing to record.
    """

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)
    compute.complete_assignment(assignment.id, a.id, artifacts=[_artifact()])

    stored = compute.assignments_for_job(job.id)[0]

    assert stored.state == "succeeded"
    assert stored.is_verified is True
    assert stored.artifacts_verified_at is not None


@integration
def test_the_artifact_is_linked_to_the_job_so_the_answer_is_readable(world) -> None:
    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)
    outcome = compute.complete_assignment(assignment.id, a.id, artifacts=[_artifact()])

    assert len(outcome.artifacts_recorded) == 1
    with compute._connection() as connection:
        linked = connection.execute(
            "SELECT count(*) AS total FROM public.job_artifacts WHERE job_id = %s", (job.id,)
        ).fetchone()["total"]
    assert linked == 1, "the artifact is not linked to the job, so nothing records what it produced"


# -- more than one node -------------------------------------------------------


@integration
def test_a_job_can_be_given_to_several_machines_at_once(world) -> None:
    """What `lease_owner` could not express."""

    compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)

    created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2)

    assert {one.node_id for one in created} == {a.id, b.id}
    assert len(compute.assignments_for_job(job.id)) == 2


@integration
def test_one_machine_is_not_given_the_same_job_twice(world) -> None:
    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)

    first = compute.assign_job(job.id, node_ids=[a.id])
    second = compute.assign_job(job.id, node_ids=[a.id])

    assert len(first) == 1
    assert len(compute.assignments_for_job(job.id)) == 1, "the same node holds the job twice"


@integration
def test_the_first_verified_machine_wins_and_the_others_are_cancelled(world) -> None:
    """Otherwise every machine on the list burns its time on a result nobody reads."""

    compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2)
    by_node = {one.node_id: one for one in created}

    compute.start_assignment(by_node[a.id].id, a.id)
    outcome = compute.complete_assignment(by_node[a.id].id, a.id, artifacts=[_artifact()])

    assert outcome.verified is True
    assert str(b.id) in outcome.superseded_nodes
    loser = next(
        one for one in compute.assignments_for_job(job.id) if one.node_id == b.id
    )
    assert loser.state == "cancelled", "the losing machine is still holding an open assignment"
    assert compute.get_job(job.id).status == COMPLETE


@integration
def test_a_machine_giving_up_does_not_disturb_one_still_working(world) -> None:
    """With a fan-out the job is not waiting for anybody.

    Requeuing or failing the job while another machine is mid-flight is how a running job
    gets handed out a second time.
    """

    compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2)
    by_node = {one.node_id: one for one in created}
    compute.start_assignment(by_node[b.id].id, b.id)
    compute.start_assignment(by_node[a.id].id, a.id)

    compute.fail_assignment(by_node[a.id].id, a.id, error="gpu caught fire")

    assert compute.get_job(job.id).status == RUNNING, (
        "the job was moved out from under the machine that is still working on it"
    )
    assert next(
        one for one in compute.assignments_for_job(job.id) if one.node_id == b.id
    ).state == "running"


@integration
def test_when_nothing_is_left_open_the_work_is_offered_again(world) -> None:
    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)

    compute.fail_assignment(assignment.id, a.id, error="out of memory")

    assert compute.get_job(job.id).status == QUEUED, (
        "the job is neither running nor waiting for anybody, so nothing will pick it up"
    )


@integration
def test_work_is_offered_again_until_the_attempts_run_out(world) -> None:
    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world, max_attempts=1)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)

    compute.fail_assignment(assignment.id, a.id, error="out of memory")

    assert compute.get_job(job.id).status == FAILED, (
        "the job keeps being re-offered after its attempts are spent, so it never ends"
    )


# -- cancelling ---------------------------------------------------------------


@integration
def test_cancelling_one_attempt_leaves_the_job_for_the_others(world) -> None:
    """Cancelling an assignment is not cancelling the job.

    Conflating them is how a job dies because one machine was withdrawn while three
    others were still entitled to finish it.
    """

    compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2)
    by_node = {one.node_id: one for one in created}

    compute.cancel_assignment(by_node[a.id].id, reason="machine withdrawn")

    assert compute.get_job(job.id).status != FAILED
    assert next(
        one for one in compute.assignments_for_job(job.id) if one.node_id == b.id
    ).state == "assigned", "the other machine lost its assignment when one was cancelled"


@integration
def test_cancelling_records_why(world) -> None:
    """A month later, withdrawn and superseded look identical unless something says so."""

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]

    compute.cancel_assignment(assignment.id, reason="machine withdrawn for reimage", actor="ops")

    events = [event["event"] for event in compute.job_events(job.id)]
    assert "cancelled" in events
    details = [
        event["detail"]
        for event in compute.job_events(job.id)
        if event["event"] == "cancelled"
    ]
    assert any("reimage" in str(detail) for detail in details), (
        f"the cancellation was recorded without its reason: {details}"
    )


@integration
def test_a_finished_job_cannot_be_assigned(world) -> None:
    """Handing completed work to a machine is duplicated effort with a real cost."""

    from umbratrans_web.errors import ConflictError

    compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)
    compute.complete_assignment(assignment.id, a.id, artifacts=[_artifact()])

    with pytest.raises(ConflictError):
        compute.assign_job(job.id, node_ids=[a.id])


@integration
def test_a_failed_verification_leaves_the_work_available(world) -> None:
    """Not lost, and not completed: available for another machine to be given.

    A node reporting success with no artifact is a failed *attempt*. If that finished the
    job, one broken node would silently end a queue.
    """

    compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"]
    job = _job(world)
    assignment = compute.assign_job(job.id, node_ids=[a.id])[0]
    compute.start_assignment(assignment.id, a.id)

    compute.complete_assignment(assignment.id, a.id, artifacts=[])

    assert compute.get_job(job.id).status == QUEUED
    retry = compute.assign_job(job.id, node_ids=[b.id])
    assert len(retry) == 1, "the work cannot be given to another machine after a failed attempt"