234 lines
7.9 KiB
Python
234 lines
7.9 KiB
Python
"""The stage cache, from the side that needs a database.
|
|
|
|
The pipeline half — what a hit restores and what a key is made of — is in
|
|
`tests/flow/test_runs.py`, which runs without one.
|
|
"""
|
|
|
|
import json
|
|
from datetime import UTC, datetime
|
|
|
|
import pytest
|
|
from sqlmodel import Session
|
|
|
|
from fluksio.core.config import settings
|
|
from fluksio.core.db import engine as db_engine
|
|
from fluksio.flow.artifacts import ArtifactStore
|
|
from fluksio.flow.messages import DType, MessageSpec
|
|
from fluksio.flow.pipeline import NodeOutcome
|
|
from fluksio.flow.runs import (
|
|
OUTPUT_CAP,
|
|
RunCache,
|
|
RunRejected,
|
|
_cacheable,
|
|
new_run_id,
|
|
resolve_references,
|
|
seed_values,
|
|
)
|
|
from fluksio.flow.schemas import FlowDef, FlowInput, NodeDef
|
|
from fluksio.models import Run, RunArtifact, RunNode
|
|
|
|
|
|
def test_a_run_cache_finds_what_an_earlier_run_recorded(tmp_path):
|
|
"""The run history is the cache; there is no second store to keep."""
|
|
store = ArtifactStore(tmp_path / "artifacts")
|
|
reference = store.put([b"payload"], name="data.bin")
|
|
plain, with_artifact, collected = "k-plain", "k-artifact", "k-collected"
|
|
with Session(db_engine) as session:
|
|
session.add(
|
|
RunNode(
|
|
run_id="cache-1",
|
|
node="study.a",
|
|
status="ok",
|
|
cache_key=plain,
|
|
outputs=json.dumps({"study.loss": 1.5}),
|
|
)
|
|
)
|
|
session.add(
|
|
RunNode(
|
|
run_id="cache-2",
|
|
node="study.b",
|
|
status="ok",
|
|
cache_key=with_artifact,
|
|
outputs=json.dumps({"study.data": reference}),
|
|
)
|
|
)
|
|
session.add(
|
|
RunNode(
|
|
run_id="cache-3",
|
|
node="study.c",
|
|
status="ok",
|
|
cache_key=collected,
|
|
outputs=json.dumps(
|
|
{"study.data": {**reference, "digest": "sha256:" + "1" * 64}}
|
|
),
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
cache = RunCache(store)
|
|
assert cache.lookup(plain) == (True, {"study.loss": 1.5})
|
|
assert cache.lookup(with_artifact) == (True, {"study.data": reference})
|
|
# Its bytes have gone from the store, so the reference names nothing a
|
|
# restored run could open. That is a miss, not a broken run.
|
|
assert cache.lookup(collected) == (False, None)
|
|
assert cache.lookup("never-seen") == (False, None)
|
|
assert cache.lookup("") == (False, None)
|
|
|
|
|
|
def test_what_may_be_stored_as_a_cache_entry():
|
|
"""A row carries a key and its outputs together, or neither."""
|
|
ok = NodeOutcome(
|
|
node="study.a", ok=True, cache_key="k", output_values={"study.loss": 1.0}
|
|
)
|
|
assert _cacheable(ok) == '{"study.loss":1.0}'
|
|
# A node that published nothing is still an answer worth reusing.
|
|
assert _cacheable(ok.model_copy(update={"output_values": None})) == "null"
|
|
# Not cacheable: it failed, it has no key, or it returned too much.
|
|
assert _cacheable(ok.model_copy(update={"ok": False})) is None
|
|
assert _cacheable(ok.model_copy(update={"cache_key": ""})) is None
|
|
big = {"study.data": "x" * (OUTPUT_CAP + 1)}
|
|
assert _cacheable(ok.model_copy(update={"output_values": big})) is None
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Naming an artifact from outside the process that made it
|
|
#
|
|
# A python caller passes the reference it holds. A shell holds nothing, so the
|
|
# same input also takes `@run:<id>.<output>` or a bare digest, resolved here
|
|
# rather than in each client.
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def artifact_flow() -> FlowDef:
|
|
"""A flow taking a dataset somebody else's run produced."""
|
|
dataset = MessageSpec(name="dataset", dtype=DType.ARTIFACT)
|
|
return FlowDef(
|
|
name="study",
|
|
mode="batch",
|
|
inputs=[FlowInput(spec=dataset)],
|
|
nodes=[NodeDef(id="train", requires=[dataset])],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def made_artifact():
|
|
"""A finished run with one artifact, as a later run would find it."""
|
|
digest = "sha256:" + "a1" * 32
|
|
reference = {
|
|
"digest": digest,
|
|
"size": 12,
|
|
"media_type": "text/csv",
|
|
"name": "cities.csv",
|
|
}
|
|
run_id = new_run_id()
|
|
with Session(db_engine) as session:
|
|
session.add(
|
|
Run(
|
|
id=run_id,
|
|
flow="prepare",
|
|
status="ok",
|
|
result={"dataset": reference},
|
|
created_at=datetime.now(UTC),
|
|
)
|
|
)
|
|
session.add(
|
|
RunArtifact(
|
|
run_id=run_id,
|
|
name="prepare.dataset",
|
|
node="load",
|
|
digest=digest,
|
|
size=12,
|
|
)
|
|
)
|
|
session.commit()
|
|
yield run_id, reference
|
|
with Session(db_engine) as session:
|
|
session.delete(session.get(RunArtifact, (run_id, "prepare.dataset")))
|
|
session.delete(session.get(Run, run_id))
|
|
session.commit()
|
|
|
|
|
|
def test_a_run_reference_resolves_to_what_that_run_produced(made_artifact):
|
|
run_id, reference = made_artifact
|
|
resolved = resolve_references(
|
|
artifact_flow(), {"dataset": f"@run:{run_id}.dataset"}
|
|
)
|
|
|
|
# The producer's own reference, file name and all — not one rebuilt from
|
|
# the row, which carries the message name instead.
|
|
assert resolved["dataset"] == reference
|
|
|
|
|
|
def test_a_bare_digest_resolves_to_the_bytes_under_it(made_artifact):
|
|
_run_id, reference = made_artifact
|
|
resolved = resolve_references(artifact_flow(), {"dataset": reference["digest"]})
|
|
|
|
assert resolved["dataset"]["digest"] == reference["digest"]
|
|
assert resolved["dataset"]["size"] == 12
|
|
|
|
|
|
def test_a_resolved_reference_passes_the_input_check(made_artifact):
|
|
run_id, _reference = made_artifact
|
|
flow = artifact_flow()
|
|
resolved = resolve_references(flow, {"dataset": f"@run:{run_id}.dataset"})
|
|
|
|
assert "study.dataset" in seed_values(flow, resolved)
|
|
|
|
|
|
def test_an_output_a_run_never_made_says_what_it_did(made_artifact):
|
|
run_id, _reference = made_artifact
|
|
with pytest.raises(RunRejected, match="prepare.dataset"):
|
|
resolve_references(artifact_flow(), {"dataset": f"@run:{run_id}.weights"})
|
|
|
|
|
|
def test_a_reference_to_no_run_at_all_is_refused():
|
|
with pytest.raises(RunRejected, match="no run"):
|
|
resolve_references(artifact_flow(), {"dataset": "@run:nothing.dataset"})
|
|
|
|
|
|
def test_an_unknown_digest_is_refused():
|
|
with pytest.raises(RunRejected, match="nothing here"):
|
|
resolve_references(artifact_flow(), {"dataset": "sha256:" + "b2" * 32})
|
|
|
|
|
|
def test_a_reference_passed_whole_is_left_alone(made_artifact):
|
|
"""A python caller already has the object, and hands it over as one."""
|
|
_run_id, reference = made_artifact
|
|
assert resolve_references(artifact_flow(), {"dataset": reference}) == {
|
|
"dataset": reference
|
|
}
|
|
|
|
|
|
def test_the_overview_counts_a_flow_the_list_page_would_not_reach(
|
|
client, superuser_token_headers
|
|
):
|
|
"""The list caps at 500 newest; the flow rail needs whole counts."""
|
|
with Session(db_engine) as session:
|
|
for index in range(3):
|
|
session.add(
|
|
Run(
|
|
id=f"ov-{index}",
|
|
flow="overviewed",
|
|
status="ok" if index else "running",
|
|
created_at=datetime(2026, 1, 1 + index, tzinfo=UTC),
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
rows = client.get(
|
|
f"{settings.API_V1_STR}/runs/overview", headers=superuser_token_headers
|
|
).json()
|
|
row = next(r for r in rows if r["flow"] == "overviewed")
|
|
|
|
assert (row["runs"], row["running"], row["queued"]) == (3, 1, 0)
|
|
|
|
|
|
def test_overview_is_not_read_as_a_run_id(client, superuser_token_headers):
|
|
"""`/overview` is declared before `/{run_id}`, which would swallow it."""
|
|
answer = client.get(
|
|
f"{settings.API_V1_STR}/runs/overview", headers=superuser_token_headers
|
|
)
|
|
|
|
assert answer.status_code == 200
|
|
assert isinstance(answer.json(), list)
|