"""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 fluksio_worker.worker_main import ARTIFACT_DIR_ENV from fluksio.flow.artifacts import ArtifactStore from fluksio.flow.controller import EmitSink, FlowController from fluksio.flow.messages import DType, MessageSpec from fluksio.flow.nodes import Node from fluksio.flow.nodes.base import NodeOutputError from fluksio.flow.pipeline import Pipeline from fluksio.flow.schemas import FlowDef, NodeDef from fluksio.flow.state import MemoryState from fluksio.flow.store import FlowStore from fluksio.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 _sink_node(node_id: str, port: str, dtype: DType) -> tuple[Node, list[dict]]: """A node declaring one streaming port, and the list its emissions land in.""" node = Node( f=lambda: None, requires=[], provides=[MessageSpec(name=port, dtype=dtype, stream=True)], name=node_id, ) node.assign_flow("demo", node_id) published: list[dict] = [] node._pipeline = type( "Recorder", (), { "publish_emission": staticmethod( lambda _node, outputs: published.append(outputs) ) }, )() return node, published 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_node_can_log_through_the_module_it_imports(pool, capsys): """The SDK exports `logger`; inside a worker `fluksio` is the reporter. Without it, `fluksio.logger.info(...)` died with AttributeError — after the training it was reporting on had already succeeded. """ result = run( pool, "import fluksio\n\n\n" "def process(value):\n" " fluksio.logger.info('tuned %s', value)\n" " return {'out': value}\n", value=7, ) assert result == {"out": 7} assert "tuned 7" 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 ": ", so both have to # survive the trip. assert type(caught.value).__name__ == "ValueError" assert str(caught.value) == "bad input" assert " 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("DATABASE_URL", "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', 'DATABASE_URL',\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_a_node_the_controller_built_reclassifies_its_last_yield(pool, tmp_path): """The sink only reaches a node if `_build_node` wrapped its function with it. Without that, a generator node's last yield is lost and its result is None — which the pieces of this checked one at a time cannot see. """ store = FlowStore(tmp_path / "flows") store.write_flow( FlowDef( name="study", mode="batch", nodes=[ NodeDef( id="train", provides=[MessageSpec(name="loss", dtype=DType.FLOAT, stream=True)], ) ], ) ) store.write_node_source( "study", "train", "def process():\n" " yield {'loss': 1.0}\n" " yield {'loss': 2.0}\n" " yield {'loss': 3.0}\n", ) controller = FlowController(store, workers=pool) nodes, _loaded, _initial, _inputs = controller._build_flows( [(store.read_flow("study"), False)] ) emitted: list[dict] = [] pipeline = Pipeline( nodes=nodes, state=MemoryState(), emission_observer=lambda _node, outputs: emitted.append(outputs), ) pipeline.run() assert [e["study.loss"] for e in emitted] == [1.0, 2.0] assert pipeline.state["study.loss"] == 3.0 def test_without_a_return_the_last_yield_is_the_result(pool): """The worker sends every yield; the engine is what holds the last one back.""" node, published = _sink_node("count", "out", DType.INT) sink = EmitSink() sink.node = node result = sink.wrap( lambda: 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=sink.handle, ) )() # The last yield is the node's output rather than an emission, so it is not # counted twice. assert result == {"out": 3} assert published == [{"demo.out": 1}, {"demo.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() def test_a_node_passes_audio_to_the_next_one(tmp_path): # The media path end to end: one node writes a clip and declares what kind # of bytes it is, the next opens it. Only the reference travels. 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", "speak", "import fluksio, io, math, struct, wave\n" "def process():\n" " buffer = io.BytesIO()\n" " with wave.open(buffer, 'wb') as out:\n" " out.setnchannels(1)\n" " out.setsampwidth(2)\n" " out.setframerate(8000)\n" " out.writeframes(b''.join(\n" " struct.pack(' 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=0, ) finally: stopper.join() def test_an_emission_on_an_undeclared_port_fails_the_call(pool): """The engine's sink raises, and that has to reach the node's author. At the *first* yield, which is where the mistake is: the worker sends each one as it happens, so a mistyped port is checked before the node goes back to work rather than when the yield after it arrives. """ node, _ = _sink_node("gen", "loss", DType.FLOAT) sink = EmitSink() sink.node = node started = time.monotonic() with pytest.raises(NodeOutputError, match="lss"): sink.wrap( lambda: pool.run( "demo", "gen", "import time\n\n\ndef process():\n" " yield {'lss': 1.0}\n" " time.sleep(30)\n" " yield {'loss': 0.5}\n", {}, "demo.gen", timeout=0, on_event=sink.handle, ) )() # It did not wait out the node: the failure stopped the call. assert time.monotonic() - started < 10 # The worker was retired rather than left mid-generator, so the slot works. assert run(pool, "def process():\n return {'out': 4}\n") == {"out": 4} # ----------------------------------------------------------------------------- # Pools per environment # # A thread cap is read when the process imports the library, so a warm worker # cannot be told a different one. A node that declared resources therefore gets # a pool started with its own environment, and nodes deriving the same one # share it rather than paying for a cold start each. # ----------------------------------------------------------------------------- def test_a_declared_environment_gets_a_pool_of_its_own(pool): child = pool.for_env({"OMP_NUM_THREADS": "2"}) assert child is not pool assert child.env["OMP_NUM_THREADS"] == "2" # The same environment is the same pool: warm workers are the point. assert pool.for_env({"OMP_NUM_THREADS": "2"}) is child assert pool.for_env({"OMP_NUM_THREADS": "3"}) is not child # Nothing to derive means the shared pool, with no second process anywhere. assert pool.for_env({}) is pool def test_a_child_worker_is_started_with_what_it_was_given(pool): child = pool.for_env({"OMP_NUM_THREADS": "2"}) seen = child.run( "demo", "env", "import os\n\n\ndef process():\n" " return {'threads': os.environ.get('OMP_NUM_THREADS', '')}\n", {}, "demo.env", timeout=10, ) assert seen == {"threads": "2"} def test_retiring_workers_reaches_the_children(pool): child = pool.for_env({"OMP_NUM_THREADS": "2"}) before = child._generation pool.respawn_all() assert child._generation > before def test_retiring_the_cards_leaves_the_other_pools_warm(pool): """A library that preallocated the card only gives it back by dying.""" card = pool.for_env({"CUDA_VISIBLE_DEVICES": "0"}) threads = pool.for_env({"OMP_NUM_THREADS": "2"}) before = (card._generation, threads._generation) pool.retire_gpu_children() assert card._generation > before[0] # Nothing to hand back, so nothing pays a cold start for it. assert threads._generation == before[1] def test_cancelling_reaches_a_node_running_in_a_child(pool): child = pool.for_env({"OMP_NUM_THREADS": "2"}) started = threading.Event() failed: list[Exception] = [] def call() -> None: try: child.run( "demo", "slow", "import time\n\n\ndef process():\n" " print('up', flush=True)\n" " time.sleep(30)\n" " return {'out': 1}\n", {}, "demo.slow", timeout=0, on_event=lambda _event: started.set(), ) except Exception as exc: # noqa: BLE001 — the point is that it stopped failed.append(exc) thread = threading.Thread(target=call) thread.start() # The node has no events, so wait for the process rather than for a frame. deadline = time.monotonic() + 10 while time.monotonic() < deadline and not child._running: time.sleep(0.01) assert pool.cancel("demo.slow"), "the parent must find a child's node" thread.join(timeout=10) assert failed def test_a_node_saves_a_frame_to_the_ring(tmp_path): """A camera writes to memory, and the next node reads it from there. The point of the ring being an ordinary store: node code says one word more, and everything downstream — the dtype, ``load_artifact``, the widget — carries on not knowing there are two of them. """ from fluksio_worker.worker_main import ARTIFACT_VOLATILE_DIR_ENV from fluksio.flow.artifacts import VolatileStore store = ArtifactStore(tmp_path / "artifacts") store.volatile = VolatileStore(tmp_path / "ring", limit_bytes=1_000_000) pool = PythonWorkerPool( python=sys.executable, size=1, env={ ARTIFACT_DIR_ENV: str(store.root), ARTIFACT_VOLATILE_DIR_ENV: str(store.volatile.root), }, ) pool.start() try: ref = pool.run( "demo", "frame", "import fluksio\n" "def process():\n" " return {'frame': fluksio.save_artifact(\n" " b'\\x89PNG' + b'p' * 64, 'f.png',\n" " media_type='image/png', volatile=True)}\n", {}, "demo.frame", timeout=10, )["frame"] assert MessageSpec(name="frame", dtype=DType.IMAGE).check(ref) is None # In the ring, and nowhere near the data volume. assert store.volatile.path(ref["digest"]) is not None assert not (store.root / ref["digest"][7:9]).exists() read = pool.run( "demo", "read", "import fluksio\n" "def process(frame):\n" " with open(fluksio.load_artifact(frame), 'rb') as f:\n" " return {'size': len(f.read())}\n", {"frame": ref}, "demo.read", timeout=10, ) assert read == {"size": 68} finally: pool.stop() def test_the_fetch_cache_is_bounded(tmp_path): """A worker downloading a media stream fills its cache with chunks nothing will ask for twice, and content addressing means nothing ever expires.""" import os from fluksio_worker.worker_main import ARTIFACT_CACHE_BYTES_ENV, _trim_cache cache = tmp_path / "cache" cache.mkdir() for index in range(5): entry = cache / f"chunk{index}" entry.write_bytes(b"c" * 1000) os.utime(entry, (index, index)) os.environ[ARTIFACT_CACHE_BYTES_ENV] = "3000" try: _trim_cache(str(cache)) finally: del os.environ[ARTIFACT_CACHE_BYTES_ENV] left = sorted(entry.name for entry in cache.iterdir()) assert left == ["chunk2", "chunk3", "chunk4"]