diff --git a/use-cases/.gitignore b/use-cases/.gitignore new file mode 100644 index 0000000..4a98218 --- /dev/null +++ b/use-cases/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +tmp.py +rhapsody.data.* +workspace/ +__pycache__/ +dump.rdb +.vscode/ \ No newline at end of file diff --git a/use-cases/dt-complete/README.md b/use-cases/dt-complete/README.md new file mode 100644 index 0000000..b86c27a --- /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 NEGATIVE_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 ---> NEGATIVE_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/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/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/dtypes.py b/use-cases/dt-complete/dtypes.py new file mode 100644 index 0000000..dbcf233 --- /dev/null +++ b/use-cases/dt-complete/dtypes.py @@ -0,0 +1,17 @@ +from digitaltwin.components import DataType, JoinDataType + +############################# +# Complete Digital Twin demo DATA_TYPES +############################# + +# use the M3DC1 sensor and prediction data types +from m3dc1.m3dc1_dtypes import * + +# use the NEGATIVE_Agent sensors and data types +from negative_agent.neg_dtypes import * + +# The JOIN output DataType +JOIN_NEG_M3DC1 = JoinDataType([M3DC1_PREDICTION, NEG_PREDICTION]) + +# use the DEMO Agent data types +from demo_agent.demo_dtypes import * 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..049a4b2 --- /dev/null +++ b/use-cases/dt-complete/m3dc1/m3dc1_investigator.py @@ -0,0 +1,363 @@ +""" +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 +import os +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 import Learner + +from .m3dc1_dtypes import * + +# 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: + 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 + + +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, + learn_backend: str | None = "learning", + ): + super().__init__(flow) + + self.learner = Learner(flow) + + # 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 + # the convergence bar. + self.metrics: dict = {} + self._metric_history: list = [] + 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 + + 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, **_learn) + async def simulation(rows, **kwargs) -> dict: + import pandas as pd + + it = int(kwargs.get("iteration", 0)) + label = str(kwargs["learner_label"]) + + if DO_PRINT: + print( + f"[M3DC1 Investigator]: [sim {label} iter={it}] waiting for {window_size} more rows", + flush=True, + ) + + if DO_PRINT: + print( + f"[M3DC1 Investigator]: [sim {label} iter={it}] Received {window_size} rows. Total: {len(rows)}", + flush=True, + ) + + 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), + } + (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.sim_task = simulation + + @self.learner.training_task(as_executable=False, **_learn) + async def training(sim_result: str, **kwargs) -> dict: + print("TRAIN ..........................") + 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 = { + # 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 + ), + "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} + + self.train_task = training + + # ── Active-learning task ────────────────────────────────────────────────── + @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"]) + 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"], + } + + self.active_learn_task = active_learn + + # ── Stop criterion ──────────────────────────────────────────────────────── + @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"]) + 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 {"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=""): + # 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 + print(f"GOT : {in_data.data}") + 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.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}) + + rows: list[dict] = [] + + 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) + + 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"]}, + ) + + self.all_data_update.clear() 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/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/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 new file mode 100644 index 0000000..2d30832 --- /dev/null +++ b/use-cases/dt-complete/out.py @@ -0,0 +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): + + 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 new file mode 100644 index 0000000..49d942a --- /dev/null +++ b/use-cases/dt-complete/run_me.py @@ -0,0 +1,133 @@ +""" +Complete Digital Twin - a DT wrapper of the M3DC1 streaming surrogate example + +Complete Digital Twin graph: + +M3DC1 Mock sensor --> M3DC1 Investigator -- + \\ + --(JOIN)--> DEMO Agent --> OUT +RAND_VAL sensor ---> NEGATIVE_Agent ------// + + +""" + +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 negative_agent.neg_agent import NEGATIVE_Agent +from demo_agent.demo_agent import DEMO_Agent + +from out import OutputSink +from dtypes import * + +import logging + +logger = logging.getLogger(__name__) + + +async def main(m3dc1_candidates, other_args): + + # Start engine + + 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, + ) + + 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_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() + + # 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/dt-complete/run_me_service.py b/use-cases/dt-complete/run_me_service.py new file mode 100644 index 0000000..51a194c --- /dev/null +++ b/use-cases/dt-complete/run_me_service.py @@ -0,0 +1,228 @@ +""" +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 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 + +# Digital Twin imports +from digitaltwin.components import NULL_DTYPE, TypedData + +# 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") + +# 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": INFERENCE_EP, "backends": [INFERENCE_BE]}, + "learning": {"endpoint_name": LEARNING_EP, "backends": [LEARNING_BE]}, + } +} + + +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. + + 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, + ) + + 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.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. + from m3dc1_mock_sensor import MockM3DC1Sensor + + probe = MockM3DC1Sensor() # samples only, no pacing + deadline = time.time() + other_args.runtime + 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) + + 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( + "--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", + 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) 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) diff --git a/use-cases/m3dc1-stream/amsc_stream.py b/use-cases/m3dc1-stream/amsc_stream.py new file mode 100644 index 0000000..a73c08b --- /dev/null +++ b/use-cases/m3dc1-stream/amsc_stream.py @@ -0,0 +1,303 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import asyncio +import json +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_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 + +N_BASE = 100 +N_STEP = 100 + +_WORKSPACE = _HERE / "workspace" + + +def _workspace_iter(iteration: int, label: str) -> Path: + 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, redis_endpoint: str +) -> list[LearnerConfig]: + configs = [] + for idx, family in enumerate(candidates): + 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( + LearnerConfig( + simulation=schedule, + training=schedule, + active_learn=schedule, + criterion=schedule, + ) + ) + return configs + + +async def run_rose_workflow( + endpoint, + *, + candidates: list[str], + max_iter: int, + r2_threshold: float, +) -> None: + import redis as _redis + + engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) + asyncflow = await WorkflowEngine.create(engine) + learner = ParallelActiveLearner(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)) + 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) + + 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 + + @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.pipeline import Pipeline + from sklearn.preprocessing import StandardScaler + + 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} + + @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"], + } + + @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 + + configs = _candidate_configs(candidates, max_iter, redis_endpoint) + 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)] + stream_len = await asyncio.to_thread(redis_client.xlen, STREAM_KEY) + 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" stream={stream_len} 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}") + + +def main() -> None: + warnings.filterwarnings("ignore", category=UserWarning) + + parser = argparse.ArgumentParser( + 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()] + + + + async def _main() -> None: + 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"Redis: {endpoint.serialize()} sensor: {args.sensor_rate} Hz" + f" candidates: {candidates} max_iter: {args.max_iter}", + flush=True, + ) + try: + await run_rose_workflow( + endpoint, + candidates=candidates, + max_iter=args.max_iter, + r2_threshold=args.r2_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/m3dc1-stream/dt/README.md b/use-cases/m3dc1-stream/dt/README.md new file mode 100644 index 0000000..71ee549 --- /dev/null +++ b/use-cases/m3dc1-stream/dt/README.md @@ -0,0 +1,32 @@ + + +# The M3DC1-stream (aka SPARC-stream) ported over to the Digital Twin framework. + +Items: +- `sensor_daemon.py` --> `dt/sensor.py` +- `amsc_stream.py` --> `dt/amsc_investigator.py` + +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: + +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/amsc_investigator.py b/use-cases/m3dc1-stream/dt/amsc_investigator.py new file mode 100644 index 0000000..aed359a --- /dev/null +++ b/use-cases/m3dc1-stream/dt/amsc_investigator.py @@ -0,0 +1,370 @@ +""" +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 dtypes import M3DC1_PREDICTION + +# 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 + + +# 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/m3dc1-stream/dt/dtypes.py b/use-cases/m3dc1-stream/dt/dtypes.py new file mode 100644 index 0000000..6789122 --- /dev/null +++ b/use-cases/m3dc1-stream/dt/dtypes.py @@ -0,0 +1,10 @@ +from digitaltwin.components import DataType + +# Sensor channel + +M3DC1_MOCK_CHANNEL = "sensors/MockM3DC1" +M3DC1_SENSOR = DataType("M3DC1") + + +# Modeling +M3DC1_PREDICTION = DataType("M3DC1_PREDICTION") 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() 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..fdefa4f --- /dev/null +++ b/use-cases/m3dc1-stream/dt/run_me.py @@ -0,0 +1,113 @@ +""" +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 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 dtypes import * + +from radical.asyncflow.logging import init_default_logger +from rhapsody.backends.data.redis import RedisDataBackend +import logging + +logger = logging.getLogger(__name__) + + +async def main(candidates, args): + max_iter = args.max_iter + r2_threshold = args.r2_threshold + max_len = args.buffer_maxlen + 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) + 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 + + m3dc1 = M3DC1_Investigator( + flow, + 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_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="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=1000, + help="Maximum observations retained in the sensor buffer.", + ) + parser.add_argument( + "--window-size", + type=int, + default=10, + help="Window size for sensor data", + ) + 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)) diff --git a/use-cases/m3dc1-stream/dt/sensor.py b/use-cases/m3dc1-stream/dt/sensor.py new file mode 100644 index 0000000..4d740a1 --- /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=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/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)