stroblme
8dcceecd95
Docs / docs (push) Successful in 33s
Playwright Tests / test-playwright (1, 2) (push) Successful in 3m45s
Playwright Tests / test-playwright (2, 2) (push) Successful in 2m47s
pre-commit / pre-commit (push) Successful in 3m44s
Test Backend / test-backend (3.12) (push) Failing after 5m55s
Test Backend / test-backend (3.14) (push) Failing after 5m31s
Compose Smoke Test / test-compose (push) Successful in 33s
Playwright Tests / merge-reports (push) Successful in 1m19s
632 lines
21 KiB
Python
632 lines
21 KiB
Python
"""A worker on another host, and the thread-to-loop bridge that reaches it.
|
|
|
|
The socket lives on an event loop; a node executes on a worker thread. These
|
|
run a real loop in a thread of its own, because that split is the whole
|
|
difficulty: what is checked is that a call handed across it comes back — with
|
|
its answer, with the reports it made on the way, or with a failure that says
|
|
what happened rather than a wait that never ends.
|
|
"""
|
|
|
|
import asyncio
|
|
import gc
|
|
import sys
|
|
import threading
|
|
import time
|
|
import warnings
|
|
from collections.abc import Iterator
|
|
|
|
import pytest
|
|
|
|
from fluksio.api.routes.flows import cancel_node
|
|
from fluksio.flow import remote
|
|
from fluksio.flow.controller import FlowController
|
|
from fluksio.flow.messages import DType, MessageSpec
|
|
from fluksio.flow.remote import (
|
|
NoWorker,
|
|
RemoteWorker,
|
|
RemoteWorkerHub,
|
|
WorkerInventory,
|
|
)
|
|
from fluksio.flow.schemas import FlowDef, NodeDef
|
|
from fluksio.flow.store import FlowStore
|
|
from fluksio.flow.workers import (
|
|
MAX_REMOTE_NAME,
|
|
NodeTimeout,
|
|
PythonWorkerPool,
|
|
RemoteError,
|
|
_remote_class,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def loop() -> Iterator[asyncio.AbstractEventLoop]:
|
|
"""An event loop running in a thread, as the server's does."""
|
|
running = asyncio.new_event_loop()
|
|
thread = threading.Thread(target=running.run_forever, daemon=True)
|
|
thread.start()
|
|
yield running
|
|
running.call_soon_threadsafe(running.stop)
|
|
thread.join(timeout=5)
|
|
running.close()
|
|
|
|
|
|
class FakeSocket:
|
|
"""Stands in for the websocket: records frames, and says when one lands."""
|
|
|
|
def __init__(self) -> None:
|
|
self.sent: list[dict] = []
|
|
self.arrived = threading.Event()
|
|
|
|
async def send_json(self, payload: dict) -> None:
|
|
self.sent.append(payload)
|
|
self.arrived.set()
|
|
|
|
|
|
def attach(hub: RemoteWorkerHub, loop: asyncio.AbstractEventLoop, name: str = "gpu1"):
|
|
socket = FakeSocket()
|
|
worker = RemoteWorker(
|
|
name=name, labels=["gpu"], send=socket.send_json, loop=loop, max_parallel=2
|
|
)
|
|
hub.attach(worker)
|
|
return worker, socket
|
|
|
|
|
|
def call_in_thread(target) -> threading.Thread:
|
|
thread = threading.Thread(target=target, daemon=True)
|
|
thread.start()
|
|
return thread
|
|
|
|
|
|
def test_a_call_crosses_to_the_thread_and_the_answer_comes_back(loop):
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
result: dict = {}
|
|
|
|
thread = call_in_thread(
|
|
lambda: result.update(
|
|
value=hub.run(
|
|
"gpu", "flow", "node", "src", {"x": 1}, "flow.node", timeout=5
|
|
)
|
|
)
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
assert socket.sent[0]["source"] == "src"
|
|
assert socket.sent[0]["kwargs"] == {"x": 1}
|
|
|
|
worker.deliver(
|
|
{"call_id": socket.sent[0]["call_id"], "ok": True, "result": {"out": 2}}
|
|
)
|
|
thread.join(timeout=5)
|
|
assert result["value"] == {"out": 2}
|
|
|
|
|
|
def test_reports_arrive_before_the_answer_and_a_heartbeat_is_not_one(loop):
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
seen: list[dict] = []
|
|
result: dict = {}
|
|
|
|
thread = call_in_thread(
|
|
lambda: result.update(
|
|
value=hub.run(
|
|
"gpu",
|
|
"flow",
|
|
"node",
|
|
"src",
|
|
{},
|
|
"flow.node",
|
|
timeout=5,
|
|
run_id="r1",
|
|
on_event=seen.append,
|
|
)
|
|
)
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
call_id = socket.sent[0]["call_id"]
|
|
# The call names its run, which is how a metric finds the run that made it.
|
|
assert call_id == "r1:flow.node"
|
|
|
|
worker.deliver(
|
|
{"call_id": call_id, "event": "metric", "name": "loss", "value": 1.0}
|
|
)
|
|
worker.deliver({"call_id": call_id, "event": "heartbeat"})
|
|
worker.deliver({"call_id": call_id, "ok": True, "result": {"done": True}})
|
|
|
|
thread.join(timeout=5)
|
|
assert result["value"] == {"done": True}
|
|
# Liveness is not a measurement; only the metric is kept.
|
|
assert [event["event"] for event in seen] == ["metric"]
|
|
|
|
|
|
def test_a_failure_keeps_its_class_across_the_socket(loop):
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
caught: list[Exception] = []
|
|
|
|
def call() -> None:
|
|
try:
|
|
hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5)
|
|
except Exception as exc:
|
|
caught.append(exc)
|
|
|
|
thread = call_in_thread(call)
|
|
assert socket.arrived.wait(5)
|
|
worker.deliver(
|
|
{
|
|
"call_id": socket.sent[0]["call_id"],
|
|
"ok": False,
|
|
"error": {"type": "ValueError", "message": "bad input", "traceback": "tb"},
|
|
}
|
|
)
|
|
thread.join(timeout=5)
|
|
assert type(caught[0]).__name__ == "ValueError"
|
|
assert str(caught[0]) == "bad input"
|
|
|
|
|
|
def test_a_class_name_off_the_wire_is_bounded_and_identifier_safe():
|
|
cls = _remote_class("Bad Name\n" + "x" * 10_000)
|
|
assert cls.__name__.isidentifier()
|
|
assert len(cls.__name__) == MAX_REMOTE_NAME
|
|
assert issubclass(cls, RemoteError)
|
|
|
|
|
|
def test_a_worker_that_goes_away_fails_the_call_rather_than_hanging(loop):
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
caught: list[Exception] = []
|
|
|
|
def call() -> None:
|
|
try:
|
|
hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=30)
|
|
except Exception as exc:
|
|
caught.append(exc)
|
|
|
|
thread = call_in_thread(call)
|
|
assert socket.arrived.wait(5)
|
|
# Pulling the cable mid-training: the node fails, and does not wait out its
|
|
# thirty-second deadline to do it.
|
|
worker.detach()
|
|
thread.join(timeout=5)
|
|
|
|
assert not thread.is_alive()
|
|
assert isinstance(caught[0], RemoteError)
|
|
assert "went away" in str(caught[0])
|
|
|
|
|
|
def test_silence_past_the_deadline_is_a_timeout(loop):
|
|
hub = RemoteWorkerHub()
|
|
attach(hub, loop)
|
|
caught: list[Exception] = []
|
|
|
|
def call() -> None:
|
|
try:
|
|
hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0.3)
|
|
except Exception as exc:
|
|
caught.append(exc)
|
|
|
|
call_in_thread(call).join(timeout=5)
|
|
assert isinstance(caught[0], NodeTimeout)
|
|
|
|
|
|
def test_a_label_nothing_carries_is_named_rather_than_waited_on(loop):
|
|
hub = RemoteWorkerHub()
|
|
attach(hub, loop)
|
|
|
|
with pytest.raises(NoWorker, match="tpu"):
|
|
hub.run("tpu", "flow", "node", "src", {}, "flow.node", timeout=5)
|
|
# Compiling against a machine that is not attached is not a broken node —
|
|
# a node importing torch is correct there and missing here.
|
|
assert hub.compile("tpu", "flow", "node", "src") is None
|
|
|
|
|
|
def test_reattaching_replaces_the_old_socket(loop):
|
|
hub = RemoteWorkerHub()
|
|
first, _ = attach(hub, loop)
|
|
second, _ = attach(hub, loop)
|
|
|
|
assert first.gone
|
|
assert hub.pick("gpu") is second
|
|
assert [worker.name for worker in hub.workers()] == ["gpu1"]
|
|
|
|
|
|
def test_cancelling_a_run_reaches_only_that_run(loop):
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
|
|
def call(run_id: str) -> None:
|
|
try:
|
|
hub.run("gpu", "flow", "node", "src", {}, "flow.node", 30, run_id=run_id)
|
|
except Exception:
|
|
pass
|
|
|
|
call_in_thread(lambda: call("run-a"))
|
|
assert socket.arrived.wait(5)
|
|
socket.arrived.clear()
|
|
call_in_thread(lambda: call("run-b"))
|
|
assert socket.arrived.wait(5)
|
|
|
|
socket.arrived.clear()
|
|
assert hub.cancel_run("run-a") == 1
|
|
# Sent by the call's own thread, once its frame is out.
|
|
assert socket.arrived.wait(5)
|
|
cancels = [frame for frame in socket.sent if frame.get("op") == "cancel"]
|
|
assert [frame["call_id"] for frame in cancels] == ["run-a:flow.node"]
|
|
|
|
|
|
class Loading(FakeSocket):
|
|
"""A worker that loads whatever it is sent and leaves each call running."""
|
|
|
|
worker: RemoteWorker
|
|
|
|
async def send_json(self, payload: dict) -> None:
|
|
await super().send_json(payload)
|
|
if payload["op"] == "compile":
|
|
self.worker.deliver({"call_id": payload["call_id"], "ok": True})
|
|
|
|
|
|
def live_voice(loop, tmp_path, pool: PythonWorkerPool, local: str = "", events=None):
|
|
"""A live flow 'voice' whose node 'gpu' runs on a worker and, given
|
|
*local*, a node 'cpu' running that source in *pool*. Built and started."""
|
|
store = FlowStore(tmp_path / "flows")
|
|
nodes = [
|
|
NodeDef(
|
|
id="gpu",
|
|
worker="gpu",
|
|
provides=[MessageSpec(name="reply", dtype=DType.STR)],
|
|
)
|
|
]
|
|
if local:
|
|
nodes.append(
|
|
NodeDef(id="cpu", provides=[MessageSpec(name="local", dtype=DType.STR)])
|
|
)
|
|
store.write_node_source("voice", "cpu", local)
|
|
store.write_flow(FlowDef(name="voice", nodes=nodes))
|
|
store.write_node_source("voice", "gpu", "def process():\n return None\n")
|
|
socket = Loading()
|
|
socket.worker = RemoteWorker(
|
|
name="gpu1", labels=["gpu"], send=socket.send_json, loop=loop, max_parallel=2
|
|
)
|
|
hub = RemoteWorkerHub()
|
|
hub.attach(socket.worker)
|
|
controller = FlowController(store, workers=pool, remote=hub, events=events)
|
|
asyncio.run(controller.reload())
|
|
socket.arrived.clear()
|
|
return controller, store, socket
|
|
|
|
|
|
class Events(list):
|
|
"""Stands in for the event bus: keeps what was published."""
|
|
|
|
def publish(self, event: dict) -> None:
|
|
self.append(event)
|
|
|
|
|
|
#: What a worker answers for a call it was told to cancel.
|
|
CANCELLED = {"ok": False, "error": {"type": "NodeCancelled", "message": "cancelled"}}
|
|
|
|
|
|
def test_deleting_a_flow_ends_its_call_and_drops_a_late_answer(loop, tmp_path):
|
|
"""A node that never returns kept its worker slot past its flow's delete,
|
|
and an answer that came anyway put the deleted flow's value back."""
|
|
pool = PythonWorkerPool(python=sys.executable, size=1)
|
|
controller, store, socket = live_voice(loop, tmp_path, pool)
|
|
worker = socket.worker
|
|
|
|
running = call_in_thread(lambda: controller.run_flow("voice"))
|
|
assert socket.arrived.wait(5)
|
|
call = socket.sent[-1]["call_id"]
|
|
assert socket.sent[-1]["op"] == "run"
|
|
|
|
socket.arrived.clear()
|
|
key = store.delete_flow("voice")
|
|
asyncio.run(controller.reload_flow("voice"))
|
|
controller.forget_flow(key)
|
|
assert socket.arrived.wait(5)
|
|
assert socket.sent[-1] == {"op": "cancel", "call_id": call}
|
|
|
|
worker.deliver({"call_id": call, "ok": True, "result": {"reply": "late"}})
|
|
running.join(timeout=5)
|
|
assert not running.is_alive()
|
|
assert worker.in_flight == 0
|
|
assert "voice.reply" not in controller.state
|
|
|
|
|
|
def test_stopping_a_flow_ends_its_calls_here_and_on_workers(loop, tmp_path):
|
|
"""A stopped flow's calls ran on, here and on the worker, and what they
|
|
returned was written to its state, or, once it was started again, taken
|
|
for a new call's: a live call's id was its node's."""
|
|
pool = PythonWorkerPool(python=sys.executable, size=1)
|
|
pool.start()
|
|
local = "import time\n\n\ndef process():\n time.sleep(60)\n return None\n"
|
|
controller, _, socket = live_voice(loop, tmp_path, pool, local=local)
|
|
worker = socket.worker
|
|
try:
|
|
running = call_in_thread(lambda: controller.run_flow("voice"))
|
|
assert socket.arrived.wait(5)
|
|
old = socket.sent[-1]["call_id"]
|
|
deadline = time.monotonic() + 10
|
|
while ("", "voice.cpu") not in pool._running:
|
|
assert time.monotonic() < deadline, "the local node never started"
|
|
time.sleep(0.01)
|
|
|
|
socket.arrived.clear()
|
|
asyncio.run(controller.set_enabled("voice", "voice", False))
|
|
assert socket.arrived.wait(5)
|
|
assert socket.sent[-1] == {"op": "cancel", "call_id": old}
|
|
deadline = time.monotonic() + 5
|
|
while ("", "voice.cpu") in pool._running:
|
|
assert time.monotonic() < deadline, "the local node was left running"
|
|
time.sleep(0.01)
|
|
|
|
# Started again before the worker has answered the old call.
|
|
asyncio.run(controller.set_enabled("voice", "voice", True))
|
|
socket.arrived.clear()
|
|
gpu = controller.get_node("voice.gpu")
|
|
assert gpu is not None
|
|
again = call_in_thread(lambda: controller.pipeline.run({}, nodes={gpu}))
|
|
assert socket.arrived.wait(5)
|
|
new = socket.sent[-1]["call_id"]
|
|
|
|
worker.deliver({"call_id": old, "ok": True, "result": {"reply": "old"}})
|
|
running.join(timeout=5)
|
|
assert not running.is_alive()
|
|
assert "voice.reply" not in controller.state
|
|
worker.deliver({"call_id": new, "ok": True, "result": {"reply": "new"}})
|
|
again.join(timeout=5)
|
|
assert controller.state.get("voice.reply") == "new"
|
|
finally:
|
|
pool.stop()
|
|
|
|
|
|
def test_a_node_run_by_hand_and_stopped_tells_the_caller_alone(loop, tmp_path):
|
|
"""Stopping a flow failed the node a person had started by hand as an
|
|
error, which the alert rules fire on."""
|
|
pool = PythonWorkerPool(python=sys.executable, size=1)
|
|
events = Events()
|
|
controller, _, socket = live_voice(loop, tmp_path, pool, events=events)
|
|
|
|
answer: list[str] = []
|
|
running = call_in_thread(
|
|
lambda: answer.append(controller.trigger_node("voice.gpu"))
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
call = socket.sent[-1]["call_id"]
|
|
socket.arrived.clear()
|
|
asyncio.run(controller.set_enabled("voice", "voice", False))
|
|
assert socket.arrived.wait(5)
|
|
socket.worker.deliver({"call_id": call, **CANCELLED})
|
|
running.join(timeout=5)
|
|
|
|
assert answer == ["NodeCancelled: cancelled"]
|
|
assert [event for event in events if event["type"] == "node_error"] == []
|
|
|
|
|
|
def test_what_a_call_emits_after_its_flow_stopped_is_dropped(loop, tmp_path):
|
|
"""An emission between a flow's stop and its call's end was still written,
|
|
though what the call returned no longer was."""
|
|
pool = PythonWorkerPool(python=sys.executable, size=1)
|
|
controller, _, socket = live_voice(loop, tmp_path, pool)
|
|
|
|
running = call_in_thread(lambda: controller.run_flow("voice"))
|
|
assert socket.arrived.wait(5)
|
|
call = socket.sent[-1]["call_id"]
|
|
socket.arrived.clear()
|
|
asyncio.run(controller.set_enabled("voice", "voice", False))
|
|
assert socket.arrived.wait(5)
|
|
socket.worker.deliver(
|
|
{"call_id": call, "event": "emit", "outputs": {"reply": "late"}}
|
|
)
|
|
socket.worker.deliver({"call_id": call, **CANCELLED})
|
|
running.join(timeout=5)
|
|
|
|
assert not running.is_alive()
|
|
assert "voice.reply" not in controller.state
|
|
|
|
|
|
def test_stopping_a_node_reaches_its_call_on_a_worker(loop, tmp_path):
|
|
"""The canvas's stop ended only a call in the local pool."""
|
|
pool = PythonWorkerPool(python=sys.executable, size=1)
|
|
controller, _, socket = live_voice(loop, tmp_path, pool)
|
|
|
|
running = call_in_thread(lambda: controller.run_flow("voice"))
|
|
assert socket.arrived.wait(5)
|
|
call = socket.sent[-1]["call_id"]
|
|
socket.arrived.clear()
|
|
assert cancel_node("voice", "gpu", controller).message == "Stopped 'gpu'"
|
|
assert socket.arrived.wait(5)
|
|
assert socket.sent[-1] == {"op": "cancel", "call_id": call}
|
|
|
|
socket.worker.deliver({"call_id": call, **CANCELLED})
|
|
running.join(timeout=5)
|
|
assert not running.is_alive()
|
|
assert socket.worker.in_flight == 0
|
|
|
|
|
|
def test_a_service_and_a_compile_leave_the_only_slot_to_calls(loop):
|
|
"""A service holding the slot for good would leave the worker blocked."""
|
|
hub = RemoteWorkerHub()
|
|
socket = FakeSocket()
|
|
worker = RemoteWorker(
|
|
name="gpu1", labels=["gpu"], send=socket.send_json, loop=loop, max_parallel=1
|
|
)
|
|
hub.attach(worker)
|
|
call_in_thread(
|
|
lambda: worker.request({"op": "serve", "call_id": "svc"}, 0, slot=False)
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
socket.arrived.clear()
|
|
result: dict = {}
|
|
thread = call_in_thread(
|
|
lambda: result.update(
|
|
value=hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5)
|
|
)
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
# Neither waits for the slot the call above holds.
|
|
compiled = call_in_thread(
|
|
lambda: result.update(problem=hub.compile("gpu", "flow", "other", "src"))
|
|
)
|
|
deadline = time.monotonic() + 5
|
|
while len(socket.sent) < 3 and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
for frame in socket.sent[1:]:
|
|
worker.deliver({"call_id": frame["call_id"], "ok": True, "result": None})
|
|
thread.join(timeout=5)
|
|
compiled.join(timeout=5)
|
|
|
|
assert [frame["op"] for frame in socket.sent] == ["serve", "run", "compile"]
|
|
assert result == {"value": None, "problem": None}
|
|
assert worker.in_flight == 1
|
|
|
|
|
|
def test_a_cancel_after_the_loop_closed_leaves_no_coroutine_behind():
|
|
closed = asyncio.new_event_loop()
|
|
closed.close()
|
|
worker = RemoteWorker(
|
|
name="gpu1", labels=[], send=FakeSocket().send_json, loop=closed
|
|
)
|
|
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
worker.cancel("run-a:flow.node")
|
|
gc.collect()
|
|
|
|
assert not [w for w in caught if "never awaited" in str(w.message)]
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Liveness, and what a heartbeat is evidence of
|
|
#
|
|
# The agent beats every ten seconds while it is executing. That says the agent
|
|
# is alive; it says nothing about the node, which is why it bounds the silence
|
|
# deadline and not the node's own timeout.
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_heartbeats_do_not_stave_off_a_nodes_own_timeout(loop):
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
caught: list[Exception] = []
|
|
|
|
def call() -> None:
|
|
try:
|
|
hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0.5)
|
|
except Exception as exc:
|
|
caught.append(exc)
|
|
|
|
thread = call_in_thread(call)
|
|
assert socket.arrived.wait(5)
|
|
call_id = socket.sent[0]["call_id"]
|
|
# Beating faster than the deadline. A node that reports its progress is
|
|
# held off; one whose agent is merely alive is not.
|
|
for _ in range(10):
|
|
worker.deliver({"call_id": call_id, "event": "heartbeat"})
|
|
time.sleep(0.1)
|
|
|
|
thread.join(timeout=5)
|
|
assert isinstance(caught[0], NodeTimeout)
|
|
|
|
|
|
def test_with_no_timeout_total_silence_is_still_bounded(loop, monkeypatch):
|
|
monkeypatch.setattr(remote, "SILENCE_S", 0.3)
|
|
hub = RemoteWorkerHub()
|
|
attach(hub, loop)
|
|
caught: list[Exception] = []
|
|
|
|
def call() -> None:
|
|
try:
|
|
hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0)
|
|
except Exception as exc:
|
|
caught.append(exc)
|
|
|
|
call_in_thread(call).join(timeout=5)
|
|
assert isinstance(caught[0], RemoteError)
|
|
assert not isinstance(caught[0], NodeTimeout)
|
|
assert "presumed gone" in str(caught[0])
|
|
|
|
|
|
def test_with_no_timeout_a_beating_worker_is_left_to_finish(loop, monkeypatch):
|
|
monkeypatch.setattr(remote, "SILENCE_S", 0.3)
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
result: dict = {}
|
|
|
|
thread = call_in_thread(
|
|
lambda: result.update(
|
|
value=hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=0)
|
|
)
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
call_id = socket.sent[0]["call_id"]
|
|
for _ in range(6):
|
|
worker.deliver({"call_id": call_id, "event": "heartbeat"})
|
|
time.sleep(0.1)
|
|
worker.deliver({"call_id": call_id, "ok": True, "result": {"done": True}})
|
|
|
|
thread.join(timeout=5)
|
|
assert result["value"] == {"done": True}
|
|
|
|
|
|
def test_an_allocation_rides_along_with_the_call(loop):
|
|
"""Thread caps and devices reach the worker on the frame that needs them.
|
|
|
|
The worker starts a process per call, so this is the only moment it can
|
|
apply them: a library reads them when it is imported and never again.
|
|
"""
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
|
|
thread = call_in_thread(
|
|
lambda: hub.run(
|
|
"gpu",
|
|
"flow",
|
|
"node",
|
|
"src",
|
|
{},
|
|
"flow.node",
|
|
timeout=5,
|
|
env={"OMP_NUM_THREADS": "4", "CUDA_VISIBLE_DEVICES": "0"},
|
|
)
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
assert socket.sent[0]["env"] == {
|
|
"OMP_NUM_THREADS": "4",
|
|
"CUDA_VISIBLE_DEVICES": "0",
|
|
}
|
|
|
|
worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
|
|
thread.join(timeout=5)
|
|
|
|
|
|
def test_a_call_with_nothing_allocated_carries_no_env(loop):
|
|
hub = RemoteWorkerHub()
|
|
worker, socket = attach(hub, loop)
|
|
|
|
thread = call_in_thread(
|
|
lambda: hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5)
|
|
)
|
|
assert socket.arrived.wait(5)
|
|
assert "env" not in socket.sent[0]
|
|
|
|
worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
|
|
thread.join(timeout=5)
|
|
|
|
|
|
def test_a_worker_that_reports_nothing_is_the_machine_it_used_to_be():
|
|
"""Inventory is read leniently: absent is a default, junk is a default.
|
|
|
|
A worker saying nothing has to keep being scheduled by its label, and one
|
|
that learns to report something new must not need this engine taught about
|
|
it first.
|
|
"""
|
|
assert WorkerInventory.from_hello({}) == WorkerInventory(cpus=1, gpus=0)
|
|
assert WorkerInventory.from_hello(None) == WorkerInventory()
|
|
|
|
read = WorkerInventory.from_hello(
|
|
{"cpus": "8", "gpus": 2, "ram_mb": 64000, "tpus": 4}
|
|
)
|
|
assert (read.cpus, read.gpus, read.ram_mb) == (8, 2, 64000)
|
|
|
|
junk = WorkerInventory.from_hello({"cpus": "many", "gpus": None, "ram_mb": 0})
|
|
assert (junk.cpus, junk.gpus, junk.ram_mb) == (1, 0, None)
|