stroblmeandClaude Opus 5
4355c917f8
A python node's settings are constants of its own function, so they are passed the way its ports are: by name. The controller binds them to the compiled function, the `params` field is gone from the worker and remote protocols, and a setting sharing a port's name is reported as a node error rather than shadowing it. The panel's scaffold follows suit and keeps the header in step with both ports and settings. The demo's `pace` moves from a flow input to a setting of the training node, which is what it always was. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NUb8YpL2s3gmN9WTACTt4q
362 lines
11 KiB
Python
362 lines
11 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, f"demo.{node}", timeout=5)
|
|
|
|
|
|
def test_a_node_returns_its_value_and_what_it_printed(pool, capsys):
|
|
result = run(
|
|
pool,
|
|
"def process(value, factor):\n"
|
|
" print('seen', value)\n"
|
|
" return {'out': value * factor}\n",
|
|
value=21,
|
|
factor=2,
|
|
)
|
|
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():\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():\n os._exit(1)\n")
|
|
|
|
assert run(pool, "def process():\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():\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():\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():\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():\n return {'out': {1, 2}}\n")
|
|
|
|
|
|
def test_compiling_reports_where_the_source_is_wrong(pool):
|
|
assert pool.compile("demo", "broken", "def process()\n return {}\n")
|
|
assert pool.compile("demo", "fine", "def process():\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():\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():\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():\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():\n return {'out': 1}\n")
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Producing values before returning: a generator node
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_a_generator_node_publishes_each_yield_and_returns_the_end(pool):
|
|
seen = []
|
|
result = pool.run(
|
|
"demo",
|
|
"train",
|
|
"def process():\n"
|
|
" for step in range(3):\n"
|
|
" yield {'loss': 1.0 / (step + 1)}\n"
|
|
" return {'weights': 'w', 'final_loss': 0.25}\n",
|
|
{},
|
|
"demo.train",
|
|
timeout=5,
|
|
run_id="r1",
|
|
on_event=seen.append,
|
|
)
|
|
|
|
# What it returned is the node's output; what it yielded went out as it
|
|
# happened, on the same ports.
|
|
assert result == {"weights": "w", "final_loss": 0.25}
|
|
assert [event["outputs"] for event in seen] == [
|
|
{"loss": 1.0},
|
|
{"loss": 0.5},
|
|
{"loss": 1 / 3},
|
|
]
|
|
# Every frame says which call it belongs to, so a sweep can tell them apart.
|
|
assert {event["call_id"] for event in seen} == {"r1:demo.train"}
|
|
|
|
|
|
def test_without_a_return_the_last_yield_is_the_result(pool):
|
|
seen = []
|
|
result = pool.run(
|
|
"demo",
|
|
"count",
|
|
"def process():\n"
|
|
" yield {'out': 1}\n"
|
|
" yield {'out': 2}\n"
|
|
" yield {'out': 3}\n",
|
|
{},
|
|
"demo.count",
|
|
timeout=5,
|
|
on_event=seen.append,
|
|
)
|
|
|
|
assert result == {"out": 3}
|
|
assert [event["outputs"] for event in seen] == [{"out": 1}, {"out": 2}]
|
|
|
|
|
|
def test_emit_reaches_the_same_ports_from_inside_a_callback(pool):
|
|
# A value produced somewhere a yield cannot reach — inside a framework's
|
|
# callback — is still an output rather than a log.
|
|
seen = []
|
|
result = pool.run(
|
|
"demo",
|
|
"fit",
|
|
"import fluksio\n"
|
|
"def process():\n"
|
|
" def on_epoch(n):\n"
|
|
" fluksio.emit(loss=1.0 / (n + 1))\n"
|
|
" for epoch in range(2):\n"
|
|
" on_epoch(epoch)\n"
|
|
" return {'done': True}\n",
|
|
{},
|
|
"demo.fit",
|
|
timeout=5,
|
|
on_event=seen.append,
|
|
)
|
|
|
|
assert result == {"done": True}
|
|
assert [event["outputs"] for event in seen] == [{"loss": 1.0}, {"loss": 0.5}]
|
|
|
|
|
|
def test_a_plain_function_still_just_returns(pool):
|
|
seen = []
|
|
assert run(pool, "def process():\n return {'out': 7}\n") == {"out": 7}
|
|
assert seen == []
|
|
|
|
|
|
def test_events_hold_off_the_timeout_but_silence_does_not(pool):
|
|
# The deadline measures silence: a node yielding 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\n"
|
|
"def process():\n"
|
|
" for step in range(12):\n"
|
|
" time.sleep(0.05)\n"
|
|
" yield {'beat': 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():\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():\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():\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):\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()
|