diff --git a/README.md b/README.md index 96e824c..7287b75 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,9 @@ Main set of features implemented: - Science Agents - Request inference API on runtime - Barrier +- Ex-situ learning (ROSE streaming learner on a second engine) + +Not yet implemented: - Split - Join - Shared SIM / subtasks running on agent, accessible by all investigators @@ -21,12 +24,16 @@ Main set of features implemented: ## Running the unit tests: -1. `pip install .[test,service]` +1. `pip install .[test,service,learn]` 2. `pytest` (or `tox` for all supported interpreters) +The `learn` extra currently needs ROSE from **PR #98** (commit +`64330d9`) -- `StreamingActiveLearner` is not in a release yet, so +`pip install ` at that commit until it merges. + The unit tests start their own stream broker on a random port; no setup. The integration tests under `test/integration` bring up a real ORBIT -broker and rhapsody endpoint and skip themselves when they cannot. +broker and two rhapsody endpoints and skip themselves when they cannot. ## Running the demos: @@ -132,12 +139,62 @@ Three contract notes: JSON-safe or `bytes`** -- ORBIT's rhapsody plugin JSON-encodes results and stringifies anything else. Return plain values from `@flow.function_task` bodies and wrap them in `TypedData` in the - component. + component. (Fixed upstream in radical.orbit `devel` after this was + written: rich results now round-trip by cloudpickle marker. Keep to + plain values until the release you deploy against contains it.) - Persistent components run inline on the service's event loop. Their bodies must be thin async glue publishing through `runtime.stream`, never `@flow.function_task`s (the service warns when it sees one). +### Ex-situ learning: the second engine + +A `StreamingLearnerInvestigator` (`digitaltwin.learn`, needs the `learn` +extra) embeds a ROSE `StreamingActiveLearner` in a model investigator: +the twin's input stream both feeds the learner and is served by the +inference task, and each window of samples retrains the model the +inference task runs with. + +That class is the *only* thing that selects an engine in v1 -- there is +no `engine=` argument. The service recognises it by subclass check and +hands it two engines: its learner tasks run on `'exsitu'`, its inference +stays on `'task'`. + +```python +dt = rt.get_plugin('broker', 'dt', config={'engines': { + 'task': {'endpoint_name': 'dt_task_ep', 'backends': ['concurrent']}, + 'exsitu': {'endpoint_name': 'dt_exsitu_ep', 'backends': ['concurrent']}, +}}) +``` + +`'exsitu'` is optional: left out, it aliases `'task'` and one endpoint +serves both. Both engines are session-shared and built once, in the +background phase of `twin_create`. + +Register the learner's training / active-learning / criterion tasks with +`as_executable=False`. ROSE's default makes them shell commands, and a +command line with local paths does not survive an endpoint that shares +no filesystem with the service; `as_executable=False` sends them as +cloudpickled function tasks instead (the component warns if it finds +executable ones). `test/10-learner/` is a complete worked example. + +### When an endpoint disappears (R8) + +`OrbitExecutionBackend` does not reconnect and components bind their +engine at construction, so a twin whose endpoint went away is stranded +and v1 cannot heal it. It is at least not silent: the plugin watches +the ORBIT topology and marks every twin that bound an engine on a lost +endpoint `failed`, with `engine endpoint lost: ` in +`twin_list`. Twins on surviving engines keep running. + +Recovery means **closing the session**, not just the twins: engines are +session-shared, so a twin created afterwards would inherit the dead one. +The session remembers the loss and refuses to hand that engine out +again, so a `twin_create` after it fails immediately with `engine +'' endpoint was lost; recreate the session` rather than coming up +`ready` and stalling. `unregister_session`, then build the session and +its twins again. + ### Binding policy for the service (R7) The plugin runs its own DT stream broker, embedded, one per plugin and diff --git a/pyproject.toml b/pyproject.toml index e28687a..1ce6652 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,15 @@ dependencies = [ service = [ "radical.orbit>=0.3", ] +# ex-situ learning (`digitaltwin.learn`). A separate extra rather than +# part of `service`: the learner runs perfectly well against a local +# engine with no ORBIT in sight, and a service host that only serves +# in-situ twins should not have to carry ROSE and its dependencies. +# NOTE: `StreamingActiveLearner` lives on ROSE PR #98 (commit 64330d9) +# and is not in any release yet -- install that branch until it merges. +learn = [ + "ROSE>=0.3", +] test = [ "pytest>=8", "pytest-asyncio>=0.24", diff --git a/src/digitaltwin/components.py b/src/digitaltwin/components.py index 71f4275..4741ab4 100644 --- a/src/digitaltwin/components.py +++ b/src/digitaltwin/components.py @@ -223,6 +223,16 @@ async def main_loop(self, runtime, *args, **kwargs) -> TypedData | None: # ------------------------------------------------------------------ + async def _on_stop(self) -> None: + """Internal teardown hook, called by `DTRuntime.stop()` just before + the runtime cancels this component's tasks. + + For components owning machinery the runtime cannot see -- a ROSE + learner loop and the source pumps it spawned, say -- winding it + down here means a cancellation never has to interrupt it + mid-flight. Not user API: best-effort, bounded, must not raise. + """ + class ModelInvestigator(_TwinComponent): """Model-oriented investigation step. ``flow`` is a diff --git a/src/digitaltwin/learn.py b/src/digitaltwin/learn.py new file mode 100644 index 0000000..7a2870b --- /dev/null +++ b/src/digitaltwin/learn.py @@ -0,0 +1,260 @@ +"""Ex-situ learning: a ROSE streaming learner inside a twin component. + +`StreamingLearnerInvestigator` packages the wiring `test/rose_streaming` +spells out by hand -- a `StreamingActiveLearner` fed from `ON_INPUT`, a +bootstrap model published up front, `on_model_ready -> +publish_new_model`, and a learner whose lifetime is the twin's. + +It is also the marker for **dual-engine injection**: the learner's +training / active-learning / criterion tasks run on the `'exsitu'` engine +(typically remote HPC hardware) while inference stays on the twin's +`'task'` engine. The service detects this class by subclass check and +passes the second engine as `learn_flow`; locally the caller passes it +(or nothing -- one engine then serves both, which is what a +single-endpoint deployment does). + +A subclass provides its learner tasks and its inference task:: + + class Fit(StreamingLearnerInvestigator): + + def __init__(self, flow, learn_flow=None): + super().__init__(flow, learn_flow, batch_size=8) + + # ex-situ, on `learn_flow`. as_executable=False makes these + # cloudpickled function tasks, which is what lets them run on + # an endpoint that shares no filesystem with the service + @self.learner.training_task(as_executable=False) + async def training(window, *args): + return {'slope': fit(window)} + + ... + + # in-situ, on `flow` + @flow.function_task + async def predict(in_data, slope=0.0): + return in_data.data * slope + + self.inference_task = ... + +A criterion task takes no dependency from ROSE's streaming loop, so +whatever it scores has to travel *with* it -- the usual pattern is a +dict mirroring the learner's state, captured by value. Mind the cost: +that mirror is re-cloudpickled with the criterion on every window and +retains every state key ever registered, so a model approaching the +~2 MiB return budget crosses the wire twice per window. Keep bulk +artifacts out of learner state and stage them instead. + +This module needs ROSE (`pip install .[learn]`); nothing else in the +package imports it. +""" + +import asyncio +import contextlib +import logging + +from typing import Any, Callable, Optional + +from radical.asyncflow import WorkflowEngine # type: ignore +from rose.al.streaming_learner import StreamingActiveLearner # type: ignore + +from .components import ModelInvestigator, TypedData +from .runtime import RuntimeAPI + +logger = logging.getLogger(__name__) + +# how long twin teardown lets the learner leave its current window before +# the runtime cancels it outright +LEARNER_STOP_TIMEOUT = 5.0 + +# the three ROSE task slots a streaming learner drives +LEARNER_TASKS = ("training", "active_learn", "criterion") + +# ROSE's own per-window bookkeeping -- state, but not model parameters +_ROSE_STATE_KEYS = ("window_size",) + + +class StreamingLearnerInvestigator(ModelInvestigator): + """A `ModelInvestigator` with a ROSE `StreamingActiveLearner` inside. + + Every item the twin routes to this investigator is fed to the learner + (`ON_INPUT`) *and* served by the inference task, so one stream drives + both retraining and prediction. Each window of `batch_size` items (or + `max_wait` seconds' worth) runs one training / active-learning / + criterion iteration on the `'exsitu'` engine; a met criterion is a + publish gate, not a terminator -- it swaps the model the in-situ + inference task runs with. + """ + + def __init__( + self, + flow: WorkflowEngine, + learn_flow: Optional[WorkflowEngine] = None, + batch_size: int = 5, + max_wait: Optional[float] = 2.0, + conflate: bool = True, + ): + super().__init__(flow) + + # Dual engine. `learn_flow` is the 'exsitu' engine the service + # injects; without one (local use, or a deployment that configured + # no 'exsitu' engine) the twin's own engine serves both roles. + self.learn_flow = flow if learn_flow is None else learn_flow + + # conflate: a stream faster than the learner drops its backlog + # rather than growing it -- a days-long twin must not queue days + # of sensor data + self.learner = StreamingActiveLearner( + self.learn_flow, + batch_size=batch_size, + max_wait=max_wait, + conflate=conflate, + ) + + # set by the subclass; the in-situ half of the pair + self.inference_task: Optional[Callable] = None + + self._started = False + self._finished = asyncio.Event() + + # -- what subclasses shape ---------------------------------------------- + + def bootstrap_model(self) -> tuple[dict, dict]: + """The model published before any training has happened. + + Inference gates on a published model, so a learner that published + only from `on_model_ready` would deadlock its twin on the very + first input -- and nothing would ever reach the learner, since the + stream feeds it through that same input. Override to bootstrap + with something better than the inference task's own defaults. + """ + + return {}, {} + + def published_model(self, state: Any) -> tuple[dict, dict]: + """`(model_kwargs, accuracy_kwargs)` for a criterion-met window. + + Whatever the learner's tasks registered as state becomes the model + -- a training task returning a dict has every key of it registered + -- and those kwargs are what the inference task is called with. + ROSE's own per-window bookkeeping is dropped. + """ + + model = { + key: value + for key, value in state.state.items() + if key not in _ROSE_STATE_KEYS + } + + return model, {"metric": state.metric_value} + + def on_window(self, state: Any) -> None: + """Called once per learning window. Default: one log line.""" + + logger.info( + "window %s (%s items): %s=%s published=%s", + state.iteration, + state.window_size, + state.metric_name, + state.metric_value, + state.should_stop, + ) + + # -- the wiring --------------------------------------------------------- + + async def main_loop(self, runtime: RuntimeAPI): + if self.inference_task is None: + raise ValueError( + f"{type(self).__name__} must set self.inference_task -- the" + " in-situ inference, on the twin's 'task' engine" + ) + + self._warn_local_learner_tasks() + + runtime.set_inference_task(self.inference_task) + runtime.subscribe_to_topic(RuntimeAPI.ON_INPUT, self._feed) + runtime.publish_new_model(*self.bootstrap_model()) + + # criterion met => this model is worth serving. In streaming mode + # the criterion is a publish gate and the loop keeps running. + self.learner.on_model_ready( + lambda state: runtime.publish_new_model(*self.published_model(state)) + ) + + self._started = True + + try: + async for state in self.learner.start(): + self.on_window(state) + + finally: + # the failure path too: no learner outlives its twin + self.learner.stop() + self._finished.set() + + async def _feed(self, in_data: TypedData) -> None: + """`ON_INPUT`: everything the twin sees also feeds the learner.""" + + await self.learner.feed(in_data.data) + + async def _on_stop(self) -> None: + """Wind the learner down before the runtime cancels the main loop. + + `learner.stop()` unblocks the window collector, so the loop leaves + its `async for` at a *window boundary* and ROSE's generator runs + its own cleanup. Cancellation alone would mostly work -- the + consumer is usually suspended in `__anext__`, so the generator's + `finally` does run -- but three things only this buys: + + - it does not kill an in-flight `await train_task` on a *shared* + engine, which a cancellation mid-window would; + - ROSE catches `Exception`, not `CancelledError`, so a cancelled + loop records `stop_reason='stream_exhausted'` to its trackers; + - it does not rely on async-generator GC finalization, which is + not something to lean on after days of running. + + Bounded: a learner parked in a remote training task is cancelled + with everything else, a moment later. + """ + + if not self._started: + return + + self.learner.stop() + + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(self._finished.wait(), LEARNER_STOP_TIMEOUT) + + def _warn_local_learner_tasks(self) -> None: + """Warn about learner tasks that cannot leave this host. + + ROSE registers tasks as *executables* by default: the task body + returns a command line, which only runs where that command exists + under that path. The `'exsitu'` engine points at other hardware, + so learner tasks belong on the cloudpickle path -- registered with + `as_executable=False` they travel as function tasks, and the + backend's Python-version guard covers the rest. + + Only when there *is* a separate ex-situ engine: a learner running + both halves on one engine is the local case, where a shell + command with local paths is a perfectly good task. + """ + + if self.learn_flow is self.flow: + return + + local = [ + name + for name in LEARNER_TASKS + if (getattr(self.learner, f"{name}_function", None) or {}).get( + "as_executable" + ) + ] + + if local: + logger.warning( + "%s registered %s as executable task(s): a shell command with" + " local paths does not survive a remote 'exsitu' endpoint." + " Register them with as_executable=False.", + type(self).__name__, + ", ".join(local), + ) diff --git a/src/digitaltwin/runtime.py b/src/digitaltwin/runtime.py index 53430a2..23bb8a1 100644 --- a/src/digitaltwin/runtime.py +++ b/src/digitaltwin/runtime.py @@ -686,6 +686,8 @@ def _teardown_done(self, task: asyncio.Task): logger.error("twin teardown failed: %s", exc, exc_info=exc) async def _teardown(self, timeout: float): + await self._quiesce(timeout) + tasks, self.running_tasks = self.running_tasks, set() for task in tasks: task.cancel() @@ -706,6 +708,42 @@ async def _teardown(self, timeout: float): except Exception as exc: self._record_error(exc) + async def _quiesce(self, timeout: float): + """Let components wind down their own machinery before cancellation. + + One shared budget, spent *concurrently*: two learners waiting five + seconds each must not add up to ten. Whoever ignores the budget + is cancelled like everything else a moment later. + """ + + async def wind_down(component: _TwinComponent): + try: + await component._on_stop() + except Exception as exc: + self._record_error(exc) + + hooks = [wind_down(ant.component) for ant in self._annotated()] + if not hooks: + return + + try: + async with asyncio.timeout(timeout): + await asyncio.gather(*hooks) + + except TimeoutError: + logger.warning("component teardown exceeded %ss", timeout) + + def _annotated(self): + """Every component in the graph, child investigators included.""" + + for ants in list(self.components.values()): + for ant in ants: + yield ant + yield from ant.investigators.values() + + async def _call_await(self, func, *args, **kwargs): + await func(*args, **kwargs) + def _to_asyncio_task(self, func, *args, **kwargs) -> Optional[asyncio.Task]: """Schedule a coroutine as an :class:`asyncio.Task` and track its completion. @@ -745,7 +783,18 @@ def _task_done(self, task: asyncio.Task): if exc is not None: self._record_error(exc) - def _record_error(self, exc: BaseException): + def fail(self, error: str): + """Route an out-of-band failure into the twin state. + + Component failures arrive through the done-callbacks; this is the + door for the ones only the host can see -- a lost engine endpoint + (R8) strands every component bound to that engine, but nothing + inside the runtime notices. + """ + + self._record_error(error) + + def _record_error(self, exc: BaseException | str): """Route a failure into the twin state, and stop the twin. A component failure is a twin failure: the other components have @@ -766,8 +815,12 @@ def _record_error(self, exc: BaseException): half-dead component -- which is logged and dropped. """ - error = f"{type(exc).__name__}: {exc}" - logger.error("twin component failed: %s", error, exc_info=exc) + error = exc if isinstance(exc, str) else f"{type(exc).__name__}: {exc}" + logger.error( + "twin component failed: %s", + error, + exc_info=exc if isinstance(exc, BaseException) else None, + ) if self.state is RuntimeState.FAILED: # the cause is already recorded, and its teardown is running @@ -1257,7 +1310,7 @@ async def _run_component( self._to_asyncio_task(self._call_await, cb, in_data) await i_select.has_published_model.wait() - answer = await i_select.inference_task(in_data, **model_kwargs) + answer = await self._infer(i_select, in_data, model_kwargs) for cb in i_select.subscriptions[RuntimeAPI.ON_FILTERED_OUTPUT]: self._to_asyncio_task(self._call_await, cb, answer) @@ -1279,6 +1332,35 @@ async def _run_component( return answer + async def _infer( + self, ant: _AnnotatedComponent, in_data: TypedData, model_kwargs: dict + ): + """Run an investigator's inference task with its published model. + + A published model is just kwargs, so publishing a key the + inference task does not accept fails as a `TypeError` deep inside + a call the user never wrote -- and for a learner that publishes + whatever its training task returned, that is an easy mistake to + make. Name it instead. + + Only the *call* is rewritten: a `TypeError` raised inside the task + body has its own frame in the traceback and is left alone. + """ + + try: + return await ant.inference_task(in_data, **model_kwargs) + + except TypeError as exc: + traceback = exc.__traceback__ + if traceback is None or traceback.tb_next is not None: + raise + + raise TypeError( + f"published model keys do not match the inference task" + f" signature of {type(ant.component).__name__}: published" + f" {sorted(model_kwargs)} -- {exc}" + ) from exc + ## flow.block async def _dtype_consumer(self, input_data: TypedData) -> None: """Consume data for a given :class:`DataType`. diff --git a/src/digitaltwin/service/client.py b/src/digitaltwin/service/client.py index 8fea2d4..33c0687 100644 --- a/src/digitaltwin/service/client.py +++ b/src/digitaltwin/service/client.py @@ -68,8 +68,14 @@ def register_session( `config` carries the engine configuration and applies at create time only:: - {"engines": {"task": {"endpoint_name": "ep1", - "backends": ["concurrent"]}}} + {"engines": {"task": {"endpoint_name": "ep1", + "backends": ["concurrent"]}, + "exsitu": {"endpoint_name": "hpc1", + "backends": ["concurrent"]}}} + + `'task'` runs the twins' components; `'exsitu'` runs the learner + tasks of a `StreamingLearnerInvestigator` and aliases `'task'` + when it is not configured. Sessions are always persistent (the service forces it), so a `lifetime` / `ttl` argument is accepted and ignored. diff --git a/src/digitaltwin/service/plugin.py b/src/digitaltwin/service/plugin.py index 42f8e84..60146f0 100644 --- a/src/digitaltwin/service/plugin.py +++ b/src/digitaltwin/service/plugin.py @@ -126,8 +126,12 @@ async def register_session(self, request: Request) -> dict: Body (all optional): `{"sid": str, "config": {...}}`. `config` carries the engine configuration and applies at create time only: - {"engines": {"task": {"endpoint_name": "ep1", - "backends": ["concurrent"]}}} + {"engines": {"task": {"endpoint_name": "ep1", + "backends": ["concurrent"]}, + "exsitu": {"endpoint_name": "hpc1", + "backends": ["concurrent"]}}} + + `'exsitu'` is optional: unconfigured, it aliases `'task'`. """ self._ensure_cleanup_task() @@ -258,6 +262,53 @@ async def admin_sessions(self, request: Request) -> dict: }, } + # -- observability ------------------------------------------------------ + + async def on_topology_change(self, participants: dict) -> None: + """Mark twins failed when an engine's endpoint is lost (risk R8). + + `OrbitExecutionBackend` has no reconnect and components bind their + engine at construction, so a twin whose endpoint went away is + stranded and cannot be healed in v1. What it must not be is + *silent*: on a days-long twin an endpoint loss would otherwise + show up as inference calls that simply never return. So this maps + the lost participants onto the sessions' engine endpoints and + turns the affected twins into `failed` + a reason in `twin_list`. + Twins on surviving engines are untouched. Recovery is the + client's, and it is the *session* that has to go: engines are + session-shared and one of them is dead, so a twin created + afterwards would inherit it. `unregister_session`, then build + the session and its twins again. + """ + + await super().on_topology_change(participants) + + lost = { + name + for name, info in (participants or {}).items() + if (info or {}).get("liveness") == "lost" + } + if not lost: + return + + for sid, session in list(self._sessions.items()): + if not isinstance(session, DTSession): + continue + + # one session's bookkeeping must not cost the others their + # notification -- this is the only announcement they get + try: + failed = session.endpoints_lost(lost) + except Exception: + log.exception("[dt] session %s: endpoint loss handling", sid) + continue + + if failed: + log.warning( + "[dt] session %s: endpoint(s) %s lost -- twins failed: %s", + sid, ", ".join(sorted(lost)), ", ".join(failed), + ) + # -- embedded stream broker -------------------------------------------- async def stream_addresses(self) -> tuple[str, str]: diff --git a/src/digitaltwin/service/session.py b/src/digitaltwin/service/session.py index dc186f3..4a69473 100644 --- a/src/digitaltwin/service/session.py +++ b/src/digitaltwin/service/session.py @@ -1,9 +1,11 @@ """Service-side session and twin instances. A `DTSession` belongs to one client and hosts many independent twins. -It owns the session-shared execution engines (M1: exactly one, `'task'`); -each `TwinInstance` owns a `DTRuntime` plus its own namespaced stream -client. Twin teardown never disturbs its siblings or the engines. +It owns the session-shared execution engines, keyed by name: `'task'` +(twin components, typically a co-located endpoint) and `'exsitu'` (ROSE +learners, typically remote HPC hardware). Each `TwinInstance` owns a +`DTRuntime` plus its own namespaced stream client. Twin teardown never +disturbs its siblings or the engines. """ import asyncio @@ -11,6 +13,7 @@ import logging import time +from collections import defaultdict from typing import Any, Optional from fastapi import HTTPException @@ -19,14 +22,27 @@ from rhapsody.backends.execution.orbit import OrbitExecutionBackend # type: ignore from ..components import DataType, TypedData -from ..runtime import DTRuntime +from ..runtime import DTRuntime, RuntimeState from ..streaming import PubSubClient, connect_stream_client from .wire import Package, check_versions, decode, encode log = logging.getLogger("radical.orbit") -# the one engine M1 knows about; M2 adds 'exsitu' as a config addition +try: + from ..learn import StreamingLearnerInvestigator + +except ImportError as _exc: # the 'learn' extra (ROSE) is optional + # logged, not swallowed: without this, a service built *with* the + # extra but with a broken ROSE looks identical to one built without + # it -- the learner twins just quietly get one engine + log.info("[dt] ex-situ learning unavailable (%s); install the 'learn'" + " extra to host StreamingLearnerInvestigator twins", _exc) + StreamingLearnerInvestigator = None + +# every twin component runs on 'task'; only a StreamingLearnerInvestigator +# also gets 'exsitu', and only when the session configured one TASK_ENGINE = "task" +EXSITU_ENGINE = "exsitu" # co-located-demo default -- 'dragon_v3' (the rhapsody default) would # break every demo on a laptop @@ -48,6 +64,9 @@ STATE_FAILED = "failed" STATE_CLOSED = "closed" +# a twin in one of these has nothing left to lose to an endpoint failure +TERMINAL_STATES = (STATE_FAILED, STATE_CLOSED, str(RuntimeState.STOPPED)) + # exactly one of these per twin_call VERBS = ( "add_task", @@ -60,6 +79,20 @@ ) +def _is_learner(cls: Any) -> bool: + """Does this shipped class want the ex-situ engine as well? + + False whenever ROSE is not installed on the service -- such a class + could not have been unpickled here in the first place. + """ + + return ( + StreamingLearnerInvestigator is not None + and isinstance(cls, type) + and issubclass(cls, StreamingLearnerInvestigator) + ) + + def _retrieve_exception(task: asyncio.Task) -> None: """Consume a background task's exception. @@ -87,6 +120,11 @@ def __init__(self, twin_id: str, config: Optional[dict] = None): self.runtime: Optional[DTRuntime] = None self.stream: Optional[PubSubClient] = None + # which session engines this twin's components actually bound to. + # Engines are session-shared, so losing one endpoint must fail the + # twins that use it and leave the others alone (R8). + self.engines: set[str] = {TASK_ENGINE} + self._state = STATE_INITIALIZING self._last_error: Optional[str] = None @@ -126,6 +164,13 @@ def fail(self, error: BaseException | str) -> None: self._last_error = ( error if isinstance(error, str) else f"{type(error).__name__}: {error}" ) + + # from `ready` on the runtime's state machine is the truth, so a + # failure the *service* saw has to be recorded there too or the + # twin keeps reporting `running` (see `state`) + if self.runtime is not None: + self.runtime.fail(self._last_error) + log.error("[dt] twin %s failed: %s", self.twin_id, self._last_error) def track(self, task: asyncio.Task) -> asyncio.Task: @@ -177,10 +222,16 @@ def __init__(self, sid: str, config: Optional[dict] = None): self.twins: dict[str, TwinInstance] = {} self._engines: dict[str, WorkflowEngine] = {} + # engine name -> the endpoint it resolved to; the R8 lookup + self._endpoints: dict[str, str] = {} # in-flight builds, owned by the session rather than by whichever - # twin asked first (see `engine`) + # twin asked first (see `engine`). Both dicts are keyed by engine + # name: a slow 'exsitu' init must not hold up a 'task' build. self._engine_tasks: dict[str, asyncio.Task] = {} - self._engine_lock = asyncio.Lock() + self._engine_locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock) + # endpoints this session's engines ran on and that ORBIT declared + # lost; nothing here is ever cleared (see `endpoints_lost`) + self._lost: set[str] = set() # -- twin lifecycle ----------------------------------------------------- @@ -383,20 +434,66 @@ async def _verb_get_inference( # -- engines ------------------------------------------------------------ + def _engine_config(self, name: str) -> dict: + """The configuration block for engine `name` (possibly empty).""" + + return (self.config.get("engines") or {}).get(name) or {} + + def configured(self, name: str) -> bool: + """Is engine `name` configured for this session? + + An unconfigured engine other than `'task'` is not built: it + aliases `'task'` instead, so adding `'exsitu'` stays a config-only + change and a single-endpoint deployment keeps working unchanged. + """ + + return bool(self._engine_config(name)) + + def _engine_endpoint(self, name: str) -> Optional[str]: + """The endpoint engine `name` runs on. + + The one the backend settled on once it has been built, the + configured one before that (`None` when neither: an auto-selecting + engine that has not resolved yet). + """ + + return self._endpoints.get(name) or self._engine_config(name).get( + "endpoint_name" + ) + async def engine(self, name: str = TASK_ENGINE) -> WorkflowEngine: """The session-shared engine `name`, created on first use. - One engine per name per session, never per twin. M1 only ever - asks for `'task'`. + One engine per name per session, never per twin. An unconfigured + name falls back to `'task'` (see `configured`). The build is a *session-owned* task that callers only ever `shield`-await: a twin whose initialization is cancelled halfway through must not take the build down with it and strand a live `OrbitExecutionBackend` that nothing holds a reference to. A build that lands after the session closed disposes of itself. + + Raises: + RuntimeError: if this engine's endpoint has been lost (R8). """ - async with self._engine_lock: + if not self.configured(name): + name = TASK_ENGINE + + # R8 arrives exactly once, but its consequences do not expire: the + # dead engine stays cached here, so without this a twin created + # *after* the loss would bind it, reach `ready`, and stall in + # silence -- the very failure mode R8 detection exists to remove. + endpoint = self._engine_endpoint(name) + if endpoint in self._lost: + raise RuntimeError( + f"engine {name!r} endpoint was lost ({endpoint});" + f" recreate the session" + ) + + # per-name lock: a 150 s 'exsitu' backend init must not serialize + # ahead of a 'task' build that another twin is waiting on + async with self._engine_locks[name]: flow = self._engines.get(name) if flow is not None: return flow @@ -438,7 +535,7 @@ async def _shutdown_engine(self, name: str, flow: WorkflowEngine) -> None: self.sid, name, exc) async def _create_engine(self, name: str) -> WorkflowEngine: - cfg = (self.config.get("engines") or {}).get(name) or {} + cfg = self._engine_config(name) log.info( "[dt] session %s building engine %r on endpoint %s", @@ -454,8 +551,53 @@ async def _create_engine(self, name: str) -> WorkflowEngine: batch_window=0, # per-call latency beats batching for in-situ ) + # the endpoint the backend *settled on* (it auto-selects when the + # config named none) -- what a topology change is matched against + self._endpoints[name] = ( + getattr(backend, "_endpoint_name", None) or cfg.get("endpoint_name") + ) + return await WorkflowEngine.create(backend=backend) + def endpoints_lost(self, lost: set[str]) -> tuple[str, ...]: + """Fail every twin bound to an engine on a lost endpoint (R8). + + Detection only: `OrbitExecutionBackend` does not reconnect and + components bind their engine at construction, so a stranded twin + cannot be healed -- but a silent stall on a days-long twin is not + acceptable either, so it becomes a `failed` state with a reason in + `twin_list`. Recovery is the client's, and it is the *session* + that has to go: its engines are shared and one of them is dead. + Close the session (`unregister_session`) and build it again. + + The loss is also remembered, because the broker announces it only + once -- see `engine`. + """ + + names = set(self._endpoints) | set(self.config.get("engines") or {}) + gone = { + name: endpoint + for name in names + if (endpoint := self._engine_endpoint(name)) in lost + } + if not gone: + return () + + self._lost.update(gone.values()) + log.warning("[dt] session %s lost endpoint(s) %s -- it must be" + " recreated", self.sid, ", ".join(sorted(gone.values()))) + + failed = [] + for twin in self.twins.values(): + if twin.state in TERMINAL_STATES: + continue + for name in sorted(twin.engines & set(gone)): + twin.fail(f"engine endpoint lost: {gone[name]}") + failed.append(twin.twin_id) + break + + return tuple(failed) + @property def broker_url(self) -> Optional[str]: """Plugin-level broker URL (`None` lets ORBIT resolve it).""" @@ -517,7 +659,16 @@ async def _init_twin(self, twin: TwinInstance) -> None: raise RuntimeError("session is not attached to a dt plugin") async with asyncio.timeout(TWIN_INIT_TIMEOUT): - flow = await self.engine(TASK_ENGINE) + # A configured 'exsitu' engine is built here as well, and + # concurrently: a learner twin must not pay a two-minute + # backend init inside `add_investigator`, which is a short + # verb like every other one. + names = [TASK_ENGINE] + if self.configured(EXSITU_ENGINE): + names.append(EXSITU_ENGINE) + + flow, *_ = await asyncio.gather(*map(self.engine, names)) + pub_addr, sub_addr = await self._plugin.stream_addresses() stream = await connect_stream_client( twin.twin_id, pub_addr, sub_addr, STREAM_CONNECT_TIMEOUT @@ -557,7 +708,13 @@ def _live_twin(self, twin_id: str) -> TwinInstance: def _instantiate( self, package: Any, twin: TwinInstance, is_persistent: bool = False ) -> Any: - """Build a component from a shipped class, injecting the engine. + """Build a component from a shipped class, injecting the engine(s). + + Dual-engine injection lives here: a `StreamingLearnerInvestigator` + additionally receives the `'exsitu'` engine as `learn_flow`, so its + learner tasks run on remote hardware while its inference stays on + `'task'`. The class *is* the marker -- there is no user-facing + engine selector in v1. Also the home of the persistent-component guard: a persistent `main_loop` runs inline on the host loop, so any `function_task` @@ -574,6 +731,15 @@ def _instantiate( ) flow = twin.runtime.flow + extra = {} + + if _is_learner(package.cls): + learn_flow = self._engines.get(EXSITU_ENGINE) + # no 'exsitu' engine configured: one engine serves both roles + extra["learn_flow"] = learn_flow or flow + if learn_flow is not None: + twin.engines.add(EXSITU_ENGINE) + registered = 0 original = flow.function_task @@ -584,7 +750,7 @@ def counting(*args, **kwargs): flow.function_task = counting try: - component = package.instantiate(flow) + component = package.instantiate(flow, **extra) finally: flow.function_task = original diff --git a/src/digitaltwin/service/wire.py b/src/digitaltwin/service/wire.py index 227d2bd..0fb7f87 100644 --- a/src/digitaltwin/service/wire.py +++ b/src/digitaltwin/service/wire.py @@ -46,8 +46,11 @@ class Package: args: tuple = () kwargs: dict = field(default_factory=dict) - def instantiate(self, flow) -> Any: - return self.cls(flow, *self.args, **self.kwargs) + def instantiate(self, flow, **engines: Any) -> Any: + """Instantiate with the twin's engine, plus any extra engine the + service injects by class (`learn_flow` for a streaming learner).""" + + return self.cls(flow, *self.args, **{**self.kwargs, **engines}) def register_user_modules(modules: list) -> None: diff --git a/test/10-learner/README.md b/test/10-learner/README.md new file mode 100644 index 0000000..d45e3f0 --- /dev/null +++ b/test/10-learner/README.md @@ -0,0 +1,95 @@ +# Learner Demo: ex-situ retraining, in-situ inference + +Digital twin application: +- sensor: a stream of raw readings (persistent component, running inline + on the service loop and publishing through its injected stream client) +- calibration learner: a `StreamingLearnerInvestigator` -- a ROSE + `StreamingActiveLearner` embedded in a model investigator +- data sink + +One stream drives both halves of the twin. Every reading feeds the +learner (through `ON_INPUT`) *and* is served by the inference task. +Each window of 8 readings runs one training / active-learning / +criterion iteration; when the resulting calibration beats the criterion +it is published, and the next prediction already uses it. + +## The dual-engine wiring + +This is the demo's point, so it is spelled out rather than defaulted: + +```python +ENGINES = {'engines': { + 'task': {'endpoint_name': TASK_ENDPOINT, 'backends': ['concurrent']}, + 'exsitu': {'endpoint_name': EXSITU_ENDPOINT, 'backends': ['concurrent']}, +}} +``` + +- **`'task'`** runs the twin's components, including the inference task. + It sits in the per-reading critical path, so it belongs on a + co-located endpoint. +- **`'exsitu'`** runs the learner's training, active-learning and + criterion tasks -- the expensive half, typically on remote HPC + hardware. + +There is no `engine=` argument anywhere. The service picks the engines +by *class*: `StreamingLearnerInvestigator` (and only it) is instantiated +with the ex-situ engine as `learn_flow` on top of the usual `flow`, and +`model.py` hands one to the learner and the other to the inference task. +Both engines are session-shared and built once, in the background phase +of `twin_create`. + +`'exsitu'` is optional. Left out of the config it aliases `'task'`, so +this demo also runs against a single endpoint -- the twin is then simply +sharing one endpoint between learning and inference. + +## Why the learner tasks are function tasks + +The learner's three tasks are registered with `as_executable=False`, so +ROSE submits them as **cloudpickled function tasks** instead of shell +commands. A command line with local paths (ROSE's default, and what +`test/rose_streaming` uses) only runs where those paths exist: it does +not survive an `'exsitu'` endpoint that shares no filesystem with the +service. The backend's Python-version guard and the wire's version +stamp cover the rest. + +Two consequences visible in `model.py`: + +- the criterion task gets its model by *value*, cloudpickled with the + task, rather than by reading a `model.json` the training task left + behind -- there is no shared filesystem to leave it on; +- the training task returns its model as a dict. Rich return values + round-trip (ORBIT cloudpickles non-JSON results back), and ROSE + registers a returned dict as learner state, which is what + `published_model()` turns into the inference task's kwargs. + +## Running it + +Four terminals. All of them need `pip install .[service,learn]` and the +ORBIT broker cert/token in `~/.radical/orbit/`. + +```sh +# 1 - the ORBIT broker, hosting the dt plugin +radical-orbit-broker.py --plugins default,dt + +# 2 - the co-located endpoint: the twin's components and inference +RADICAL_ORBIT_RHAPSODY_NOTIFY_WINDOW=0 \ +RADICAL_ORBIT_RHAPSODY_BACKEND=concurrent \ +radical-orbit-endpoint.py -n dt_task_ep + +# 3 - the ex-situ endpoint: the learner's tasks. It keeps the default +# notify window -- 250 ms is noise under a training task +RADICAL_ORBIT_RHAPSODY_BACKEND=concurrent \ +radical-orbit-endpoint.py -n dt_exsitu_ep + +# 4 - the client +cd test/10-learner +DT_TASK_ENDPOINT=dt_task_ep DT_EXSITU_ENDPOINT=dt_exsitu_ep python run_me.py +``` + +## What to watch + +`run_me.py` asks the twin for the same reading (`4.0`) every few +seconds. The uncalibrated bootstrap model answers `0.000`; as windows +are learned the answer climbs to the true calibration +(`2.5 * 4.0 + 1.0 = 11.0`). Nothing about the *request* changes -- only +the model behind it. diff --git a/test/10-learner/data_sink.py b/test/10-learner/data_sink.py new file mode 100644 index 0000000..a376804 --- /dev/null +++ b/test/10-learner/data_sink.py @@ -0,0 +1,12 @@ +import logging + +from digitaltwin.components import UtilityTask + +logger = logging.getLogger(__name__) + + +class MySink(UtilityTask): + """Terminal component: prints what the twin predicted.""" + + async def main_loop(self, runtime, in_data): + print(f"prediction: {in_data.data:.3f}") diff --git a/test/10-learner/dtypes.py b/test/10-learner/dtypes.py new file mode 100644 index 0000000..6adeefa --- /dev/null +++ b/test/10-learner/dtypes.py @@ -0,0 +1,4 @@ +from digitaltwin.components import DataType + +SENSOR_DTYPE = DataType("sensor") +PREDICTION_DTYPE = DataType("prediction") diff --git a/test/10-learner/model.py b/test/10-learner/model.py new file mode 100644 index 0000000..d516586 --- /dev/null +++ b/test/10-learner/model.py @@ -0,0 +1,141 @@ +"""The ex-situ learner: a calibration refitted from the live stream. + +The twin serves `prediction = slope * reading + intercept` in-situ while +a ROSE streaming learner refits `(slope, intercept)` ex-situ from the +same stream. Two engines, stated explicitly: + +- the **learner** tasks (training / active learning / criterion) run on + the `'exsitu'` engine -- typically remote HPC hardware. They are + registered `as_executable=False`, which makes them *cloudpickled + function tasks*: a shell command with local paths would not survive an + endpoint that shares no filesystem with the service. +- the **inference** task runs on the twin's `'task'` engine, co-located + with the service, because it sits in the per-reading critical path. + +`StreamingLearnerInvestigator` is what tells the service to inject both: +it takes the ex-situ engine as `learn_flow`, and passes the ordinary +`flow` to the inference task. +""" + +import logging +import math +import random + +from digitaltwin.components import TypedData +from digitaltwin.learn import StreamingLearnerInvestigator + +from dtypes import PREDICTION_DTYPE + +logger = logging.getLogger(__name__) + +# the calibration the learner has to discover. It lives in the training +# task's world (a stand-in for the reference instrument / simulation a +# real twin would run ex-situ), never in the inference task's. +TRUE_SLOPE = 2.5 +TRUE_INTERCEPT = 1.0 +NOISE = 0.4 + +# how much of the previous window's model a new one keeps. ROSE feeds +# the last active-learning result back in as a training dependency; that +# is the warm start, and it is what makes the fit converge across +# windows instead of hopping around with the noise. +MEMORY = 0.6 + +# readings the criterion scores the model on, and the error at which the +# model is good enough to serve +HOLDOUT = [0.0, 2.5, 5.0, 7.5, 10.0] +PUBLISH_RMSE = 0.25 + + +def fit(xs: list, ys: list) -> tuple: + """Ordinary least squares through `(xs, ys)`.""" + + n = len(xs) + mean_x = sum(xs) / n + mean_y = sum(ys) / n + + cov = sum((x - mean_x) * (y - mean_y) for x, y in zip(xs, ys)) + var = sum((x - mean_x) ** 2 for x in xs) + + slope = cov / var if var else 0.0 + + return slope, mean_y - slope * mean_x + + +def rmse(slope: float, intercept: float, xs: list) -> float: + """Error of a calibration against the truth, on `xs`.""" + + errors = [ + (slope * x + intercept - (TRUE_SLOPE * x + TRUE_INTERCEPT)) ** 2 for x in xs + ] + + return math.sqrt(sum(errors) / len(errors)) + + +class CalibrationLearner(StreamingLearnerInvestigator): + """Refits the sensor calibration from every window of readings.""" + + def __init__(self, flow, learn_flow=None, batch_size: int = 8): + super().__init__(flow, learn_flow, batch_size=batch_size, max_wait=5.0) + + # The criterion task takes no dependency, so the model it scores + # has to travel *with* it: ROSE registers a task's returned dict + # as state, this mirror collects it service-side, and it is + # cloudpickled by value on every submission. On a remote + # endpoint there is no `model.json` to read. + latest: dict = {} + self.learner.on_state_update(latest.__setitem__) + + # -- ex-situ, on `learn_flow` --------------------------------------- + + @self.learner.training_task(as_executable=False) + async def training(window, previous=None, *args): + # labelling the window is the expensive, ex-situ half: this + # stands in for the reference instrument or the simulation + xs = [float(x) for x in window] + ys = [TRUE_SLOPE * x + TRUE_INTERCEPT + random.gauss(0, NOISE) + for x in xs] + + slope, intercept = fit(xs, ys) + + # warm start: `previous` is the last window's model, handed + # back by the active-learning task + if isinstance(previous, dict): + slope = MEMORY * previous["slope"] + (1 - MEMORY) * slope + intercept = MEMORY * previous["intercept"] + (1 - MEMORY) * intercept + + return {"slope": slope, "intercept": intercept} + + @self.learner.active_learn_task(as_executable=False) + async def active_learn(model, *args): + # a real one picks the next samples to label; this one just + # carries the model forward as the next window's warm start + return model + + @self.learner.as_stop_criterion( + metric_name="rmse", + threshold=PUBLISH_RMSE, + operator="<", + as_executable=False, + ) + async def criterion(*args, model=latest, holdout=HOLDOUT): + return rmse(model.get("slope", 0.0), model.get("intercept", 0.0), + holdout) + + # -- in-situ, on `flow` --------------------------------------------- + + @flow.function_task + async def predict(in_data: TypedData, slope=0.0, intercept=0.0): + return slope * in_data.data + intercept + + async def infer(in_data: TypedData, slope=0.0, intercept=0.0): + value = await predict(in_data, slope=slope, intercept=intercept) + return TypedData(PREDICTION_DTYPE, value) + + self.inference_task = infer + + def bootstrap_model(self) -> tuple: + """An uncalibrated sensor: everything reads zero until the first + window has been learned.""" + + return {"slope": 0.0, "intercept": 0.0}, {} diff --git a/test/10-learner/run_me.py b/test/10-learner/run_me.py new file mode 100644 index 0000000..ae650ea --- /dev/null +++ b/test/10-learner/run_me.py @@ -0,0 +1,114 @@ +"""Service demo: a twin that retrains ex-situ while it serves in-situ. + + sensor --> calibration learner --> data sink + +One stream drives both halves. Every reading feeds the ROSE streaming +learner, whose training / active-learning / criterion tasks run on the +`'exsitu'` engine; every reading is also served by the inference task, +which runs on the co-located `'task'` engine. When a window's model +beats the criterion it is published, and the very next prediction uses +it -- which is what the before/after inference below shows. + +See README.md for the services this needs. +""" + +import json +import logging +import os +import time + +from radical.orbit import EndpointRuntime + +from digitaltwin.components import NULL_DTYPE, TRUTHY, TypedData +from digitaltwin.service import register_user_modules + +from dtypes import PREDICTION_DTYPE, SENSOR_DTYPE +from model import CalibrationLearner +from sensor import MySensor +from data_sink import MySink + +# the service has none of this code -- ship it by value +import data_sink +import dtypes +import model +import sensor + +register_user_modules([dtypes, sensor, model, data_sink]) + +logger = logging.getLogger(__name__) + +# where the `dt` plugin is hosted ('broker', or an endpoint name) +DT_HOST = os.environ.get("DT_SERVICE_HOST", "broker") + +# the two endpoints. `DT_TASK_ENDPOINT` unset: ORBIT picks one +# advertising rhapsody. +TASK_ENDPOINT = os.environ.get("DT_TASK_ENDPOINT") or None +EXSITU_ENDPOINT = os.environ.get("DT_EXSITU_ENDPOINT") or None + +# Engine wiring, stated explicitly. 'task' runs the twin's components +# (co-located: it is in the per-reading critical path); 'exsitu' runs the +# learner's tasks (typically remote HPC hardware). One engine of each +# per session, shared by every twin in it. +ENGINES = { + "engines": { + "task": {"endpoint_name": TASK_ENDPOINT, "backends": ["concurrent"]}, + } +} + +# Without `DT_EXSITU_ENDPOINT` the key is left out entirely rather than +# configured with a `None` endpoint: that takes the documented alias path +# ('exsitu' resolves to 'task') instead of quietly building a second +# backend against whichever endpoint ORBIT happens to pick. +if EXSITU_ENDPOINT: + ENGINES["engines"]["exsitu"] = { + "endpoint_name": EXSITU_ENDPOINT, "backends": ["concurrent"] + } + +RUN_TIME = 40.0 +PROBE = 4.0 + + +def main(): + logging.basicConfig(level=logging.INFO) + logging.getLogger("radical.orbit").setLevel(logging.WARNING) + + runtime = EndpointRuntime() + runtime.start(wait=True) + + try: + dt = runtime.get_plugin(DT_HOST, "dt", config=ENGINES) + print(f"session: {dt.sid} (reattach with this sid)") + + twin = dt.create_twin() + print(f"twin: {twin}") + + dt.add_task(twin, dt.package(MySensor), TRUTHY, SENSOR_DTYPE, + is_persistent=True) + dt.add_investigator(twin, dt.package(CalibrationLearner), + SENSOR_DTYPE, PREDICTION_DTYPE) + dt.add_task(twin, dt.package(MySink), PREDICTION_DTYPE, NULL_DTYPE) + + print(json.dumps(dt.describe(twin), indent=2)) + + dt.start(twin) + + # the same reading, over and over: the answer changes only + # because the learner keeps publishing better calibrations + probe = TypedData(SENSOR_DTYPE, 4.0) + deadline = time.time() + RUN_TIME + + while time.time() < deadline: + answer = dt.get_inference(twin, probe, PREDICTION_DTYPE) + print(f"calibrated reading of 4.0: {answer.data:.3f}") + time.sleep(PROBE) + + print(json.dumps(dt.twin(twin), indent=2)) + + dt.twin_close(twin) + + finally: + runtime.stop() + + +if __name__ == "__main__": + main() diff --git a/test/10-learner/sensor.py b/test/10-learner/sensor.py new file mode 100644 index 0000000..0ec4e12 --- /dev/null +++ b/test/10-learner/sensor.py @@ -0,0 +1,29 @@ +import asyncio +import logging +import random + +from digitaltwin.components import UtilityTask + +from dtypes import SENSOR_DTYPE + +logger = logging.getLogger(__name__) + + +class MySensor(UtilityTask): + """Persistent source: a stream of raw readings. + + Plain async code on the service loop, publishing through the injected + stream client -- the persistent-component contract. Each reading is + both a learning sample (it feeds the learner through `ON_INPUT`) and + an inference input. + """ + + def __init__(self, flow, count: int = 200, interval: float = 0.2): + super().__init__(flow) + self.count = count + self.interval = interval + + async def main_loop(self, runtime, in_data): + for _ in range(self.count): + await asyncio.sleep(self.interval) + await runtime.stream.publish(SENSOR_DTYPE, random.uniform(0.0, 10.0)) diff --git a/test/integration/conftest.py b/test/integration/conftest.py index 6c43465..e9844fa 100644 --- a/test/integration/conftest.py +++ b/test/integration/conftest.py @@ -6,12 +6,19 @@ - a co-located rhapsody endpoint with `backends=['concurrent']` and the notification window at 0 (P2 -- otherwise every sequential task pays 250 ms), +- a second rhapsody endpoint standing in for remote HPC hardware, which + is where the `'exsitu'` engine sends learner tasks, - a consumer runtime the tests get `DTClient`s from. +Every endpoint carries a `DT_TEST_ENDPOINT_TAG` in its environment, so a +task can report where it ran -- which is how the dual-engine tests prove +that learning and inference really landed on different endpoints. + Everything is skipped when the broker cannot be started (no certs, no token, port taken), so the suite stays runnable without a deployment. """ +import contextlib import os import shutil import socket @@ -40,17 +47,33 @@ BROKER_URL = f"https://{BROKER_HOST}:{BROKER_PORT}" TASK_ENDPOINT = "dt_test_task_ep" +EXSITU_ENDPOINT = "dt_test_exsitu_ep" # stands in for remote HPC hardware +DOOMED_ENDPOINT = "dt_test_doomed_ep" # started to be killed (R8) DT_ENDPOINT = "dt_test_dt_ep" # endpoint-hosted `dt`, for the smoke test STARTUP_TIMEOUT = 60.0 LOGS = Path(os.environ.get("DT_TEST_LOG_DIR", "/tmp")) / "dt-integration-logs" -# engine wiring every test uses: the co-located endpoint, concurrent backend -ENGINES = { - "engines": { - "task": {"endpoint_name": TASK_ENDPOINT, "backends": ["concurrent"]} + +def engines(**endpoints: str) -> dict: + """Session config for the named engines, concurrent backend each.""" + + return { + "engines": { + name: {"endpoint_name": endpoint, "backends": ["concurrent"]} + for name, endpoint in endpoints.items() + } } -} + + +# the single-engine wiring most tests use: the co-located endpoint only +ENGINES = engines(task=TASK_ENDPOINT) + +# dual-engine wiring: learner tasks ex-situ, everything else co-located +ENGINES_DUAL = engines(task=TASK_ENDPOINT, exsitu=EXSITU_ENDPOINT) + +# same, but with an ex-situ engine on an endpoint the test will kill +ENGINES_DOOMED = engines(task=TASK_ENDPOINT, exsitu=DOOMED_ENDPOINT) def _port_free(port: int) -> bool: @@ -182,27 +205,70 @@ def _await_broker(proc: subprocess.Popen) -> None: pytest.skip(f"broker did not come up -- see {LOGS / 'broker.log'}") -@pytest.fixture(scope="session") -def task_endpoint(broker): - """A co-located rhapsody endpoint: where the twins' tasks execute.""" +@contextlib.contextmanager +def _rhapsody_endpoint(name: str, broker, **env: str): + """A rhapsody endpoint, up and advertised, torn down on exit. + + `DT_TEST_ENDPOINT_TAG` is the endpoint's name to a task running on + it: `os.environ` is the only channel a cloudpickled function body has + for finding out where it landed. + """ argv = _orbit_script("radical-orbit-endpoint.py") + [ - "-n", TASK_ENDPOINT, "-u", broker.url, "-p", "default", + "-n", name, "-u", broker.url, "-p", "default", ] proc = _spawn( - TASK_ENDPOINT, + name, argv, RADICAL_ORBIT_RHAPSODY_BACKEND="concurrent", - RADICAL_ORBIT_RHAPSODY_NOTIFY_WINDOW="0", + DT_TEST_ENDPOINT_TAG=name, + **env, ) try: - _await_plugin(TASK_ENDPOINT, "rhapsody", proc) - yield TASK_ENDPOINT + _await_plugin(name, "rhapsody", proc) + yield proc finally: _terminate(proc) +@pytest.fixture(scope="session") +def task_endpoint(broker): + """A co-located rhapsody endpoint: where the twins' tasks execute.""" + + # notify window 0 (P2): every sequential in-situ prediction would + # otherwise pay 250 ms + with _rhapsody_endpoint(TASK_ENDPOINT, broker, + RADICAL_ORBIT_RHAPSODY_NOTIFY_WINDOW="0"): + yield TASK_ENDPOINT + + +@pytest.fixture(scope="session") +def exsitu_endpoint(broker): + """A second rhapsody endpoint: where learner tasks execute. + + Distinct from `task_endpoint` on purpose -- that is the whole point + of the `'exsitu'` engine. It keeps the default notify window: 250 ms + is noise under a training task. + """ + + with _rhapsody_endpoint(EXSITU_ENDPOINT, broker): + yield EXSITU_ENDPOINT + + +@pytest.fixture +def doomed_endpoint(broker): + """A disposable rhapsody endpoint the test is expected to kill (R8). + + Its own endpoint rather than a shared one: the R8 test asserts on + what an endpoint *loss* does, and the rest of the suite still needs + somewhere to run. + """ + + with _rhapsody_endpoint(DOOMED_ENDPOINT, broker) as proc: + yield proc + + @pytest.fixture(scope="session") def dt_endpoint(broker): """An endpoint hosting the `dt` plugin itself (endpoint-hosted mode). @@ -275,14 +341,32 @@ def runtime(stack): @pytest.fixture -def dt(runtime): - """A `DTClient` on a fresh session, torn down with the test.""" +def dt_client(runtime): + """Factory for `DTClient`s with an explicit engine configuration. + + Every session it hands out is closed with the test -- sessions are + immortal, so nothing else would ever reclaim them. + """ + + clients = [] + + def make(config: dict = ENGINES): + client = runtime.get_plugin("broker", "dt", config=config) + clients.append(client) + return client - client = runtime.get_plugin("broker", "dt", config=ENGINES) try: - yield client + yield make finally: - _drop_session(client) + for client in clients: + _drop_session(client) + + +@pytest.fixture +def dt(dt_client): + """A `DTClient` on a fresh single-engine session.""" + + return dt_client() def _drop_session(client) -> None: diff --git a/test/integration/learner_components.py b/test/integration/learner_components.py new file mode 100644 index 0000000..a94c3b8 --- /dev/null +++ b/test/integration/learner_components.py @@ -0,0 +1,107 @@ +"""The dual-engine twin components the M2 integration tests ship. + +Cloudpickled by value (the service has no copy of this module), so it +must stay importable on its own -- no test imports, no fixtures. + +Separate from `twin_components` because importing it needs ROSE: a +service without the `learn` extra still hosts every M1 twin. +""" + +import os + +from digitaltwin.components import DataType, TypedData +from digitaltwin.learn import StreamingLearnerInvestigator + +SENSOR_DTYPE = DataType("sensor") +INFERENCE_DTYPE = DataType("inference") + +# the calibration the learner has to recover. `CountingSensor` streams +# 0, 1, 2, ... so a least-squares fit through the origin lands on it +# exactly -- no tolerance games in the assertions. +SLOPE = 10.0 + +# batch_size items per window, or `max_wait` seconds' worth +BATCH_SIZE = 4 +MAX_WAIT = 10.0 + + +def _tag() -> str: + """Which endpoint is this task running on? + + `os.environ` is the only channel a cloudpickled function body has for + finding that out; the fixtures stamp every endpoint with its name. + """ + + return os.environ.get("DT_TEST_ENDPOINT_TAG", "?") + + +class LinearLearner(StreamingLearnerInvestigator): + """Fits `y = slope * x` ex-situ and serves it in-situ. + + Both halves report the endpoint they ran on, so a test can assert + that the learner tasks and the inference really went to different + engines. + """ + + def __init__(self, flow, learn_flow=None): + super().__init__(flow, learn_flow, batch_size=BATCH_SIZE, + max_wait=MAX_WAIT) + + # the criterion task takes no dependency, so the model it scores + # travels with it: this mirror is filled service-side from each + # training result and cloudpickled by value on every submission + latest: dict = {} + self.learner.on_state_update(latest.__setitem__) + + # -- ex-situ, on `learn_flow` --------------------------------------- + + @self.learner.training_task(as_executable=False) + async def training(window, *args): + # labelling the window stands in for the simulation an + # ex-situ learner would run + xs = [float(x) for x in window] + ys = [SLOPE * x for x in xs] + den = sum(x * x for x in xs) or 1.0 + + return { + "slope": sum(x * y for x, y in zip(xs, ys)) / den, + "trained_on": _tag(), + } + + @self.learner.active_learn_task(as_executable=False) + async def active_learn(model, *args): + # a non-dict result: nothing here belongs in the model + return len(model) + + @self.learner.as_stop_criterion( + metric_name="fit_error", + threshold=1e-6, + operator="<", + as_executable=False, + ) + async def criterion(*args, model=latest): + error = model.get("slope", 0.0) - SLOPE + + return error * error + + # -- in-situ, on `flow` --------------------------------------------- + + @flow.function_task + async def predict(in_data: TypedData, slope=0.0, trained_on=""): + # a dict return value: rich results round-trip through ORBIT + return { + "value": slope * in_data.data, + "served_by": _tag(), + "trained_on": trained_on, + } + + async def infer(in_data: TypedData, slope=0.0, trained_on=""): + answer = await predict(in_data, slope=slope, trained_on=trained_on) + return TypedData(INFERENCE_DTYPE, answer) + + self.inference_task = infer + + def bootstrap_model(self) -> tuple: + """Nothing learned yet: every reading predicts zero.""" + + return {"slope": 0.0}, {} diff --git a/test/integration/test_dt_learner.py b/test/integration/test_dt_learner.py new file mode 100644 index 0000000..bb1e8bc --- /dev/null +++ b/test/integration/test_dt_learner.py @@ -0,0 +1,211 @@ +"""Integration tests for M2: ex-situ learning on a second endpoint. + +Covers the DTaaS plan's M2 item 11 against a live stack with *distinct* +`'task'` and `'exsitu'` endpoints: a twin whose +`StreamingLearnerInvestigator` retrains on streamed windows ex-situ +while serving inference in-situ, a model update actually propagating to +the next prediction, and (risk R8) an endpoint loss failing exactly the +twins that used it. +""" + +import time +import uuid + +import pytest + +from digitaltwin.components import TRUTHY, TypedData + +from digitaltwin.service import register_user_modules + +import learner_components +import twin_components + +from conftest import ( + DOOMED_ENDPOINT, + ENGINES_DOOMED, + ENGINES_DUAL, + EXSITU_ENDPOINT, + TASK_ENDPOINT, + _terminate, +) +from learner_components import ( + INFERENCE_DTYPE, + SENSOR_DTYPE, + SLOPE, + LinearLearner, +) +from test_dt_service import await_state +from twin_components import CountingSensor, OffsetModel + +pytestmark = pytest.mark.integration + +register_user_modules([learner_components, twin_components]) + +# the learner needs a few windows off the sensor stream, and each window +# is a round trip to a second endpoint +LEARN_TIMEOUT = 120.0 +INFER_TIMEOUT = 120.0 + +# the broker's suspect -> lost grace is 10 s by default +LOST_TIMEOUT = 120.0 + +PROBE = 3.0 + + +def build_learner_twin(dt, twin, interval: float = 0.2): + """sensor -> streaming learner, the standard dual-engine test twin.""" + + dt.create_twin(twin) + dt.add_task(twin, dt.package(CountingSensor, interval=interval), + TRUTHY, SENSOR_DTYPE, is_persistent=True) + dt.add_investigator(twin, dt.package(LinearLearner), SENSOR_DTYPE, + INFERENCE_DTYPE) + dt.start(twin) + + +def infer(dt, twin, value=PROBE): + """One inference through the twin, as the plain dict the task built.""" + + answer = dt.get_inference(twin, TypedData(SENSOR_DTYPE, value), + INFERENCE_DTYPE, timeout=INFER_TIMEOUT) + + return answer.data + + +def await_learned(dt, twin, timeout=LEARN_TIMEOUT): + """Poll inference until the published model is no longer the + bootstrap one -- the twin has to *serve* the update, not merely + compute it.""" + + deadline = time.time() + timeout + + while True: + answer = infer(dt, twin) + if answer["value"]: + return answer + if time.time() > deadline: + pytest.fail(f"twin {twin} never published a learned model: " + f"{dt.twin(twin)}") + time.sleep(1) + + +# --------------------------------------------------------------------------- +# two engines, two endpoints +# --------------------------------------------------------------------------- + +def test_learned_model_propagates_across_two_endpoints( + dt_client, task_endpoint, exsitu_endpoint, twin_id +): + """The M2 acceptance test. + + One stream feeds both halves of the twin: the learner retrains on + windows of it via the `'exsitu'` endpoint, the inference task serves + it from the `'task'` endpoint, and a published model changes what + the *next* prediction answers. + """ + + dt = dt_client(ENGINES_DUAL) + build_learner_twin(dt, twin_id) + + # the bootstrap model: published up front, or the first input would + # deadlock on `has_published_model` and nothing would ever be learned + before = infer(dt, twin_id) + assert before["value"] == 0.0 + assert before["served_by"] == TASK_ENDPOINT + + after = await_learned(dt, twin_id) + + # the same request, a different answer -- only the model changed + assert after["value"] == pytest.approx(SLOPE * PROBE) + + # and the two halves really did run on different endpoints + assert after["served_by"] == TASK_ENDPOINT + assert after["trained_on"] == EXSITU_ENDPOINT + + session = next(s for s in dt.admin_sessions()["sessions"] + if s["sid"] == dt.sid) + assert session["engines"] == ["exsitu", "task"] + + +def test_a_learner_twin_stops_cleanly(dt_client, task_endpoint, + exsitu_endpoint, twin_id): + """The learner's lifetime is the twin's. + + `stop` is terminal and must not hang on a learner parked in a + window, and the twin must not land in `failed` on the way out. + """ + + dt = dt_client(ENGINES_DUAL) + build_learner_twin(dt, twin_id) + await_learned(dt, twin_id) + + t0 = time.time() + assert dt.stop(twin_id) == "stopped" + assert time.time() - t0 < 60, "stop waited for the learner" + + assert dt.twin(twin_id)["last_error"] is None + assert dt.twin_close(twin_id) == "closed" + + +def test_an_unconfigured_exsitu_engine_aliases_task(dt, task_endpoint, + twin_id): + """Adding `'exsitu'` is a config-only change: a single-endpoint + deployment keeps working, with one engine serving both roles.""" + + build_learner_twin(dt, twin_id) + + answer = await_learned(dt, twin_id) + assert answer["served_by"] == TASK_ENDPOINT + assert answer["trained_on"] == TASK_ENDPOINT + + session = next(s for s in dt.admin_sessions()["sessions"] + if s["sid"] == dt.sid) + assert session["engines"] == ["task"] + + +# --------------------------------------------------------------------------- +# R8: a lost endpoint is visible, not silent +# --------------------------------------------------------------------------- + +def test_a_lost_endpoint_fails_only_the_twins_that_used_it( + dt_client, task_endpoint, doomed_endpoint, twin_id +): + """Killing the ex-situ endpoint strands the learner twin -- which + must show up as `failed` with a readable reason, while its + task-only sibling keeps serving.""" + + dt = dt_client(ENGINES_DOOMED) + + plain = str(uuid.uuid4()) + dt.create_twin(plain) + dt.add_investigator(plain, dt.package(OffsetModel, offset=7), + SENSOR_DTYPE, INFERENCE_DTYPE) + dt.start(plain) + + build_learner_twin(dt, twin_id) + await_learned(dt, twin_id) + + _terminate(doomed_endpoint) + + entry = await_state(dt, twin_id, "failed", timeout=LOST_TIMEOUT) + assert entry["last_error"] == f"engine endpoint lost: {DOOMED_ENDPOINT}" + + # the sibling never touched that engine + sibling = dt.twin(plain) + assert sibling["state"] == "running", sibling + assert sibling["last_error"] is None + + answer = dt.get_inference(plain, TypedData(SENSOR_DTYPE, 5), + INFERENCE_DTYPE, timeout=INFER_TIMEOUT) + assert answer.data == 12 + + # The loss is announced once, but the dead engine is still cached in + # the session. A twin created now must not inherit it and come up + # `ready` only to stall: the session is what has to be recreated. + doomed = str(uuid.uuid4()) + with pytest.raises(RuntimeError, match="recreate the session"): + dt.create_twin(doomed) + + entry = dt.twin(doomed) + assert entry["state"] == "failed" + assert "endpoint was lost" in entry["last_error"] diff --git a/test/unit/test_learn.py b/test/unit/test_learn.py new file mode 100644 index 0000000..a79fa71 --- /dev/null +++ b/test/unit/test_learn.py @@ -0,0 +1,336 @@ +"""M2 -- `StreamingLearnerInvestigator`: wiring, propagation, teardown. + +Against a real (local, thread-backed) engine: ROSE typechecks its +`WorkflowEngine` argument, so there is no useful fake to substitute. +""" + +import asyncio +import logging + +from concurrent.futures import ThreadPoolExecutor + +import pytest + +pytest.importorskip("rose") + +from radical.asyncflow import WorkflowEngine # noqa: E402 +from rhapsody.backends import ConcurrentExecutionBackend # noqa: E402 + +from digitaltwin import ( # noqa: E402 + TRUTHY, + DTRuntime, + DataType, + RuntimeState, + TypedData, + UtilityTask, +) +from digitaltwin.learn import StreamingLearnerInvestigator # noqa: E402 + +X = DataType("x") +Y = DataType("y") + +SLOPE = 10.0 +BATCH = 4 + + +@pytest.fixture +async def engines(): + """Factory for local engines, all shut down with the test. + + Threads, not processes: the learner's tasks are closures, and a + process pool would have to pickle them. + """ + + made = [] + + async def make(): + backend = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + made.append(await WorkflowEngine.create(backend=backend)) + return made[-1] + + try: + yield make + finally: + for engine in made: + await engine.shutdown() + + +@pytest.fixture +async def flow(engines): + """The twin's `'task'` engine.""" + + return await engines() + + +class Counter(UtilityTask): + """Persistent source: 0, 1, 2, ... into the twin's stream.""" + + async def main_loop(self, runtime, in_data): + for value in range(1000): + await asyncio.sleep(0.05) + await runtime.stream.publish(X, float(value)) + + +class LinearLearner(StreamingLearnerInvestigator): + """Fits `y = slope * x` on each window; serves `slope * x`.""" + + def __init__(self, flow, learn_flow=None): + super().__init__(flow, learn_flow, batch_size=BATCH, max_wait=5.0) + + latest: dict = {} + self.learner.on_state_update(latest.__setitem__) + + @self.learner.training_task(as_executable=False) + async def training(window, *args): + xs = [float(x) for x in window] + den = sum(x * x for x in xs) or 1.0 + return {"slope": sum(SLOPE * x * x for x in xs) / den} + + @self.learner.active_learn_task(as_executable=False) + async def active_learn(model, *args): + return len(model) + + @self.learner.as_stop_criterion( + metric_name="fit_error", threshold=1e-6, operator="<", + as_executable=False) + async def criterion(*args, model=latest): + return (model.get("slope", 0.0) - SLOPE) ** 2 + + @flow.function_task + async def predict(in_data, slope=0.0): + return slope * in_data.data + + async def infer(in_data, slope=0.0): + return TypedData(Y, await predict(in_data, slope=slope)) + + self.inference_task = infer + + def bootstrap_model(self): + return {"slope": 0.0}, {} + + +async def _twin(flow, stream, learn_flow=None): + """A started twin: counter -> learner.""" + + runtime = DTRuntime(flow, stream) + learner = LinearLearner(flow, learn_flow) + + runtime.add_task(Counter(flow), TRUTHY, X, is_persistent=True) + runtime.add_investigator(learner, X, Y) + runtime.start() + + return runtime, learner + + +async def _await_learned(runtime, timeout=30.0): + """Poll inference until a learned model is being served.""" + + deadline = asyncio.get_running_loop().time() + timeout + + while True: + answer = await runtime.get_inference(TypedData(X, 3.0), Y) + if answer.data: + return answer.data + if asyncio.get_running_loop().time() > deadline: + pytest.fail(f"no model published: {runtime.state} {runtime.last_error}") + await asyncio.sleep(0.25) + + +# --------------------------------------------------------------------------- +# the wiring +# --------------------------------------------------------------------------- + +async def test_a_published_model_changes_the_next_prediction(flow, + stream_clients): + runtime, _ = await _twin(flow, await stream_clients("twin-learn")) + + try: + # the bootstrap model, published before any input -- without it + # inference would block on `has_published_model` and the stream + # that feeds the learner would never get past the first item + first = await runtime.get_inference(TypedData(X, 3.0), Y) + assert first.data == 0.0 + + assert await _await_learned(runtime) == pytest.approx(3.0 * SLOPE) + assert runtime.state is RuntimeState.RUNNING + + finally: + await runtime.stop() + + +async def test_the_learner_uses_the_engine_it_was_given(flow, engines): + """Dual-engine: the learner's tasks go to `learn_flow`, the + inference task stays on `flow`.""" + + exsitu = await engines() + learner = LinearLearner(flow, learn_flow=exsitu) + + assert learner.learn_flow is exsitu + assert learner.learner.asyncflow is exsitu + assert learner.flow is flow + + +async def test_an_absent_exsitu_engine_falls_back_to_the_twins(flow): + learner = LinearLearner(flow) + + assert learner.learn_flow is flow + assert learner.learner.asyncflow is flow + + +# --------------------------------------------------------------------------- +# lifetime +# --------------------------------------------------------------------------- + +async def test_stop_winds_the_learner_down(flow, stream_clients, + no_task_leaks): + """The learner's lifetime is the twin's: `stop()` leaves no loop, no + task and no error behind.""" + + runtime, learner = await _twin(flow, await stream_clients("twin-stop")) + await _await_learned(runtime) + + await runtime.stop() + + assert learner.learner.is_stopped + assert learner._finished.is_set() + assert runtime.state is RuntimeState.STOPPED + assert runtime.last_error is None + + +async def test_stop_before_start_does_not_wait_for_the_learner(flow, + stream_clients): + """A twin torn down before its learner ever ran must not sit out the + stop hook's timeout.""" + + runtime = DTRuntime(flow, await stream_clients("twin-idle")) + runtime.add_investigator(LinearLearner(flow), X, Y) + + loop = asyncio.get_running_loop() + t0 = loop.time() + await runtime.stop() + + assert loop.time() - t0 < 2.0 + + +async def test_a_missing_inference_task_is_a_clear_error(flow, + stream_clients): + class NoInference(StreamingLearnerInvestigator): + pass + + runtime = DTRuntime(flow, await stream_clients("twin-broken")) + runtime.add_investigator(NoInference(flow), X, Y) + runtime.start() + + try: + await asyncio.sleep(0.2) + assert runtime.state is RuntimeState.FAILED + assert "inference_task" in runtime.last_error + + finally: + await runtime.stop() + + +# --------------------------------------------------------------------------- +# the remote-executability guard +# --------------------------------------------------------------------------- + +class Shelling(LinearLearner): + """A learner left on ROSE's executable default.""" + + def __init__(self, flow, learn_flow=None): + super().__init__(flow, learn_flow) + + @self.learner.training_task + async def training(window, *args, task_description={"shell": True}): + return "/bin/true" + + +async def test_executable_learner_tasks_warn(flow, engines, stream_clients, + caplog): + """ROSE's default is a shell command with local paths, which cannot + reach a remote 'exsitu' endpoint.""" + + runtime = DTRuntime(flow, await stream_clients("twin-shell")) + runtime.add_investigator(Shelling(flow, await engines()), X, Y) + + with caplog.at_level(logging.WARNING): + runtime.start() + await asyncio.sleep(0.2) + + assert "as executable task(s)" in caplog.text + assert "training" in caplog.text + + await runtime.stop() + + +async def test_a_purely_local_learner_does_not_warn(flow, stream_clients, + caplog): + """One engine for both halves is the local case, where a shell + command with local paths is a perfectly good task.""" + + runtime = DTRuntime(flow, await stream_clients("twin-local")) + runtime.add_investigator(Shelling(flow), X, Y) + + with caplog.at_level(logging.WARNING): + runtime.start() + await asyncio.sleep(0.2) + + assert "as executable task(s)" not in caplog.text + + await runtime.stop() + + +# --------------------------------------------------------------------------- +# a published model the inference task cannot take +# --------------------------------------------------------------------------- + +async def test_a_model_the_inference_task_rejects_is_named(flow, + stream_clients): + """A learner publishes whatever its training task returned, so a key + the inference task does not accept is an easy mistake -- and it must + not surface as a bare `TypeError` from a call the user never wrote.""" + + class Mismatched(LinearLearner): + def published_model(self, state): + return {"nonesuch": 1}, {} + + runtime = DTRuntime(flow, await stream_clients("twin-mismatch")) + runtime.add_task(Counter(flow), TRUTHY, X, is_persistent=True) + runtime.add_investigator(Mismatched(flow), X, Y) + runtime.start() + + try: + for _ in range(120): + await asyncio.sleep(0.25) + if runtime.state is RuntimeState.FAILED: + break + + assert runtime.state is RuntimeState.FAILED + assert "published model keys do not match" in runtime.last_error + assert "nonesuch" in runtime.last_error + + finally: + await runtime.stop() + + +async def test_the_task_bodys_own_TypeError_is_left_alone(flow, + stream_clients): + class Exploding(LinearLearner): + def __init__(self, flow, learn_flow=None): + super().__init__(flow, learn_flow) + + async def infer(in_data, slope=0.0): + raise TypeError("the body's own complaint") + + self.inference_task = infer + + runtime = DTRuntime(flow, await stream_clients("twin-boom")) + runtime.add_investigator(Exploding(flow), X, Y) + runtime.start() + + try: + with pytest.raises(TypeError, match="the body's own complaint"): + await runtime.get_inference(TypedData(X, 1.0), Y) + + finally: + await runtime.stop() diff --git a/test/unit/test_service.py b/test/unit/test_service.py index 8fb8908..611cbd5 100644 --- a/test/unit/test_service.py +++ b/test/unit/test_service.py @@ -17,6 +17,7 @@ from starlette.testclient import TestClient # noqa: E402 from digitaltwin.components import UtilityTask # noqa: E402 +from digitaltwin.runtime import DTRuntime # noqa: E402 from digitaltwin.service.plugin import PluginDT # noqa: E402 from digitaltwin.service.session import DTSession, TwinInstance # noqa: E402 from digitaltwin.service.wire import ( # noqa: E402 @@ -275,6 +276,13 @@ class _Plain(UtilityTask): pass +class _FakeStream: + """Just enough stream client for a `DTRuntime` that never streams.""" + + namespace = "twin" + on_error = None + + def _twin_with(flow): twin = TwinInstance("t1") twin.runtime = type("R", (), {"flow": flow})() @@ -401,6 +409,179 @@ async def test_close_shuts_down_a_built_engine(): assert session._engines == {} +# --------------------------------------------------------------------------- +# the 'exsitu' engine (M2) +# --------------------------------------------------------------------------- + +def _dual(**endpoints: str) -> dict: + return {"engines": {name: {"endpoint_name": endpoint} + for name, endpoint in endpoints.items()}} + + +async def test_an_unconfigured_engine_aliases_task(): + """Adding `'exsitu'` must stay a config-only change: without one, a + learner twin runs both halves on the twin's own engine.""" + + session = DTSession("s1", _dual(task="ep1")) + built = _slow_build(session, 0) + + assert await session.engine("exsitu") is await session.engine("task") + assert len(built) == 1 + assert sorted(session._engines) == ["task"] + + +async def test_a_slow_exsitu_build_does_not_hold_up_task(): + """Per-name build tasks and locks: a remote backend taking minutes + must not serialize ahead of the engine a sibling twin needs.""" + + session = DTSession("s1", _dual(task="ep1", exsitu="hpc1")) + delays = {"task": 0.0, "exsitu": 5.0} + + async def build(name): + await asyncio.sleep(delays[name]) + return _FakeFlow() + + session._create_engine = build + + slow = asyncio.create_task(session.engine("exsitu")) + await asyncio.sleep(0.05) + + assert await asyncio.wait_for(session.engine("task"), 1.0) + + slow.cancel() + + +async def test_a_learner_gets_the_exsitu_engine(): + """Dual-engine injection, by subclass check -- there is no + user-facing engine selector in v1.""" + + learn = pytest.importorskip("digitaltwin.learn") + + class _Learner(learn.StreamingLearnerInvestigator): + def __init__(self, flow, learn_flow=None): + self.flow, self.learn_flow = flow, learn_flow + + session = DTSession("s1", _dual(task="ep1", exsitu="hpc1")) + exsitu = session._engines["exsitu"] = _FakeFlow() + + twin = _twin_with(_FakeFlow()) + component = session._instantiate(Package(_Learner), twin) + + assert component.learn_flow is exsitu + assert component.flow is twin.runtime.flow + assert twin.engines == {"task", "exsitu"} + + +async def test_a_plain_component_never_sees_a_second_engine(): + session = DTSession("s1", _dual(task="ep1", exsitu="hpc1")) + session._engines["exsitu"] = _FakeFlow() + + twin = _twin_with(_FakeFlow()) + session._instantiate(Package(_Plain), twin) + + assert twin.engines == {"task"} + + +# --------------------------------------------------------------------------- +# R8: a lost endpoint is visible, not silent +# --------------------------------------------------------------------------- + +def _running_twin(session, twin_id, *engines: str): + twin = session.twins[twin_id] = TwinInstance(twin_id) + twin.engines.update(engines) + twin.runtime = DTRuntime(_FakeFlow(), _FakeStream()) + twin.runtime.start() + + return twin + + +async def test_a_lost_endpoint_fails_only_the_twins_that_used_it(): + session = DTSession("s1") + session._endpoints = {"task": "ep1", "exsitu": "hpc1"} + + learner = _running_twin(session, "learner", "exsitu") + plain = _running_twin(session, "plain") + + assert session.endpoints_lost({"hpc1"}) == ("learner",) + + assert learner.state == "failed" + assert learner.last_error == "engine endpoint lost: hpc1" + + # a session-shared engine the twin never bound to is not its problem + assert plain.state == "running" + assert plain.last_error is None + + +async def test_a_lost_endpoint_is_never_handed_out_again(): + """R8 is announced once, but the dead engine stays cached: a twin + created afterwards must fail fast instead of binding it and + stalling.""" + + session = DTSession("s1", _dual(task="ep1", exsitu="hpc1")) + _slow_build(session, 0) + + task_flow = await session.engine("task") + await session.engine("exsitu") + session._endpoints = {"task": "ep1", "exsitu": "hpc1"} + + session.endpoints_lost({"hpc1"}) + + with pytest.raises(RuntimeError, match="recreate the session"): + await session.engine("exsitu") + + # the surviving engine is still handed out + assert await session.engine("task") is task_flow + + +async def test_a_loss_before_the_build_is_remembered_too(): + """The engine need not have been built for its endpoint to be + known: the configuration named it.""" + + session = DTSession("s1", _dual(task="ep1", exsitu="hpc1")) + _slow_build(session, 0) + + session.endpoints_lost({"hpc1"}) + + with pytest.raises(RuntimeError, match="recreate the session"): + await session.engine("exsitu") + + +async def test_a_surviving_topology_change_fails_nothing(): + session = DTSession("s1") + session._endpoints = {"task": "ep1"} + twin = _running_twin(session, "t1") + + assert session.endpoints_lost({"someone-else"}) == () + assert twin.state == "running" + + +async def test_the_plugin_routes_lost_participants_to_its_sessions(plugin): + session = plugin._sessions["s1"] = DTSession("s1") + session._endpoints = {"task": "ep1"} + twin = _running_twin(session, "t1") + + await plugin.on_topology_change({ + "ep1": {"liveness": "lost"}, + "ep2": {"liveness": "present"}, + }) + + assert twin.state == "failed" + assert twin.last_error == "engine endpoint lost: ep1" + + +async def test_a_suspect_participant_is_not_a_loss(plugin): + """A transient blip reaches `suspect` at most and must never fail a + twin -- the broker's grace timer decides.""" + + session = plugin._sessions["s1"] = DTSession("s1") + session._endpoints = {"task": "ep1"} + twin = _running_twin(session, "t1") + + await plugin.on_topology_change({"ep1": {"liveness": "suspect"}}) + + assert twin.state == "running" + + # --------------------------------------------------------------------------- # the embedded stream broker and its supervisor # ---------------------------------------------------------------------------