From 7f0ecb32f3b9f1e02e4dffed356827d06382b7e8 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 26 Aug 2026 06:50:03 +0200 Subject: [PATCH 01/22] adding stream version of m3dc1 --- use-cases/m3dc1-stream/amsc_stream.py | 330 ++++++++++++++++++++++++ use-cases/m3dc1-stream/sensor.py | 99 +++++++ use-cases/m3dc1-stream/sensor_buffer.py | 122 +++++++++ 3 files changed, 551 insertions(+) create mode 100644 use-cases/m3dc1-stream/amsc_stream.py create mode 100644 use-cases/m3dc1-stream/sensor.py create mode 100644 use-cases/m3dc1-stream/sensor_buffer.py diff --git a/use-cases/m3dc1-stream/amsc_stream.py b/use-cases/m3dc1-stream/amsc_stream.py new file mode 100644 index 0000000..e72c7d0 --- /dev/null +++ b/use-cases/m3dc1-stream/amsc_stream.py @@ -0,0 +1,330 @@ +#!/usr/bin/env python3 +""" +M3DC1 streaming surrogate — sensor-based input version. + +Replaces the predefined dataset pool from amsc.py with a continuous sensor +stream (mocked by MockM3DC1Sensor). The simulation task no longer slices +from a pre-loaded PKL; instead it calls buffer.wait_for(n) and consumes +however many observations the sensor has delivered so far. + +Architecture +──────────── + SensorStream (sensor.py) ← implement this for a real sensor + │ + SensorBuffer (sensor_buffer.py) ← background asyncio task, accumulates rows + │ + simulation task ← waits for N rows, writes parquet + │ + training task ← fits sklearn surrogate on that parquet + │ + active_learn task ← records AL decision + │ + stop_on_r2 criterion ← stops when val_r2 ≥ threshold + +To swap in a real sensor: + class MyRealSensor(SensorStream): + async def read_one(self) -> dict[str, float]: + ... # read from your hardware / REST API / Kafka / etc. + + sensor = MyRealSensor() + # everything else in this file stays the same +""" +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +import warnings +from pathlib import Path + +import numpy as np +import rhapsody + +_HERE = Path(__file__).resolve().parent +if str(_HERE) not in sys.path: + sys.path.insert(0, str(_HERE)) + +from sensor import MockM3DC1Sensor, SensorStream # noqa: E402 +from sensor_buffer import SensorBuffer # noqa: E402 + +from radical.asyncflow import WorkflowEngine # noqa: E402 +from rose.al import ParallelActiveLearner # noqa: E402 +from rose.learner import LearnerConfig, TaskConfig # noqa: E402 + +# ── Row-budget policy (mirrors amsc.py's growing-pool logic) ───────────────── +# At iteration i, the simulation task requests N_BASE + i * N_STEP rows from +# the buffer. Increase N_STEP to consume more data per iteration. +N_BASE = 100 +N_STEP = 100 + +# Workspace for iteration artefacts (parquet snapshots, metric JSON) +_WORKSPACE = _HERE / "workspace" + + +def _workspace_iter(iteration: int, label: str) -> Path: + d = _WORKSPACE / f"{label}" / f"iter_{iteration:03d}" + d.mkdir(parents=True, exist_ok=True) + return d + + +def _candidate_configs(candidates: list[str], max_iter: int) -> list[LearnerConfig]: + configs = [] + for idx, family in enumerate(candidates): + label = f"{idx}_{family}" + kwargs = {"learner_label": label, "model_family": family} + schedule = {i: TaskConfig(kwargs={**kwargs, "iteration": i}) for i in range(max_iter + 1)} + schedule[-1] = TaskConfig(kwargs={**kwargs, "iteration": max_iter}) + configs.append( + LearnerConfig( + simulation=schedule, + training=schedule, + active_learn=schedule, + criterion=schedule, + ) + ) + return configs + + +# ── ROSE workflow ───────────────────────────────────────────────────────────── + +async def run_rose_workflow( + buffer: SensorBuffer, + *, + candidates: list[str], + max_iter: int, + r2_threshold: float, +) -> None: + engine = await rhapsody.get_backend("concurrent") + asyncflow = await WorkflowEngine.create(engine) + learner = ParallelActiveLearner(asyncflow) + + # ── Simulation task ─────────────────────────────────────────────────────── + # KEY CHANGE vs amsc.py: reads from sensor buffer, not from the PKL pool. + # The buffer is captured by closure; swapping the sensor changes nothing here. + @learner.simulation_task(as_executable=False) + async def simulation(*args, **kwargs) -> dict: + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + + n_rows = N_BASE + it * N_STEP + out_dir = _workspace_iter(it, label) + parquet = out_dir / "sensor_snapshot.parquet" + + print(f" [sim {label} iter={it}] waiting for {n_rows} sensor rows …", flush=True) + n_written = await buffer.write_snapshot(n_rows, parquet, timeout=600.0) + + meta = { + "iteration" : it, + "learner_label": label, + "dataset" : str(parquet), + "n_rows" : n_written, + "source" : "sensor_stream", + } + (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) + return meta + + # ── Training task ───────────────────────────────────────────────────────── + # Fits a surrogate model locally using sklearn. + # Replace with subprocess to surge_train.py if running on HPC. + @learner.training_task(as_executable=False) + async def training(sim_result: dict, **kwargs) -> dict: + import pandas as pd + from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor + from sklearn.linear_model import Ridge + from sklearn.metrics import mean_squared_error, r2_score + from sklearn.model_selection import train_test_split + from sklearn.neural_network import MLPRegressor + from sklearn.preprocessing import StandardScaler + from sklearn.pipeline import Pipeline + + it = int(kwargs.get("iteration", sim_result["iteration"])) + label = str(kwargs["learner_label"]) + family = str(kwargs.get("model_family", "rf")) + + df = pd.read_parquet(sim_result["dataset"]) + X = df.drop(columns=["output_gamma"]).values + y = df["output_gamma"].values + + X_train, X_val, y_train, y_val = train_test_split( + X, y, test_size=0.2, random_state=it + ) + + _models = { + "rf": RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=-1), + "mlp": Pipeline([ + ("scaler", StandardScaler()), + ("net", MLPRegressor(hidden_layer_sizes=(64, 64), max_iter=500, random_state=42)), + ]), + "gbr": GradientBoostingRegressor(n_estimators=100, random_state=42), + "ridge": Pipeline([ + ("scaler", StandardScaler()), + ("reg", Ridge()), + ]), + } + model = _models.get(family, _models["rf"]) + model.fit(X_train, y_train) + + y_pred = model.predict(X_val) + val_r2 = float(r2_score(y_val, y_pred)) + val_rmse = float(np.sqrt(mean_squared_error(y_val, y_pred))) + + metrics = { + "iteration" : it, + "learner_label": label, + "model_family" : family, + "val_r2" : val_r2, + "val_rmse" : val_rmse, + "n_train" : int(len(X_train)), + "n_val" : int(len(X_val)), + } + out_dir = _workspace_iter(it, label) + (out_dir / "surge_metrics.json").write_text(json.dumps(metrics, indent=2)) + return {"simulation": sim_result, "surge": metrics} + + # ── Active-learning task ────────────────────────────────────────────────── + @learner.active_learn_task(as_executable=False) + async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: + it = int(kwargs.get("iteration", train_bundle["simulation"]["iteration"])) + label = str(kwargs["learner_label"]) + surge = train_bundle["surge"] + + decision = { + "iteration" : it, + "learner_label": label, + "policy" : "monitor_best_val_r2", + "val_r2" : surge["val_r2"], + "val_rmse" : surge["val_rmse"], + } + out_dir = _workspace_iter(it, label) + (out_dir / "active.json").write_text(json.dumps(decision, indent=2)) + return { + "iteration" : it, + "learner_label": label, + "train" : train_bundle, + "val_r2" : surge["val_r2"], + "val_rmse" : surge["val_rmse"], + } + + # ── Stop criterion ──────────────────────────────────────────────────────── + @learner.as_stop_criterion( + metric_name="val_r2", + threshold=r2_threshold, + operator=">=", + as_executable=False, + ) + async def stop_on_r2(*args, **kwargs) -> float: + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + path = _workspace_iter(it, label) / "surge_metrics.json" + meta = json.loads(path.read_text()) + r2 = float(meta["val_r2"]) + forced = it >= max_iter - 1 and r2 < r2_threshold + return r2_threshold if forced else r2 + + # ── Run ─────────────────────────────────────────────────────────────────── + configs = _candidate_configs(candidates, max_iter) + rows: list[dict] = [] + + try: + async for state in learner.start( + parallel_learners=len(candidates), + max_iter=max_iter, + learner_configs=configs, + ): + label = candidates[int(state.learner_id)] + rows.append({ + "learner" : label, + "iter" : state.iteration, + "val_r2" : state.val_r2, + "val_rmse": state.val_rmse, + }) + print( + f" learner={label} iter={state.iteration}" + f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}" + f" buffer={len(buffer)} obs", + flush=True, + ) + if len(rows) >= len(candidates) * max_iter: + learner.stop() + break + finally: + await asyncflow.shutdown() + + rows.sort(key=lambda r: float(r["val_r2"]), reverse=True) + print("\n── Summary ──────────────────────────────────────────────────────") + print(f"{'rank':>4} {'learner':<8} {'val_r2':>9} {'val_rmse':>10}") + for rank, row in enumerate(rows, 1): + print( + f"{rank:>4} {row['learner']:<8} " + f"{float(row['val_r2']):>9.5f} {float(row['val_rmse']):>10.6f}" + ) + print(f"Workspace: {_WORKSPACE}") + + +# ── Entry point ─────────────────────────────────────────────────────────────── + +def main() -> None: + warnings.filterwarnings("ignore", category=UserWarning) + + parser = argparse.ArgumentParser( + description="M3DC1 streaming surrogate — sensor input version." + ) + parser.add_argument( + "--candidates", + default="rf,mlp", + help="Comma-separated model families: rf, mlp, gbr, ridge.", + ) + parser.add_argument("--max-iter", type=int, default=3) + parser.add_argument("--r2-threshold", type=float, default=0.80) + parser.add_argument( + "--sensor-rate", + type=float, + default=10.0, + help="Mock sensor emission rate in observations/second (default: 10).", + ) + parser.add_argument( + "--sensor-seed", + type=int, + default=42, + help="RNG seed for the mock sensor.", + ) + parser.add_argument( + "--buffer-maxlen", + type=int, + default=10_000, + help="Maximum observations retained in the sensor buffer.", + ) + args = parser.parse_args() + candidates = [x.strip() for x in args.candidates.split(",") if x.strip()] + + if len(candidates) < 2: + parser.error("Need at least two candidates for ParallelActiveLearner.") + + sensor = MockM3DC1Sensor(rate_hz=args.sensor_rate, seed=args.sensor_seed) + buffer = SensorBuffer(sensor, maxlen=args.buffer_maxlen) + + async def _main() -> None: + await buffer.start() + print( + f"Sensor stream started rate={args.sensor_rate} Hz" + f" candidates={candidates} max_iter={args.max_iter}", + flush=True, + ) + try: + await run_rose_workflow( + buffer, + candidates=candidates, + max_iter=args.max_iter, + r2_threshold=args.r2_threshold, + ) + finally: + await buffer.stop() + print("Sensor stream stopped.", flush=True) + + asyncio.run(_main()) + + +if __name__ == "__main__": + main() diff --git a/use-cases/m3dc1-stream/sensor.py b/use-cases/m3dc1-stream/sensor.py new file mode 100644 index 0000000..eb09069 --- /dev/null +++ b/use-cases/m3dc1-stream/sensor.py @@ -0,0 +1,99 @@ +""" +Sensor stream interface and mock implementation for M3DC1 streaming workflows. + +To plug in a real sensor, subclass SensorStream and implement read_one(). +Everything else (SensorBuffer, amsc_stream.py) works unchanged. +""" +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator + +import numpy as np + + +class SensorStream(ABC): + """Abstract sensor that emits one M3DC1 observation at a time. + + Implement read_one() for any real sensor, REST endpoint, Kafka consumer, + shared-memory ring buffer, etc. The default stream() calls read_one() in + a loop; override it for push-based sources where data arrives on its own. + """ + + @abstractmethod + async def read_one(self) -> dict[str, float]: + """Return a single observation as a column-name → value mapping.""" + + async def stream(self) -> AsyncIterator[dict[str, float]]: + """Yield observations indefinitely. Override for push-based sensors.""" + while True: + yield await self.read_one() + + +# ── Physical parameter ranges for SPARC M3DC1 D1 ──────────────────────────── +# Derived from sparc_m3dc1_D1_metadata.yaml (inputs + output_gamma). +# Replace bounds with real calibration data when integrating a live source. +_M3DC1_RANGES: dict[str, tuple[float, float]] = { + "input_batemanscale": (0.5, 2.0), + "input_ntor": (1.0, 15.0), # toroidal mode number + "input_pscale": (0.5, 2.0), + "eq_a": (0.1, 0.5), # minor radius [m] + "eq_R0": (1.5, 6.0), # major radius [m] + "eq_kappa": (1.0, 2.5), # elongation + "eq_delta": (0.0, 0.8), # triangularity + "eq_simag": (-2.0, 0.0), # poloidal flux at magnetic axis + "eq_sibry": (-5.0, -0.5), # poloidal flux at boundary + "eq_current": (0.5, 15.0), # plasma current [MA] + "q0": (0.8, 2.5), # safety factor on axis + "q95": (3.0, 8.0), # safety factor at 95 % flux + "p0": (1e4, 1e6), # peak pressure [Pa] +} + +COLUMNS: list[str] = list(_M3DC1_RANGES) + ["output_gamma"] + + +class MockM3DC1Sensor(SensorStream): + """Simulates a real-time M3DC1 physics sensor at a configurable rate. + + Observations are drawn from the realistic parameter ranges in _M3DC1_RANGES. + output_gamma is a nonlinear surrogate of the MHD stability growth rate plus + Gaussian noise — non-trivial enough to make the surrogate task meaningful. + + Args: + rate_hz: Target emission rate in observations per second. + seed: RNG seed for reproducibility. + noise_std: Std-dev of Gaussian noise on output_gamma. + """ + + def __init__( + self, + rate_hz: float = 2.0, + seed: int = 42, + noise_std: float = 0.005, + ) -> None: + self._delay = 1.0 / max(rate_hz, 1e-6) + self._rng = np.random.default_rng(seed) + self._noise_std = noise_std + + async def read_one(self) -> dict[str, float]: + await asyncio.sleep(self._delay) + return self._sample() + + def _sample(self) -> dict[str, float]: + rng = self._rng + obs: dict[str, float] = { + col: float(rng.uniform(lo, hi)) + for col, (lo, hi) in _M3DC1_RANGES.items() + } + # Surrogate physics: gamma grows with mode number and pressure scale, + # falls with safety factor — rough but nonlinear enough for surrogates. + obs["output_gamma"] = float(max( + 0.0, + 0.08 * obs["input_batemanscale"] * (obs["input_ntor"] / 8.0) + + 0.06 * obs["input_pscale"] / max(obs["q95"], 0.1) + + 0.04 * obs["eq_kappa"] * obs["eq_delta"] + - 0.02 * obs["q0"] + + float(rng.normal(0.0, self._noise_std)), + )) + return obs diff --git a/use-cases/m3dc1-stream/sensor_buffer.py b/use-cases/m3dc1-stream/sensor_buffer.py new file mode 100644 index 0000000..676ce79 --- /dev/null +++ b/use-cases/m3dc1-stream/sensor_buffer.py @@ -0,0 +1,122 @@ +""" +SensorBuffer: accumulates observations from a SensorStream into a bounded +in-memory deque and surfaces snapshots as pandas DataFrames. + +The buffer runs a background asyncio task that calls sensor.stream() and +appends each observation as it arrives. Workflow tasks call wait_for(n) to +block until enough data is available, then take a snapshot for training. +""" +from __future__ import annotations + +import asyncio +from collections import deque +from pathlib import Path + +import pandas as pd + +from sensor import SensorStream + + +class SensorBuffer: + """Thread-of-execution-safe accumulator for a continuous sensor stream. + + Args: + sensor: Any SensorStream implementation (mock or real). + maxlen: Maximum observations to retain (oldest discarded when full). + Set to None for an unbounded buffer. + """ + + def __init__(self, sensor: SensorStream, maxlen: int | None = 10_000) -> None: + self._sensor = sensor + self._deque: deque[dict] = deque(maxlen=maxlen) + self._event: asyncio.Event | None = None + self._loop: asyncio.AbstractEventLoop | None = None # loop that owns the Event + self._task: asyncio.Task | None = None + + # ── Lifecycle ──────────────────────────────────────────────────────────── + + async def start(self) -> None: + """Start the background ingestion task.""" + self._loop = asyncio.get_running_loop() + self._event = asyncio.Event() + self._task = asyncio.create_task(self._ingest(), name="sensor-ingest") + + async def stop(self) -> None: + """Cancel the background ingestion task and wait for it to exit.""" + if self._task is not None: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + # ── Internal ───────────────────────────────────────────────────────────── + + async def _ingest(self) -> None: + async for obs in self._sensor.stream(): + self._deque.append(obs) + self._event.set() # wake any waiter; they will clear it themselves + + # ── Public API ─────────────────────────────────────────────────────────── + + def __len__(self) -> int: + return len(self._deque) + + def snapshot(self, n: int | None = None) -> pd.DataFrame: + """Return the most recent n observations as a DataFrame (no blocking).""" + rows = list(self._deque) + if n is not None: + rows = rows[-n:] + return pd.DataFrame(rows) + + async def wait_for(self, n: int, timeout: float | None = None) -> pd.DataFrame: + """Block until at least n observations are buffered, then return them. + + Args: + n: Minimum number of observations required. + timeout: Seconds to wait before raising TimeoutError (None = forever). + + Returns: + DataFrame of the most recent n observations. + + Raises: + TimeoutError: If timeout elapses before n observations arrive. + """ + async def _wait() -> pd.DataFrame: + while True: + if len(self._deque) >= n: + return self.snapshot(n) + # Double-check after clearing to avoid missing an arrival that + # landed between the size check above and the clear below. + self._event.clear() + if len(self._deque) >= n: + return self.snapshot(n) + # Always schedule event.wait() on the loop that owns the Event. + # If the caller is on a different loop (e.g. WorkflowEngine's + # internal loop), bridge via run_coroutine_threadsafe so the + # wait runs on self._loop and the result is surfaced back here. + if asyncio.get_running_loop() is self._loop: + await self._event.wait() + else: + fut = asyncio.run_coroutine_threadsafe( + self._event.wait(), self._loop + ) + await asyncio.wrap_future(fut) + + if timeout is not None: + return await asyncio.wait_for(_wait(), timeout=timeout) + return await _wait() + + async def write_snapshot( + self, + n: int, + path: Path, + *, + timeout: float | None = None, + ) -> int: + """Wait for n observations, write them to a Parquet file, return row count.""" + df = await self.wait_for(n, timeout=timeout) + path.parent.mkdir(parents=True, exist_ok=True) + df.to_parquet(path, index=False) + return len(df) From 9f895920afab2a73f0decd538ee2fd0847a38d43 Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:45:56 -0700 Subject: [PATCH 02/22] DT Framework example foundations --- use-cases/m3dc1-stream/dt/.gitignore | 1 + use-cases/m3dc1-stream/dt/README.md | 12 + .../m3dc1-stream/dt/amsc_investigator.py | 341 ++++++++++++++++++ use-cases/m3dc1-stream/dt/buffer.py | 56 +++ use-cases/m3dc1-stream/dt/dtypes.py | 14 + use-cases/m3dc1-stream/dt/run_me.py | 110 ++++++ use-cases/m3dc1-stream/dt/sensor.py | 116 ++++++ 7 files changed, 650 insertions(+) create mode 100644 use-cases/m3dc1-stream/dt/.gitignore create mode 100644 use-cases/m3dc1-stream/dt/README.md create mode 100644 use-cases/m3dc1-stream/dt/amsc_investigator.py create mode 100644 use-cases/m3dc1-stream/dt/buffer.py create mode 100644 use-cases/m3dc1-stream/dt/dtypes.py create mode 100644 use-cases/m3dc1-stream/dt/run_me.py create mode 100644 use-cases/m3dc1-stream/dt/sensor.py diff --git a/use-cases/m3dc1-stream/dt/.gitignore b/use-cases/m3dc1-stream/dt/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/use-cases/m3dc1-stream/dt/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/use-cases/m3dc1-stream/dt/README.md b/use-cases/m3dc1-stream/dt/README.md new file mode 100644 index 0000000..b9ea574 --- /dev/null +++ b/use-cases/m3dc1-stream/dt/README.md @@ -0,0 +1,12 @@ + + +# The M3DC1 (spark) ported over to the Digital Twin framework. + +Items: +- `sensor.py` --> `dt/sensor.py` +- `sensor_buffer.py` --> `dt/buffer.py` +- `amsc_stream.py` --> `dt/amsc_investigator.py` + +Other: +`dt/run_me.py` +`dt/dtypes.py` diff --git a/use-cases/m3dc1-stream/dt/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py new file mode 100644 index 0000000..2310666 --- /dev/null +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -0,0 +1,341 @@ +""" +M3DC1 streaming digital twin. + +This is a version of the M3DC1 streaming surrogate that uses the digital twin +framework. + +Because we are using one active learner and only care about one physics +property, it's simplest to use just a single DT model investigator (no need for +a ScienceAgent. The point of a science agent is in the event you have various +surrogates with separate ALs) + +Also, in the M3DC1 streaming example, it assumes a ThreadedPoolExecutor, and all +tasks running on the same machine. This is due to the use of "buffer" inside the +simulation task. However, we want to demonstrate this code working across +machines, so this restriction must be lifted. + +A key point here is that the simulation task "waits" for new data. As the +ParallelActiveLearner is used, we don't have the flexibility to trigger the +simulation when we want to ourselves. + +Currently, the purpose of the simulation is to simply wait until the data is +available. So, + +""" + +import asyncio +import json +from pathlib import Path +import shlex + +import cloudpickle +from digitaltwin import ( + ModelInvestigator, + RuntimeAPI, + TypedData, + UtilityTask, + WindowedTypeData, +) +import numpy as np +import pandas as pd +from radical.asyncflow import WorkflowEngine +from rose import LearnerConfig, TaskConfig +from rose.al import ParallelActiveLearner + +from dtypes import M3DC1_PREDICTION + +# ── Row-budget policy (mirrors amsc.py's growing-pool logic) ───────────────── +# At iteration i, the simulation task requests N_BASE + i * N_STEP rows from +# the buffer. Increase N_STEP to consume more data per iteration. +N_BASE = 100 +N_STEP = 100 + + +# Workspace for iteration artefacts (parquet snapshots, metric JSON) +_HERE = Path(__file__).resolve().parent +_WORKSPACE = _HERE / "workspace" + + +def _workspace_iter(iteration: int, label: str) -> Path: + d = _WORKSPACE / f"{label}" / f"iter_{iteration:03d}" + d.mkdir(parents=True, exist_ok=True) + return d + + +def _candidate_configs(candidates: list[str], max_iter: int) -> list[LearnerConfig]: + configs = [] + for idx, family in enumerate(candidates): + label = f"{idx}_{family}" + kwargs = {"learner_label": label, "model_family": family} + schedule = { + i: TaskConfig(kwargs={**kwargs, "iteration": i}) + for i in range(max_iter + 1) + } + schedule[-1] = TaskConfig(kwargs={**kwargs, "iteration": max_iter}) + configs.append( + LearnerConfig( + simulation=schedule, + training=schedule, + active_learn=schedule, + criterion=schedule, + ) + ) + return configs + + +class M3DC1_Investigator(ModelInvestigator): + def __init__( + self, + flow: WorkflowEngine, + candidates: list[str], + max_iter: int, + r2_threshold: float, + ): + super().__init__(flow) + + self.learner = ParallelActiveLearner(flow) + self.candidates = candidates + self.max_iter = max_iter + self.r2_threshold = r2_threshold + + # each surrogate gets their own event + self.data_event: dict[str, asyncio.Event] = {} + + for candidate in self.candidates: + self.data_event[candidate] = asyncio.Event() + + # ── Simulation task ─────────────────────────────────────────────────────── + # KEY CHANGE vs amsc_stream. Buffered inputs come in from the DT. + + @self.learner.simulation_task + async def simulation(*args, **kwargs) -> str: + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + family = str(kwargs["model_family"]) + + n_rows = N_BASE + it * N_STEP + out_dir = _workspace_iter(it, label) + parquet = out_dir / "sensor_snapshot.parquet" + + print( + f" [sim {label} iter={it}] Now has {n_rows} sensor rows …", flush=True + ) + + await self.data_event[family].wait() + n_written = len(self.all_data) + + meta = { + "iteration": it, + "learner_label": label, + "dataset": str(parquet), + "n_rows": n_written, + "source": "sensor_stream", + } + (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) + + self.data_event[family].clear() + + # workaround: I know, looks a little goofy. Read the notes at top of + # file. Parameter is passed via `out_dir/simulation.json` + # + # Consequence of this workaround: on shutdown, AsyncFlow will + # complain that the future is missing some attributes. + # This is because AsyncFlow stamps the attributes when the cmdline + # is returned. It doesn't hurt accuracy in any way, just doesn't + # look as clean. + return f"echo {shlex.quote(str(out_dir / "simulation.json"))}" + + # ── Training task ───────────────────────────────────────────────────────── + # Fits a surrogate model locally using sklearn. + # Replace with subprocess to surge_train.py if running on HPC. + @self.learner.training_task(as_executable=False) + async def training(sim_result_path: str, **kwargs) -> dict: + import pandas as pd + from sklearn.ensemble import ( + GradientBoostingRegressor, + RandomForestRegressor, + ) + from sklearn.linear_model import Ridge + from sklearn.metrics import mean_squared_error, r2_score + from sklearn.model_selection import train_test_split + from sklearn.neural_network import MLPRegressor + from sklearn.preprocessing import StandardScaler + from sklearn.pipeline import Pipeline + + with open(sim_result_path, "r") as f: + sim_result = json.load(f) + + it = int(kwargs.get("iteration", sim_result["iteration"])) + label = str(kwargs["learner_label"]) + family = str(kwargs.get("model_family", "rf")) + + df = pd.read_parquet(sim_result["dataset"]) + X = df.drop(columns=["output_gamma"]).values + y = df["output_gamma"].values + + X_train, X_val, y_train, y_val = train_test_split( + X, y, test_size=0.2, random_state=it + ) + + _models = { + "rf": RandomForestRegressor( + n_estimators=100, random_state=42, n_jobs=-1 + ), + "mlp": Pipeline( + [ + ("scaler", StandardScaler()), + ( + "net", + MLPRegressor( + hidden_layer_sizes=(64, 64), + max_iter=500, + random_state=42, + ), + ), + ] + ), + "gbr": GradientBoostingRegressor(n_estimators=100, random_state=42), + "ridge": Pipeline( + [ + ("scaler", StandardScaler()), + ("reg", Ridge()), + ] + ), + } + model = _models.get(family, _models["rf"]) + model.fit(X_train, y_train) + + y_pred = model.predict(X_val) + val_r2 = float(r2_score(y_val, y_pred)) + val_rmse = float(np.sqrt(mean_squared_error(y_val, y_pred))) + + metrics = { + "iteration": it, + "learner_label": label, + "model_family": family, + "val_r2": val_r2, + "val_rmse": val_rmse, + "n_train": int(len(X_train)), + "n_val": int(len(X_val)), + } + out_dir = _workspace_iter(it, label) + (out_dir / "surge_metrics.json").write_text(json.dumps(metrics, indent=2)) + + # save + model_path = str((out_dir / "model.pkl")) + with open(model_path, "wb") as f: + cloudpickle.dump(model, f) + + return {"simulation": sim_result, "surge": metrics, "model": model_path} + + # ── Active-learning task ────────────────────────────────────────────────── + @self.learner.active_learn_task(as_executable=False) + async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: + it = int(kwargs.get("iteration", train_bundle["simulation"]["iteration"])) + label = str(kwargs["learner_label"]) + surge = train_bundle["surge"] + + decision = { + "iteration": it, + "learner_label": label, + "policy": "monitor_best_val_r2", + "val_r2": surge["val_r2"], + "val_rmse": surge["val_rmse"], + "model": train_bundle["model"], + } + out_dir = _workspace_iter(it, label) + (out_dir / "active.json").write_text(json.dumps(decision, indent=2)) + return { + "iteration": it, + "learner_label": label, + "train": train_bundle, + "val_r2": surge["val_r2"], + "val_rmse": surge["val_rmse"], + "model": train_bundle["model"], + } + + # ── Stop criterion ──────────────────────────────────────────────────────── + @self.learner.as_stop_criterion( + metric_name="val_r2", + threshold=r2_threshold, + operator=">=", + as_executable=False, + ) + async def stop_on_r2(*args, **kwargs) -> float: + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + path = _workspace_iter(it, label) / "surge_metrics.json" + meta = json.loads(path.read_text()) + r2 = float(meta["val_r2"]) + forced = it >= max_iter - 1 and r2 < r2_threshold + return r2_threshold if forced else r2 + + async def do_inference(in_data: WindowedTypeData, model=None): + # the ASMC_stream.py demo doesn't tackle streaming inference. + # Put streaming inference code here. + if model is None: + return TypedData(M3DC1_PREDICTION, None) + + with open(model, "rb") as f: + model_obj = cloudpickle.load(f) + + df = pd.DataFrame(in_data.sequence) + + X = df.drop(columns=["output_gamma"]).values + out = model_obj.predict(X) + + return TypedData(M3DC1_PREDICTION, out) + + self.inference = do_inference + + async def input_callback(self, in_data: WindowedTypeData): + # add the data to large database + self.all_data += in_data.sequence + # trigger event + for v in self.data_event.values(): + v.set() + + async def main_loop(self, runtime: RuntimeAPI): + # call the pipeline + + runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) + runtime.publish_new_model({"model": None}) + + configs = _candidate_configs(self.candidates, self.max_iter) + rows: list[dict] = [] + + async for state in self.learner.start( + parallel_learners=len(self.candidates), + max_iter=self.max_iter, + learner_configs=configs, + ): + label = self.candidates[int(state.learner_id)] + rows.append( + { + "learner": label, + "iter": state.iteration, + "val_r2": state.val_r2, + "val_rmse": state.val_rmse, + } + ) + print( + f" learner={label} iter={state.iteration}" + f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}" + f" buffer={len(self.all_data)} obs", + flush=True, + ) + + # publish model with stats + runtime.publish_new_model({"model": state.model}, rows[-1]) + + if len(rows) >= len(self.candidates) * self.max_iter: + break + + +# this is needed to +class OutputSink(UtilityTask): + def __init__(self, flow): + super().__init__(flow) + + async def main_loop(self, runtime, in_data: TypedData): + print("Received: ", in_data.data) diff --git a/use-cases/m3dc1-stream/dt/buffer.py b/use-cases/m3dc1-stream/dt/buffer.py new file mode 100644 index 0000000..4660d5d --- /dev/null +++ b/use-cases/m3dc1-stream/dt/buffer.py @@ -0,0 +1,56 @@ +""" +Buffer: generic data buffer. Batches inputs and only emits when done. + +The digital twin framework provides a stream-processing dataflow paradigm. It already +comes with a Windowing feature via its BARRIER. However, it leaves the logic up to the user of +how their windows should be defined. + +So, you need a split task that generates an event of when to trigger the barrier. + +Dataflow Digital Twin graph: + +SENSOR --> BUFFER_EVENT_EMIT --> BARRIER --> SINK + | | + +--------------> BARRIER --> goes nowhere + +From this, the BUFFER_EVENT_EMIT is a SPLIT task. It takes in one stream and +generates two. + +The BARRIER is provided by the DT framework. This file simply implements the +buffer event split task, and a null sink. + +The null sink is needed to drain the queues. + +""" + +import time + +from digitaltwin import SplitTask, TypedData, UtilityTask +from dtypes import BUFFER_EVENT, SYNC_SENSOR + + +class BufferEventEmit(SplitTask): + def __init__(self, flow, buffer_length): + super().__init__(flow) + self.buffer_length = buffer_length + self.counter = 0 + + async def main_loop(self, runtime, in_data: TypedData): + self.counter += 1 + out_data = TypedData( + SYNC_SENSOR, + ) + if self.counter >= self.buffer_length: + # emit event + return in_data, TypedData(BUFFER_EVENT, time.time()) + + return in_data, None + + +# this is needed to +class DeadSink(UtilityTask): + def __init__(self, flow): + super().__init__(flow) + + async def main_loop(self, runtime, in_data): + pass # do nothing. diff --git a/use-cases/m3dc1-stream/dt/dtypes.py b/use-cases/m3dc1-stream/dt/dtypes.py new file mode 100644 index 0000000..fc9c828 --- /dev/null +++ b/use-cases/m3dc1-stream/dt/dtypes.py @@ -0,0 +1,14 @@ +from digitaltwin.components import DataType + +# Sensor channel + +M3DC1_MOCK_CHANNEL = "sensors/MockM3DC1" +M3DC1_SENSOR = DataType("M3DC1") + +# Buffer event +BUFFER_EVENT = DataType("BUFFER_EVENT") +SYNC_SENSOR = DataType("M3DC1_SYNC") + + +# Modeling +M3DC1_PREDICTION = DataType("M3DC1_PREDICTION") diff --git a/use-cases/m3dc1-stream/dt/run_me.py b/use-cases/m3dc1-stream/dt/run_me.py new file mode 100644 index 0000000..e72e1f7 --- /dev/null +++ b/use-cases/m3dc1-stream/dt/run_me.py @@ -0,0 +1,110 @@ +""" +M3DC1 Digital Twin - a DT wrapper of the M3DC1 streaming surrogate example + +Complete Digital Twin graph: + +MOCK_SENSOR --> BUFFER_EVENT_TASK --> BARRIER --> M3DC1_Investigator --> OUTPUT TASK + | ||| + +-------------> BARRIER --> DEAD SINK + + +""" + +import argparse +import asyncio +from concurrent.futures import ProcessPoolExecutor +from digitaltwin import Barrier +from radical.asyncflow import WorkflowEngine +from rhapsody.backends import ConcurrentExecutionBackend + +from digitaltwin.runtime import DTRuntime +from digitaltwin.streaming import connect_stream_client +from digitaltwin.components import NULL_DTYPE + +from amsc_investigator import M3DC1_Investigator, OutputSink +from buffer import BufferEventEmit, DeadSink +from dtypes import * + +from radical.asyncflow.logging import init_default_logger +import logging + +logger = logging.getLogger(__name__) + +# put it all together +# sensor channel --> model --> data_sink +# +# The sensor is external: run sensor.py in its own terminal. + + +async def main(candidates, max_iter, r2_threshold, buffer_length): + init_default_logger(logging.INFO) + logging.getLogger("radical.asyncflow").setLevel(logging.WARNING) + logging.getLogger("rhapsody").setLevel(logging.WARNING) + + # create engine + exe = await ConcurrentExecutionBackend(ProcessPoolExecutor()) + flow = await WorkflowEngine.create(backend=exe) + + # create the twin's namespaced stream client + pubsub_client = await connect_stream_client("M3DC1-Demo") + + runtime = DTRuntime(flow, pubsub_client) + + ################### + # create tasks and investigators + + # for buffer: + buffer_event_task = BufferEventEmit(flow, buffer_length) + barrier = Barrier("sensor_windowing_buffer") + WINDOW_DTYPE = barrier.add_dtype(SYNC_SENSOR, hard=False) + dead_sink = DeadSink(flow) + + barrier.add_dtype(BUFFER_EVENT, hard=True) + + # for investigator + m3dc1 = M3DC1_Investigator(flow, candidates, max_iter, r2_threshold) + output_sink = OutputSink(flow) + + # create graph + runtime.add_input(M3DC1_SENSOR, M3DC1_MOCK_CHANNEL) + runtime.add_barrier(barrier) + runtime.add_data_split_task( + buffer_event_task, M3DC1_SENSOR, [SYNC_SENSOR, BUFFER_EVENT] + ) + runtime.add_task(dead_sink, BUFFER_EVENT, NULL_DTYPE) + runtime.add_investigator(m3dc1, SYNC_SENSOR, M3DC1_PREDICTION) + runtime.add_task(output_sink, M3DC1_PREDICTION, NULL_DTYPE) + + runtime.print_graph() + # runtime.start() + + # let it run + await asyncio.sleep(5) + await runtime.stop() + await flow.shutdown() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="M3DC1 streaming surrogate — sensor input version." + ) + parser.add_argument( + "--candidates", + default="rf,mlp", + help="Comma-separated model families: rf, mlp, gbr, ridge.", + ) + parser.add_argument("--max-iter", type=int, default=3) + parser.add_argument("--r2-threshold", type=float, default=0.80) + parser.add_argument( + "--buffer-maxlen", + type=int, + default=10_000, + help="Maximum observations retained in the sensor buffer.", + ) + args = parser.parse_args() + candidates = [x.strip() for x in args.candidates.split(",") if x.strip()] + + if len(candidates) < 2: + parser.error("Need at least two candidates for ParallelActiveLearner.") + + asyncio.run(main(candidates, args.max_iter, args.r2_threshold, args.buffer_maxlen)) diff --git a/use-cases/m3dc1-stream/dt/sensor.py b/use-cases/m3dc1-stream/dt/sensor.py new file mode 100644 index 0000000..1ceff08 --- /dev/null +++ b/use-cases/m3dc1-stream/dt/sensor.py @@ -0,0 +1,116 @@ +""" +Sensor stream interface and mock implementation for M3DC1 streaming workflows. + +To plug in a real sensor, subclass SensorStream and implement read_one(). +Everything else (SensorBuffer, amsc_stream.py) works unchanged. +""" + +from __future__ import annotations + +import argparse +import asyncio + +from digitaltwin import ChannelPublisher +import numpy as np +from dtypes import * + +# ── Physical parameter ranges for SPARC M3DC1 D1 ──────────────────────────── +# Derived from sparc_m3dc1_D1_metadata.yaml (inputs + output_gamma). +# Replace bounds with real calibration data when integrating a live source. +_M3DC1_RANGES: dict[str, tuple[float, float]] = { + "input_batemanscale": (0.5, 2.0), + "input_ntor": (1.0, 15.0), # toroidal mode number + "input_pscale": (0.5, 2.0), + "eq_a": (0.1, 0.5), # minor radius [m] + "eq_R0": (1.5, 6.0), # major radius [m] + "eq_kappa": (1.0, 2.5), # elongation + "eq_delta": (0.0, 0.8), # triangularity + "eq_simag": (-2.0, 0.0), # poloidal flux at magnetic axis + "eq_sibry": (-5.0, -0.5), # poloidal flux at boundary + "eq_current": (0.5, 15.0), # plasma current [MA] + "q0": (0.8, 2.5), # safety factor on axis + "q95": (3.0, 8.0), # safety factor at 95 % flux + "p0": (1e4, 1e6), # peak pressure [Pa] +} + +COLUMNS: list[str] = list(_M3DC1_RANGES) + ["output_gamma"] + + +class MockM3DC1Sensor: + """Simulates a real-time M3DC1 physics sensor at a configurable rate. + + Observations are drawn from the realistic parameter ranges in _M3DC1_RANGES. + output_gamma is a nonlinear surrogate of the MHD stability growth rate plus + Gaussian noise — non-trivial enough to make the surrogate task meaningful. + + Args: + rate_hz: Target emission rate in observations per second. + seed: RNG seed for reproducibility. + noise_std: Std-dev of Gaussian noise on output_gamma. + """ + + def __init__( + self, + rate_hz: float = 2.0, + seed: int = 42, + noise_std: float = 0.005, + ) -> None: + self._delay = 1.0 / max(rate_hz, 1e-6) + self._rng = np.random.default_rng(seed) + self._noise_std = noise_std + + async def read_one(self) -> dict[str, float]: + await asyncio.sleep(self._delay) + return self._sample() + + def _sample(self) -> dict[str, float]: + rng = self._rng + obs: dict[str, float] = { + col: float(rng.uniform(lo, hi)) for col, (lo, hi) in _M3DC1_RANGES.items() + } + # Surrogate physics: gamma grows with mode number and pressure scale, + # falls with safety factor — rough but nonlinear enough for surrogates. + obs["output_gamma"] = float( + max( + 0.0, + 0.08 * obs["input_batemanscale"] * (obs["input_ntor"] / 8.0) + + 0.06 * obs["input_pscale"] / max(obs["q95"], 0.1) + + 0.04 * obs["eq_kappa"] * obs["eq_delta"] + - 0.02 * obs["q0"] + + float(rng.normal(0.0, self._noise_std)), + ) + ) + return obs + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="M3DC1 mock sensor.") + parser.add_argument( + "--sensor-rate", + type=float, + default=10.0, + help="Mock sensor emission rate in observations/second (default: 10).", + ) + parser.add_argument( + "--sensor-seed", + type=int, + default=42, + help="RNG seed for the mock sensor.", + ) + + args = parser.parse_args() + + async def main(): + publisher = await ChannelPublisher.open(M3DC1_MOCK_CHANNEL) + + sensor = MockM3DC1Sensor(rate_hz=args.sensor_rate, seed=args.sensor_seed) + try: + while True: + val = await sensor.read_one() + await publisher.publish(val) + finally: + await publisher.close() + + if __name__ == "__main__": + asyncio.run(main()) From 027033742795433f4180b83acbf49c58308e5c5f Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 26 Aug 2026 19:43:26 +0200 Subject: [PATCH 03/22] use rhapsody databackends --- use-cases/m3dc1-stream/amsc_stream.py | 219 +++++++++++------------- use-cases/m3dc1-stream/sensor_daemon.py | 97 +++++++++++ 2 files changed, 193 insertions(+), 123 deletions(-) create mode 100644 use-cases/m3dc1-stream/sensor_daemon.py diff --git a/use-cases/m3dc1-stream/amsc_stream.py b/use-cases/m3dc1-stream/amsc_stream.py index e72c7d0..a73c08b 100644 --- a/use-cases/m3dc1-stream/amsc_stream.py +++ b/use-cases/m3dc1-stream/amsc_stream.py @@ -1,79 +1,54 @@ #!/usr/bin/env python3 -""" -M3DC1 streaming surrogate — sensor-based input version. - -Replaces the predefined dataset pool from amsc.py with a continuous sensor -stream (mocked by MockM3DC1Sensor). The simulation task no longer slices -from a pre-loaded PKL; instead it calls buffer.wait_for(n) and consumes -however many observations the sensor has delivered so far. - -Architecture -──────────── - SensorStream (sensor.py) ← implement this for a real sensor - │ - SensorBuffer (sensor_buffer.py) ← background asyncio task, accumulates rows - │ - simulation task ← waits for N rows, writes parquet - │ - training task ← fits sklearn surrogate on that parquet - │ - active_learn task ← records AL decision - │ - stop_on_r2 criterion ← stops when val_r2 ≥ threshold - -To swap in a real sensor: - class MyRealSensor(SensorStream): - async def read_one(self) -> dict[str, float]: - ... # read from your hardware / REST API / Kafka / etc. - - sensor = MyRealSensor() - # everything else in this file stays the same -""" from __future__ import annotations import argparse import asyncio import json -import os import sys import warnings from pathlib import Path import numpy as np import rhapsody +from rhapsody.backends.data.redis import RedisDataBackend _HERE = Path(__file__).resolve().parent if str(_HERE) not in sys.path: sys.path.insert(0, str(_HERE)) -from sensor import MockM3DC1Sensor, SensorStream # noqa: E402 -from sensor_buffer import SensorBuffer # noqa: E402 +from sensor_daemon import STREAM_KEY, MockM3DC1Sensor, SensorDaemon # noqa: E402 + +from concurrent.futures import ProcessPoolExecutor + +from rhapsody.backends import ConcurrentExecutionBackend from radical.asyncflow import WorkflowEngine # noqa: E402 from rose.al import ParallelActiveLearner # noqa: E402 from rose.learner import LearnerConfig, TaskConfig # noqa: E402 -# ── Row-budget policy (mirrors amsc.py's growing-pool logic) ───────────────── -# At iteration i, the simulation task requests N_BASE + i * N_STEP rows from -# the buffer. Increase N_STEP to consume more data per iteration. N_BASE = 100 N_STEP = 100 -# Workspace for iteration artefacts (parquet snapshots, metric JSON) _WORKSPACE = _HERE / "workspace" def _workspace_iter(iteration: int, label: str) -> Path: - d = _WORKSPACE / f"{label}" / f"iter_{iteration:03d}" + d = _WORKSPACE / label / f"iter_{iteration:03d}" d.mkdir(parents=True, exist_ok=True) return d -def _candidate_configs(candidates: list[str], max_iter: int) -> list[LearnerConfig]: +def _candidate_configs( + candidates: list[str], max_iter: int, redis_endpoint: str +) -> list[LearnerConfig]: configs = [] for idx, family in enumerate(candidates): - label = f"{idx}_{family}" - kwargs = {"learner_label": label, "model_family": family} + label = f"{idx}_{family}" + kwargs = { + "learner_label" : label, + "model_family" : family, + "redis_endpoint": redis_endpoint, + } schedule = {i: TaskConfig(kwargs={**kwargs, "iteration": i}) for i in range(max_iter + 1)} schedule[-1] = TaskConfig(kwargs={**kwargs, "iteration": max_iter}) configs.append( @@ -87,47 +62,63 @@ def _candidate_configs(candidates: list[str], max_iter: int) -> list[LearnerConf return configs -# ── ROSE workflow ───────────────────────────────────────────────────────────── - async def run_rose_workflow( - buffer: SensorBuffer, + endpoint, *, candidates: list[str], max_iter: int, r2_threshold: float, ) -> None: - engine = await rhapsody.get_backend("concurrent") + import redis as _redis + + engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) learner = ParallelActiveLearner(asyncflow) - # ── Simulation task ─────────────────────────────────────────────────────── - # KEY CHANGE vs amsc.py: reads from sensor buffer, not from the PKL pool. - # The buffer is captured by closure; swapping the sensor changes nothing here. + redis_endpoint = endpoint.serialize() + redis_client = _redis.Redis(host=endpoint.host, port=endpoint.port, decode_responses=True) + @learner.simulation_task(as_executable=False) async def simulation(*args, **kwargs) -> dict: - it = int(kwargs.get("iteration", 0)) - label = str(kwargs["learner_label"]) + import time + import redis as _redis + import pandas as pd + + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + redis_ep = str(kwargs["redis_endpoint"]) + n_rows = N_BASE + it * N_STEP + + host, port_str = redis_ep.rsplit(":", 1) + redis_client = _redis.Redis(host=host, port=int(port_str), decode_responses=True) + + print(f" [sim {label} iter={it}] waiting for {n_rows} rows …", flush=True) + deadline = time.monotonic() + 600.0 + while redis_client.xlen(STREAM_KEY) < n_rows: + if time.monotonic() > deadline: + raise TimeoutError(f"Timeout waiting for {n_rows} sensor rows") + time.sleep(0.5) + + entries = redis_client.xrevrange(STREAM_KEY, count=n_rows) + entries.reverse() + + rows = [{key: float(val) for key, val in fields.items()} for _, fields in entries] + df = pd.DataFrame(rows) - n_rows = N_BASE + it * N_STEP out_dir = _workspace_iter(it, label) parquet = out_dir / "sensor_snapshot.parquet" - - print(f" [sim {label} iter={it}] waiting for {n_rows} sensor rows …", flush=True) - n_written = await buffer.write_snapshot(n_rows, parquet, timeout=600.0) + df.to_parquet(parquet, index=False) meta = { "iteration" : it, "learner_label": label, "dataset" : str(parquet), - "n_rows" : n_written, - "source" : "sensor_stream", + "n_rows" : len(df), + "source" : "redis_stream", } (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) return meta - # ── Training task ───────────────────────────────────────────────────────── - # Fits a surrogate model locally using sklearn. - # Replace with subprocess to surge_train.py if running on HPC. @learner.training_task(as_executable=False) async def training(sim_result: dict, **kwargs) -> dict: import pandas as pd @@ -136,8 +127,8 @@ async def training(sim_result: dict, **kwargs) -> dict: from sklearn.metrics import mean_squared_error, r2_score from sklearn.model_selection import train_test_split from sklearn.neural_network import MLPRegressor - from sklearn.preprocessing import StandardScaler from sklearn.pipeline import Pipeline + from sklearn.preprocessing import StandardScaler it = int(kwargs.get("iteration", sim_result["iteration"])) label = str(kwargs["learner_label"]) @@ -147,26 +138,24 @@ async def training(sim_result: dict, **kwargs) -> dict: X = df.drop(columns=["output_gamma"]).values y = df["output_gamma"].values - X_train, X_val, y_train, y_val = train_test_split( - X, y, test_size=0.2, random_state=it - ) + X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=it) _models = { - "rf": RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=-1), - "mlp": Pipeline([ - ("scaler", StandardScaler()), - ("net", MLPRegressor(hidden_layer_sizes=(64, 64), max_iter=500, random_state=42)), - ]), - "gbr": GradientBoostingRegressor(n_estimators=100, random_state=42), + "rf": RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=-1), + "mlp": Pipeline([ + ("scaler", StandardScaler()), + ("net", MLPRegressor(hidden_layer_sizes=(64, 64), max_iter=500, random_state=42)), + ]), + "gbr": GradientBoostingRegressor(n_estimators=100, random_state=42), "ridge": Pipeline([ - ("scaler", StandardScaler()), - ("reg", Ridge()), - ]), + ("scaler", StandardScaler()), + ("reg", Ridge()), + ]), } model = _models.get(family, _models["rf"]) model.fit(X_train, y_train) - y_pred = model.predict(X_val) + y_pred = model.predict(X_val) val_r2 = float(r2_score(y_val, y_pred)) val_rmse = float(np.sqrt(mean_squared_error(y_val, y_pred))) @@ -183,7 +172,6 @@ async def training(sim_result: dict, **kwargs) -> dict: (out_dir / "surge_metrics.json").write_text(json.dumps(metrics, indent=2)) return {"simulation": sim_result, "surge": metrics} - # ── Active-learning task ────────────────────────────────────────────────── @learner.active_learn_task(as_executable=False) async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: it = int(kwargs.get("iteration", train_bundle["simulation"]["iteration"])) @@ -207,7 +195,6 @@ async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: "val_rmse" : surge["val_rmse"], } - # ── Stop criterion ──────────────────────────────────────────────────────── @learner.as_stop_criterion( metric_name="val_r2", threshold=r2_threshold, @@ -215,16 +202,15 @@ async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: as_executable=False, ) async def stop_on_r2(*args, **kwargs) -> float: - it = int(kwargs.get("iteration", 0)) - label = str(kwargs["learner_label"]) - path = _workspace_iter(it, label) / "surge_metrics.json" - meta = json.loads(path.read_text()) - r2 = float(meta["val_r2"]) + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + path = _workspace_iter(it, label) / "surge_metrics.json" + meta = json.loads(path.read_text()) + r2 = float(meta["val_r2"]) forced = it >= max_iter - 1 and r2 < r2_threshold return r2_threshold if forced else r2 - # ── Run ─────────────────────────────────────────────────────────────────── - configs = _candidate_configs(candidates, max_iter) + configs = _candidate_configs(candidates, max_iter, redis_endpoint) rows: list[dict] = [] try: @@ -233,7 +219,8 @@ async def stop_on_r2(*args, **kwargs) -> float: max_iter=max_iter, learner_configs=configs, ): - label = candidates[int(state.learner_id)] + label = candidates[int(state.learner_id)] + stream_len = await asyncio.to_thread(redis_client.xlen, STREAM_KEY) rows.append({ "learner" : label, "iter" : state.iteration, @@ -243,7 +230,7 @@ async def stop_on_r2(*args, **kwargs) -> float: print( f" learner={label} iter={state.iteration}" f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}" - f" buffer={len(buffer)} obs", + f" stream={stream_len} obs", flush=True, ) if len(rows) >= len(candidates) * max_iter: @@ -263,65 +250,51 @@ async def stop_on_r2(*args, **kwargs) -> float: print(f"Workspace: {_WORKSPACE}") -# ── Entry point ─────────────────────────────────────────────────────────────── - def main() -> None: warnings.filterwarnings("ignore", category=UserWarning) parser = argparse.ArgumentParser( - description="M3DC1 streaming surrogate — sensor input version." - ) - parser.add_argument( - "--candidates", - default="rf,mlp", - help="Comma-separated model families: rf, mlp, gbr, ridge.", - ) - parser.add_argument("--max-iter", type=int, default=3) - parser.add_argument("--r2-threshold", type=float, default=0.80) - parser.add_argument( - "--sensor-rate", - type=float, - default=10.0, - help="Mock sensor emission rate in observations/second (default: 10).", - ) - parser.add_argument( - "--sensor-seed", - type=int, - default=42, - help="RNG seed for the mock sensor.", - ) - parser.add_argument( - "--buffer-maxlen", - type=int, - default=10_000, - help="Maximum observations retained in the sensor buffer.", + description="M3DC1 streaming surrogate — RedisDataBackend version." ) + parser.add_argument("--candidates", default="rf,mlp", + help="Comma-separated model families: rf, mlp, gbr, ridge.") + parser.add_argument("--max-iter", type=int, default=3) + parser.add_argument("--r2-threshold", type=float, default=0.80) + parser.add_argument("--sensor-rate", type=float, default=10.0, + help="Mock sensor rate in obs/s.") + parser.add_argument("--sensor-seed", type=int, default=42) + parser.add_argument("--buffer-maxlen", type=int, default=10_000, + help="Max stream length retained in Redis.") + args = parser.parse_args() candidates = [x.strip() for x in args.candidates.split(",") if x.strip()] - if len(candidates) < 2: - parser.error("Need at least two candidates for ParallelActiveLearner.") - - sensor = MockM3DC1Sensor(rate_hz=args.sensor_rate, seed=args.sensor_seed) - buffer = SensorBuffer(sensor, maxlen=args.buffer_maxlen) + async def _main() -> None: - await buffer.start() + redis_backend = await RedisDataBackend() + endpoint = redis_backend.endpoints[0] + + sensor = MockM3DC1Sensor(rate_hz=args.sensor_rate, seed=args.sensor_seed) + daemon = SensorDaemon(sensor, endpoint.serialize(), maxlen=args.buffer_maxlen) + await daemon.start() + print( - f"Sensor stream started rate={args.sensor_rate} Hz" - f" candidates={candidates} max_iter={args.max_iter}", + f"Redis: {endpoint.serialize()} sensor: {args.sensor_rate} Hz" + f" candidates: {candidates} max_iter: {args.max_iter}", flush=True, ) try: await run_rose_workflow( - buffer, + endpoint, candidates=candidates, max_iter=args.max_iter, r2_threshold=args.r2_threshold, ) finally: - await buffer.stop() - print("Sensor stream stopped.", flush=True) + await daemon.stop() + await redis_backend.shutdown() + print("Sensor daemon and Redis stopped.", flush=True) asyncio.run(_main()) diff --git a/use-cases/m3dc1-stream/sensor_daemon.py b/use-cases/m3dc1-stream/sensor_daemon.py new file mode 100644 index 0000000..49f13f3 --- /dev/null +++ b/use-cases/m3dc1-stream/sensor_daemon.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import asyncio +import time +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator +from pathlib import Path +import sys + +_HERE = Path(__file__).resolve().parent +if str(_HERE) not in sys.path: + sys.path.insert(0, str(_HERE)) + +import numpy as np +import redis + +STREAM_KEY = "m3dc1:sensor" + +_M3DC1_RANGES: dict[str, tuple[float, float]] = { + "input_batemanscale": (0.5, 2.0), + "input_ntor": (1.0, 15.0), + "input_pscale": (0.5, 2.0), + "eq_a": (0.1, 0.5), + "eq_R0": (1.5, 6.0), + "eq_kappa": (1.0, 2.5), + "eq_delta": (0.0, 0.8), + "eq_simag": (-2.0, 0.0), + "eq_sibry": (-5.0, -0.5), + "eq_current": (0.5, 15.0), + "q0": (0.8, 2.5), + "q95": (3.0, 8.0), + "p0": (1e4, 1e6), +} + +COLUMNS: list[str] = list(_M3DC1_RANGES) + ["output_gamma"] + + +class SensorStream(ABC): + @abstractmethod + async def read_one(self) -> dict[str, float]: ... + + async def stream(self) -> AsyncIterator[dict[str, float]]: + while True: + yield await self.read_one() + + +class MockM3DC1Sensor(SensorStream): + def __init__(self, rate_hz: float = 2.0, seed: int = 42, noise_std: float = 0.005) -> None: + self._delay = 1.0 / max(rate_hz, 1e-6) + self._rng = np.random.default_rng(seed) + self._noise_std = noise_std + + async def read_one(self) -> dict[str, float]: + await asyncio.sleep(self._delay) + return self._sample() + + def _sample(self) -> dict[str, float]: + rng = self._rng + obs: dict[str, float] = { + col: float(rng.uniform(lo, hi)) + for col, (lo, hi) in _M3DC1_RANGES.items() + } + obs["output_gamma"] = float(max( + 0.0, + 0.08 * obs["input_batemanscale"] * (obs["input_ntor"] / 8.0) + + 0.06 * obs["input_pscale"] / max(obs["q95"], 0.1) + + 0.04 * obs["eq_kappa"] * obs["eq_delta"] + - 0.02 * obs["q0"] + + float(rng.normal(0.0, self._noise_std)), + )) + return obs + + +class SensorDaemon: + def __init__(self, sensor: SensorStream, redis_endpoint: str, maxlen: int = 10_000) -> None: + host, port = redis_endpoint.rsplit(":", 1) + self._sensor = sensor + self._r = redis.Redis(host=host, port=int(port), decode_responses=True) + self._maxlen = maxlen + self._task: asyncio.Task | None = None + + async def start(self) -> None: + self._task = asyncio.create_task(self._run(), name="sensor-daemon") + + async def stop(self) -> None: + if self._task is not None: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + async def _run(self) -> None: + async for obs in self._sensor.stream(): + fields = {k: str(v) for k, v in obs.items()} + await asyncio.to_thread(self._r.xadd, STREAM_KEY, fields, maxlen=self._maxlen) From 9cecd4794cae79b61c795241d9ddadeefe507709 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 26 Aug 2026 19:44:46 +0200 Subject: [PATCH 04/22] remove old files --- use-cases/m3dc1-stream/sensor.py | 99 ------------------- use-cases/m3dc1-stream/sensor_buffer.py | 122 ------------------------ 2 files changed, 221 deletions(-) delete mode 100644 use-cases/m3dc1-stream/sensor.py delete mode 100644 use-cases/m3dc1-stream/sensor_buffer.py diff --git a/use-cases/m3dc1-stream/sensor.py b/use-cases/m3dc1-stream/sensor.py deleted file mode 100644 index eb09069..0000000 --- a/use-cases/m3dc1-stream/sensor.py +++ /dev/null @@ -1,99 +0,0 @@ -""" -Sensor stream interface and mock implementation for M3DC1 streaming workflows. - -To plug in a real sensor, subclass SensorStream and implement read_one(). -Everything else (SensorBuffer, amsc_stream.py) works unchanged. -""" -from __future__ import annotations - -import asyncio -from abc import ABC, abstractmethod -from collections.abc import AsyncIterator - -import numpy as np - - -class SensorStream(ABC): - """Abstract sensor that emits one M3DC1 observation at a time. - - Implement read_one() for any real sensor, REST endpoint, Kafka consumer, - shared-memory ring buffer, etc. The default stream() calls read_one() in - a loop; override it for push-based sources where data arrives on its own. - """ - - @abstractmethod - async def read_one(self) -> dict[str, float]: - """Return a single observation as a column-name → value mapping.""" - - async def stream(self) -> AsyncIterator[dict[str, float]]: - """Yield observations indefinitely. Override for push-based sensors.""" - while True: - yield await self.read_one() - - -# ── Physical parameter ranges for SPARC M3DC1 D1 ──────────────────────────── -# Derived from sparc_m3dc1_D1_metadata.yaml (inputs + output_gamma). -# Replace bounds with real calibration data when integrating a live source. -_M3DC1_RANGES: dict[str, tuple[float, float]] = { - "input_batemanscale": (0.5, 2.0), - "input_ntor": (1.0, 15.0), # toroidal mode number - "input_pscale": (0.5, 2.0), - "eq_a": (0.1, 0.5), # minor radius [m] - "eq_R0": (1.5, 6.0), # major radius [m] - "eq_kappa": (1.0, 2.5), # elongation - "eq_delta": (0.0, 0.8), # triangularity - "eq_simag": (-2.0, 0.0), # poloidal flux at magnetic axis - "eq_sibry": (-5.0, -0.5), # poloidal flux at boundary - "eq_current": (0.5, 15.0), # plasma current [MA] - "q0": (0.8, 2.5), # safety factor on axis - "q95": (3.0, 8.0), # safety factor at 95 % flux - "p0": (1e4, 1e6), # peak pressure [Pa] -} - -COLUMNS: list[str] = list(_M3DC1_RANGES) + ["output_gamma"] - - -class MockM3DC1Sensor(SensorStream): - """Simulates a real-time M3DC1 physics sensor at a configurable rate. - - Observations are drawn from the realistic parameter ranges in _M3DC1_RANGES. - output_gamma is a nonlinear surrogate of the MHD stability growth rate plus - Gaussian noise — non-trivial enough to make the surrogate task meaningful. - - Args: - rate_hz: Target emission rate in observations per second. - seed: RNG seed for reproducibility. - noise_std: Std-dev of Gaussian noise on output_gamma. - """ - - def __init__( - self, - rate_hz: float = 2.0, - seed: int = 42, - noise_std: float = 0.005, - ) -> None: - self._delay = 1.0 / max(rate_hz, 1e-6) - self._rng = np.random.default_rng(seed) - self._noise_std = noise_std - - async def read_one(self) -> dict[str, float]: - await asyncio.sleep(self._delay) - return self._sample() - - def _sample(self) -> dict[str, float]: - rng = self._rng - obs: dict[str, float] = { - col: float(rng.uniform(lo, hi)) - for col, (lo, hi) in _M3DC1_RANGES.items() - } - # Surrogate physics: gamma grows with mode number and pressure scale, - # falls with safety factor — rough but nonlinear enough for surrogates. - obs["output_gamma"] = float(max( - 0.0, - 0.08 * obs["input_batemanscale"] * (obs["input_ntor"] / 8.0) - + 0.06 * obs["input_pscale"] / max(obs["q95"], 0.1) - + 0.04 * obs["eq_kappa"] * obs["eq_delta"] - - 0.02 * obs["q0"] - + float(rng.normal(0.0, self._noise_std)), - )) - return obs diff --git a/use-cases/m3dc1-stream/sensor_buffer.py b/use-cases/m3dc1-stream/sensor_buffer.py deleted file mode 100644 index 676ce79..0000000 --- a/use-cases/m3dc1-stream/sensor_buffer.py +++ /dev/null @@ -1,122 +0,0 @@ -""" -SensorBuffer: accumulates observations from a SensorStream into a bounded -in-memory deque and surfaces snapshots as pandas DataFrames. - -The buffer runs a background asyncio task that calls sensor.stream() and -appends each observation as it arrives. Workflow tasks call wait_for(n) to -block until enough data is available, then take a snapshot for training. -""" -from __future__ import annotations - -import asyncio -from collections import deque -from pathlib import Path - -import pandas as pd - -from sensor import SensorStream - - -class SensorBuffer: - """Thread-of-execution-safe accumulator for a continuous sensor stream. - - Args: - sensor: Any SensorStream implementation (mock or real). - maxlen: Maximum observations to retain (oldest discarded when full). - Set to None for an unbounded buffer. - """ - - def __init__(self, sensor: SensorStream, maxlen: int | None = 10_000) -> None: - self._sensor = sensor - self._deque: deque[dict] = deque(maxlen=maxlen) - self._event: asyncio.Event | None = None - self._loop: asyncio.AbstractEventLoop | None = None # loop that owns the Event - self._task: asyncio.Task | None = None - - # ── Lifecycle ──────────────────────────────────────────────────────────── - - async def start(self) -> None: - """Start the background ingestion task.""" - self._loop = asyncio.get_running_loop() - self._event = asyncio.Event() - self._task = asyncio.create_task(self._ingest(), name="sensor-ingest") - - async def stop(self) -> None: - """Cancel the background ingestion task and wait for it to exit.""" - if self._task is not None: - self._task.cancel() - try: - await self._task - except asyncio.CancelledError: - pass - self._task = None - - # ── Internal ───────────────────────────────────────────────────────────── - - async def _ingest(self) -> None: - async for obs in self._sensor.stream(): - self._deque.append(obs) - self._event.set() # wake any waiter; they will clear it themselves - - # ── Public API ─────────────────────────────────────────────────────────── - - def __len__(self) -> int: - return len(self._deque) - - def snapshot(self, n: int | None = None) -> pd.DataFrame: - """Return the most recent n observations as a DataFrame (no blocking).""" - rows = list(self._deque) - if n is not None: - rows = rows[-n:] - return pd.DataFrame(rows) - - async def wait_for(self, n: int, timeout: float | None = None) -> pd.DataFrame: - """Block until at least n observations are buffered, then return them. - - Args: - n: Minimum number of observations required. - timeout: Seconds to wait before raising TimeoutError (None = forever). - - Returns: - DataFrame of the most recent n observations. - - Raises: - TimeoutError: If timeout elapses before n observations arrive. - """ - async def _wait() -> pd.DataFrame: - while True: - if len(self._deque) >= n: - return self.snapshot(n) - # Double-check after clearing to avoid missing an arrival that - # landed between the size check above and the clear below. - self._event.clear() - if len(self._deque) >= n: - return self.snapshot(n) - # Always schedule event.wait() on the loop that owns the Event. - # If the caller is on a different loop (e.g. WorkflowEngine's - # internal loop), bridge via run_coroutine_threadsafe so the - # wait runs on self._loop and the result is surfaced back here. - if asyncio.get_running_loop() is self._loop: - await self._event.wait() - else: - fut = asyncio.run_coroutine_threadsafe( - self._event.wait(), self._loop - ) - await asyncio.wrap_future(fut) - - if timeout is not None: - return await asyncio.wait_for(_wait(), timeout=timeout) - return await _wait() - - async def write_snapshot( - self, - n: int, - path: Path, - *, - timeout: float | None = None, - ) -> int: - """Wait for n observations, write them to a Parquet file, return row count.""" - df = await self.wait_for(n, timeout=timeout) - path.parent.mkdir(parents=True, exist_ok=True) - df.to_parquet(path, index=False) - return len(df) From 66fdb190a82dd242f4127e3084bfea517bc264f9 Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:16:09 -0700 Subject: [PATCH 05/22] Switch to use REDIS for investigator --> sim comm --- .../m3dc1-stream/dt/amsc_investigator.py | 106 +++++++++++------- use-cases/m3dc1-stream/dt/buffer.py | 11 +- use-cases/m3dc1-stream/dt/dtypes.py | 2 +- use-cases/m3dc1-stream/dt/run_me.py | 47 ++++++-- 4 files changed, 110 insertions(+), 56 deletions(-) diff --git a/use-cases/m3dc1-stream/dt/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py index 2310666..9a81dfe 100644 --- a/use-cases/m3dc1-stream/dt/amsc_investigator.py +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -41,7 +41,7 @@ from radical.asyncflow import WorkflowEngine from rose import LearnerConfig, TaskConfig from rose.al import ParallelActiveLearner - +import redis from dtypes import M3DC1_PREDICTION # ── Row-budget policy (mirrors amsc.py's growing-pool logic) ───────────────── @@ -89,7 +89,10 @@ def __init__( flow: WorkflowEngine, candidates: list[str], max_iter: int, + buffer_max: int, r2_threshold: float, + redis_endpoint: str, + redis_key: str, ): super().__init__(flow) @@ -97,59 +100,76 @@ def __init__( self.candidates = candidates self.max_iter = max_iter self.r2_threshold = r2_threshold - - # each surrogate gets their own event - self.data_event: dict[str, asyncio.Event] = {} - + self.buffer_max = buffer_max + + self.all_data: list[dict] = [] + + # use REDIS for communication from the investigator to the Simulation. + # see note at top of file. This is required as the simulation task + # itself it waiting for data. (Other DT examples have it where the + # simulation task is fired after receiving the data.) + host, port = redis_endpoint.rsplit(":", 1) + self.redis = redis.Redis(host=host, port=int(port), decode_responses=True) + self.redis_key = redis_key + # ensure start clear for candidate in self.candidates: - self.data_event[candidate] = asyncio.Event() + self.redis.delete(f"{redis_key}/{candidate}") # ── Simulation task ─────────────────────────────────────────────────────── # KEY CHANGE vs amsc_stream. Buffered inputs come in from the DT. - @self.learner.simulation_task - async def simulation(*args, **kwargs) -> str: + @self.learner.simulation_task(as_executable=False) + async def simulation(*args, **kwargs) -> dict: + import time + import redis as _redis + import pandas as pd + it = int(kwargs.get("iteration", 0)) label = str(kwargs["learner_label"]) - family = str(kwargs["model_family"]) - n_rows = N_BASE + it * N_STEP - out_dir = _workspace_iter(it, label) - parquet = out_dir / "sensor_snapshot.parquet" + family = str(kwargs["model_family"]) - print( - f" [sim {label} iter={it}] Now has {n_rows} sensor rows …", flush=True + host, port_str = redis_endpoint.rsplit(":", 1) + redis_client = _redis.Redis( + host=host, port=int(port_str), decode_responses=True ) - await self.data_event[family].wait() - n_written = len(self.all_data) + print(f" [sim {label} iter={it}] waiting for {n_rows} rows …", flush=True) + deadline = time.monotonic() + 600.0 + while not redis_client.exists( + redis_key + "/MAIN" + ) or not redis_client.exists(f"{redis_key}/{family}"): + time.sleep(0.5) + if time.monotonic() > deadline: + raise TimeoutError(f"Timeout waiting for {n_rows} sensor rows") + + resp = redis_client.get(redis_key + "/MAIN") + assert resp is not None + rows = json.loads(resp) + redis_client.delete(f"{redis_key}/{family}") + + df = pd.DataFrame(rows) + + out_dir = _workspace_iter(it, label) + parquet = out_dir / "sensor_snapshot.parquet" + df.to_parquet(parquet, index=False) meta = { "iteration": it, "learner_label": label, "dataset": str(parquet), - "n_rows": n_written, - "source": "sensor_stream", + "n_rows": len(df), + "source": "redis_stream", } (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) - - self.data_event[family].clear() - - # workaround: I know, looks a little goofy. Read the notes at top of - # file. Parameter is passed via `out_dir/simulation.json` - # - # Consequence of this workaround: on shutdown, AsyncFlow will - # complain that the future is missing some attributes. - # This is because AsyncFlow stamps the attributes when the cmdline - # is returned. It doesn't hurt accuracy in any way, just doesn't - # look as clean. - return f"echo {shlex.quote(str(out_dir / "simulation.json"))}" + return meta # ── Training task ───────────────────────────────────────────────────────── # Fits a surrogate model locally using sklearn. # Replace with subprocess to surge_train.py if running on HPC. + @self.learner.training_task(as_executable=False) - async def training(sim_result_path: str, **kwargs) -> dict: + async def training(sim_result: str, **kwargs) -> dict: import pandas as pd from sklearn.ensemble import ( GradientBoostingRegressor, @@ -162,9 +182,6 @@ async def training(sim_result_path: str, **kwargs) -> dict: from sklearn.preprocessing import StandardScaler from sklearn.pipeline import Pipeline - with open(sim_result_path, "r") as f: - sim_result = json.load(f) - it = int(kwargs.get("iteration", sim_result["iteration"])) label = str(kwargs["learner_label"]) family = str(kwargs.get("model_family", "rf")) @@ -225,7 +242,6 @@ async def training(sim_result_path: str, **kwargs) -> dict: model_path = str((out_dir / "model.pkl")) with open(model_path, "wb") as f: cloudpickle.dump(model, f) - return {"simulation": sim_result, "surge": metrics, "model": model_path} # ── Active-learning task ────────────────────────────────────────────────── @@ -270,6 +286,7 @@ async def stop_on_r2(*args, **kwargs) -> float: forced = it >= max_iter - 1 and r2 < r2_threshold return r2_threshold if forced else r2 + @self.flow.function_task async def do_inference(in_data: WindowedTypeData, model=None): # the ASMC_stream.py demo doesn't tackle streaming inference. # Put streaming inference code here. @@ -292,13 +309,20 @@ async def input_callback(self, in_data: WindowedTypeData): # add the data to large database self.all_data += in_data.sequence # trigger event - for v in self.data_event.values(): - v.set() + # put all_data onto redis + + if len(self.all_data) > self.buffer_max: + self.all_data = self.all_data[-self.buffer_max :] + + self.redis.set(self.redis_key + "/MAIN", json.dumps(self.all_data)) + for c in self.candidates: + self.redis.set(f"{self.redis_key}/{c}", 1) async def main_loop(self, runtime: RuntimeAPI): # call the pipeline runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) + runtime.set_inference_task(self.inference) runtime.publish_new_model({"model": None}) configs = _candidate_configs(self.candidates, self.max_iter) @@ -319,9 +343,11 @@ async def main_loop(self, runtime: RuntimeAPI): } ) print( - f" learner={label} iter={state.iteration}" - f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}" - f" buffer={len(self.all_data)} obs", + "\nMODEL PUBLISHED -----------------------------------" + f" learner={label} iter={state.iteration}\n" + f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}\n" + f" buffer={len(self.all_data)} obs\n" + f" ---------------------------------------------------\n", flush=True, ) diff --git a/use-cases/m3dc1-stream/dt/buffer.py b/use-cases/m3dc1-stream/dt/buffer.py index 4660d5d..8ec561a 100644 --- a/use-cases/m3dc1-stream/dt/buffer.py +++ b/use-cases/m3dc1-stream/dt/buffer.py @@ -37,14 +37,15 @@ def __init__(self, flow, buffer_length): async def main_loop(self, runtime, in_data: TypedData): self.counter += 1 - out_data = TypedData( - SYNC_SENSOR, - ) + out_data = TypedData(SYNC_SENSOR, in_data.data) if self.counter >= self.buffer_length: # emit event - return in_data, TypedData(BUFFER_EVENT, time.time()) + print("EMIT!!") + self.counter = 0 + return out_data, TypedData(BUFFER_EVENT, time.time()) - return in_data, None + print(f"Buffer counter: {self.counter} / {self.buffer_length}") + return out_data, None # this is needed to diff --git a/use-cases/m3dc1-stream/dt/dtypes.py b/use-cases/m3dc1-stream/dt/dtypes.py index fc9c828..e08b427 100644 --- a/use-cases/m3dc1-stream/dt/dtypes.py +++ b/use-cases/m3dc1-stream/dt/dtypes.py @@ -7,7 +7,7 @@ # Buffer event BUFFER_EVENT = DataType("BUFFER_EVENT") -SYNC_SENSOR = DataType("M3DC1_SYNC") +SYNC_SENSOR = DataType("SYNC_SENSOR") # Modeling diff --git a/use-cases/m3dc1-stream/dt/run_me.py b/use-cases/m3dc1-stream/dt/run_me.py index e72e1f7..28dc76d 100644 --- a/use-cases/m3dc1-stream/dt/run_me.py +++ b/use-cases/m3dc1-stream/dt/run_me.py @@ -36,8 +36,14 @@ # The sensor is external: run sensor.py in its own terminal. -async def main(candidates, max_iter, r2_threshold, buffer_length): - init_default_logger(logging.INFO) +async def main(candidates, args): + max_iter = args.max_iter + r2_threshold = args.r2_threshold + max_len = args.buffer_maxlen + window_size = args.buffer_size + redis_endpoint = args.redis_endpoint + + init_default_logger(logging.WARNING) logging.getLogger("radical.asyncflow").setLevel(logging.WARNING) logging.getLogger("rhapsody").setLevel(logging.WARNING) @@ -54,15 +60,24 @@ async def main(candidates, max_iter, r2_threshold, buffer_length): # create tasks and investigators # for buffer: - buffer_event_task = BufferEventEmit(flow, buffer_length) + buffer_event_task = BufferEventEmit(flow, window_size) barrier = Barrier("sensor_windowing_buffer") - WINDOW_DTYPE = barrier.add_dtype(SYNC_SENSOR, hard=False) + SENSOR_WINDOW = barrier.add_dtype(SYNC_SENSOR, hard=False) + dead_sink = DeadSink(flow) barrier.add_dtype(BUFFER_EVENT, hard=True) # for investigator - m3dc1 = M3DC1_Investigator(flow, candidates, max_iter, r2_threshold) + m3dc1 = M3DC1_Investigator( + flow, + candidates, + max_iter, + max_len, + r2_threshold, + redis_endpoint, + "M3DC1", + ) output_sink = OutputSink(flow) # create graph @@ -72,14 +87,14 @@ async def main(candidates, max_iter, r2_threshold, buffer_length): buffer_event_task, M3DC1_SENSOR, [SYNC_SENSOR, BUFFER_EVENT] ) runtime.add_task(dead_sink, BUFFER_EVENT, NULL_DTYPE) - runtime.add_investigator(m3dc1, SYNC_SENSOR, M3DC1_PREDICTION) + runtime.add_investigator(m3dc1, SENSOR_WINDOW, M3DC1_PREDICTION) runtime.add_task(output_sink, M3DC1_PREDICTION, NULL_DTYPE) runtime.print_graph() - # runtime.start() + runtime.start() # let it run - await asyncio.sleep(5) + await asyncio.sleep(45) await runtime.stop() await flow.shutdown() @@ -98,13 +113,25 @@ async def main(candidates, max_iter, r2_threshold, buffer_length): parser.add_argument( "--buffer-maxlen", type=int, - default=10_000, + default=100, help="Maximum observations retained in the sensor buffer.", ) + parser.add_argument( + "--buffer-size", + type=int, + default=100, + help="Window size for sensor data", + ) + parser.add_argument( + "--redis-endpoint", + type=str, + default="localhost:6379", + help="Redis endpoint used in M3D1C investigator.", + ) args = parser.parse_args() candidates = [x.strip() for x in args.candidates.split(",") if x.strip()] if len(candidates) < 2: parser.error("Need at least two candidates for ParallelActiveLearner.") - asyncio.run(main(candidates, args.max_iter, args.r2_threshold, args.buffer_maxlen)) + asyncio.run(main(candidates, args)) From d4fe9c658fd298e9483ba2e63fd5078ca7f7f866 Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:46:26 -0700 Subject: [PATCH 06/22] Simply to do buffering in investigator. Switch to use RedisDataBackend --- use-cases/m3dc1-stream/dt/.gitignore | 5 ++ use-cases/m3dc1-stream/dt/README.md | 5 +- .../m3dc1-stream/dt/amsc_investigator.py | 49 ++++++++-------- use-cases/m3dc1-stream/dt/buffer.py | 57 ------------------ use-cases/m3dc1-stream/dt/dtypes.py | 4 -- use-cases/m3dc1-stream/dt/run_me.py | 58 ++++++------------- 6 files changed, 48 insertions(+), 130 deletions(-) delete mode 100644 use-cases/m3dc1-stream/dt/buffer.py diff --git a/use-cases/m3dc1-stream/dt/.gitignore b/use-cases/m3dc1-stream/dt/.gitignore index c18dd8d..db0f920 100644 --- a/use-cases/m3dc1-stream/dt/.gitignore +++ b/use-cases/m3dc1-stream/dt/.gitignore @@ -1 +1,6 @@ __pycache__/ +tmp.py +rhapsody.data.* +workspace/ +__pycache__/ +dump.rdb \ No newline at end of file diff --git a/use-cases/m3dc1-stream/dt/README.md b/use-cases/m3dc1-stream/dt/README.md index b9ea574..7da461b 100644 --- a/use-cases/m3dc1-stream/dt/README.md +++ b/use-cases/m3dc1-stream/dt/README.md @@ -1,10 +1,9 @@ -# The M3DC1 (spark) ported over to the Digital Twin framework. +# The M3DC1-stream (aka spark-stream) ported over to the Digital Twin framework. Items: -- `sensor.py` --> `dt/sensor.py` -- `sensor_buffer.py` --> `dt/buffer.py` +- `sensor_daemon.py` --> `dt/sensor.py` - `amsc_stream.py` --> `dt/amsc_investigator.py` Other: diff --git a/use-cases/m3dc1-stream/dt/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py index 9a81dfe..65c064b 100644 --- a/use-cases/m3dc1-stream/dt/amsc_investigator.py +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -9,24 +9,17 @@ a ScienceAgent. The point of a science agent is in the event you have various surrogates with separate ALs) -Also, in the M3DC1 streaming example, it assumes a ThreadedPoolExecutor, and all -tasks running on the same machine. This is due to the use of "buffer" inside the -simulation task. However, we want to demonstrate this code working across -machines, so this restriction must be lifted. - -A key point here is that the simulation task "waits" for new data. As the -ParallelActiveLearner is used, we don't have the flexibility to trigger the -simulation when we want to ourselves. - -Currently, the purpose of the simulation is to simply wait until the data is -available. So, +One tricky part is that the simulation task itself is waiting for streaming data +(opposed to the pipeline waiting and then launching the sim.). This requires a +way to transfer data from the investigator to the simulation task as the +simulation task is running. This example uses REDIS from the Rhapsody Data +Backend. """ import asyncio import json from pathlib import Path -import shlex import cloudpickle from digitaltwin import ( @@ -34,7 +27,6 @@ RuntimeAPI, TypedData, UtilityTask, - WindowedTypeData, ) import numpy as np import pandas as pd @@ -87,9 +79,11 @@ class M3DC1_Investigator(ModelInvestigator): def __init__( self, flow: WorkflowEngine, + *, candidates: list[str], max_iter: int, buffer_max: int, + window_size: int, r2_threshold: float, redis_endpoint: str, redis_key: str, @@ -101,8 +95,9 @@ def __init__( self.max_iter = max_iter self.r2_threshold = r2_threshold self.buffer_max = buffer_max - + self.window_size = window_size self.all_data: list[dict] = [] + self.input_counter = 0 # use REDIS for communication from the investigator to the Simulation. # see note at top of file. This is required as the simulation task @@ -287,7 +282,7 @@ async def stop_on_r2(*args, **kwargs) -> float: return r2_threshold if forced else r2 @self.flow.function_task - async def do_inference(in_data: WindowedTypeData, model=None): + async def do_inference(in_data: TypedData, model=None): # the ASMC_stream.py demo doesn't tackle streaming inference. # Put streaming inference code here. if model is None: @@ -296,7 +291,7 @@ async def do_inference(in_data: WindowedTypeData, model=None): with open(model, "rb") as f: model_obj = cloudpickle.load(f) - df = pd.DataFrame(in_data.sequence) + df = pd.DataFrame([in_data.data]) X = df.drop(columns=["output_gamma"]).values out = model_obj.predict(X) @@ -305,18 +300,20 @@ async def do_inference(in_data: WindowedTypeData, model=None): self.inference = do_inference - async def input_callback(self, in_data: WindowedTypeData): + async def input_callback(self, in_data: TypedData): # add the data to large database - self.all_data += in_data.sequence - # trigger event - # put all_data onto redis + self.all_data.append(in_data.data) if len(self.all_data) > self.buffer_max: - self.all_data = self.all_data[-self.buffer_max :] + self.all_data.pop(0) + + self.input_counter += 1 - self.redis.set(self.redis_key + "/MAIN", json.dumps(self.all_data)) - for c in self.candidates: - self.redis.set(f"{self.redis_key}/{c}", 1) + if self.input_counter > self.window_size: + self.input_counter = 0 + self.redis.set(self.redis_key + "/MAIN", json.dumps(self.all_data)) + for c in self.candidates: + self.redis.set(f"{self.redis_key}/{c}", 1) async def main_loop(self, runtime: RuntimeAPI): # call the pipeline @@ -343,7 +340,7 @@ async def main_loop(self, runtime: RuntimeAPI): } ) print( - "\nMODEL PUBLISHED -----------------------------------" + "\nMODEL PUBLISHED -----------------------------------\n" f" learner={label} iter={state.iteration}\n" f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}\n" f" buffer={len(self.all_data)} obs\n" @@ -364,4 +361,6 @@ def __init__(self, flow): super().__init__(flow) async def main_loop(self, runtime, in_data: TypedData): + if in_data.data is None: + return # don't print out None... that means there wasn't a model ready yet print("Received: ", in_data.data) diff --git a/use-cases/m3dc1-stream/dt/buffer.py b/use-cases/m3dc1-stream/dt/buffer.py deleted file mode 100644 index 8ec561a..0000000 --- a/use-cases/m3dc1-stream/dt/buffer.py +++ /dev/null @@ -1,57 +0,0 @@ -""" -Buffer: generic data buffer. Batches inputs and only emits when done. - -The digital twin framework provides a stream-processing dataflow paradigm. It already -comes with a Windowing feature via its BARRIER. However, it leaves the logic up to the user of -how their windows should be defined. - -So, you need a split task that generates an event of when to trigger the barrier. - -Dataflow Digital Twin graph: - -SENSOR --> BUFFER_EVENT_EMIT --> BARRIER --> SINK - | | - +--------------> BARRIER --> goes nowhere - -From this, the BUFFER_EVENT_EMIT is a SPLIT task. It takes in one stream and -generates two. - -The BARRIER is provided by the DT framework. This file simply implements the -buffer event split task, and a null sink. - -The null sink is needed to drain the queues. - -""" - -import time - -from digitaltwin import SplitTask, TypedData, UtilityTask -from dtypes import BUFFER_EVENT, SYNC_SENSOR - - -class BufferEventEmit(SplitTask): - def __init__(self, flow, buffer_length): - super().__init__(flow) - self.buffer_length = buffer_length - self.counter = 0 - - async def main_loop(self, runtime, in_data: TypedData): - self.counter += 1 - out_data = TypedData(SYNC_SENSOR, in_data.data) - if self.counter >= self.buffer_length: - # emit event - print("EMIT!!") - self.counter = 0 - return out_data, TypedData(BUFFER_EVENT, time.time()) - - print(f"Buffer counter: {self.counter} / {self.buffer_length}") - return out_data, None - - -# this is needed to -class DeadSink(UtilityTask): - def __init__(self, flow): - super().__init__(flow) - - async def main_loop(self, runtime, in_data): - pass # do nothing. diff --git a/use-cases/m3dc1-stream/dt/dtypes.py b/use-cases/m3dc1-stream/dt/dtypes.py index e08b427..6789122 100644 --- a/use-cases/m3dc1-stream/dt/dtypes.py +++ b/use-cases/m3dc1-stream/dt/dtypes.py @@ -5,10 +5,6 @@ M3DC1_MOCK_CHANNEL = "sensors/MockM3DC1" M3DC1_SENSOR = DataType("M3DC1") -# Buffer event -BUFFER_EVENT = DataType("BUFFER_EVENT") -SYNC_SENSOR = DataType("SYNC_SENSOR") - # Modeling M3DC1_PREDICTION = DataType("M3DC1_PREDICTION") diff --git a/use-cases/m3dc1-stream/dt/run_me.py b/use-cases/m3dc1-stream/dt/run_me.py index 28dc76d..0d209e6 100644 --- a/use-cases/m3dc1-stream/dt/run_me.py +++ b/use-cases/m3dc1-stream/dt/run_me.py @@ -3,9 +3,7 @@ Complete Digital Twin graph: -MOCK_SENSOR --> BUFFER_EVENT_TASK --> BARRIER --> M3DC1_Investigator --> OUTPUT TASK - | ||| - +-------------> BARRIER --> DEAD SINK +MOCK_SENSOR --> M3DC1_Investigator --> OUTPUT TASK """ @@ -13,7 +11,6 @@ import argparse import asyncio from concurrent.futures import ProcessPoolExecutor -from digitaltwin import Barrier from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend @@ -22,26 +19,24 @@ from digitaltwin.components import NULL_DTYPE from amsc_investigator import M3DC1_Investigator, OutputSink -from buffer import BufferEventEmit, DeadSink from dtypes import * from radical.asyncflow.logging import init_default_logger +from rhapsody.backends.data.redis import RedisDataBackend import logging logger = logging.getLogger(__name__) -# put it all together -# sensor channel --> model --> data_sink -# -# The sensor is external: run sensor.py in its own terminal. - async def main(candidates, args): max_iter = args.max_iter r2_threshold = args.r2_threshold max_len = args.buffer_maxlen - window_size = args.buffer_size - redis_endpoint = args.redis_endpoint + window_size = args.window_size + + redis_backend = await RedisDataBackend() + endpoint = redis_backend.endpoints[0] + redis_endpoint = endpoint.serialize() init_default_logger(logging.WARNING) logging.getLogger("radical.asyncflow").setLevel(logging.WARNING) @@ -59,35 +54,21 @@ async def main(candidates, args): ################### # create tasks and investigators - # for buffer: - buffer_event_task = BufferEventEmit(flow, window_size) - barrier = Barrier("sensor_windowing_buffer") - SENSOR_WINDOW = barrier.add_dtype(SYNC_SENSOR, hard=False) - - dead_sink = DeadSink(flow) - - barrier.add_dtype(BUFFER_EVENT, hard=True) - - # for investigator m3dc1 = M3DC1_Investigator( flow, - candidates, - max_iter, - max_len, - r2_threshold, - redis_endpoint, - "M3DC1", + candidates=candidates, + max_iter=max_iter, + buffer_max=max_len, + window_size=window_size, + r2_threshold=r2_threshold, + redis_endpoint=redis_endpoint, + redis_key="M3DC1", ) output_sink = OutputSink(flow) # create graph runtime.add_input(M3DC1_SENSOR, M3DC1_MOCK_CHANNEL) - runtime.add_barrier(barrier) - runtime.add_data_split_task( - buffer_event_task, M3DC1_SENSOR, [SYNC_SENSOR, BUFFER_EVENT] - ) - runtime.add_task(dead_sink, BUFFER_EVENT, NULL_DTYPE) - runtime.add_investigator(m3dc1, SENSOR_WINDOW, M3DC1_PREDICTION) + runtime.add_investigator(m3dc1, M3DC1_SENSOR, M3DC1_PREDICTION) runtime.add_task(output_sink, M3DC1_PREDICTION, NULL_DTYPE) runtime.print_graph() @@ -95,6 +76,7 @@ async def main(candidates, args): # let it run await asyncio.sleep(45) + print("SHUTDOWN") await runtime.stop() await flow.shutdown() @@ -117,17 +99,11 @@ async def main(candidates, args): help="Maximum observations retained in the sensor buffer.", ) parser.add_argument( - "--buffer-size", + "--window-size", type=int, default=100, help="Window size for sensor data", ) - parser.add_argument( - "--redis-endpoint", - type=str, - default="localhost:6379", - help="Redis endpoint used in M3D1C investigator.", - ) args = parser.parse_args() candidates = [x.strip() for x in args.candidates.split(",") if x.strip()] From 3b2b2a4e64b36c359e9480f6b896d8d90ad0ae7f Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:55:28 -0700 Subject: [PATCH 07/22] Add run instructions to README --- use-cases/m3dc1-stream/dt/README.md | 17 +++++++++++++++++ use-cases/m3dc1-stream/dt/local_broker.py | 20 ++++++++++++++++++++ 2 files changed, 37 insertions(+) create mode 100644 use-cases/m3dc1-stream/dt/local_broker.py diff --git a/use-cases/m3dc1-stream/dt/README.md b/use-cases/m3dc1-stream/dt/README.md index 7da461b..5e82ea2 100644 --- a/use-cases/m3dc1-stream/dt/README.md +++ b/use-cases/m3dc1-stream/dt/README.md @@ -9,3 +9,20 @@ Items: Other: `dt/run_me.py` `dt/dtypes.py` + + +## To run: + +1. Install the digital twins library: + +``` bash +git clone https://github.com/radical-cybertools/digital.twins + +# this is for the plain DT framework without all the as-a-service changes +git checkout release/vanilla-framework +pip install . +``` + +2. Run in one terminal `python3 sensor.py` +3. Run in a second terminal `python3 local_broker.py` +4. Finally, in a third terminal, run `python3 run_me.py` diff --git a/use-cases/m3dc1-stream/dt/local_broker.py b/use-cases/m3dc1-stream/dt/local_broker.py new file mode 100644 index 0000000..8587ffe --- /dev/null +++ b/use-cases/m3dc1-stream/dt/local_broker.py @@ -0,0 +1,20 @@ +"""Standalone stream broker for the two-terminal demos. + +Addresses come from configuration (`DT_STREAM_PUB_ADDR` / +`DT_STREAM_SUB_ADDR`, loopback defaults) -- the same resolution the +demos use, so both terminals agree without any literal in the code. +""" + +from digitaltwin.config import stream_addresses +from digitaltwin.streaming import ZMQ_Broker + +if __name__ == "__main__": + broker = ZMQ_Broker(*stream_addresses()) + + publish_addr, subscribe_addr = broker.bind() + print( + f"stream broker: publish to {publish_addr}, subscribe on {subscribe_addr}", + flush=True, + ) + + broker.run() From 72a80f73e1c5aaf3b9aa9b7d52963a36e09c8f80 Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:00:45 -0700 Subject: [PATCH 08/22] Switch to window size --- use-cases/m3dc1-stream/dt/amsc_investigator.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/use-cases/m3dc1-stream/dt/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py index 65c064b..9afd967 100644 --- a/use-cases/m3dc1-stream/dt/amsc_investigator.py +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -36,13 +36,6 @@ import redis from dtypes import M3DC1_PREDICTION -# ── Row-budget policy (mirrors amsc.py's growing-pool logic) ───────────────── -# At iteration i, the simulation task requests N_BASE + i * N_STEP rows from -# the buffer. Increase N_STEP to consume more data per iteration. -N_BASE = 100 -N_STEP = 100 - - # Workspace for iteration artefacts (parquet snapshots, metric JSON) _HERE = Path(__file__).resolve().parent _WORKSPACE = _HERE / "workspace" @@ -121,7 +114,6 @@ async def simulation(*args, **kwargs) -> dict: it = int(kwargs.get("iteration", 0)) label = str(kwargs["learner_label"]) - n_rows = N_BASE + it * N_STEP family = str(kwargs["model_family"]) host, port_str = redis_endpoint.rsplit(":", 1) @@ -129,14 +121,16 @@ async def simulation(*args, **kwargs) -> dict: host=host, port=int(port_str), decode_responses=True ) - print(f" [sim {label} iter={it}] waiting for {n_rows} rows …", flush=True) + print( + f" [sim {label} iter={it}] waiting for {window_size} rows", flush=True + ) deadline = time.monotonic() + 600.0 while not redis_client.exists( redis_key + "/MAIN" ) or not redis_client.exists(f"{redis_key}/{family}"): time.sleep(0.5) if time.monotonic() > deadline: - raise TimeoutError(f"Timeout waiting for {n_rows} sensor rows") + raise TimeoutError(f"Timeout waiting for {window_size} sensor rows") resp = redis_client.get(redis_key + "/MAIN") assert resp is not None From a852991dd202e9425a4fd4c814b8310f91f201c8 Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:03:39 -0700 Subject: [PATCH 09/22] Add note --- use-cases/m3dc1-stream/{dt => }/.gitignore | 3 ++- use-cases/m3dc1-stream/dt/README.md | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) rename use-cases/m3dc1-stream/{dt => }/.gitignore (77%) diff --git a/use-cases/m3dc1-stream/dt/.gitignore b/use-cases/m3dc1-stream/.gitignore similarity index 77% rename from use-cases/m3dc1-stream/dt/.gitignore rename to use-cases/m3dc1-stream/.gitignore index db0f920..4a98218 100644 --- a/use-cases/m3dc1-stream/dt/.gitignore +++ b/use-cases/m3dc1-stream/.gitignore @@ -3,4 +3,5 @@ tmp.py rhapsody.data.* workspace/ __pycache__/ -dump.rdb \ No newline at end of file +dump.rdb +.vscode/ \ No newline at end of file diff --git a/use-cases/m3dc1-stream/dt/README.md b/use-cases/m3dc1-stream/dt/README.md index 5e82ea2..0e39cae 100644 --- a/use-cases/m3dc1-stream/dt/README.md +++ b/use-cases/m3dc1-stream/dt/README.md @@ -10,6 +10,10 @@ Other: `dt/run_me.py` `dt/dtypes.py` +The digital twin framework handles sensor streams directly, so the demo defers +in-stream data movement to the digital twin framework. (Data is small enough where +this works). + ## To run: From d0b8e50338283199c7efe65e13d15dde0acd8504 Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Wed, 26 Aug 2026 20:25:59 -0700 Subject: [PATCH 10/22] Typo --- use-cases/m3dc1-stream/dt/amsc_investigator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/use-cases/m3dc1-stream/dt/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py index 9afd967..680451f 100644 --- a/use-cases/m3dc1-stream/dt/amsc_investigator.py +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -310,7 +310,7 @@ async def input_callback(self, in_data: TypedData): self.redis.set(f"{self.redis_key}/{c}", 1) async def main_loop(self, runtime: RuntimeAPI): - # call the pipeline + # run the pipeline runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) runtime.set_inference_task(self.inference) From 35bbc1dee6c0b78d8e4f604ca4aafa12b11b38ba Mon Sep 17 00:00:00 2001 From: Benjamin Carter <59350660+BenCarter44@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:54:29 -0400 Subject: [PATCH 11/22] Update comment Clarified comment regarding buffered inputs in simulation task. --- use-cases/m3dc1-stream/dt/amsc_investigator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/use-cases/m3dc1-stream/dt/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py index 680451f..86c5e5f 100644 --- a/use-cases/m3dc1-stream/dt/amsc_investigator.py +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -104,7 +104,8 @@ def __init__( self.redis.delete(f"{redis_key}/{candidate}") # ── Simulation task ─────────────────────────────────────────────────────── - # KEY CHANGE vs amsc_stream. Buffered inputs come in from the DT. + # KEY CHANGE vs amsc_stream. Buffered inputs come in from the investigator's + # input callback @self.learner.simulation_task(as_executable=False) async def simulation(*args, **kwargs) -> dict: From f522dd171e87eebe7159a7382f77cf82df2044ab Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:03:57 -0700 Subject: [PATCH 12/22] Foundations of the DT-Complete demo --- use-cases/{m3dc1-stream => }/.gitignore | 0 use-cases/dt-complete/README.md | 47 +++ use-cases/dt-complete/dtypes.py | 22 ++ use-cases/dt-complete/local_broker.py | 20 + use-cases/dt-complete/m3dc1/m3dc1_dtypes.py | 6 + .../dt-complete/m3dc1/m3dc1_investigator.py | 360 ++++++++++++++++++ .../dt-complete/m3dc1/m3dc1_mock_sensor.py | 117 ++++++ use-cases/dt-complete/out.py | 12 + use-cases/dt-complete/run_me.py | 115 ++++++ use-cases/m3dc1-stream/dt/README.md | 2 +- .../m3dc1-stream/dt/amsc_investigator.py | 19 +- use-cases/m3dc1-stream/dt/run_me.py | 4 +- use-cases/m3dc1-stream/dt/sensor.py | 4 +- 13 files changed, 718 insertions(+), 10 deletions(-) rename use-cases/{m3dc1-stream => }/.gitignore (100%) create mode 100644 use-cases/dt-complete/README.md create mode 100644 use-cases/dt-complete/dtypes.py create mode 100644 use-cases/dt-complete/local_broker.py create mode 100644 use-cases/dt-complete/m3dc1/m3dc1_dtypes.py create mode 100644 use-cases/dt-complete/m3dc1/m3dc1_investigator.py create mode 100644 use-cases/dt-complete/m3dc1/m3dc1_mock_sensor.py create mode 100644 use-cases/dt-complete/out.py create mode 100644 use-cases/dt-complete/run_me.py diff --git a/use-cases/m3dc1-stream/.gitignore b/use-cases/.gitignore similarity index 100% rename from use-cases/m3dc1-stream/.gitignore rename to use-cases/.gitignore diff --git a/use-cases/dt-complete/README.md b/use-cases/dt-complete/README.md new file mode 100644 index 0000000..7b8090a --- /dev/null +++ b/use-cases/dt-complete/README.md @@ -0,0 +1,47 @@ + +# DT-Complete + +A demonstration of a complete digital twin: + +2 sensors + 3 surrogates* / physics entities + +Sensors: +- M3DC1 Mock sensor +- Random Value sensor + +Three physical entities: +- M3DC1 Investigator +- Runs a M3DC1 Investigator +- Runs a DEMO_AGENT Agent (a simple pass through) +- Runs a NEG_AGENT Agent (simply computes the negative of sensor input) + +> *Technically, the M3DC1 trains two surrogates and then picks the best one. + + +**Digital Twin Description Graph:** +``` + +M3DC1 Mock sensor --> M3DC1 Investigator -- + \ + --(JOIN)--> DEMO Agent --> OUT +RAND_VAL sensor ---> PN_AGENT -----------/ + +``` + + + +## To run: + +1. Install the digital twins library: + +``` bash +git clone https://github.com/radical-cybertools/digital.twins + +# this is for the plain DT framework without all the as-a-service changes +git checkout release/vanilla-framework +pip install . +``` + +2. Start up your sensors: `python3 m3dc1_mock_sensor.py` and `python3 rand_sensor.py` +3. Start up the PUB/SUB streaming broker: `python3 local_broker.py` +4. Finally, run `python3 run_me.py` diff --git a/use-cases/dt-complete/dtypes.py b/use-cases/dt-complete/dtypes.py new file mode 100644 index 0000000..6d38dde --- /dev/null +++ b/use-cases/dt-complete/dtypes.py @@ -0,0 +1,22 @@ +from digitaltwin.components import DataType + +############################# +# Complete Digital Twin demo DATA_TYPES +############################# + +# use the M3DC1 sensor and prediction data types +from m3dc1.m3dc1_dtypes import * + +RANDOM_CHANNEL = "sensors/RANDOM" +RAND_SENSOR = DataType("RAND") + +# Physical entities: + +# M3DC1 Prediction + + +# Demo Agent Prediction +DEMO_PREDICTION = DataType("DEMO_PREDICTION") + +# Negate Agent Prediction +NEG_PREDICTION = DataType("NEG_PREDICTION") diff --git a/use-cases/dt-complete/local_broker.py b/use-cases/dt-complete/local_broker.py new file mode 100644 index 0000000..8587ffe --- /dev/null +++ b/use-cases/dt-complete/local_broker.py @@ -0,0 +1,20 @@ +"""Standalone stream broker for the two-terminal demos. + +Addresses come from configuration (`DT_STREAM_PUB_ADDR` / +`DT_STREAM_SUB_ADDR`, loopback defaults) -- the same resolution the +demos use, so both terminals agree without any literal in the code. +""" + +from digitaltwin.config import stream_addresses +from digitaltwin.streaming import ZMQ_Broker + +if __name__ == "__main__": + broker = ZMQ_Broker(*stream_addresses()) + + publish_addr, subscribe_addr = broker.bind() + print( + f"stream broker: publish to {publish_addr}, subscribe on {subscribe_addr}", + flush=True, + ) + + broker.run() diff --git a/use-cases/dt-complete/m3dc1/m3dc1_dtypes.py b/use-cases/dt-complete/m3dc1/m3dc1_dtypes.py new file mode 100644 index 0000000..3de3cbf --- /dev/null +++ b/use-cases/dt-complete/m3dc1/m3dc1_dtypes.py @@ -0,0 +1,6 @@ +from digitaltwin.components import DataType + +M3DC1_MOCK_CHANNEL = "sensors/MockM3DC1" +M3DC1_SENSOR = DataType("M3DC1") + +M3DC1_PREDICTION = DataType("M3DC1_PREDICTION") diff --git a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py new file mode 100644 index 0000000..1c0475b --- /dev/null +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -0,0 +1,360 @@ +""" +M3DC1 streaming digital twin. + +This is a version of the M3DC1 streaming surrogate that uses the digital twin +framework. + +Because we are using one active learner and only care about one physics +property, it's simplest to use just a single DT model investigator (no need for +a ScienceAgent. The point of a science agent is in the event you have various +surrogates with separate ALs) + +One tricky part is that the simulation task itself is waiting for streaming data +(opposed to the pipeline waiting and then launching the sim.). This requires a +way to transfer data from the investigator to the simulation task as the +simulation task is running. This example uses REDIS from the Rhapsody Data +Backend. + +""" + +import asyncio +import json +from pathlib import Path + +import cloudpickle +from digitaltwin import ( + ModelInvestigator, + RuntimeAPI, + TypedData, + UtilityTask, +) +import numpy as np +import pandas as pd +from radical.asyncflow import WorkflowEngine +from rose import LearnerConfig, TaskConfig +from rose.al import ParallelActiveLearner +import redis + +from .m3dc1_dtypes import * + +# Workspace for iteration artefacts (parquet snapshots, metric JSON) +_HERE = Path(__file__).resolve().parent +_WORKSPACE = _HERE / "workspace" + + +def _workspace_iter(iteration: int, label: str) -> Path: + d = _WORKSPACE / f"{label}" / f"iter_{iteration:03d}" + d.mkdir(parents=True, exist_ok=True) + return d + + +def _candidate_configs(candidates: list[str], max_iter: int) -> list[LearnerConfig]: + configs = [] + for idx, family in enumerate(candidates): + label = f"{idx}_{family}" + kwargs = {"learner_label": label, "model_family": family} + schedule = { + i: TaskConfig(kwargs={**kwargs, "iteration": i}) + for i in range(max_iter + 1) + } + schedule[-1] = TaskConfig(kwargs={**kwargs, "iteration": max_iter}) + configs.append( + LearnerConfig( + simulation=schedule, + training=schedule, + active_learn=schedule, + criterion=schedule, + ) + ) + return configs + + +class M3DC1_Investigator(ModelInvestigator): + def __init__( + self, + flow: WorkflowEngine, + *, + candidates: list[str], + max_iter: int, + buffer_max: int, + window_size: int, + r2_threshold: float, + redis_endpoint: str, + redis_key: str, + ): + super().__init__(flow) + + self.learner = ParallelActiveLearner(flow) + self.candidates = candidates + self.max_iter = max_iter + self.r2_threshold = r2_threshold + self.buffer_max = buffer_max + self.window_size = window_size + self.all_data: list[dict] = [] + self.input_counter = 0 + + # use REDIS for communication from the investigator to the Simulation. + # see note at top of file. This is required as the simulation task + # itself it waiting for data. (Other DT examples have it where the + # simulation task is fired after receiving the data.) + host, port = redis_endpoint.rsplit(":", 1) + self.redis = redis.Redis(host=host, port=int(port), decode_responses=True) + self.redis_key = redis_key + # ensure start clear + for candidate in self.candidates: + self.redis.delete(f"{redis_key}/{candidate}") + + # ── Simulation task ─────────────────────────────────────────────────────── + # KEY CHANGE vs amsc_stream. Buffered inputs come in from the investigator's + # input callback + + @self.learner.simulation_task(as_executable=False) + async def simulation(*args, **kwargs) -> dict: + import time + import redis as _redis + import pandas as pd + + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + family = str(kwargs["model_family"]) + + host, port_str = redis_endpoint.rsplit(":", 1) + redis_client = _redis.Redis( + host=host, port=int(port_str), decode_responses=True + ) + + print( + f" [sim {label} iter={it}] waiting for {window_size} more rows", + flush=True, + ) + deadline = time.monotonic() + 600.0 + while not redis_client.exists( + redis_key + "/MAIN" + ) or not redis_client.exists(f"{redis_key}/{family}"): + time.sleep(0.5) + if time.monotonic() > deadline: + raise TimeoutError(f"Timeout waiting for {window_size} sensor rows") + + resp = redis_client.get(redis_key + "/MAIN") + assert resp is not None + rows = json.loads(resp) + print( + f" [sim {label} iter={it}] Received {window_size} rows. Total: {len(rows)}", + flush=True, + ) + redis_client.delete(f"{redis_key}/{family}") + + df = pd.DataFrame(rows) + + out_dir = _workspace_iter(it, label) + parquet = out_dir / "sensor_snapshot.parquet" + df.to_parquet(parquet, index=False) + + meta = { + "iteration": it, + "learner_label": label, + "dataset": str(parquet), + "n_rows": len(df), + "source": "redis_stream", + } + (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) + return meta + + # ── Training task ───────────────────────────────────────────────────────── + # Fits a surrogate model locally using sklearn. + # Replace with subprocess to surge_train.py if running on HPC. + + @self.learner.training_task(as_executable=False) + async def training(sim_result: str, **kwargs) -> dict: + import pandas as pd + from sklearn.ensemble import ( + GradientBoostingRegressor, + RandomForestRegressor, + ) + from sklearn.linear_model import Ridge + from sklearn.metrics import mean_squared_error, r2_score + from sklearn.model_selection import train_test_split + from sklearn.neural_network import MLPRegressor + from sklearn.preprocessing import StandardScaler + from sklearn.pipeline import Pipeline + + it = int(kwargs.get("iteration", sim_result["iteration"])) + label = str(kwargs["learner_label"]) + family = str(kwargs.get("model_family", "rf")) + + df = pd.read_parquet(sim_result["dataset"]) + X = df.drop(columns=["output_gamma"]).values + y = df["output_gamma"].values + + X_train, X_val, y_train, y_val = train_test_split( + X, y, test_size=0.2, random_state=it + ) + + _models = { + "rf": RandomForestRegressor( + n_estimators=100, random_state=42, n_jobs=-1 + ), + "mlp": Pipeline( + [ + ("scaler", StandardScaler()), + ( + "net", + MLPRegressor( + hidden_layer_sizes=(64, 64), + max_iter=500, + random_state=42, + ), + ), + ] + ), + "gbr": GradientBoostingRegressor(n_estimators=100, random_state=42), + "ridge": Pipeline( + [ + ("scaler", StandardScaler()), + ("reg", Ridge()), + ] + ), + } + model = _models.get(family, _models["rf"]) + model.fit(X_train, y_train) + + y_pred = model.predict(X_val) + val_r2 = float(r2_score(y_val, y_pred)) + val_rmse = float(np.sqrt(mean_squared_error(y_val, y_pred))) + + metrics = { + "iteration": it, + "learner_label": label, + "model_family": family, + "val_r2": val_r2, + "val_rmse": val_rmse, + "n_train": int(len(X_train)), + "n_val": int(len(X_val)), + } + out_dir = _workspace_iter(it, label) + (out_dir / "surge_metrics.json").write_text(json.dumps(metrics, indent=2)) + + # save + model_path = str((out_dir / "model.pkl")) + with open(model_path, "wb") as f: + cloudpickle.dump(model, f) + return {"simulation": sim_result, "surge": metrics, "model": model_path} + + # ── Active-learning task ────────────────────────────────────────────────── + @self.learner.active_learn_task(as_executable=False) + async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: + it = int(kwargs.get("iteration", train_bundle["simulation"]["iteration"])) + label = str(kwargs["learner_label"]) + surge = train_bundle["surge"] + + decision = { + "iteration": it, + "learner_label": label, + "policy": "monitor_best_val_r2", + "val_r2": surge["val_r2"], + "val_rmse": surge["val_rmse"], + "model": train_bundle["model"], + } + out_dir = _workspace_iter(it, label) + (out_dir / "active.json").write_text(json.dumps(decision, indent=2)) + return { + "iteration": it, + "learner_label": label, + "train": train_bundle, + "val_r2": surge["val_r2"], + "val_rmse": surge["val_rmse"], + "model": train_bundle["model"], + } + + # ── Stop criterion ──────────────────────────────────────────────────────── + @self.learner.as_stop_criterion( + metric_name="val_r2", + threshold=r2_threshold, + operator=">=", + as_executable=False, + ) + async def stop_on_r2(*args, **kwargs) -> float: + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + path = _workspace_iter(it, label) / "surge_metrics.json" + meta = json.loads(path.read_text()) + r2 = float(meta["val_r2"]) + forced = it >= max_iter - 1 and r2 < r2_threshold + return r2_threshold if forced else r2 + + @self.flow.function_task + async def do_inference(in_data: TypedData, model=None, iter=0, label=""): + # the ASMC_stream.py demo doesn't tackle streaming inference. + # Put streaming inference code here. + if model is None: + return TypedData(M3DC1_PREDICTION, None) + + print(f"Using model: {label}-iter:{iter}") + with open(model, "rb") as f: + model_obj = cloudpickle.load(f) + + df = pd.DataFrame([in_data.data]) + + X = df.drop(columns=["output_gamma"]).values + out = model_obj.predict(X) + + return TypedData(M3DC1_PREDICTION, out) + + self.inference = do_inference + + async def input_callback(self, in_data: TypedData): + # add the data to large database + self.all_data.append(in_data.data) + + if len(self.all_data) > self.buffer_max: + self.all_data.pop(0) + + self.input_counter += 1 + + if self.input_counter >= self.window_size: + self.input_counter = 0 + self.redis.set(self.redis_key + "/MAIN", json.dumps(self.all_data)) + for c in self.candidates: + self.redis.set(f"{self.redis_key}/{c}", 1) + + async def main_loop(self, runtime: RuntimeAPI): + # run the pipeline + + runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) + runtime.set_inference_task(self.inference) + runtime.publish_new_model({"model": None}) + + configs = _candidate_configs(self.candidates, self.max_iter) + rows: list[dict] = [] + + async for state in self.learner.start( + parallel_learners=len(self.candidates), + max_iter=self.max_iter, + learner_configs=configs, + ): + label = self.candidates[int(state.learner_id)] + rows.append( + { + "learner": label, + "iter": state.iteration, + "val_r2": state.val_r2, + "val_rmse": state.val_rmse, + } + ) + print( + "\nMODEL PUBLISHED -----------------------------------\n" + f" learner={label} iter={state.iteration}\n" + f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}\n" + f" buffer={len(self.all_data)} obs\n" + f" ---------------------------------------------------\n", + flush=True, + ) + + # publish model with stats + runtime.publish_new_model( + {"model": state.model, "iter": state.iteration, "label": label}, + rows[-1], + ) + + if len(rows) >= len(self.candidates) * self.max_iter: + break diff --git a/use-cases/dt-complete/m3dc1/m3dc1_mock_sensor.py b/use-cases/dt-complete/m3dc1/m3dc1_mock_sensor.py new file mode 100644 index 0000000..f4a2628 --- /dev/null +++ b/use-cases/dt-complete/m3dc1/m3dc1_mock_sensor.py @@ -0,0 +1,117 @@ +""" +Sensor stream interface and mock implementation for M3DC1 streaming workflows. + +To plug in a real sensor, subclass SensorStream and implement read_one(). +Everything else (SensorBuffer, amsc_stream.py) works unchanged. +""" + +from __future__ import annotations + +import argparse +import asyncio + +from digitaltwin import ChannelPublisher +import numpy as np + +from m3dc1_dtypes import * + +# ── Physical parameter ranges for SPARC M3DC1 D1 ──────────────────────────── +# Derived from sparc_m3dc1_D1_metadata.yaml (inputs + output_gamma). +# Replace bounds with real calibration data when integrating a live source. +_M3DC1_RANGES: dict[str, tuple[float, float]] = { + "input_batemanscale": (0.5, 2.0), + "input_ntor": (1.0, 15.0), # toroidal mode number + "input_pscale": (0.5, 2.0), + "eq_a": (0.1, 0.5), # minor radius [m] + "eq_R0": (1.5, 6.0), # major radius [m] + "eq_kappa": (1.0, 2.5), # elongation + "eq_delta": (0.0, 0.8), # triangularity + "eq_simag": (-2.0, 0.0), # poloidal flux at magnetic axis + "eq_sibry": (-5.0, -0.5), # poloidal flux at boundary + "eq_current": (0.5, 15.0), # plasma current [MA] + "q0": (0.8, 2.5), # safety factor on axis + "q95": (3.0, 8.0), # safety factor at 95 % flux + "p0": (1e4, 1e6), # peak pressure [Pa] +} + +COLUMNS: list[str] = list(_M3DC1_RANGES) + ["output_gamma"] + + +class MockM3DC1Sensor: + """Simulates a real-time M3DC1 physics sensor at a configurable rate. + + Observations are drawn from the realistic parameter ranges in _M3DC1_RANGES. + output_gamma is a nonlinear surrogate of the MHD stability growth rate plus + Gaussian noise — non-trivial enough to make the surrogate task meaningful. + + Args: + rate_hz: Target emission rate in observations per second. + seed: RNG seed for reproducibility. + noise_std: Std-dev of Gaussian noise on output_gamma. + """ + + def __init__( + self, + rate_hz: float = 2.0, + seed: int = 42, + noise_std: float = 0.005, + ) -> None: + self._delay = 1.0 / max(rate_hz, 1e-6) + self._rng = np.random.default_rng(seed) + self._noise_std = noise_std + + async def read_one(self) -> dict[str, float]: + await asyncio.sleep(self._delay) + return self._sample() + + def _sample(self) -> dict[str, float]: + rng = self._rng + obs: dict[str, float] = { + col: float(rng.uniform(lo, hi)) for col, (lo, hi) in _M3DC1_RANGES.items() + } + # Surrogate physics: gamma grows with mode number and pressure scale, + # falls with safety factor — rough but nonlinear enough for surrogates. + obs["output_gamma"] = float( + max( + 0.0, + 0.08 * obs["input_batemanscale"] * (obs["input_ntor"] / 8.0) + + 0.06 * obs["input_pscale"] / max(obs["q95"], 0.1) + + 0.04 * obs["eq_kappa"] * obs["eq_delta"] + - 0.02 * obs["q0"] + + float(rng.normal(0.0, self._noise_std)), + ) + ) + return obs + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="M3DC1 mock sensor.") + parser.add_argument( + "--sensor-rate", + type=float, + default=2.0, + help="Mock sensor emission rate in observations/second (default: 2).", + ) + parser.add_argument( + "--sensor-seed", + type=int, + default=42, + help="RNG seed for the mock sensor.", + ) + + args = parser.parse_args() + + async def main(): + publisher = await ChannelPublisher.open(M3DC1_MOCK_CHANNEL) + + sensor = MockM3DC1Sensor(rate_hz=args.sensor_rate, seed=args.sensor_seed) + try: + while True: + val = await sensor.read_one() + await publisher.publish(val) + finally: + await publisher.close() + + if __name__ == "__main__": + asyncio.run(main()) diff --git a/use-cases/dt-complete/out.py b/use-cases/dt-complete/out.py new file mode 100644 index 0000000..2a7c4b0 --- /dev/null +++ b/use-cases/dt-complete/out.py @@ -0,0 +1,12 @@ + +from digitaltwin.components import UtilityTask, TypedData + +# this is needed to +class OutputSink(UtilityTask): + def __init__(self, flow): + super().__init__(flow) + + async def main_loop(self, runtime, in_data: TypedData): + if in_data.data is None: + return # don't print out None... that means there wasn't a model ready yet + print("Received: ", in_data.data) diff --git a/use-cases/dt-complete/run_me.py b/use-cases/dt-complete/run_me.py new file mode 100644 index 0000000..00f8b71 --- /dev/null +++ b/use-cases/dt-complete/run_me.py @@ -0,0 +1,115 @@ +""" +M3DC1 Digital Twin - a DT wrapper of the M3DC1 streaming surrogate example + +Complete Digital Twin graph: + +MOCK_SENSOR --> M3DC1_Investigator --> OUTPUT TASK + + +""" + +import argparse +import asyncio +from concurrent.futures import ProcessPoolExecutor +from radical.asyncflow import WorkflowEngine +from rhapsody.backends import ConcurrentExecutionBackend +from radical.asyncflow.logging import init_default_logger +from rhapsody.backends.data.redis import RedisDataBackend + +# Digital Twin imports +from digitaltwin.runtime import DTRuntime +from digitaltwin.streaming import connect_stream_client +from digitaltwin.components import NULL_DTYPE + +# User code imports +from m3dc1.m3dc1_investigator import M3DC1_Investigator +from out import OutputSink +from dtypes import * + +import logging + +logger = logging.getLogger(__name__) + + +async def main(m3dc1_candidates, other_args): + + # Start engine + redis_backend = await RedisDataBackend() + endpoint = redis_backend.endpoints[0] + redis_endpoint = endpoint.serialize() + + init_default_logger(logging.WARNING) + logging.getLogger("radical.asyncflow").setLevel(logging.WARNING) + logging.getLogger("rhapsody").setLevel(logging.WARNING) + + # create engine + exe = await ConcurrentExecutionBackend(ProcessPoolExecutor()) + flow = await WorkflowEngine.create(backend=exe) + + # Connect to the namespaced stream client + # create the twin's namespaced stream client + pubsub_client = await connect_stream_client("Complete-DT") + runtime = DTRuntime(flow, pubsub_client) + + ############################ + # Create the tasks and investigators + + m3dc1 = M3DC1_Investigator( + flow, + candidates=m3dc1_candidates, + max_iter=other_args.m3dc1_max_iter, + buffer_max=other_args.m3dc1_buffer_maxlen, + window_size=other_args.m3dc1_window_size, + r2_threshold=other_args.m3dc1_r2_threshold, + redis_endpoint=redis_endpoint, + redis_key="M3DC1", + ) + output_sink = OutputSink(flow) + + ########################## + # Create Digital Twin description graph + + runtime.add_input(M3DC1_SENSOR, M3DC1_MOCK_CHANNEL) + runtime.add_investigator(m3dc1, M3DC1_SENSOR, M3DC1_PREDICTION) + runtime.add_task(output_sink, M3DC1_PREDICTION, NULL_DTYPE) + + runtime.print_graph() + runtime.start() + + # let it run + await asyncio.sleep(45) + print("SHUTDOWN") + await runtime.stop() + await flow.shutdown() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Complete Digital Twin Run") + parser.add_argument( + "--m3dc1-candidates", + default="rf,mlp", + help="Comma-separated model families: rf, mlp, gbr, ridge.", + ) + parser.add_argument("--m3dc1-max-iter", type=int, default=3) + parser.add_argument("--m3dc1-r2-threshold", type=float, default=0.80) + parser.add_argument( + "--m3dc1-buffer-maxlen", + type=int, + default=1000, + help="Maximum observations retained in the sensor buffer.", + ) + parser.add_argument( + "--m3dc1-window-size", + type=int, + default=10, + help="Window size for sensor data", + ) + args = parser.parse_args() + m3dc1_candidates = [ + x.strip() for x in args.m3dc1_candidates.split(",") if x.strip() + ] + + if len(m3dc1_candidates) < 2: + parser.error("Need at least two candidates for ParallelActiveLearner.") + + asyncio.run(main(m3dc1_candidates, args)) diff --git a/use-cases/m3dc1-stream/dt/README.md b/use-cases/m3dc1-stream/dt/README.md index 0e39cae..71ee549 100644 --- a/use-cases/m3dc1-stream/dt/README.md +++ b/use-cases/m3dc1-stream/dt/README.md @@ -1,6 +1,6 @@ -# The M3DC1-stream (aka spark-stream) ported over to the Digital Twin framework. +# The M3DC1-stream (aka SPARC-stream) ported over to the Digital Twin framework. Items: - `sensor_daemon.py` --> `dt/sensor.py` diff --git a/use-cases/m3dc1-stream/dt/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py index 86c5e5f..aed359a 100644 --- a/use-cases/m3dc1-stream/dt/amsc_investigator.py +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -104,7 +104,7 @@ def __init__( self.redis.delete(f"{redis_key}/{candidate}") # ── Simulation task ─────────────────────────────────────────────────────── - # KEY CHANGE vs amsc_stream. Buffered inputs come in from the investigator's + # KEY CHANGE vs amsc_stream. Buffered inputs come in from the investigator's # input callback @self.learner.simulation_task(as_executable=False) @@ -123,7 +123,8 @@ async def simulation(*args, **kwargs) -> dict: ) print( - f" [sim {label} iter={it}] waiting for {window_size} rows", flush=True + f" [sim {label} iter={it}] waiting for {window_size} more rows", + flush=True, ) deadline = time.monotonic() + 600.0 while not redis_client.exists( @@ -136,6 +137,10 @@ async def simulation(*args, **kwargs) -> dict: resp = redis_client.get(redis_key + "/MAIN") assert resp is not None rows = json.loads(resp) + print( + f" [sim {label} iter={it}] Received {window_size} rows. Total: {len(rows)}", + flush=True, + ) redis_client.delete(f"{redis_key}/{family}") df = pd.DataFrame(rows) @@ -277,12 +282,13 @@ async def stop_on_r2(*args, **kwargs) -> float: return r2_threshold if forced else r2 @self.flow.function_task - async def do_inference(in_data: TypedData, model=None): + async def do_inference(in_data: TypedData, model=None, iter=0, label=""): # the ASMC_stream.py demo doesn't tackle streaming inference. # Put streaming inference code here. if model is None: return TypedData(M3DC1_PREDICTION, None) + print(f"Using model: {label}-iter:{iter}") with open(model, "rb") as f: model_obj = cloudpickle.load(f) @@ -304,7 +310,7 @@ async def input_callback(self, in_data: TypedData): self.input_counter += 1 - if self.input_counter > self.window_size: + if self.input_counter >= self.window_size: self.input_counter = 0 self.redis.set(self.redis_key + "/MAIN", json.dumps(self.all_data)) for c in self.candidates: @@ -344,7 +350,10 @@ async def main_loop(self, runtime: RuntimeAPI): ) # publish model with stats - runtime.publish_new_model({"model": state.model}, rows[-1]) + runtime.publish_new_model( + {"model": state.model, "iter": state.iteration, "label": label}, + rows[-1], + ) if len(rows) >= len(self.candidates) * self.max_iter: break diff --git a/use-cases/m3dc1-stream/dt/run_me.py b/use-cases/m3dc1-stream/dt/run_me.py index 0d209e6..fdefa4f 100644 --- a/use-cases/m3dc1-stream/dt/run_me.py +++ b/use-cases/m3dc1-stream/dt/run_me.py @@ -95,13 +95,13 @@ async def main(candidates, args): parser.add_argument( "--buffer-maxlen", type=int, - default=100, + default=1000, help="Maximum observations retained in the sensor buffer.", ) parser.add_argument( "--window-size", type=int, - default=100, + default=10, help="Window size for sensor data", ) args = parser.parse_args() diff --git a/use-cases/m3dc1-stream/dt/sensor.py b/use-cases/m3dc1-stream/dt/sensor.py index 1ceff08..4d740a1 100644 --- a/use-cases/m3dc1-stream/dt/sensor.py +++ b/use-cases/m3dc1-stream/dt/sensor.py @@ -89,8 +89,8 @@ def _sample(self) -> dict[str, float]: parser.add_argument( "--sensor-rate", type=float, - default=10.0, - help="Mock sensor emission rate in observations/second (default: 10).", + default=2.0, + help="Mock sensor emission rate in observations/second (default: 2).", ) parser.add_argument( "--sensor-seed", From a6171a9a828c679af19551e05ecdc2cdf4c623ad Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:53:37 -0700 Subject: [PATCH 13/22] Complete three-physics two-sensor DT demo --- use-cases/dt-complete/README.md | 6 +- .../dt-complete/demo_agent/demo_agent.py | 103 ++++++++++++++++ .../dt-complete/demo_agent/demo_dtypes.py | 4 + .../demo_agent/demo_investigator1.py | 113 +++++++++++++++++ .../demo_agent/demo_investigator2.py | 114 ++++++++++++++++++ use-cases/dt-complete/dtypes.py | 19 ++- .../dt-complete/m3dc1/m3dc1_investigator.py | 34 +++--- .../inference_only_investigator.py | 45 +++++++ .../dt-complete/negative_agent/neg_agent.py | 67 ++++++++++ .../dt-complete/negative_agent/neg_dtypes.py | 6 + .../dt-complete/negative_agent/rand_sensor.py | 45 +++++++ use-cases/dt-complete/out.py | 17 ++- use-cases/dt-complete/run_me.py | 29 ++++- 13 files changed, 563 insertions(+), 39 deletions(-) create mode 100644 use-cases/dt-complete/demo_agent/demo_agent.py create mode 100644 use-cases/dt-complete/demo_agent/demo_dtypes.py create mode 100644 use-cases/dt-complete/demo_agent/demo_investigator1.py create mode 100644 use-cases/dt-complete/demo_agent/demo_investigator2.py create mode 100644 use-cases/dt-complete/negative_agent/inference_only_investigator.py create mode 100644 use-cases/dt-complete/negative_agent/neg_agent.py create mode 100644 use-cases/dt-complete/negative_agent/neg_dtypes.py create mode 100644 use-cases/dt-complete/negative_agent/rand_sensor.py diff --git a/use-cases/dt-complete/README.md b/use-cases/dt-complete/README.md index 7b8090a..b86c27a 100644 --- a/use-cases/dt-complete/README.md +++ b/use-cases/dt-complete/README.md @@ -13,7 +13,7 @@ Three physical entities: - M3DC1 Investigator - Runs a M3DC1 Investigator - Runs a DEMO_AGENT Agent (a simple pass through) -- Runs a NEG_AGENT Agent (simply computes the negative of sensor input) +- Runs a NEGATIVE_Agent Agent (simply computes the negative of sensor input) > *Technically, the M3DC1 trains two surrogates and then picks the best one. @@ -23,8 +23,8 @@ Three physical entities: M3DC1 Mock sensor --> M3DC1 Investigator -- \ - --(JOIN)--> DEMO Agent --> OUT -RAND_VAL sensor ---> PN_AGENT -----------/ + --(JOIN)--> DEMO Agent --> OUT +RAND_VAL sensor ---> NEGATIVE_Agent ------/ ``` diff --git a/use-cases/dt-complete/demo_agent/demo_agent.py b/use-cases/dt-complete/demo_agent/demo_agent.py new file mode 100644 index 0000000..97ad60b --- /dev/null +++ b/use-cases/dt-complete/demo_agent/demo_agent.py @@ -0,0 +1,103 @@ +""" +This agent is a demo of the "SciAgents" abstraction. + +This Demo_Agent stores all the models and their accuracies generated by both +investigators. It then updates the model selector to always use the most +accurate model. + +--- +More about SciAgents: + +A SciAgent is used to group together multiple investigators that operate on the +same input / output DataTypes under one roof. It also has a "model selector" +task that runs in-stream, deciding what investigator and model to run for +inference. + +The purpose of the Science Agent is to contain one physics property. The +investigator then provides the implementation. +This implementation can have an Active Learner, and publishes one surrogate. + + +The alternative is to have only an investigator, and put all the various +surrogates inside one active learning loop. This is absolutely acceptable (see +the m3dc1 investigator), though the SciAgent format is more generalizable and scalable. +It separates the concerns from training a specific surrogate architecture from +the decision making of what surrogate to train/run when. +""" + +DO_PRINT = False +import asyncio +import json + +import pandas as pd +from radical.asyncflow import WorkflowEngine +from digitaltwin.components import ModelInvestigator, TypedData, SciAgent +from digitaltwin.runtime import RuntimeAPI + +from .demo_investigator1 import Demo_Investigator_1 +from .demo_investigator2 import Demo_Investigator_2 + +from dtypes import * + +import logging + +logger = logging.getLogger(__name__) + + +class DEMO_Agent(SciAgent): + def __init__(self, flow: WorkflowEngine): + super().__init__(flow) + self.flow = flow + + # no learning. Simple investigator + self.investigator_1 = Demo_Investigator_1(flow) + self.investigator_2 = Demo_Investigator_2(flow) + + self.model_inventory: list[dict] = [] + self.update_event = asyncio.Event() + + @self.flow.function_task + async def model_select(in_data: TypedData, i_id=0, model_kwargs={}): + return i_id, model_kwargs + + self.model_selector = model_select + + async def model_publish_cb( + self, investigator: ModelInvestigator, model_args, acc_metrics + ): + # when a new model is published, add it to my model inventory + model = { + "investigator": investigator.get_id(), + "model_args": json.dumps(model_args), + "acc": acc_metrics.get("acc", 0), + } + self.model_inventory.append(model) + self.update_event.set() + + async def main_loop(self, runtime: RuntimeAPI): + # Start up the investigator + runtime.start_investigator(self.investigator_1) + runtime.start_investigator(self.investigator_2) + runtime.subscribe_to_topic(runtime.ON_MODEL_PUBLISH, self.model_publish_cb) + + runtime.set_model_selection_task(self.model_selector) + + # default to investigator 1 + runtime.update_model_selector(i_id=self.investigator_1.get_id()) + + while True: + await self.update_event.wait() + model_df = pd.DataFrame(self.model_inventory) + + # select the model with the best accuracy. + sorted_df = model_df.sort_values(by="acc", ascending=False) + investigator_out = int(sorted_df.iloc[0]["investigator"]) + + m_args = json.loads(sorted_df.iloc[0]["model_args"]) + if DO_PRINT: + print( + f"[Demo Agent]: Winner is Investigator {investigator_out + 1}, model: {m_args}" + ) + runtime.update_model_selector(i_id=investigator_out, model_kwargs=m_args) + del model_df + self.update_event.clear() diff --git a/use-cases/dt-complete/demo_agent/demo_dtypes.py b/use-cases/dt-complete/demo_agent/demo_dtypes.py new file mode 100644 index 0000000..04eea27 --- /dev/null +++ b/use-cases/dt-complete/demo_agent/demo_dtypes.py @@ -0,0 +1,4 @@ +from digitaltwin.components import DataType + +# Demo Agent Prediction +DEMO_PREDICTION = DataType("DEMO_PREDICTION") diff --git a/use-cases/dt-complete/demo_agent/demo_investigator1.py b/use-cases/dt-complete/demo_agent/demo_investigator1.py new file mode 100644 index 0000000..5311aaf --- /dev/null +++ b/use-cases/dt-complete/demo_agent/demo_investigator1.py @@ -0,0 +1,113 @@ +""" +The Demo Investigator is a simple investigator that triggers an active learning +workflow per input on callback. It batches the input when the active learning +workflow is running, so there is at most one workflow running at once. + +This is a simple example of how to interact with ROSE's AL inside the digital +twin framework. This example shows how to have an active learner where the workflow is launched +from the data stream. + +Though Demo Investigator 1 & 2 are identical, this is to demonstrate that I can +have different "implementations" of a physics property. I can have each +investigator focus on a single surrogate, have a single active learner loop, or +other custom logic / lifecycle management. The DEMO_AGENT selects the +investigator / surrogate to run. + +Note: Compare this to the M3DC1 investigator which instructs the simulation itself to wait for +data. Therefore, the M3DC1 investigator requires a side-channel as the +simulation is fetching the data. The approach here does not require REDIS or +some side-channel method for sending data as the workflow is built already +knowing the input data. + +""" + +DO_PRINT = False + +import asyncio +import random +from typing import Any +import cloudpickle +from digitaltwin import ( + ModelInvestigator, + RuntimeAPI, + TypedData, +) +from radical.asyncflow import WorkflowEngine + +from rose.al.active_learner import Learner +from .demo_dtypes import DEMO_PREDICTION + + +class Demo_Investigator_1(ModelInvestigator): + def __init__(self, flow: WorkflowEngine): + super().__init__(flow) + + # Learners + self.acl = Learner(flow) + + self.data_update = asyncio.Event() + self.dataset: list[Any] = [] + self.new_values: list[Any] = [] + + # Learning tasks.............. + @self.acl.simulation_task(as_executable=False) + async def simulation(*args): + import time + + time.sleep(1) + return time.time() + + self.simulation = simulation + + @self.acl.training_task(as_executable=False) + async def training(*args): + return random.random() + + self.training = training + + @self.flow.function_task + async def do_inference(in_data: TypedData, model=None): + # gamma = in_data.data[0].data + # neg = in_data.data[1].data + + # out = [gamma, neg] + # if gamma is None: + # out[0] = None + + return TypedData(DEMO_PREDICTION, in_data.data) + + self.inference = do_inference + + async def input_callback(self, in_data: TypedData): + # only trigger update for ~10% of inputs + if random.random() > 0.1: + return + self.new_values.append(in_data) + self.data_update.set() + + async def main_loop(self, runtime: RuntimeAPI): + # run the pipeline + runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) + runtime.set_inference_task(self.inference) + runtime.publish_new_model() + counter = 0 + while True: + await self.data_update.wait() + + self.dataset += self.new_values + self.new_values = [] + + # Start the active learning workflow on the dataset. + if DO_PRINT: + print("[Demo Agent / Investigator 1]: Start AL Workflow") + model = await self.training(self.simulation(self.dataset)) + + # publish model and accuracy metrics. + acc = random.random() + if DO_PRINT: + print( + f"[Demo Agent / Investigator 1]: Publish model {counter}. Acc: {acc}" + ) + runtime.publish_new_model({"model": counter}, {"acc": acc}) + self.data_update.clear() + counter += 1 diff --git a/use-cases/dt-complete/demo_agent/demo_investigator2.py b/use-cases/dt-complete/demo_agent/demo_investigator2.py new file mode 100644 index 0000000..a82ae56 --- /dev/null +++ b/use-cases/dt-complete/demo_agent/demo_investigator2.py @@ -0,0 +1,114 @@ +""" +The Demo Investigator is a simple investigator that triggers an active learning +workflow per input on callback. It batches the input when the active learning +workflow is running, so there is at most one workflow running at once. + +This is a simple example of how to interact with ROSE's AL inside the digital +twin framework. This example shows how to have an active learner where the workflow is launched +from the data stream. + +Though Demo Investigator 1 & 2 are identical, this is to demonstrate that I can +have different "implementations" of a physics property. I can have each +investigator focus on a single surrogate, have a single active learner loop, or +other custom logic / lifecycle management. The DEMO_AGENT selects the +investigator / surrogate to run. + +Note: Compare this to the M3DC1 investigator which instructs the simulation itself to wait for +data. Therefore, the M3DC1 investigator requires a side-channel as the +simulation is fetching the data. The approach here does not require REDIS or +some side-channel method for sending data as the workflow is built already +knowing the input data. + +""" + +DO_PRINT = False + +import asyncio +import random +from typing import Any +import cloudpickle +from digitaltwin import ( + ModelInvestigator, + RuntimeAPI, + TypedData, +) +from radical.asyncflow import WorkflowEngine + +from rose.al.active_learner import Learner +from .demo_dtypes import DEMO_PREDICTION + + +class Demo_Investigator_2(ModelInvestigator): + def __init__(self, flow: WorkflowEngine): + super().__init__(flow) + + # Learners + self.acl = Learner(flow) + + self.data_update = asyncio.Event() + self.dataset: list[Any] = [] + self.new_values: list[Any] = [] + + # Learning tasks.............. + @self.acl.simulation_task(as_executable=False) + async def simulation(*args): + import time + + time.sleep(1) + return time.time() + + self.simulation = simulation + + @self.acl.training_task(as_executable=False) + async def training(*args): + return random.random() + + self.training = training + + @self.flow.function_task + async def do_inference(in_data: TypedData, model=None): + # gamma = in_data.data[0].data + # neg = in_data.data[1].data + + # out = [gamma, neg] + # if gamma is None: + # out[0] = None + + return TypedData(DEMO_PREDICTION, in_data.data) + + self.inference = do_inference + + async def input_callback(self, in_data: TypedData): + self.new_values.append(in_data) + + # only trigger update for ~10% of inputs + if random.random() > 0.1: + return + self.data_update.set() + + async def main_loop(self, runtime: RuntimeAPI): + # run the pipeline + runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) + runtime.set_inference_task(self.inference) + runtime.publish_new_model() + counter = 0 + while True: + await self.data_update.wait() + + self.dataset += self.new_values + self.new_values = [] + + if DO_PRINT: + print("[Demo Agent / Investigator 2]: Start AL Workflow") + # Start the active learning workflow on the dataset. + model = await self.training(self.simulation(self.dataset)) + + # publish model and accuracy metrics. + acc = random.random() + if DO_PRINT: + print( + f"[Demo Agent / Investigator 2]: Publish model {counter}. Acc: {acc}" + ) + runtime.publish_new_model({"model": counter}, {"acc": acc}) + self.data_update.clear() + counter += 1 diff --git a/use-cases/dt-complete/dtypes.py b/use-cases/dt-complete/dtypes.py index 6d38dde..dbcf233 100644 --- a/use-cases/dt-complete/dtypes.py +++ b/use-cases/dt-complete/dtypes.py @@ -1,4 +1,4 @@ -from digitaltwin.components import DataType +from digitaltwin.components import DataType, JoinDataType ############################# # Complete Digital Twin demo DATA_TYPES @@ -7,16 +7,11 @@ # use the M3DC1 sensor and prediction data types from m3dc1.m3dc1_dtypes import * -RANDOM_CHANNEL = "sensors/RANDOM" -RAND_SENSOR = DataType("RAND") +# use the NEGATIVE_Agent sensors and data types +from negative_agent.neg_dtypes import * -# Physical entities: +# The JOIN output DataType +JOIN_NEG_M3DC1 = JoinDataType([M3DC1_PREDICTION, NEG_PREDICTION]) -# M3DC1 Prediction - - -# Demo Agent Prediction -DEMO_PREDICTION = DataType("DEMO_PREDICTION") - -# Negate Agent Prediction -NEG_PREDICTION = DataType("NEG_PREDICTION") +# use the DEMO Agent data types +from demo_agent.demo_dtypes import * diff --git a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py index 1c0475b..234fff5 100644 --- a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -17,6 +17,8 @@ """ +DO_PRINT = False + import asyncio import json from pathlib import Path @@ -123,10 +125,11 @@ async def simulation(*args, **kwargs) -> dict: host=host, port=int(port_str), decode_responses=True ) - print( - f" [sim {label} iter={it}] waiting for {window_size} more rows", - flush=True, - ) + if DO_PRINT: + print( + f"[M3DC1 Investigator]: [sim {label} iter={it}] waiting for {window_size} more rows", + flush=True, + ) deadline = time.monotonic() + 600.0 while not redis_client.exists( redis_key + "/MAIN" @@ -138,10 +141,11 @@ async def simulation(*args, **kwargs) -> dict: resp = redis_client.get(redis_key + "/MAIN") assert resp is not None rows = json.loads(resp) - print( - f" [sim {label} iter={it}] Received {window_size} rows. Total: {len(rows)}", - flush=True, - ) + if DO_PRINT: + print( + f"[M3DC1 Investigator]: [sim {label} iter={it}] Received {window_size} rows. Total: {len(rows)}", + flush=True, + ) redis_client.delete(f"{redis_key}/{family}") df = pd.DataFrame(rows) @@ -289,7 +293,6 @@ async def do_inference(in_data: TypedData, model=None, iter=0, label=""): if model is None: return TypedData(M3DC1_PREDICTION, None) - print(f"Using model: {label}-iter:{iter}") with open(model, "rb") as f: model_obj = cloudpickle.load(f) @@ -341,14 +344,11 @@ async def main_loop(self, runtime: RuntimeAPI): "val_rmse": state.val_rmse, } ) - print( - "\nMODEL PUBLISHED -----------------------------------\n" - f" learner={label} iter={state.iteration}\n" - f" val_r2={state.val_r2:.5f} val_rmse={state.val_rmse:.5f}\n" - f" buffer={len(self.all_data)} obs\n" - f" ---------------------------------------------------\n", - flush=True, - ) + if DO_PRINT: + print( + f"[M3DC1 Investigator]: Model Publish: {label}-{state.iteration}", + flush=True, + ) # publish model with stats runtime.publish_new_model( diff --git a/use-cases/dt-complete/negative_agent/inference_only_investigator.py b/use-cases/dt-complete/negative_agent/inference_only_investigator.py new file mode 100644 index 0000000..37ceb65 --- /dev/null +++ b/use-cases/dt-complete/negative_agent/inference_only_investigator.py @@ -0,0 +1,45 @@ +""" +M3DC1 streaming digital twin. + +This is a version of the M3DC1 streaming surrogate that uses the digital twin +framework. + +Because we are using one active learner and only care about one physics +property, it's simplest to use just a single DT model investigator (no need for +a ScienceAgent. The point of a science agent is in the event you have various +surrogates with separate ALs) + +One tricky part is that the simulation task itself is waiting for streaming data +(opposed to the pipeline waiting and then launching the sim.). This requires a +way to transfer data from the investigator to the simulation task as the +simulation task is running. This example uses REDIS from the Rhapsody Data +Backend. + +""" + +from digitaltwin import ( + ModelInvestigator, + RuntimeAPI, + TypedData, +) +from radical.asyncflow import WorkflowEngine +from .neg_dtypes import NEG_PREDICTION + + +class Neg_Inference_Only_Investigator(ModelInvestigator): + def __init__(self, flow: WorkflowEngine): + super().__init__(flow) + + @self.flow.function_task + async def do_inference(in_data: TypedData): + val = in_data.data + + return TypedData(NEG_PREDICTION, -1 * val) + + self.inference = do_inference + + async def main_loop(self, runtime: RuntimeAPI): + # run the pipeline + + runtime.set_inference_task(self.inference) + runtime.publish_new_model() diff --git a/use-cases/dt-complete/negative_agent/neg_agent.py b/use-cases/dt-complete/negative_agent/neg_agent.py new file mode 100644 index 0000000..ec80ea8 --- /dev/null +++ b/use-cases/dt-complete/negative_agent/neg_agent.py @@ -0,0 +1,67 @@ +""" +This agent is a demo of the "SciAgents" abstraction. + +This NEGATIVE_Agent calls the "Neg_Inference_Only_Investigator" which is an +investigator that only does inference. Since there is only one investigator, the +NEGATIVE_Agent is a very light wrapper that merely passes through all requests to the +investigator. + +--- +More about SciAgents: + +A SciAgent is used to group together multiple investigators that operate on the +same input / output DataTypes under one roof. It also has a "model selector" +task that runs in-stream, deciding what investigator and model to run for +inference. + +The purpose of the Science Agent is to contain one physics property. The +investigator then provides the implementation. +This implementation can have an Active Learner, and publishes one surrogate. + + +The alternative is to have only an investigator, and put all the various +surrogates inside one active learning loop. This is absolutely acceptable (see +the m3dc1 investigator), though the SciAgent format is more generalizable and scalable. +It separates the concerns from training a specific surrogate architecture from +the decision making of what surrogate to train/run when. +""" + +import asyncio + +from radical.asyncflow import WorkflowEngine +from digitaltwin.components import ModelInvestigator, TypedData, SciAgent +from digitaltwin.runtime import RuntimeAPI + +from .inference_only_investigator import Neg_Inference_Only_Investigator + +from dtypes import * + +import logging + +logger = logging.getLogger(__name__) + + +class NEGATIVE_Agent(SciAgent): + def __init__(self, flow: WorkflowEngine): + super().__init__(flow) + self.flow = flow + + # no learning. Simple investigator + self.investigator = Neg_Inference_Only_Investigator(flow) + + @self.flow.function_task + async def model_select( + in_data: TypedData, i_id=self.investigator.get_id(), model_kwargs={} + ): + return i_id # default to latest model + + self.model_selector = model_select + + async def main_loop(self, runtime: RuntimeAPI): + # Start up the investigator + runtime.start_investigator(self.investigator) + + runtime.set_model_selection_task(self.model_selector) + + # set the investigator for primary inference + runtime.update_model_selector(i_id=self.investigator.get_id()) diff --git a/use-cases/dt-complete/negative_agent/neg_dtypes.py b/use-cases/dt-complete/negative_agent/neg_dtypes.py new file mode 100644 index 0000000..276d262 --- /dev/null +++ b/use-cases/dt-complete/negative_agent/neg_dtypes.py @@ -0,0 +1,6 @@ +from digitaltwin.components import DataType + +RAND_SENSOR_CHANNEL = "sensors/RAND" +RAND_SENSOR = DataType("RAND") + +NEG_PREDICTION = DataType("NEG_PREDICTION") diff --git a/use-cases/dt-complete/negative_agent/rand_sensor.py b/use-cases/dt-complete/negative_agent/rand_sensor.py new file mode 100644 index 0000000..2a7b123 --- /dev/null +++ b/use-cases/dt-complete/negative_agent/rand_sensor.py @@ -0,0 +1,45 @@ +""" +Random value sensor stream +""" + +from __future__ import annotations + +import argparse +import asyncio +import random + +from digitaltwin import ChannelPublisher + +from neg_dtypes import * + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="RAND mock sensor.") + parser.add_argument( + "--sensor-rate", + type=float, + default=2.0, + help="Mock sensor emission rate in observations/second (default: 2).", + ) + parser.add_argument( + "--sensor-seed", + type=int, + default=42, + help="RNG seed for the rand sensor.", + ) + + args = parser.parse_args() + + random.seed(args.sensor_seed) + + async def main(): + publisher = await ChannelPublisher.open(RAND_SENSOR_CHANNEL) + + while True: + val = random.random() + await publisher.publish(val) + + await asyncio.sleep(1 / args.sensor_rate) + + if __name__ == "__main__": + asyncio.run(main()) diff --git a/use-cases/dt-complete/out.py b/use-cases/dt-complete/out.py index 2a7c4b0..2d30832 100644 --- a/use-cases/dt-complete/out.py +++ b/use-cases/dt-complete/out.py @@ -1,12 +1,21 @@ - from digitaltwin.components import UtilityTask, TypedData +GREEN = "\033[92m" +RESET = "\033[0m" + + # this is needed to class OutputSink(UtilityTask): def __init__(self, flow): super().__init__(flow) async def main_loop(self, runtime, in_data: TypedData): - if in_data.data is None: - return # don't print out None... that means there wasn't a model ready yet - print("Received: ", in_data.data) + + prediction = in_data.data[0].data + neg_val = in_data.data[1].data + + if prediction is not None: + prediction = prediction[0] + print(f"{GREEN}[OUT]: Gamma: {prediction}. NEGATIVE: {neg_val}{RESET}") + else: + print(f"{GREEN}[OUT]: Gamma model not ready. NEGATIVE: {neg_val}{RESET}") diff --git a/use-cases/dt-complete/run_me.py b/use-cases/dt-complete/run_me.py index 00f8b71..cd98ec0 100644 --- a/use-cases/dt-complete/run_me.py +++ b/use-cases/dt-complete/run_me.py @@ -1,9 +1,12 @@ """ -M3DC1 Digital Twin - a DT wrapper of the M3DC1 streaming surrogate example +Complete Digital Twin - a DT wrapper of the M3DC1 streaming surrogate example Complete Digital Twin graph: -MOCK_SENSOR --> M3DC1_Investigator --> OUTPUT TASK +M3DC1 Mock sensor --> M3DC1 Investigator -- + \\ + --(JOIN)--> DEMO Agent --> OUT +RAND_VAL sensor ---> NEGATIVE_Agent ------// """ @@ -23,6 +26,9 @@ # User code imports from m3dc1.m3dc1_investigator import M3DC1_Investigator +from negative_agent.neg_agent import NEGATIVE_Agent +from demo_agent.demo_agent import DEMO_Agent + from out import OutputSink from dtypes import * @@ -64,14 +70,31 @@ async def main(m3dc1_candidates, other_args): redis_endpoint=redis_endpoint, redis_key="M3DC1", ) + + neg_agent = NEGATIVE_Agent(flow) + + demo_agent = DEMO_Agent(flow) + output_sink = OutputSink(flow) ########################## # Create Digital Twin description graph + # sensors runtime.add_input(M3DC1_SENSOR, M3DC1_MOCK_CHANNEL) + runtime.add_input(RAND_SENSOR, RAND_SENSOR_CHANNEL) + + # investigator and agents runtime.add_investigator(m3dc1, M3DC1_SENSOR, M3DC1_PREDICTION) - runtime.add_task(output_sink, M3DC1_PREDICTION, NULL_DTYPE) + runtime.add_agent(neg_agent, RAND_SENSOR, NEG_PREDICTION) + + # JOIN + runtime.add_data_join(JOIN_NEG_M3DC1) + + runtime.add_agent(demo_agent, JOIN_NEG_M3DC1, DEMO_PREDICTION) + + # output + runtime.add_task(output_sink, DEMO_PREDICTION, NULL_DTYPE) runtime.print_graph() runtime.start() From dffcc73731d285bf36cef8a6c0dbb09713b7c472 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Thu, 27 Aug 2026 21:37:29 +0200 Subject: [PATCH 14/22] adding heat stream use case --- use-cases/heat-stream/amsc_stream.py | 239 +++++++++++++++++++++++++ use-cases/heat-stream/sensor_daemon.py | 92 ++++++++++ 2 files changed, 331 insertions(+) create mode 100644 use-cases/heat-stream/amsc_stream.py create mode 100644 use-cases/heat-stream/sensor_daemon.py diff --git a/use-cases/heat-stream/amsc_stream.py b/use-cases/heat-stream/amsc_stream.py new file mode 100644 index 0000000..49b3496 --- /dev/null +++ b/use-cases/heat-stream/amsc_stream.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +import warnings +from pathlib import Path + +import rhapsody +from rhapsody.backends.data.redis import RedisDataBackend + +_HERE = Path(__file__).resolve().parent +if str(_HERE) not in sys.path: + sys.path.insert(0, str(_HERE)) + +from sensor_daemon import STREAM_KEY, MockHeatSensor, SensorDaemon # noqa: E402 + +from radical.asyncflow import WorkflowEngine # noqa: E402 +from rose.al import SequentialActiveLearner # noqa: E402 +from rose.learner import LearnerConfig, TaskConfig # noqa: E402 + +N_BASE = 50 +N_STEP = 25 + +CONVERGENCE_THRESHOLD = 0.05 + +_WORKSPACE = _HERE / "workspace" + + +def _workspace_iter(iteration: int) -> Path: + directory = _WORKSPACE / f"iter_{iteration:03d}" + directory.mkdir(parents=True, exist_ok=True) + return directory + + +def _build_learner_config(max_iter: int, redis_endpoint: str) -> LearnerConfig: + kwargs = {"redis_endpoint": redis_endpoint} + schedule = {i: TaskConfig(kwargs={**kwargs, "iteration": i}) for i in range(max_iter + 1)} + schedule[-1] = TaskConfig(kwargs={**kwargs, "iteration": max_iter}) + return LearnerConfig( + simulation=schedule, + training=schedule, + active_learn=schedule, + criterion=schedule, + ) + + +async def run_rose_workflow( + endpoint, + *, + max_iter: int, + convergence_threshold: float, +) -> None: + import redis as _redis + + engine = await rhapsody.get_backend("concurrent") + asyncflow = await WorkflowEngine.create(engine) + learner = SequentialActiveLearner(asyncflow) + + redis_endpoint = endpoint.serialize() + redis_client = _redis.Redis(host=endpoint.host, port=endpoint.port, decode_responses=True) + + @learner.simulation_task(as_executable=False) + async def simulation(*args, **kwargs) -> dict: + import time + import redis as _redis + import pandas as pd + + it = int(kwargs.get("iteration", 0)) + redis_ep = str(kwargs["redis_endpoint"]) + n_rows = N_BASE + it * N_STEP + + host, port_str = redis_ep.rsplit(":", 1) + redis_client = _redis.Redis(host=host, port=int(port_str), decode_responses=True) + + print(f" [sim iter={it}] waiting for {n_rows} rows …", flush=True) + deadline = time.monotonic() + 600.0 + while redis_client.xlen(STREAM_KEY) < n_rows: + if time.monotonic() > deadline: + raise TimeoutError(f"Timeout waiting for {n_rows} sensor rows") + time.sleep(0.5) + + entries = redis_client.xrevrange(STREAM_KEY, count=n_rows) + entries.reverse() + + rows = [{key: float(val) for key, val in fields.items()} for _, fields in entries] + df = pd.DataFrame(rows) + out_dir = _workspace_iter(it) + parquet = out_dir / "sensor_snapshot.parquet" + df.to_parquet(parquet, index=False) + + meta = { + "iteration": it, + "dataset" : str(parquet), + "n_rows" : len(df), + "source" : "redis_stream", + } + (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) + return meta + + @learner.training_task(as_executable=False) + async def training(sim_result: dict, **kwargs) -> dict: + import numpy as np + import pandas as pd + from sklearn.gaussian_process import GaussianProcessRegressor + from sklearn.gaussian_process.kernels import Matern + from sklearn.model_selection import train_test_split + from sklearn.preprocessing import StandardScaler + + it = int(kwargs.get("iteration", sim_result["iteration"])) + + df = pd.read_parquet(sim_result["dataset"]) + X = df.drop(columns=["q_max"]).values + y = df["q_max"].values + + X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=it) + + scaler = StandardScaler().fit(X_train) + gp = GaussianProcessRegressor( + kernel=Matern(nu=2.5), + n_restarts_optimizer=3, + normalize_y=True, + ) + gp.fit(scaler.transform(X_train), y_train) + + y_pred, y_std = gp.predict(scaler.transform(X_val), return_std=True) + mean_uncertainty = float(y_std.mean() / max(y_train.std(), 1e-6)) + + metrics = { + "iteration" : it, + "n_train" : int(len(X_train)), + "n_val" : int(len(X_val)), + "mean_uncertainty": mean_uncertainty, + } + out_dir = _workspace_iter(it) + (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2)) + return {"simulation": sim_result, "metrics": metrics} + + @learner.active_learn_task(as_executable=False) + async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: + it = int(kwargs.get("iteration", train_bundle["simulation"]["iteration"])) + metrics = train_bundle["metrics"] + + decision = { + "iteration" : it, + "policy" : "stream_consume", + "mean_uncertainty": metrics["mean_uncertainty"], + } + out_dir = _workspace_iter(it) + (out_dir / "active.json").write_text(json.dumps(decision, indent=2)) + return { + "iteration" : it, + "train" : train_bundle, + "mean_uncertainty": metrics["mean_uncertainty"], + } + + @learner.as_stop_criterion( + metric_name="mean_uncertainty", + threshold=convergence_threshold, + operator="<", + as_executable=False, + ) + async def check_convergence(*args, **kwargs) -> float: + it = int(kwargs.get("iteration", 0)) + path = _workspace_iter(it) / "metrics.json" + meta = json.loads(path.read_text()) + return float(meta["mean_uncertainty"]) + + initial_config = _build_learner_config(max_iter, redis_endpoint) + print("\nStarting HEAT surrogate stream loop\n" + "─" * 60, flush=True) + + try: + async for state in learner.start(max_iter=max_iter, initial_config=initial_config): + stream_len = await asyncio.to_thread(redis_client.xlen, STREAM_KEY) + print( + f"[iter {state.iteration}]" + f" uncertainty={state.metric_value:.4f}" + f" (target <{convergence_threshold})" + f" stream={stream_len} obs", + flush=True, + ) + finally: + await asyncflow.shutdown() + + +def main() -> None: + warnings.filterwarnings("ignore", category=UserWarning) + + parser = argparse.ArgumentParser( + description="HEAT streaming surrogate — RedisDataBackend version." + ) + parser.add_argument("--max-iter", type=int, default=10) + parser.add_argument("--convergence-threshold", type=float, default=CONVERGENCE_THRESHOLD) + parser.add_argument("--sensor-rate", type=float, default=10.0, + help="Mock sensor rate in obs/s.") + parser.add_argument("--sensor-seed", type=int, default=42) + parser.add_argument("--buffer-maxlen", type=int, default=10_000) + # HPC: set --redis-port and --redis-cmd for remote Redis launch + parser.add_argument("--redis-port", type=int, default=None) + parser.add_argument("--redis-cmd", default=None, + help='e.g. "srun --nodelist={host} redis-server --port {port}"') + + args = parser.parse_args() + + redis_backend = RedisDataBackend( + **({"cmd": args.redis_cmd, "port": args.redis_port} if args.redis_cmd else {}) + ) + + async def _main() -> None: + await redis_backend.start() + endpoint = redis_backend.endpoints[0] + + sensor = MockHeatSensor(rate_hz=args.sensor_rate, seed=args.sensor_seed) + daemon = SensorDaemon(sensor, endpoint.serialize(), maxlen=args.buffer_maxlen) + await daemon.start() + + print( + f"Redis: {endpoint.serialize()} sensor: {args.sensor_rate} Hz" + f" max_iter: {args.max_iter} threshold: {args.convergence_threshold}", + flush=True, + ) + try: + await run_rose_workflow( + endpoint, + max_iter=args.max_iter, + convergence_threshold=args.convergence_threshold, + ) + finally: + await daemon.stop() + await redis_backend.shutdown() + print("Sensor daemon and Redis stopped.", flush=True) + + asyncio.run(_main()) + + +if __name__ == "__main__": + main() diff --git a/use-cases/heat-stream/sensor_daemon.py b/use-cases/heat-stream/sensor_daemon.py new file mode 100644 index 0000000..e937e12 --- /dev/null +++ b/use-cases/heat-stream/sensor_daemon.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator +from pathlib import Path +import sys + +_HERE = Path(__file__).resolve().parent +if str(_HERE) not in sys.path: + sys.path.insert(0, str(_HERE)) + +import numpy as np +import redis + +STREAM_KEY = "heat:sensor" + +# Eich optical heat flux model — parameter ranges from NSTX-U operational space +_HEAT_RANGES: dict[str, tuple[float, float]] = { + "lqCN" : (0.5, 5.0), # near-side decay length [mm] + "lqCF" : (2.0, 15.0), # far-side decay length [mm] + "S" : (0.5, 5.0), # spreading factor [mm] + "P" : (5.0, 20.0), # input power [MW] + "radFrac": (0.1, 0.8), # radiated power fraction + "fracCN" : (0.4, 0.8), # near-side power fraction + "fracCF" : (0.1, 0.6), # far-side power fraction +} + +COLUMNS: list[str] = list(_HEAT_RANGES) + ["q_max"] + + +class SensorStream(ABC): + @abstractmethod + async def read_one(self) -> dict[str, float]: ... + + async def stream(self) -> AsyncIterator[dict[str, float]]: + while True: + yield await self.read_one() + + +class MockHeatSensor(SensorStream): + def __init__(self, rate_hz: float = 2.0, seed: int = 42, noise_std: float = 0.1) -> None: + self._delay = 1.0 / max(rate_hz, 1e-6) + self._rng = np.random.default_rng(seed) + self._noise_std = noise_std + + async def read_one(self) -> dict[str, float]: + await asyncio.sleep(self._delay) + return self._sample() + + def _sample(self) -> dict[str, float]: + rng = self._rng + obs: dict[str, float] = { + col: float(rng.uniform(lo, hi)) + for col, (lo, hi) in _HEAT_RANGES.items() + } + # Eich-inspired approximation: q_max ~ P_net * (fracCN/lqCN + fracCF/lqCF) + # Scaled to realistic NSTX-U range of ~2–50 MW/m² + p_net = obs["P"] * (1.0 - obs["radFrac"]) + obs["q_max"] = float(max( + 0.0, + 2.0 * p_net * (obs["fracCN"] / max(obs["lqCN"], 0.01) + + obs["fracCF"] / max(obs["lqCF"], 0.01)) + + float(rng.normal(0.0, self._noise_std)), + )) + return obs + + +class SensorDaemon: + def __init__(self, sensor: SensorStream, redis_endpoint: str, maxlen: int = 10_000) -> None: + host, port = redis_endpoint.rsplit(":", 1) + self._sensor = sensor + self._redis = redis.Redis(host=host, port=int(port), decode_responses=True) + self._maxlen = maxlen + self._task: asyncio.Task | None = None + + async def start(self) -> None: + self._task = asyncio.create_task(self._run(), name="sensor-daemon") + + async def stop(self) -> None: + if self._task is not None: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + async def _run(self) -> None: + async for obs in self._sensor.stream(): + fields = {k: str(v) for k, v in obs.items()} + await asyncio.to_thread(self._redis.xadd, STREAM_KEY, fields, maxlen=self._maxlen) From 1df65c9a26e5b7084643aa8bcb5de2cda59814ff Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Thu, 27 Aug 2026 13:07:57 -0700 Subject: [PATCH 15/22] Add DT-as-a-Service script. Awaiting missing DTClient apis. --- use-cases/dt-complete/run_me_service.py | 189 ++++++++++++++++++++++++ 1 file changed, 189 insertions(+) create mode 100644 use-cases/dt-complete/run_me_service.py diff --git a/use-cases/dt-complete/run_me_service.py b/use-cases/dt-complete/run_me_service.py new file mode 100644 index 0000000..56fe1d4 --- /dev/null +++ b/use-cases/dt-complete/run_me_service.py @@ -0,0 +1,189 @@ +""" +Complete Digital Twin - a DT wrapper of the M3DC1 streaming surrogate example + +A Digital Twin as a Service implementation. + +Not fully ready to run yet. The DTaaS is missing some features still: +- data join (missing .add_data_join() API in client) +- input channels (missing .add_channel() API in client) + +Complete Digital Twin graph: + +M3DC1 Mock sensor --> M3DC1 Investigator -- + \\ + --(JOIN)--> DEMO Agent --> OUT +RAND_VAL sensor ---> NEGATIVE_Agent ------// + + +""" + +import argparse +import asyncio +import os +import time + +from radical.orbit import EndpointRuntime +from digitaltwin.service import register_user_modules +from rhapsody.backends.data.redis import RedisDataBackend + +# Digital Twin imports +from digitaltwin.components import NULL_DTYPE + +# User code imports +from m3dc1.m3dc1_investigator import M3DC1_Investigator +from negative_agent.neg_agent import NEGATIVE_Agent +from demo_agent.demo_agent import DEMO_Agent + +from out import OutputSink +from dtypes import * + +import logging + +############################################# +# register user modules that twin will run +import demo_agent.demo_agent +import demo_agent.demo_dtypes +import demo_agent.demo_investigator1 +import demo_agent.demo_investigator2 +import m3dc1.m3dc1_dtypes +import m3dc1.m3dc1_investigator +import negative_agent.inference_only_investigator +import negative_agent.neg_agent +import negative_agent.neg_dtypes +import dtypes +import out + +register_user_modules( + [ + demo_agent.demo_agent, + demo_agent.demo_dtypes, + demo_agent.demo_investigator1, + demo_agent.demo_investigator2, + m3dc1.m3dc1_dtypes, + m3dc1.m3dc1_investigator, + negative_agent.inference_only_investigator, + negative_agent.neg_agent, + negative_agent.neg_dtypes, + dtypes, + out, + ] +) +############################################ + + +logger = logging.getLogger(__name__) + +DT_HOST = os.environ.get("DT_SERVICE_HOST", "broker") +TASK_ENDPOINT = os.environ.get("DT_INFERENCE_ENDPOINT") or None + +ENGINES = { + "engines": { + "inference": {"endpoint_name": "small", "backends": ["concurrent"]}, + "learning": {"endpoint_name": "hpc", "backends": ["dragon"]}, + } +} + + +def main(m3dc1_candidates, other_args): + + logging.basicConfig(level=logging.INFO) + logging.getLogger("radical.orbit").setLevel(logging.WARNING) + + runtime = EndpointRuntime() + runtime.start(wait=True) + + # Start redis -- needed by M3DC1. This later would be moved to more of an + # as-a-service approach. + redis_backend = RedisDataBackend() + endpoint = redis_backend.endpoints[0] + redis_endpoint = endpoint.serialize() + + try: + dt = runtime.get_plugin(DT_HOST, "dt", config=ENGINES) + print(f"[ORBIT Client]: session: {dt.sid} (reattach with this sid)") + + twin = dt.create_twin() + print(f"twin: {twin}") + + ############################ + # Package the tasks and investigators + + m3dc1 = dt.package( + M3DC1_Investigator, + candidates=m3dc1_candidates, + max_iter=other_args.m3dc1_max_iter, + buffer_max=other_args.m3dc1_buffer_maxlen, + window_size=other_args.m3dc1_window_size, + r2_threshold=other_args.m3dc1_r2_threshold, + redis_endpoint=redis_endpoint, + redis_key="M3DC1", + ) + + neg_agent = dt.package(NEGATIVE_Agent) + + demo_agent = dt.package(DEMO_Agent) + + output_sink = dt.package(OutputSink) + + ########################## + # Create Digital Twin description graph + + # sensors + dt.add_input(twin, M3DC1_SENSOR, M3DC1_MOCK_CHANNEL) + dt.add_input(twin, RAND_SENSOR, RAND_SENSOR_CHANNEL) + + # investigator and agents + dt.add_investigator(twin, m3dc1, M3DC1_SENSOR, M3DC1_PREDICTION) + dt.add_agent(twin, neg_agent, RAND_SENSOR, NEG_PREDICTION) + + # JOIN + dt.add_data_join(twin, JOIN_NEG_M3DC1) + + dt.add_agent(twin, demo_agent, JOIN_NEG_M3DC1, DEMO_PREDICTION) + + # output + dt.add_task(twin, output_sink, DEMO_PREDICTION, NULL_DTYPE) + + dt.print_graph() + + dt.start(twin) + + # let it run + time.sleep(45) + print("SHUTDOWN") + dt.twin_close(twin) + + finally: + runtime.stop() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Complete Digital Twin Run") + parser.add_argument( + "--m3dc1-candidates", + default="rf,mlp", + help="Comma-separated model families: rf, mlp, gbr, ridge.", + ) + parser.add_argument("--m3dc1-max-iter", type=int, default=3) + parser.add_argument("--m3dc1-r2-threshold", type=float, default=0.80) + parser.add_argument( + "--m3dc1-buffer-maxlen", + type=int, + default=1000, + help="Maximum observations retained in the sensor buffer.", + ) + parser.add_argument( + "--m3dc1-window-size", + type=int, + default=10, + help="Window size for sensor data", + ) + args = parser.parse_args() + m3dc1_candidates = [ + x.strip() for x in args.m3dc1_candidates.split(",") if x.strip() + ] + + if len(m3dc1_candidates) < 2: + parser.error("Need at least two candidates for ParallelActiveLearner.") + + main(m3dc1_candidates, args) From 48a80511f29283e8913d232092be9d533a92b3e9 Mon Sep 17 00:00:00 2001 From: Benjamin C Carter <59350660+BenCarter44@users.noreply.github.com> Date: Mon, 31 Aug 2026 07:35:50 -0700 Subject: [PATCH 16/22] Quick service fix: Simplify and remove Redis --- .../dt-complete/m3dc1/m3dc1_investigator.py | 111 ++---- .../m3dc1/redis_m3dc1_investigator.py | 360 ++++++++++++++++++ use-cases/dt-complete/run_me.py | 5 - use-cases/dt-complete/run_me_service.py | 14 +- 4 files changed, 402 insertions(+), 88 deletions(-) create mode 100644 use-cases/dt-complete/m3dc1/redis_m3dc1_investigator.py diff --git a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py index 234fff5..6ddd054 100644 --- a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -34,8 +34,7 @@ import pandas as pd from radical.asyncflow import WorkflowEngine from rose import LearnerConfig, TaskConfig -from rose.al import ParallelActiveLearner -import redis +from rose import Learner from .m3dc1_dtypes import * @@ -81,12 +80,10 @@ def __init__( buffer_max: int, window_size: int, r2_threshold: float, - redis_endpoint: str, - redis_key: str, ): super().__init__(flow) - self.learner = ParallelActiveLearner(flow) + self.learner = Learner(flow) self.candidates = candidates self.max_iter = max_iter self.r2_threshold = r2_threshold @@ -95,58 +92,30 @@ def __init__( self.all_data: list[dict] = [] self.input_counter = 0 - # use REDIS for communication from the investigator to the Simulation. - # see note at top of file. This is required as the simulation task - # itself it waiting for data. (Other DT examples have it where the - # simulation task is fired after receiving the data.) - host, port = redis_endpoint.rsplit(":", 1) - self.redis = redis.Redis(host=host, port=int(port), decode_responses=True) - self.redis_key = redis_key - # ensure start clear - for candidate in self.candidates: - self.redis.delete(f"{redis_key}/{candidate}") + self.all_data_update = asyncio.Event() # ── Simulation task ─────────────────────────────────────────────────────── # KEY CHANGE vs amsc_stream. Buffered inputs come in from the investigator's # input callback @self.learner.simulation_task(as_executable=False) - async def simulation(*args, **kwargs) -> dict: - import time - import redis as _redis + async def simulation(rows, **kwargs) -> dict: import pandas as pd it = int(kwargs.get("iteration", 0)) label = str(kwargs["learner_label"]) - family = str(kwargs["model_family"]) - - host, port_str = redis_endpoint.rsplit(":", 1) - redis_client = _redis.Redis( - host=host, port=int(port_str), decode_responses=True - ) if DO_PRINT: print( f"[M3DC1 Investigator]: [sim {label} iter={it}] waiting for {window_size} more rows", flush=True, ) - deadline = time.monotonic() + 600.0 - while not redis_client.exists( - redis_key + "/MAIN" - ) or not redis_client.exists(f"{redis_key}/{family}"): - time.sleep(0.5) - if time.monotonic() > deadline: - raise TimeoutError(f"Timeout waiting for {window_size} sensor rows") - - resp = redis_client.get(redis_key + "/MAIN") - assert resp is not None - rows = json.loads(resp) + if DO_PRINT: print( f"[M3DC1 Investigator]: [sim {label} iter={it}] Received {window_size} rows. Total: {len(rows)}", flush=True, ) - redis_client.delete(f"{redis_key}/{family}") df = pd.DataFrame(rows) @@ -159,7 +128,6 @@ async def simulation(*args, **kwargs) -> dict: "learner_label": label, "dataset": str(parquet), "n_rows": len(df), - "source": "redis_stream", } (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) return meta @@ -168,8 +136,11 @@ async def simulation(*args, **kwargs) -> dict: # Fits a surrogate model locally using sklearn. # Replace with subprocess to surge_train.py if running on HPC. + self.sim_task = simulation + @self.learner.training_task(as_executable=False) async def training(sim_result: str, **kwargs) -> dict: + print("TRAIN ..........................") import pandas as pd from sklearn.ensemble import ( GradientBoostingRegressor, @@ -244,6 +215,8 @@ async def training(sim_result: str, **kwargs) -> dict: cloudpickle.dump(model, f) return {"simulation": sim_result, "surge": metrics, "model": model_path} + self.train_task = training + # ── Active-learning task ────────────────────────────────────────────────── @self.learner.active_learn_task(as_executable=False) async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: @@ -270,21 +243,20 @@ async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: "model": train_bundle["model"], } + self.active_learn_task = active_learn + # ── Stop criterion ──────────────────────────────────────────────────────── - @self.learner.as_stop_criterion( - metric_name="val_r2", - threshold=r2_threshold, - operator=">=", - as_executable=False, - ) - async def stop_on_r2(*args, **kwargs) -> float: + @self.learner.utility_task(as_executable=False) + async def stop_on_r2(*args, **kwargs) -> dict: it = int(kwargs.get("iteration", 0)) label = str(kwargs["learner_label"]) path = _workspace_iter(it, label) / "surge_metrics.json" meta = json.loads(path.read_text()) r2 = float(meta["val_r2"]) forced = it >= max_iter - 1 and r2 < r2_threshold - return r2_threshold if forced else r2 + return {"val_r2": r2_threshold if forced else r2} + + self.stop_criterion = stop_on_r2 @self.flow.function_task async def do_inference(in_data: TypedData, model=None, iter=0, label=""): @@ -307,6 +279,7 @@ async def do_inference(in_data: TypedData, model=None, iter=0, label=""): async def input_callback(self, in_data: TypedData): # add the data to large database + print(f"GOT : {in_data.data}") self.all_data.append(in_data.data) if len(self.all_data) > self.buffer_max: @@ -316,45 +289,37 @@ async def input_callback(self, in_data: TypedData): if self.input_counter >= self.window_size: self.input_counter = 0 - self.redis.set(self.redis_key + "/MAIN", json.dumps(self.all_data)) - for c in self.candidates: - self.redis.set(f"{self.redis_key}/{c}", 1) + self.all_data_update.set() async def main_loop(self, runtime: RuntimeAPI): # run the pipeline + print("START ..........................") runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) runtime.set_inference_task(self.inference) runtime.publish_new_model({"model": None}) - configs = _candidate_configs(self.candidates, self.max_iter) rows: list[dict] = [] - async for state in self.learner.start( - parallel_learners=len(self.candidates), - max_iter=self.max_iter, - learner_configs=configs, - ): - label = self.candidates[int(state.learner_id)] - rows.append( - { - "learner": label, - "iter": state.iteration, - "val_r2": state.val_r2, - "val_rmse": state.val_rmse, - } - ) - if DO_PRINT: - print( - f"[M3DC1 Investigator]: Model Publish: {label}-{state.iteration}", - flush=True, - ) + iteration = 0 + while True: + await self.all_data_update.wait() + rows = self.all_data + # do pipeline + kwargs = { + "iteration": iteration, + "learner_label": self.candidates[0], + "model_family": self.candidates[0], + } + sim = self.sim_task(rows, **kwargs) + + model = self.train_task(sim, **kwargs) + + out = await self.active_learn_task(sim, model, **kwargs) - # publish model with stats runtime.publish_new_model( - {"model": state.model, "iter": state.iteration, "label": label}, - rows[-1], + {"model": out["model"]}, + {"acc": out["val_r2"]}, ) - if len(rows) >= len(self.candidates) * self.max_iter: - break + self.all_data_update.clear() diff --git a/use-cases/dt-complete/m3dc1/redis_m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/redis_m3dc1_investigator.py new file mode 100644 index 0000000..234fff5 --- /dev/null +++ b/use-cases/dt-complete/m3dc1/redis_m3dc1_investigator.py @@ -0,0 +1,360 @@ +""" +M3DC1 streaming digital twin. + +This is a version of the M3DC1 streaming surrogate that uses the digital twin +framework. + +Because we are using one active learner and only care about one physics +property, it's simplest to use just a single DT model investigator (no need for +a ScienceAgent. The point of a science agent is in the event you have various +surrogates with separate ALs) + +One tricky part is that the simulation task itself is waiting for streaming data +(opposed to the pipeline waiting and then launching the sim.). This requires a +way to transfer data from the investigator to the simulation task as the +simulation task is running. This example uses REDIS from the Rhapsody Data +Backend. + +""" + +DO_PRINT = False + +import asyncio +import json +from pathlib import Path + +import cloudpickle +from digitaltwin import ( + ModelInvestigator, + RuntimeAPI, + TypedData, + UtilityTask, +) +import numpy as np +import pandas as pd +from radical.asyncflow import WorkflowEngine +from rose import LearnerConfig, TaskConfig +from rose.al import ParallelActiveLearner +import redis + +from .m3dc1_dtypes import * + +# Workspace for iteration artefacts (parquet snapshots, metric JSON) +_HERE = Path(__file__).resolve().parent +_WORKSPACE = _HERE / "workspace" + + +def _workspace_iter(iteration: int, label: str) -> Path: + d = _WORKSPACE / f"{label}" / f"iter_{iteration:03d}" + d.mkdir(parents=True, exist_ok=True) + return d + + +def _candidate_configs(candidates: list[str], max_iter: int) -> list[LearnerConfig]: + configs = [] + for idx, family in enumerate(candidates): + label = f"{idx}_{family}" + kwargs = {"learner_label": label, "model_family": family} + schedule = { + i: TaskConfig(kwargs={**kwargs, "iteration": i}) + for i in range(max_iter + 1) + } + schedule[-1] = TaskConfig(kwargs={**kwargs, "iteration": max_iter}) + configs.append( + LearnerConfig( + simulation=schedule, + training=schedule, + active_learn=schedule, + criterion=schedule, + ) + ) + return configs + + +class M3DC1_Investigator(ModelInvestigator): + def __init__( + self, + flow: WorkflowEngine, + *, + candidates: list[str], + max_iter: int, + buffer_max: int, + window_size: int, + r2_threshold: float, + redis_endpoint: str, + redis_key: str, + ): + super().__init__(flow) + + self.learner = ParallelActiveLearner(flow) + self.candidates = candidates + self.max_iter = max_iter + self.r2_threshold = r2_threshold + self.buffer_max = buffer_max + self.window_size = window_size + self.all_data: list[dict] = [] + self.input_counter = 0 + + # use REDIS for communication from the investigator to the Simulation. + # see note at top of file. This is required as the simulation task + # itself it waiting for data. (Other DT examples have it where the + # simulation task is fired after receiving the data.) + host, port = redis_endpoint.rsplit(":", 1) + self.redis = redis.Redis(host=host, port=int(port), decode_responses=True) + self.redis_key = redis_key + # ensure start clear + for candidate in self.candidates: + self.redis.delete(f"{redis_key}/{candidate}") + + # ── Simulation task ─────────────────────────────────────────────────────── + # KEY CHANGE vs amsc_stream. Buffered inputs come in from the investigator's + # input callback + + @self.learner.simulation_task(as_executable=False) + async def simulation(*args, **kwargs) -> dict: + import time + import redis as _redis + import pandas as pd + + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + family = str(kwargs["model_family"]) + + host, port_str = redis_endpoint.rsplit(":", 1) + redis_client = _redis.Redis( + host=host, port=int(port_str), decode_responses=True + ) + + if DO_PRINT: + print( + f"[M3DC1 Investigator]: [sim {label} iter={it}] waiting for {window_size} more rows", + flush=True, + ) + deadline = time.monotonic() + 600.0 + while not redis_client.exists( + redis_key + "/MAIN" + ) or not redis_client.exists(f"{redis_key}/{family}"): + time.sleep(0.5) + if time.monotonic() > deadline: + raise TimeoutError(f"Timeout waiting for {window_size} sensor rows") + + resp = redis_client.get(redis_key + "/MAIN") + assert resp is not None + rows = json.loads(resp) + if DO_PRINT: + print( + f"[M3DC1 Investigator]: [sim {label} iter={it}] Received {window_size} rows. Total: {len(rows)}", + flush=True, + ) + redis_client.delete(f"{redis_key}/{family}") + + df = pd.DataFrame(rows) + + out_dir = _workspace_iter(it, label) + parquet = out_dir / "sensor_snapshot.parquet" + df.to_parquet(parquet, index=False) + + meta = { + "iteration": it, + "learner_label": label, + "dataset": str(parquet), + "n_rows": len(df), + "source": "redis_stream", + } + (out_dir / "simulation.json").write_text(json.dumps(meta, indent=2)) + return meta + + # ── Training task ───────────────────────────────────────────────────────── + # Fits a surrogate model locally using sklearn. + # Replace with subprocess to surge_train.py if running on HPC. + + @self.learner.training_task(as_executable=False) + async def training(sim_result: str, **kwargs) -> dict: + import pandas as pd + from sklearn.ensemble import ( + GradientBoostingRegressor, + RandomForestRegressor, + ) + from sklearn.linear_model import Ridge + from sklearn.metrics import mean_squared_error, r2_score + from sklearn.model_selection import train_test_split + from sklearn.neural_network import MLPRegressor + from sklearn.preprocessing import StandardScaler + from sklearn.pipeline import Pipeline + + it = int(kwargs.get("iteration", sim_result["iteration"])) + label = str(kwargs["learner_label"]) + family = str(kwargs.get("model_family", "rf")) + + df = pd.read_parquet(sim_result["dataset"]) + X = df.drop(columns=["output_gamma"]).values + y = df["output_gamma"].values + + X_train, X_val, y_train, y_val = train_test_split( + X, y, test_size=0.2, random_state=it + ) + + _models = { + "rf": RandomForestRegressor( + n_estimators=100, random_state=42, n_jobs=-1 + ), + "mlp": Pipeline( + [ + ("scaler", StandardScaler()), + ( + "net", + MLPRegressor( + hidden_layer_sizes=(64, 64), + max_iter=500, + random_state=42, + ), + ), + ] + ), + "gbr": GradientBoostingRegressor(n_estimators=100, random_state=42), + "ridge": Pipeline( + [ + ("scaler", StandardScaler()), + ("reg", Ridge()), + ] + ), + } + model = _models.get(family, _models["rf"]) + model.fit(X_train, y_train) + + y_pred = model.predict(X_val) + val_r2 = float(r2_score(y_val, y_pred)) + val_rmse = float(np.sqrt(mean_squared_error(y_val, y_pred))) + + metrics = { + "iteration": it, + "learner_label": label, + "model_family": family, + "val_r2": val_r2, + "val_rmse": val_rmse, + "n_train": int(len(X_train)), + "n_val": int(len(X_val)), + } + out_dir = _workspace_iter(it, label) + (out_dir / "surge_metrics.json").write_text(json.dumps(metrics, indent=2)) + + # save + model_path = str((out_dir / "model.pkl")) + with open(model_path, "wb") as f: + cloudpickle.dump(model, f) + return {"simulation": sim_result, "surge": metrics, "model": model_path} + + # ── Active-learning task ────────────────────────────────────────────────── + @self.learner.active_learn_task(as_executable=False) + async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: + it = int(kwargs.get("iteration", train_bundle["simulation"]["iteration"])) + label = str(kwargs["learner_label"]) + surge = train_bundle["surge"] + + decision = { + "iteration": it, + "learner_label": label, + "policy": "monitor_best_val_r2", + "val_r2": surge["val_r2"], + "val_rmse": surge["val_rmse"], + "model": train_bundle["model"], + } + out_dir = _workspace_iter(it, label) + (out_dir / "active.json").write_text(json.dumps(decision, indent=2)) + return { + "iteration": it, + "learner_label": label, + "train": train_bundle, + "val_r2": surge["val_r2"], + "val_rmse": surge["val_rmse"], + "model": train_bundle["model"], + } + + # ── Stop criterion ──────────────────────────────────────────────────────── + @self.learner.as_stop_criterion( + metric_name="val_r2", + threshold=r2_threshold, + operator=">=", + as_executable=False, + ) + async def stop_on_r2(*args, **kwargs) -> float: + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + path = _workspace_iter(it, label) / "surge_metrics.json" + meta = json.loads(path.read_text()) + r2 = float(meta["val_r2"]) + forced = it >= max_iter - 1 and r2 < r2_threshold + return r2_threshold if forced else r2 + + @self.flow.function_task + async def do_inference(in_data: TypedData, model=None, iter=0, label=""): + # the ASMC_stream.py demo doesn't tackle streaming inference. + # Put streaming inference code here. + if model is None: + return TypedData(M3DC1_PREDICTION, None) + + with open(model, "rb") as f: + model_obj = cloudpickle.load(f) + + df = pd.DataFrame([in_data.data]) + + X = df.drop(columns=["output_gamma"]).values + out = model_obj.predict(X) + + return TypedData(M3DC1_PREDICTION, out) + + self.inference = do_inference + + async def input_callback(self, in_data: TypedData): + # add the data to large database + self.all_data.append(in_data.data) + + if len(self.all_data) > self.buffer_max: + self.all_data.pop(0) + + self.input_counter += 1 + + if self.input_counter >= self.window_size: + self.input_counter = 0 + self.redis.set(self.redis_key + "/MAIN", json.dumps(self.all_data)) + for c in self.candidates: + self.redis.set(f"{self.redis_key}/{c}", 1) + + async def main_loop(self, runtime: RuntimeAPI): + # run the pipeline + + runtime.subscribe_to_topic(runtime.ON_INPUT, self.input_callback) + runtime.set_inference_task(self.inference) + runtime.publish_new_model({"model": None}) + + configs = _candidate_configs(self.candidates, self.max_iter) + rows: list[dict] = [] + + async for state in self.learner.start( + parallel_learners=len(self.candidates), + max_iter=self.max_iter, + learner_configs=configs, + ): + label = self.candidates[int(state.learner_id)] + rows.append( + { + "learner": label, + "iter": state.iteration, + "val_r2": state.val_r2, + "val_rmse": state.val_rmse, + } + ) + if DO_PRINT: + print( + f"[M3DC1 Investigator]: Model Publish: {label}-{state.iteration}", + flush=True, + ) + + # publish model with stats + runtime.publish_new_model( + {"model": state.model, "iter": state.iteration, "label": label}, + rows[-1], + ) + + if len(rows) >= len(self.candidates) * self.max_iter: + break diff --git a/use-cases/dt-complete/run_me.py b/use-cases/dt-complete/run_me.py index cd98ec0..49d942a 100644 --- a/use-cases/dt-complete/run_me.py +++ b/use-cases/dt-complete/run_me.py @@ -40,9 +40,6 @@ async def main(m3dc1_candidates, other_args): # Start engine - redis_backend = await RedisDataBackend() - endpoint = redis_backend.endpoints[0] - redis_endpoint = endpoint.serialize() init_default_logger(logging.WARNING) logging.getLogger("radical.asyncflow").setLevel(logging.WARNING) @@ -67,8 +64,6 @@ async def main(m3dc1_candidates, other_args): buffer_max=other_args.m3dc1_buffer_maxlen, window_size=other_args.m3dc1_window_size, r2_threshold=other_args.m3dc1_r2_threshold, - redis_endpoint=redis_endpoint, - redis_key="M3DC1", ) neg_agent = NEGATIVE_Agent(flow) diff --git a/use-cases/dt-complete/run_me_service.py b/use-cases/dt-complete/run_me_service.py index 56fe1d4..c3f425f 100644 --- a/use-cases/dt-complete/run_me_service.py +++ b/use-cases/dt-complete/run_me_service.py @@ -24,7 +24,6 @@ from radical.orbit import EndpointRuntime from digitaltwin.service import register_user_modules -from rhapsody.backends.data.redis import RedisDataBackend # Digital Twin imports from digitaltwin.components import NULL_DTYPE @@ -78,8 +77,8 @@ ENGINES = { "engines": { - "inference": {"endpoint_name": "small", "backends": ["concurrent"]}, - "learning": {"endpoint_name": "hpc", "backends": ["dragon"]}, + "inference": {"endpoint_name": "hpc", "backends": ["dragon"]}, + # "learning": {"endpoint_name": "hpc", "backends": ["dragon"]}, } } @@ -94,9 +93,6 @@ def main(m3dc1_candidates, other_args): # Start redis -- needed by M3DC1. This later would be moved to more of an # as-a-service approach. - redis_backend = RedisDataBackend() - endpoint = redis_backend.endpoints[0] - redis_endpoint = endpoint.serialize() try: dt = runtime.get_plugin(DT_HOST, "dt", config=ENGINES) @@ -115,8 +111,6 @@ def main(m3dc1_candidates, other_args): buffer_max=other_args.m3dc1_buffer_maxlen, window_size=other_args.m3dc1_window_size, r2_threshold=other_args.m3dc1_r2_threshold, - redis_endpoint=redis_endpoint, - redis_key="M3DC1", ) neg_agent = dt.package(NEGATIVE_Agent) @@ -144,12 +138,12 @@ def main(m3dc1_candidates, other_args): # output dt.add_task(twin, output_sink, DEMO_PREDICTION, NULL_DTYPE) - dt.print_graph() + # dt.print_graph() dt.start(twin) # let it run - time.sleep(45) + time.sleep(240) print("SHUTDOWN") dt.twin_close(twin) From 93b1623842847296ac5c8ed8089a66d1a41229d1 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 31 Aug 2026 17:23:32 +0200 Subject: [PATCH 17/22] service demo: the client terminal shows lifecycle and predictions The demo is stream driven, so all component output lands on the service and the client was silent for its whole four-minute run -- a stuck twin looked identical to a working one. The wait loop now polls the twin (state, verb counts, metrics) every 10s and probes get_inference with a fresh mock observation, printing true vs predicted gamma. Verified against a local broker: state/calls lines plus converging predictions from the first probe on. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016Npyz3Hbnwos12ESsdJ2YU --- use-cases/dt-complete/run_me_service.py | 36 ++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/use-cases/dt-complete/run_me_service.py b/use-cases/dt-complete/run_me_service.py index c3f425f..c36bb12 100644 --- a/use-cases/dt-complete/run_me_service.py +++ b/use-cases/dt-complete/run_me_service.py @@ -26,7 +26,7 @@ from digitaltwin.service import register_user_modules # Digital Twin imports -from digitaltwin.components import NULL_DTYPE +from digitaltwin.components import NULL_DTYPE, TypedData # User code imports from m3dc1.m3dc1_investigator import M3DC1_Investigator @@ -142,8 +142,38 @@ def main(m3dc1_candidates, other_args): dt.start(twin) - # let it run - time.sleep(240) + # Client-side feedback while the twin runs: the demo is stream + # driven, so all component output lands on the service. Poll the + # twin and probe inference so the client terminal shows lifecycle + # and predictions too -- and a stuck twin is visible immediately. + import sys + + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "m3dc1")) + from m3dc1_mock_sensor import MockM3DC1Sensor + + probe = MockM3DC1Sensor() # samples only, no pacing + deadline = time.time() + 240 + while time.time() < deadline: + time.sleep(10) + + info = dt.twin(twin) + print( + f"[ORBIT Client]: state={info['state']}" + f" calls={info.get('calls') or {}}" + f" metrics={list((info.get('metrics') or {}).keys())}" + ) + + obs = probe._sample() + answer = dt.get_inference( + twin, TypedData(M3DC1_SENSOR, obs), M3DC1_PREDICTION, + timeout=30, + ) + print( + f"[ORBIT Client]: inference" + f" gamma_true={obs['output_gamma']:.4f}" + f" -> prediction={answer.data}" + ) + print("SHUTDOWN") dt.twin_close(twin) From 34de9e57dd1430f4b94ea4528903072a82f14839 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 31 Aug 2026 17:51:23 +0200 Subject: [PATCH 18/22] service demo: rhapsody names the backend dragon_v3, not dragon First remote run said so exactly: "Backend 'dragon' not found. Available: [... 'dragon_v1', 'dragon_v2', 'dragon_v3', ...]". Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016Npyz3Hbnwos12ESsdJ2YU --- use-cases/dt-complete/run_me_service.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/use-cases/dt-complete/run_me_service.py b/use-cases/dt-complete/run_me_service.py index c36bb12..0882f89 100644 --- a/use-cases/dt-complete/run_me_service.py +++ b/use-cases/dt-complete/run_me_service.py @@ -77,7 +77,7 @@ ENGINES = { "engines": { - "inference": {"endpoint_name": "hpc", "backends": ["dragon"]}, + "inference": {"endpoint_name": "hpc", "backends": ["dragon_v3"]}, # "learning": {"endpoint_name": "hpc", "backends": ["dragon"]}, } } From 1ec999423b6a118000fff35c28063ac6749e0b7c Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 31 Aug 2026 18:03:33 +0200 Subject: [PATCH 19/22] m3dc1: resolve the task workspace at task runtime, not client import The module ships by value and its tasks run on the remote endpoint; the module-global workspace path (from __file__ on the client) named a directory that does not exist on Perlmutter, the training task failed its first write, and active_learn died with DependencyFailureError. The base now resolves inside the function on the executing host -- M3DC1_WORKSPACE overrides (e.g. $SCRATCH), default is the host's home. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016Npyz3Hbnwos12ESsdJ2YU --- use-cases/dt-complete/m3dc1/m3dc1_investigator.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py index 6ddd054..ad072be 100644 --- a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -21,6 +21,7 @@ import asyncio import json +import os from pathlib import Path import cloudpickle @@ -38,13 +39,18 @@ from .m3dc1_dtypes import * -# Workspace for iteration artefacts (parquet snapshots, metric JSON) -_HERE = Path(__file__).resolve().parent -_WORKSPACE = _HERE / "workspace" +# Workspace for iteration artefacts (parquet snapshots, metric JSON). +# Resolved INSIDE the function, at task runtime: this module ships by +# value to the service and its tasks run on the remote endpoint, so a +# module-global path (evaluated on the client) would name a directory +# that does not exist there. M3DC1_WORKSPACE overrides (e.g. $SCRATCH +# on an HPC endpoint); the default lands in the executing host's home. def _workspace_iter(iteration: int, label: str) -> Path: - d = _WORKSPACE / f"{label}" / f"iter_{iteration:03d}" + base = Path(os.environ.get("M3DC1_WORKSPACE", "") + or Path.home() / "m3dc1_workspace") + d = base / f"{label}" / f"iter_{iteration:03d}" d.mkdir(parents=True, exist_ok=True) return d From 5cc67084e86939836bac4bbaa1ef38e679d67005 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 31 Aug 2026 18:29:11 +0200 Subject: [PATCH 20/22] m3dc1: sequential forest fit -- joblib pools break under Dragon RandomForestRegressor(n_jobs=-1) makes joblib create a stdlib ThreadPool, which Dragon's mpbridge monkeypatching sends through DragonPool.__init__ -- super(type, obj) TypeError, task failed, active_learn dead. n_jobs=1 uses joblib's sequential backend and never creates a pool. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016Npyz3Hbnwos12ESsdJ2YU --- use-cases/dt-complete/m3dc1/m3dc1_investigator.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py index ad072be..173c881 100644 --- a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -172,8 +172,12 @@ async def training(sim_result: str, **kwargs) -> dict: ) _models = { + # n_jobs=1: joblib's ThreadPool is broken under Dragon's + # multiprocessing bridge (dragon.mpbridge patches the pool + # classes; stdlib ThreadPool then dies in DragonPool's + # super().__init__). Sequential fit is fine at window size. "rf": RandomForestRegressor( - n_estimators=100, random_state=42, n_jobs=-1 + n_estimators=100, random_state=42, n_jobs=1 ), "mlp": Pipeline( [ From 00a0169a1a4717cc47ac925992a6586c036e9682 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 31 Aug 2026 18:34:57 +0200 Subject: [PATCH 21/22] service demo: the learning lane and convergence bar show the training - ROSE window tasks carry backend='learning' (the pre-#25 label seam, installed before the decorators run); the driver declares a learning engine on the same endpoint with the concurrent executor -- its own dashboard lane, no second dragon runtime, and the training tasks stay clear of dragon's mp bridge - the investigator reports val_r2 through the runtime's duck-typed component metrics, which feeds the dashboard's convergence bar Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016Npyz3Hbnwos12ESsdJ2YU --- .../dt-complete/m3dc1/m3dc1_investigator.py | 35 +++++++++++++++++++ use-cases/dt-complete/run_me_service.py | 5 ++- 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py index 173c881..70665ab 100644 --- a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -90,6 +90,28 @@ def __init__( super().__init__(flow) self.learner = Learner(flow) + + # Route the ROSE window tasks (simulation/training/active_learn/ + # criterion) over the session's 'learning' engine, so the dashboard's + # learning lane shows the training pipeline. Same seam the framework + # used pre-#25: labels ride as decor_kwargs into function_task; a + # session without a learning engine aliases them back to inference, + # so this is safe either way. Installed BEFORE the task decorators + # below run. + inner_register = self.learner._register_task + + def _labeled(task_obj, *args, **kwargs): + decor = task_obj.setdefault("decor_kwargs", {}) + decor.setdefault("backend", "learning") + return inner_register(task_obj, *args, **kwargs) + + self.learner._register_task = _labeled + + # Convergence reporting: the runtime duck-types a `metrics` dict off + # any component (see DTRuntime.metrics); the dashboard renders it as + # the convergence bar. + self.metrics: dict = {} + self._metric_history: list = [] self.candidates = candidates self.max_iter = max_iter self.r2_threshold = r2_threshold @@ -327,6 +349,19 @@ async def main_loop(self, runtime: RuntimeAPI): out = await self.active_learn_task(sim, model, **kwargs) + val_r2 = float(out["val_r2"]) + self._metric_history.append(val_r2) + self.metrics = { + "val_r2": { + "value": val_r2, + "threshold": float(self.r2_threshold), + "operator": ">", + "should_stop": val_r2 >= float(self.r2_threshold), + "windows": iteration + 1, + "history": self._metric_history[-24:], + } + } + runtime.publish_new_model( {"model": out["model"]}, {"acc": out["val_r2"]}, diff --git a/use-cases/dt-complete/run_me_service.py b/use-cases/dt-complete/run_me_service.py index 0882f89..17b0b33 100644 --- a/use-cases/dt-complete/run_me_service.py +++ b/use-cases/dt-complete/run_me_service.py @@ -78,7 +78,10 @@ ENGINES = { "engines": { "inference": {"endpoint_name": "hpc", "backends": ["dragon_v3"]}, - # "learning": {"endpoint_name": "hpc", "backends": ["dragon"]}, + # learning on the same endpoint, concurrent executor: the lane gets + # its own engine (visible in the dashboard) without a second dragon + # runtime, and the training tasks stay clear of dragon's mp bridge. + "learning": {"endpoint_name": "hpc", "backends": ["concurrent"]}, } } From e7ede9879346000cce48468b3ba7f6e7376940ee Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Tue, 1 Sep 2026 09:23:37 +0200 Subject: [PATCH 22/22] dt-complete: post-demo consolidation -- generic placement, deploy kit - ROSE task labels ride the decorators natively (backend= is a plain decor kwarg); the _register_task wrapper is gone. learn_backend is a constructor knob, None drops the label. - engine placement comes from the environment (DT_INFERENCE_ENDPOINT/ _BACKEND, DT_LEARNING_ENDPOINT/_BACKEND) with the HPC demo defaults; a laptop run overrides four variables instead of editing code. - the run window is --runtime (default 240s). - deploy/: setup + run scripts for the three tiers (broker host, HPC endpoint under the dragon launcher, client env) and a README carrying the hard-won constraints (dragon launcher requirement, python >= 3.12.1, SLURM_EXPORT_ENV, joblib-under-dragon, rhapsody branch pin). Validated end to end against a local broker: val_r2 metric live, predictions tracking, clean shutdown. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_016Npyz3Hbnwos12ESsdJ2YU --- use-cases/dt-complete/deploy/README.md | 66 +++++++++++++++++++ use-cases/dt-complete/deploy/client-env.sh | 16 +++++ .../dt-complete/deploy/run-hpc-endpoint.sh | 27 ++++++++ use-cases/dt-complete/deploy/setup-broker.sh | 18 +++++ .../dt-complete/deploy/setup-hpc-endpoint.sh | 32 +++++++++ .../dt-complete/m3dc1/m3dc1_investigator.py | 31 ++++----- use-cases/dt-complete/run_me_service.py | 36 ++++++---- 7 files changed, 195 insertions(+), 31 deletions(-) create mode 100644 use-cases/dt-complete/deploy/README.md create mode 100755 use-cases/dt-complete/deploy/client-env.sh create mode 100755 use-cases/dt-complete/deploy/run-hpc-endpoint.sh create mode 100755 use-cases/dt-complete/deploy/setup-broker.sh create mode 100755 use-cases/dt-complete/deploy/setup-hpc-endpoint.sh diff --git a/use-cases/dt-complete/deploy/README.md b/use-cases/dt-complete/deploy/README.md new file mode 100644 index 0000000..051ab40 --- /dev/null +++ b/use-cases/dt-complete/deploy/README.md @@ -0,0 +1,66 @@ +# Deploying the dt-complete demo as a service + +Three tiers, one pinned stack: a broker host runs the ORBIT broker with +the `dt` plugin, an HPC host runs the rhapsody endpoint (dragon), and +the client machine drives the twin and the sensors. Every host installs +via digital.twins' `deploy/install.sh`, which pins the same commit and +Python minor everywhere -- the wire rejects skew, and cloudpickle does +not survive it. + +## One-time setup + + # broker host + ./setup-broker.sh + + # HPC login node (fetches the broker's cert + token via scp) + ./setup-hpc-endpoint.sh + + # client machine + git clone https://github.com/radical-cybertools/digital.twins.git + (cd digital_twins && ./deploy/install.sh client && + ./ve.demo/bin/pip install pandas scikit-learn pyarrow) + scp :.radical/orbit/broker_cert.pem \ + :.radical/orbit/broker.token ~/.radical/orbit/ + +## Running (in this order) + + # 1. broker host + cd ~/digital_twins && ./deploy/run-broker.sh $PWD/ve.demo + + # 2. HPC: get an allocation, then on the compute node + salloc -N1 -C cpu -q interactive -t 2:00:00 -A + ./run-hpc-endpoint.sh # watch for: registered as 'hpc' + + # 3. client: three terminals, each sourced + source deploy/client-env.sh + /python m3dc1/m3dc1_mock_sensor.py # terminal 1 + /python negative_agent/rand_sensor.py # terminal 2 + /python run_me_service.py # terminal 3 + +Dashboard: `https://:8000/broker/dt/ui?live=1` (broker +token at the prompt). The learning lane shows the ROSE window tasks; +the `val_r2` convergence bar fills as windows complete. + +## Placement knobs (client env) + + DT_INFERENCE_ENDPOINT (hpc) DT_INFERENCE_BACKEND (dragon_v3) + DT_LEARNING_ENDPOINT (=inference) DT_LEARNING_BACKEND (concurrent) + +A laptop-only run: point both endpoints at a local one and both +backends at `concurrent`. + +## Hard-won constraints (do not relax casually) + +- The endpoint MUST be launched via the `dragon` launcher (the run + script does): rhapsody's dragon backend uses Dragon's Batch API, + which only exists inside a Dragon-launched process tree. +- Python >= 3.12.1 on the endpoint: exactly 3.12.0 breaks dragon's + transport import (CPython gh-112358). +- `SLURM_EXPORT_ENV=ALL`: dragon's inner sruns scrub their env + otherwise and lose the venv PATH. +- joblib/sklearn `n_jobs != 1` breaks under Dragon's multiprocessing + bridge (stdlib ThreadPool lands in DragonPool.__init__); the demo + trains sequentially. +- The rhapsody install pins the `fix/dragon-cancel-idempotent` branch + until its cancel-idempotency and traceback-logging fixes are merged + upstream. diff --git a/use-cases/dt-complete/deploy/client-env.sh b/use-cases/dt-complete/deploy/client-env.sh new file mode 100755 index 0000000..3f1de6b --- /dev/null +++ b/use-cases/dt-complete/deploy/client-env.sh @@ -0,0 +1,16 @@ +# Client-side environment for the dt-complete service demo -- source me +# in EVERY client terminal (driver and both sensors): +# +# source deploy/client-env.sh +# +# The client venv comes from digital.twins: ./deploy/install.sh client +# (plus: pip install pandas scikit-learn pyarrow). +# +# DT_BROKER_CERT overrides the pinned-cert path -- needed when this +# machine runs a broker of its own and ~/.radical/orbit/broker_cert.pem +# is that one, not the demo broker's. +BROKER="${1:?usage: source client-env.sh }" + +export RADICAL_ORBIT_BROKER_URL="wss://$BROKER:8000" +export RADICAL_ORBIT_BROKER_CERT="${DT_BROKER_CERT:-$HOME/.radical/orbit/broker_cert.pem}" +export DT_STREAM_BACKEND=orbit diff --git a/use-cases/dt-complete/deploy/run-hpc-endpoint.sh b/use-cases/dt-complete/deploy/run-hpc-endpoint.sh new file mode 100755 index 0000000..a164bad --- /dev/null +++ b/use-cases/dt-complete/deploy/run-hpc-endpoint.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Run the demo's HPC endpoint -- INSIDE a compute allocation. +# +# ./run-hpc-endpoint.sh +# +# The endpoint is launched via the ``dragon`` launcher: rhapsody's dragon +# backend drives Dragon's Batch API, which only works inside a +# Dragon-launched process tree. DT_DIR as in setup-hpc-endpoint.sh. +set -euo pipefail +BROKER="${1:?usage: $0 }" +DT_DIR="${DT_DIR:-$HOME/digital_twins}" +VENV="$DT_DIR/ve.demo" + +# dragon resolves its helpers BY NAME through srun on the task side +export PATH="$VENV/bin:$PATH" +export RADICAL_ORBIT_BROKER_URL="wss://$BROKER:8000" +export RADICAL_ORBIT_BROKER_CERT="$HOME/.radical/orbit/broker_cert.pem" +export RADICAL_ORBIT_RHAPSODY_BACKEND=dragon_v3 +export RADICAL_ORBIT_RHAPSODY_NOTIFY_WINDOW=0 +export SLURM_EXPORT_ENV=ALL # inner sruns must not scrub the env +export DT_STREAM_BACKEND=orbit +export DT_ENDPOINT_TAG=hpc +# training snapshots belong on scratch where available +export M3DC1_WORKSPACE="${M3DC1_WORKSPACE:-${SCRATCH:-$HOME}/m3dc1_workspace}" + +exec "$VENV/bin/dragon" "$VENV/bin/radical-orbit-endpoint.py" -n hpc \ + 2>&1 | tee endpoint.log diff --git a/use-cases/dt-complete/deploy/setup-broker.sh b/use-cases/dt-complete/deploy/setup-broker.sh new file mode 100755 index 0000000..8f34766 --- /dev/null +++ b/use-cases/dt-complete/deploy/setup-broker.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# DTaaS broker host for the dt-complete demo (e.g. radical.3). Run there. +# +# Once per host, before the first run: broker_cert.pem, broker_key.pem and +# broker.token in ~/.radical/orbit/. DT_DIR overrides where the +# digital.twins checkout + venv live (default: ~/digital_twins). +set -euo pipefail +DT_DIR="${DT_DIR:-$HOME/digital_twins}" + +[ -d "$DT_DIR" ] || git clone https://github.com/radical-cybertools/digital.twins.git "$DT_DIR" +cd "$DT_DIR" && git checkout devel && git pull + +./deploy/install.sh broker # pinned stack -> ./ve.demo +# the M3DC1 investigator instantiates on the broker and imports these +./ve.demo/bin/pip install -q pandas scikit-learn pyarrow + +echo "done. start the broker with:" +echo " cd $DT_DIR && ./deploy/run-broker.sh \$PWD/ve.demo" diff --git a/use-cases/dt-complete/deploy/setup-hpc-endpoint.sh b/use-cases/dt-complete/deploy/setup-hpc-endpoint.sh new file mode 100755 index 0000000..7091bf3 --- /dev/null +++ b/use-cases/dt-complete/deploy/setup-hpc-endpoint.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# DTaaS HPC-endpoint venv for the dt-complete demo (e.g. Perlmutter). +# Run on a login node. +# +# ./setup-hpc-endpoint.sh +# +# DT_DIR overrides where the digital.twins checkout + venv live +# (default: ~/digital_twins). +set -euo pipefail +BROKER="${1:?usage: $0 }" +DT_DIR="${DT_DIR:-$HOME/digital_twins}" + +# same Python minor as every other host -- the service rejects skew, and +# exactly 3.12.0 breaks dragon's transport import (needs >= 3.12.1) +module load python/3.12 2>/dev/null || true + +[ -d "$DT_DIR" ] || git clone https://github.com/radical-cybertools/digital.twins.git "$DT_DIR" +cd "$DT_DIR" && git checkout devel && git pull + +./deploy/install.sh endpoint # pinned stack -> ./ve.demo +# sklearn/parquet task bodies unpickle and run here +./ve.demo/bin/pip install -q pandas scikit-learn pyarrow +# dragon backend; the branch carries the idempotent-cancel and +# failure-traceback fixes (pending upstream merge) +./ve.demo/bin/pip install -q --force-reinstall --no-deps \ + "rhapsody-py[telemetry,dragon] @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" + +mkdir -p ~/.radical/orbit +scp "$BROKER:.radical/orbit/broker_cert.pem" "$BROKER:.radical/orbit/broker.token" ~/.radical/orbit/ + +echo "done. get an allocation (e.g. salloc -N1 -C cpu -q interactive -t 2:00:00 -A )," +echo "then run: run-hpc-endpoint.sh $BROKER" diff --git a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py index 70665ab..049a4b2 100644 --- a/use-cases/dt-complete/m3dc1/m3dc1_investigator.py +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -86,26 +86,19 @@ def __init__( buffer_max: int, window_size: int, r2_threshold: float, + learn_backend: str | None = "learning", ): super().__init__(flow) self.learner = Learner(flow) - # Route the ROSE window tasks (simulation/training/active_learn/ - # criterion) over the session's 'learning' engine, so the dashboard's - # learning lane shows the training pipeline. Same seam the framework - # used pre-#25: labels ride as decor_kwargs into function_task; a - # session without a learning engine aliases them back to inference, - # so this is safe either way. Installed BEFORE the task decorators - # below run. - inner_register = self.learner._register_task - - def _labeled(task_obj, *args, **kwargs): - decor = task_obj.setdefault("decor_kwargs", {}) - decor.setdefault("backend", "learning") - return inner_register(task_obj, *args, **kwargs) - - self.learner._register_task = _labeled + # The ROSE window tasks (simulation/training/active_learn/criterion) + # carry this engine-role label, so a session with a 'learning' + # engine runs them there and the dashboard's learning lane shows + # the training pipeline. A session without one aliases the label + # back to inference, so the default is safe either way; None drops + # the label entirely. + _learn = {"backend": learn_backend} if learn_backend else {} # Convergence reporting: the runtime duck-types a `metrics` dict off # any component (see DTRuntime.metrics); the dashboard renders it as @@ -126,7 +119,7 @@ def _labeled(task_obj, *args, **kwargs): # KEY CHANGE vs amsc_stream. Buffered inputs come in from the investigator's # input callback - @self.learner.simulation_task(as_executable=False) + @self.learner.simulation_task(as_executable=False, **_learn) async def simulation(rows, **kwargs) -> dict: import pandas as pd @@ -166,7 +159,7 @@ async def simulation(rows, **kwargs) -> dict: self.sim_task = simulation - @self.learner.training_task(as_executable=False) + @self.learner.training_task(as_executable=False, **_learn) async def training(sim_result: str, **kwargs) -> dict: print("TRAIN ..........................") import pandas as pd @@ -250,7 +243,7 @@ async def training(sim_result: str, **kwargs) -> dict: self.train_task = training # ── Active-learning task ────────────────────────────────────────────────── - @self.learner.active_learn_task(as_executable=False) + @self.learner.active_learn_task(as_executable=False, **_learn) async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: it = int(kwargs.get("iteration", train_bundle["simulation"]["iteration"])) label = str(kwargs["learner_label"]) @@ -278,7 +271,7 @@ async def active_learn(sim_result: dict, train_bundle: dict, **kwargs) -> dict: self.active_learn_task = active_learn # ── Stop criterion ──────────────────────────────────────────────────────── - @self.learner.utility_task(as_executable=False) + @self.learner.utility_task(as_executable=False, **_learn) async def stop_on_r2(*args, **kwargs) -> dict: it = int(kwargs.get("iteration", 0)) label = str(kwargs["learner_label"]) diff --git a/use-cases/dt-complete/run_me_service.py b/use-cases/dt-complete/run_me_service.py index 17b0b33..51a194c 100644 --- a/use-cases/dt-complete/run_me_service.py +++ b/use-cases/dt-complete/run_me_service.py @@ -20,8 +20,13 @@ import argparse import asyncio import os +import sys import time +# the sensor modules are script-style (flat imports); make them importable +# from here for the client-side inference probe +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3dc1")) + from radical.orbit import EndpointRuntime from digitaltwin.service import register_user_modules @@ -73,15 +78,21 @@ logger = logging.getLogger(__name__) DT_HOST = os.environ.get("DT_SERVICE_HOST", "broker") -TASK_ENDPOINT = os.environ.get("DT_INFERENCE_ENDPOINT") or None + +# Engine placement, overridable per deployment. Defaults match the HPC +# demo: inference on the dragon-launched endpoint, learning on the same +# endpoint with the concurrent executor -- its own dashboard lane and no +# second dragon runtime, and the training tasks stay clear of dragon's +# multiprocessing bridge. +INFERENCE_EP = os.environ.get("DT_INFERENCE_ENDPOINT", "hpc") +INFERENCE_BE = os.environ.get("DT_INFERENCE_BACKEND", "dragon_v3") +LEARNING_EP = os.environ.get("DT_LEARNING_ENDPOINT", INFERENCE_EP) +LEARNING_BE = os.environ.get("DT_LEARNING_BACKEND", "concurrent") ENGINES = { "engines": { - "inference": {"endpoint_name": "hpc", "backends": ["dragon_v3"]}, - # learning on the same endpoint, concurrent executor: the lane gets - # its own engine (visible in the dashboard) without a second dragon - # runtime, and the training tasks stay clear of dragon's mp bridge. - "learning": {"endpoint_name": "hpc", "backends": ["concurrent"]}, + "inference": {"endpoint_name": INFERENCE_EP, "backends": [INFERENCE_BE]}, + "learning": {"endpoint_name": LEARNING_EP, "backends": [LEARNING_BE]}, } } @@ -141,21 +152,16 @@ def main(m3dc1_candidates, other_args): # output dt.add_task(twin, output_sink, DEMO_PREDICTION, NULL_DTYPE) - # dt.print_graph() - dt.start(twin) # Client-side feedback while the twin runs: the demo is stream # driven, so all component output lands on the service. Poll the # twin and probe inference so the client terminal shows lifecycle # and predictions too -- and a stuck twin is visible immediately. - import sys - - sys.path.insert(0, os.path.join(os.path.dirname(__file__), "m3dc1")) from m3dc1_mock_sensor import MockM3DC1Sensor probe = MockM3DC1Sensor() # samples only, no pacing - deadline = time.time() + 240 + deadline = time.time() + other_args.runtime while time.time() < deadline: time.sleep(10) @@ -192,6 +198,12 @@ def main(m3dc1_candidates, other_args): help="Comma-separated model families: rf, mlp, gbr, ridge.", ) parser.add_argument("--m3dc1-max-iter", type=int, default=3) + parser.add_argument( + "--runtime", + type=int, + default=240, + help="Seconds to keep the twin running before teardown.", + ) parser.add_argument("--m3dc1-r2-threshold", type=float, default=0.80) parser.add_argument( "--m3dc1-buffer-maxlen",