stroblmeandClaude Fable 5
d27704a2bc
A checkpoint is not a message. DType.ARTIFACT carries a reference — digest, size, media type, name — so everything on the wire stays JSON and thirty megabytes never sit in Redis, which answers the vision's open binary-payload question by narrowing it: inline codecs would only serve payloads too small to be worth a round trip, and nothing asks for that. The store is content-addressed rather than per-run, for three reasons that all pay later: a sweep whose fifty configs share one preprocessed input stores it once, a reference stays valid however it is passed around because it names content instead of a location, and the digest is what a stage cache will compare — so building it in now is what keeps that from being a change to the message contract. Node code calls fluksio.save_artifact/load_artifact and cannot tell whether it is writing the engine's own directory or putting bytes over HTTP, which is what will let the same flow run on a remote worker unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01AD8SfVhzXBG2nAfFcVh3iD
328 lines
10 KiB
Python
328 lines
10 KiB
Python
"""Python nodes run in a worker process, and stay there when things go wrong."""
|
|
|
|
import sys
|
|
import threading
|
|
import time
|
|
from collections.abc import Iterator
|
|
|
|
import pytest
|
|
|
|
from app.flow.artifacts import ArtifactStore
|
|
from app.flow.messages import DType, MessageSpec
|
|
from app.flow.worker_main import ARTIFACT_DIR_ENV
|
|
from app.flow.workers import NodeCancelled, NodeTimeout, PythonWorkerPool
|
|
|
|
|
|
@pytest.fixture
|
|
def pool() -> Iterator[PythonWorkerPool]:
|
|
# One worker: a respawn is then provably the same slot coming back.
|
|
worker_pool = PythonWorkerPool(python=sys.executable, size=1)
|
|
worker_pool.start()
|
|
yield worker_pool
|
|
worker_pool.stop()
|
|
|
|
|
|
def run(pool: PythonWorkerPool, code: str, node: str = "demo", **kwargs):
|
|
return pool.run(
|
|
"demo", node, code, kwargs, {"factor": 2}, f"demo.{node}", timeout=5
|
|
)
|
|
|
|
|
|
def test_a_node_returns_its_value_and_what_it_printed(pool, capsys):
|
|
result = run(
|
|
pool,
|
|
"def process(value, params):\n"
|
|
" print('seen', value)\n"
|
|
" return {'out': value * params['factor']}\n",
|
|
value=21,
|
|
)
|
|
assert result == {"out": 42}
|
|
# The proxy writes them to stdout, which is where the engine's tee is.
|
|
assert "seen 21" in capsys.readouterr().out
|
|
|
|
|
|
def test_a_failure_keeps_its_class_and_points_at_the_node(pool):
|
|
with pytest.raises(Exception) as caught:
|
|
run(pool, "def process(params):\n raise ValueError('bad input')\n")
|
|
|
|
# The engine renders a node error as "<class>: <message>", so both have to
|
|
# survive the trip.
|
|
assert type(caught.value).__name__ == "ValueError"
|
|
assert str(caught.value) == "bad input"
|
|
assert "<node demo." in caught.value.remote_traceback
|
|
assert "ValueError: bad input" in caught.value.remote_traceback
|
|
|
|
|
|
def test_a_node_that_kills_its_worker_is_an_ordinary_error(pool):
|
|
with pytest.raises(Exception, match="worker died"):
|
|
run(pool, "import os\n\n\ndef process(params):\n os._exit(1)\n")
|
|
|
|
assert run(pool, "def process(params):\n return {'out': 1}\n") == {"out": 1}
|
|
|
|
|
|
def test_a_node_that_runs_too_long_is_killed_and_the_pool_recovers(pool):
|
|
started = time.monotonic()
|
|
with pytest.raises(NodeTimeout):
|
|
pool.run(
|
|
"demo",
|
|
"slow",
|
|
"import time\n\n\ndef process(params):\n time.sleep(30)\n",
|
|
{},
|
|
{},
|
|
"demo.slow",
|
|
timeout=1,
|
|
)
|
|
assert time.monotonic() - started < 10
|
|
|
|
# The killed worker's slot is refilled on the next call.
|
|
assert run(pool, "def process(params):\n return {'out': 2}\n") == {"out": 2}
|
|
|
|
|
|
def test_a_running_node_can_be_cancelled(pool):
|
|
def stop_it() -> None:
|
|
for _ in range(100):
|
|
if pool.cancel("demo.slow"):
|
|
return
|
|
time.sleep(0.05)
|
|
|
|
stopper = threading.Thread(target=stop_it)
|
|
stopper.start()
|
|
try:
|
|
with pytest.raises(NodeCancelled):
|
|
pool.run(
|
|
"demo",
|
|
"slow",
|
|
"import time\n\n\ndef process(params):\n time.sleep(30)\n",
|
|
{},
|
|
{},
|
|
"demo.slow",
|
|
timeout=30,
|
|
)
|
|
finally:
|
|
stopper.join()
|
|
|
|
|
|
def test_a_result_that_is_not_json_is_refused(pool):
|
|
with pytest.raises(Exception, match="cannot be sent back as JSON"):
|
|
run(pool, "def process(params):\n return {'out': {1, 2}}\n")
|
|
|
|
|
|
def test_compiling_reports_where_the_source_is_wrong(pool):
|
|
assert pool.compile("demo", "broken", "def process(params)\n return {}\n")
|
|
assert pool.compile("demo", "fine", "def process(params):\n return {}\n") is None
|
|
|
|
|
|
def test_a_node_imports_the_standard_library_not_the_engines_own_modules(pool):
|
|
# The worker script lives in app/flow, which holds queue.py, secrets.py and
|
|
# more; the interpreter would put that directory first on sys.path.
|
|
result = run(
|
|
pool,
|
|
"import queue\nimport secrets\n\n\n"
|
|
"def process(params):\n"
|
|
" return {'out': [queue.Queue().qsize(), len(secrets.token_hex(4))]}\n",
|
|
)
|
|
assert result == {"out": [0, 8]}
|
|
|
|
|
|
def test_the_engines_secrets_are_not_in_a_workers_environment(pool, monkeypatch):
|
|
monkeypatch.setenv("SECRET_KEY", "not-for-nodes")
|
|
monkeypatch.setenv("POSTGRES_PASSWORD", "not-for-nodes")
|
|
monkeypatch.setenv("FLUKSIO_HARMLESS", "fine")
|
|
# A fresh process, so it is built from the environment set just now.
|
|
pool.respawn_all()
|
|
|
|
result = run(
|
|
pool,
|
|
"import os\n\n\n"
|
|
"def process(params):\n"
|
|
" return {'out': [k for k in ('SECRET_KEY', 'POSTGRES_PASSWORD',\n"
|
|
" 'FLUKSIO_HARMLESS') if k in os.environ]}\n",
|
|
)
|
|
assert result == {"out": ["FLUKSIO_HARMLESS"]}
|
|
|
|
|
|
def test_a_pool_can_stop_while_a_node_is_running(pool):
|
|
# One slot, taken by a node that will not finish on its own, and a second
|
|
# call queued behind it. The engine's node threads are not daemons, so a
|
|
# wait here is a shutdown that never completes.
|
|
outcomes: list[str] = []
|
|
|
|
def call(node: str) -> None:
|
|
try:
|
|
pool.run(
|
|
"demo",
|
|
node,
|
|
"import time\n\n\ndef process(params):\n time.sleep(60)\n",
|
|
{},
|
|
{},
|
|
f"demo.{node}",
|
|
timeout=60,
|
|
)
|
|
outcomes.append("returned")
|
|
except Exception as exc:
|
|
outcomes.append(type(exc).__name__)
|
|
|
|
busy = threading.Thread(target=call, args=("busy",))
|
|
busy.start()
|
|
# Let the first one take the slot, so the second is blocked acquiring it.
|
|
time.sleep(1)
|
|
waiting = threading.Thread(target=call, args=("waiting",))
|
|
waiting.start()
|
|
time.sleep(0.2)
|
|
|
|
pool.stop()
|
|
for thread in (busy, waiting):
|
|
thread.join(timeout=10)
|
|
assert not thread.is_alive()
|
|
assert len(outcomes) == 2
|
|
|
|
with pytest.raises(Exception, match="shutting down"):
|
|
run(pool, "def process(params):\n return {'out': 1}\n")
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Reporting from inside a node that has not returned yet
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_a_node_reports_metrics_while_it_is_still_running(pool):
|
|
seen = []
|
|
result = pool.run(
|
|
"demo",
|
|
"train",
|
|
"import fluksio\n"
|
|
"def process(params):\n"
|
|
" for step in range(3):\n"
|
|
" fluksio.log_metric('loss', 1.0 / (step + 1), step)\n"
|
|
" fluksio.progress(0.5, 'halfway')\n"
|
|
" return {'out': 1}\n",
|
|
{},
|
|
{},
|
|
"demo.train",
|
|
timeout=5,
|
|
run_id="r1",
|
|
on_event=seen.append,
|
|
)
|
|
|
|
assert result == {"out": 1}
|
|
metrics = [event for event in seen if event["event"] == "metric"]
|
|
assert [(m["name"], m["step"]) for m in metrics] == [
|
|
("loss", 0),
|
|
("loss", 1),
|
|
("loss", 2),
|
|
]
|
|
assert metrics[0]["value"] == 1.0
|
|
# Every event says which call it belongs to, so a sweep can tell them apart.
|
|
assert {m["call_id"] for m in metrics} == {"r1:demo.train"}
|
|
assert [event["event"] for event in seen if event["event"] == "progress"] == [
|
|
"progress"
|
|
]
|
|
|
|
|
|
def test_events_hold_off_the_timeout_but_silence_does_not(pool):
|
|
# The deadline measures silence: a node reporting every 0.05s stays alive
|
|
# well past a 0.3s timeout, which is what a two-hour training needs.
|
|
result = pool.run(
|
|
"demo",
|
|
"slow",
|
|
"import time, fluksio\n"
|
|
"def process(params):\n"
|
|
" for step in range(12):\n"
|
|
" time.sleep(0.05)\n"
|
|
" fluksio.log_metric('beat', step, step)\n"
|
|
" return {'done': True}\n",
|
|
{},
|
|
{},
|
|
"demo.slow",
|
|
timeout=0.3,
|
|
run_id="r2",
|
|
on_event=lambda _event: None,
|
|
)
|
|
assert result == {"done": True}
|
|
|
|
with pytest.raises(NodeTimeout):
|
|
pool.run(
|
|
"demo",
|
|
"quiet",
|
|
"import time\ndef process(params):\n time.sleep(2)\n return {}\n",
|
|
{},
|
|
{},
|
|
"demo.quiet",
|
|
timeout=0.3,
|
|
)
|
|
|
|
|
|
def test_cancelling_one_run_leaves_the_same_node_in_another_alone(pool):
|
|
# Keyed by (run, node): cancelling a config of a sweep must not kill the
|
|
# rest of it. With one slot the second run is not executing, so the check
|
|
# is that the pool refuses to find it rather than killing the wrong worker.
|
|
started = threading.Event()
|
|
|
|
def hold():
|
|
try:
|
|
pool.run(
|
|
"demo",
|
|
"hold",
|
|
"import time\ndef process(params):\n time.sleep(5)\n return {}\n",
|
|
{},
|
|
{},
|
|
"demo.hold",
|
|
timeout=10,
|
|
run_id="run-a",
|
|
)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
started.set()
|
|
|
|
thread = threading.Thread(target=hold, daemon=True)
|
|
thread.start()
|
|
time.sleep(0.5)
|
|
|
|
assert pool.cancel("demo.hold", run_id="run-b") is False
|
|
assert pool.cancel("demo.hold", run_id="run-a") is True
|
|
started.wait(timeout=5)
|
|
thread.join(timeout=5)
|
|
|
|
|
|
def test_a_node_saves_and_loads_an_artifact(tmp_path):
|
|
# Bytes never travel as a message: the node stores them and returns a
|
|
# reference, which the next node opens.
|
|
store = ArtifactStore(tmp_path / "artifacts")
|
|
pool = PythonWorkerPool(
|
|
python=sys.executable, size=1, env={ARTIFACT_DIR_ENV: str(store.root)}
|
|
)
|
|
pool.start()
|
|
try:
|
|
ref = pool.run(
|
|
"demo",
|
|
"save",
|
|
"import fluksio\n"
|
|
"def process(params):\n"
|
|
" return {'weights': fluksio.save_artifact(b'x' * 2048, 'w.npz')}\n",
|
|
{},
|
|
{},
|
|
"demo.save",
|
|
timeout=10,
|
|
)["weights"]
|
|
|
|
assert ref["size"] == 2048
|
|
assert MessageSpec(name="weights", dtype=DType.ARTIFACT).check(ref) is None
|
|
assert store.path(ref["digest"]) is not None
|
|
|
|
loaded = pool.run(
|
|
"demo",
|
|
"load",
|
|
"import fluksio\n"
|
|
"def process(weights, params):\n"
|
|
" with open(fluksio.load_artifact(weights), 'rb') as f:\n"
|
|
" return {'size': len(f.read())}\n",
|
|
{"weights": ref},
|
|
{},
|
|
"demo.load",
|
|
timeout=10,
|
|
)
|
|
assert loaded == {"size": 2048}
|
|
finally:
|
|
pool.stop()
|