diff --git a/service/README.md b/service/README.md new file mode 100644 index 0000000..c1d17c7 --- /dev/null +++ b/service/README.md @@ -0,0 +1,106 @@ +# xGFabric twin, service mode + +`../twin_service.py` runs twin.py's graph on a DTaaS broker: the twin +lives in the ORBIT `dt` plugin, its tasks run on a rhapsody endpoint, +the sensor is an external channel publisher. Physics is faked at the +same seams twin.py already fakes (sensor records, simulation); the +selection story is real — three surrogate architectures with different +costs, ranked by profiler-predicted Pi runtime. + +## Local run (three terminals + stack) + +Stack (from a digital.twins checkout with `./deploy/install.sh` done, +venv `ve3`/`ve.demo`): + + ./deploy/run-broker.sh $PWD/ + ./deploy/run-endpoint.sh dt_inference_ep localhost $PWD/ + +Client terminals (each): + + export RADICAL_ORBIT_BROKER_URL=wss://localhost:8000 + export DT_STREAM_BACKEND=orbit + + /bin/python service/sensor_publisher.py # terminal 1 + /bin/python twin_service.py --runtime 240 # terminal 2 + +Dashboard: `https://localhost:8000/broker/dt/ui?live=1` (broker token). +Heatmaps land in `$XGF_WORKSPACE` (default `~/xgf_twin/`) on the host +running the endpoint tasks. + +Placement: `DT_INFERENCE_ENDPOINT` / `DT_INFERENCE_BACKEND` override +the defaults (`dt_inference_ep` / `concurrent`). + +## Remote run (broker on radical.3, endpoint on Perlmutter) + +`service/deploy/` adapts the AmSC dt-complete deploy kit (same debugged +constraints: dragon launcher requirement, python >= 3.12.1, +SLURM_EXPORT_ENV, cert staging), pinning rhapsody's +`fix/dragon-cancel-idempotent` branch (e491cd2-based) on every tier: it +carries the dragon cancel + traceback fixes and stays compatible with +the pinned dragonhpc 0.14.1 (rhapsody main's dragon backend passes +`task_logs=` to `Batch()`, which that dragon rejects). + + # broker host, once + service/deploy/setup-broker.sh + cd ~/digital_twins && ./deploy/run-broker.sh $PWD/ve.demo + + # Perlmutter login node, once + service/deploy/setup-hpc-endpoint.sh + # then inside salloc -N1 -C cpu -q interactive -t 2:00:00 -A : + service/deploy/run-hpc-endpoint.sh # registers as 'hpc' + + # client terminals (driver + sensor), each: + source service/deploy/client-env.sh remote + /bin/python service/sensor_publisher.py # terminal 1 + /bin/python twin_service.py --runtime 240 # terminal 2 + +The client venv must be the same Python minor (digital.twins +`./deploy/install.sh client` + `pip install numpy`); a 3.13 venv is +rejected at the first verb. + +## Real workload (Level A: real training, faked simulation) + +The fake `service/*` components run the DTaaS mechanics end to end. The +real path swaps in the actual FNO/PINN/PCR investigators (TensorFlow / +scikit-learn) and the real profiler, still on precalc simulation data +(`/global/cfs/cdirs/m5290/precalc_sims`) -- real OpenFOAM (Level B, +`cups_structure.zip`) is a later step. + +Only the **endpoint** changes: the trainings need TensorFlow, so it runs +in a clone of Ben's `cfdaai` conda env (which carries the stack) with our +runtime installed into the clone -- built by `setup-hpc-endpoint-real.sh`, +then launched via `run-hpc-endpoint.sh` with `DT_VENV`/`XGF_DIR` set (the +setup script prints the exact line). Broker and client stay on `ve.demo` +and need no TensorFlow. + + # Perlmutter login node, once (clones cfdaai, installs our runtime, + # checks out the tasks tree, stages the profiler dataset) + service/deploy/setup-hpc-endpoint-real.sh + # then in an allocation, per the line it prints: + DT_VENV= XGF_DIR= \ + service/deploy/run-hpc-endpoint.sh + +Two things are NOT done yet and gate an actual real run: + +- **Lazy TensorFlow imports.** The real investigator wrappers + `import tensorflow` at module load, which would drag TF onto the + client and broker just to package/instantiate the class. The wrappers + need their heavy imports deferred into the task bodies (which run on + the endpoint, where TF lives). Until that lands, only the endpoint + tier is provisioned; the driver still uses the fake components. +- **cloudpickle parity.** The cloned conda env and `ve.demo` may carry + different cloudpickle versions; align them (the endpoint install pulls + our pinned stack, but verify against the client) or unpickling the + shipped classes can fail. + +## What maps to what + +| twin.py (standalone) | twin_service.py (DTaaS) | +|----------------------------------|--------------------------------------| +| local WorkflowEngine + backend | session engine on the broker, tasks on the endpoint (ENGINES config) | +| DavisWind persistent component | external `ChannelPublisher` + `add_input` binding | +| WindFieldAgent + FNO/PINN/PCR | `ServiceWindFieldAgent` + fake `SurrogateInvestigator`s (same selection logic) | +| profiler subprocess + Pi learner | `ServiceProfiler` (inline timed run) + `ServicePiPredictor` | +| CUPS_Sink | `ServiceSink` (runtime-resolved workspace) | +| ZMQ stream | ORBIT data plane | +| asyncflow telemetry + reports | not yet — see the telemetry branch | diff --git a/service/__init__.py b/service/__init__.py new file mode 100644 index 0000000..32e0bd1 --- /dev/null +++ b/service/__init__.py @@ -0,0 +1,9 @@ +# Service-mode (DTaaS / ORBIT) variant of the xGFabric twin. +# +# Same graph and the same model-selection story as ../twin.py -- a wind +# sensor feeds a field agent whose three competing surrogates are ranked +# by profiler-predicted Pi runtime -- but the components here are +# service-safe: they ship to the broker by value, run their tasks on a +# rhapsody endpoint, and fake the physics (as twin.py already does for +# the sensor and the simulation). The real FNO/PINN/PCR training stacks +# stay in ../tasks and are NOT imported here. diff --git a/service/deploy/client-env.sh b/service/deploy/client-env.sh new file mode 100755 index 0000000..9b3a812 --- /dev/null +++ b/service/deploy/client-env.sh @@ -0,0 +1,23 @@ +# Client-side environment for the xGFabric service twin -- source me in +# EVERY client terminal (driver and sensor): +# +# source service/deploy/client-env.sh [remote] +# +# With `remote`, task placement targets the HPC endpoint ('hpc' on +# dragon_v3, as run-hpc-endpoint.sh launches it); without it the local +# defaults apply (dt_inference_ep / concurrent). +# +# The client venv is digital.twins' `./deploy/install.sh client` plus +# `pip install numpy` -- same Python minor as broker and endpoint. +# DT_BROKER_CERT overrides the pinned-cert path when this machine runs +# a broker of its own. +BROKER="${1:?usage: source client-env.sh [remote]}" + +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 + +if [ "${2:-}" = "remote" ]; then + export DT_INFERENCE_ENDPOINT=hpc + export DT_INFERENCE_BACKEND=dragon_v3 +fi diff --git a/service/deploy/run-hpc-endpoint.sh b/service/deploy/run-hpc-endpoint.sh new file mode 100755 index 0000000..cf8fd0e --- /dev/null +++ b/service/deploy/run-hpc-endpoint.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +# Run the xGFabric twin's HPC endpoint -- INSIDE a compute allocation. +# Adapted from the AmSC dt-complete deploy kit; the launch constraints +# (dragon launcher, PATH-by-name helpers, SLURM_EXPORT_ENV) are the +# debugged ones from there. +# +# ./run-hpc-endpoint.sh +# +# The endpoint registers as 'hpc' -- what the client's remote profile +# (client-env.sh remote) asks for. DT_DIR as in setup-hpc-endpoint.sh. +# +# Two modes, same script: +# faked demo -- DT_VENV unset: uses $DT_DIR/ve.demo (setup-hpc-endpoint.sh) +# real workload-- DT_VENV=, XGF_DIR=: the +# real trainings need TensorFlow (Ben's cloned cfdaai env) +# and the profiler/task bodies import from the tasks tree +# (setup-hpc-endpoint-real.sh prints the exact invocation). +set -euo pipefail +BROKER="${1:?usage: $0 }" +DT_DIR="${DT_DIR:-${SCRATCH:-$HOME}/digital_twins}" +VENV="${DT_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 +# heatmaps land on scratch where available +export XGF_WORKSPACE="${XGF_WORKSPACE:-${SCRATCH:-$HOME}/xgf_twin}" + +# real workload: the profiler shells out to a script in the xGFabric tasks +# tree and the (lazy-imported) task bodies import tasks.* -- put the +# checkout on PYTHONPATH when XGF_DIR names one +if [ -n "${XGF_DIR:-}" ]; then + export PYTHONPATH="$XGF_DIR:${PYTHONPATH:-}" +fi + +"$VENV/bin/dragon" "$VENV/bin/radical-orbit-endpoint.py" -n hpc \ + 2>&1 | tee endpoint.log diff --git a/service/deploy/setup-broker.sh b/service/deploy/setup-broker.sh new file mode 100755 index 0000000..a91d375 --- /dev/null +++ b/service/deploy/setup-broker.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# DTaaS broker host for the xGFabric service twin (e.g. radical.3). +# Run there. Adapted from the AmSC dt-complete deploy kit, which +# carries the debugged install/launch lore; delta here: numpy for the +# by-value components. +# +# Once per host: broker_cert.pem / broker_key.pem / broker.token in +# ~/.radical/orbit/. DT_DIR overrides the checkout+venv location. +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 +./ve.demo/bin/pip install -q numpy +# match the endpoint's rhapsody exactly (e491cd2-based): main breaks the +# endpoint's dragon backend (task_logs= to Batch), so both tiers pin the +# proven branch. The dragon cancel fix on it is a no-op broker-side; the +# orbit-backend cancel fix lives only on the main-based #91 branch, so a +# harmless KeyError-cancel line may appear on teardown here. +./ve.demo/bin/pip install -q --force-reinstall --no-deps \ + "rhapsody-py[telemetry] @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" + +echo "done. start the broker with:" +echo " cd $DT_DIR && ./deploy/run-broker.sh \$PWD/ve.demo" diff --git a/service/deploy/setup-hpc-endpoint-real.sh b/service/deploy/setup-hpc-endpoint-real.sh new file mode 100755 index 0000000..7ec27cc --- /dev/null +++ b/service/deploy/setup-hpc-endpoint-real.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# Real-workload HPC endpoint for the xGFabric service twin (Perlmutter). +# Run on a login node. +# +# ./setup-hpc-endpoint-real.sh +# +# Unlike setup-hpc-endpoint.sh (faked components, plain ve.demo), the real +# FNO/PINN/PCR trainings need TensorFlow + the CFD/ML stack. Rather than +# build that, we CLONE Ben's cfdaai conda env (which has it) and install +# our runtime into the clone -- never into his shared env. +# +# Prerequisites (read access): /global/common/software/m5290 (Ben's envs +# + cups_structure.zip) and /global/cfs/cdirs/m5290/precalc_sims. +# `conda` must be on PATH (module load, or source Ben's mconda). +# +# Overrides: DT_DIR (our checkout), XGF_DIR (xGFabric checkout), +# DT_ENV (clone location), DT_BASE_ENV (base env name, default cfdaai). +set -euo pipefail +BROKER="${1:?usage: $0 }" +DT_DIR="${DT_DIR:-${SCRATCH:-$HOME}/digital_twins}" +XGF_DIR="${XGF_DIR:-${SCRATCH:-$HOME}/xGFabric}" +ENV_PREFIX="${DT_ENV:-${SCRATCH:-$HOME}/dt-endpoint-env}" +BEN_ENVS="/global/common/software/m5290/bcarter/mconda/envs" +BASE_ENV="${DT_BASE_ENV:-cfdaai}" + +command -v conda >/dev/null || { + echo "ERROR: conda not on PATH -- module load python, or source Ben's" >&2 + echo " mconda, then re-run." >&2; exit 1; } + +# clone Ben's env (never install into his shared copy) +conda config --add envs_dirs "$BEN_ENVS" 2>/dev/null || true +if [ ! -d "$ENV_PREFIX" ]; then + echo "==> cloning $BASE_ENV -> $ENV_PREFIX (this is large; once)" + conda create -y -p "$ENV_PREFIX" --clone "$BEN_ENVS/$BASE_ENV" +fi +PY="$ENV_PREFIX/bin/python" + +# wire contract: the same Python minor on every tier +ver="$("$PY" -c 'import sys; print("%d.%d" % sys.version_info[:2])')" +if [ "$ver" != "3.12" ]; then + echo "ERROR: $BASE_ENV is Python $ver; the DT wire contract pins 3.12" >&2 + echo " on every host. Set DT_BASE_ENV to a 3.12 env (try" >&2 + echo " xgfabric), or rebuild from its environment.yml on 3.12." >&2 + exit 1 +fi + +# our runtime into the clone -- git-pinned deps first (as deploy/install.sh +# does), then digitaltwin resolves the rest (asyncflow, orbit) from PyPI +[ -d "$DT_DIR" ] || git clone https://github.com/radical-cybertools/digital.twins.git "$DT_DIR" +( cd "$DT_DIR" && git checkout devel && git pull ) +"$PY" -m pip install -q \ + "rose @ git+https://github.com/radical-cybertools/ROSE@64330d9cb43c3e13ca67daf0d8ae84a2ae6c3f17" +"$PY" -m pip install -q --force-reinstall --no-deps \ + "rhapsody-py[telemetry,dragon] @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" +"$PY" -m pip install -q "$DT_DIR" + +# xGFabric tasks tree: the profiler shells out to a script here, and the +# real task bodies import tasks.* at run time (run-hpc-endpoint.sh puts +# XGF_DIR on PYTHONPATH) +[ -d "$XGF_DIR" ] || git clone https://github.com/radical-collaboration/xGFabric.git "$XGF_DIR" +( cd "$XGF_DIR" && git checkout feature/dtaas-twin && git pull ) +# the Pi predictor's clean-slate dataset (Ben: the sample IS the real one) +cp -n "$XGF_DIR/tasks/profiler/pi_profiler/data.csv.sample" \ + "$XGF_DIR/tasks/profiler/pi_profiler/data.csv" 2>/dev/null || true + +# broker cert + token +mkdir -p ~/.radical/orbit +scp "$BROKER:.radical/orbit/broker_cert.pem" "$BROKER:.radical/orbit/broker.token" ~/.radical/orbit/ + +cat < +# +# DT_DIR overrides the checkout+venv location; the default lands on +# $SCRATCH -- the venv is far too big for a Perlmutter home quota. +set -euo pipefail +BROKER="${1:?usage: $0 }" +DT_DIR="${DT_DIR:-${SCRATCH:-$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 +# the surrogate/sink task bodies unpickle and run here +./ve.demo/bin/pip install -q numpy matplotlib +# dragon backend + idempotent cancel + failure-traceback logging. +# e491cd2-based (NOT main): main's dragon backend passes task_logs= to +# Batch(), which the pinned dragonhpc 0.14.1 does not accept. This is +# the branch the AmSC demo proved on the same dragon. +./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 (salloc -N1 -C cpu -q interactive -t 2:00:00 -A )," +echo "then run: service/deploy/run-hpc-endpoint.sh $BROKER" diff --git a/service/investigators.py b/service/investigators.py new file mode 100644 index 0000000..1c56476 --- /dev/null +++ b/service/investigators.py @@ -0,0 +1,91 @@ +"""Service-safe surrogate investigators: FNO / PINN / PCR stand-ins. + +Each keeps the real investigators' shape -- batch sensor windows, run a +"training" flow task on the endpoint, publish the model, serve field +inference -- with fake physics: training sleeps an architecture-typical +time and inference synthesizes a wind field. What stays real is the +part the demo is about: three architectures with different costs +competing for selection, and the shared simulation subtask feeding all +of them (memoised once per sensor window, see SIM_MASTER). +""" + +import asyncio +import random + +import numpy as np + +from digitaltwin.components import ModelInvestigator, TypedData +from digitaltwin.runtime import RuntimeAPI + +from tasks.common.dtypes import SIM_MASTER, WIND_FIELD + +# architecture-typical (train_seconds, inference_seconds) -- the spread is +# what makes profiler-driven selection visible +ARCH_COST = { + "fno": (6.0, 0.20), + "pinn": (9.0, 0.60), + "pcr": (3.0, 0.05), +} + +WINDOW = 4 # sensor points per training window + + +class SurrogateInvestigator(ModelInvestigator): + """One fake surrogate; ``arch`` picks its cost profile.""" + + def __init__(self, flow, arch: str, learn_backend: str | None = None): + super().__init__(flow) + self.flow = flow + self.arch = arch + self.batch: list = [] + train_s, infer_s = ARCH_COST[arch] + + # Route retraining to the 'learning' engine when the session has + # one, so it shows in its own dashboard lane; inference stays on + # the default (inference) engine. backend=None is the plain + # default -- asyncflow has no aliasing, so a label is set only + # when the caller declares the matching engine. + @flow.function_task(backend=learn_backend) + async def train(arch, points, sim): + # endpoint-side "training": cost is the architecture's + await asyncio.sleep(train_s) + speeds = [p["wind_speed"] for p in points] + return {"model": arch, "w": float(np.mean(speeds)) / 20.0, + "sim": sim[1]} + + @flow.function_task + async def infer(in_data, model="na", w=0.0, **_): + await asyncio.sleep(infer_s) + if model == "na": + return TypedData(WIND_FIELD, + {"arch": "na", "w": w, "result": None}) + # synthesized field: smooth bump scaled by the model weight + x = np.linspace(-2, 2, 32) + xx, yy = np.meshgrid(x, x) + field = (2.5 * w) * np.exp(-(xx ** 2 + yy ** 2)) + field += np.random.default_rng().normal(0, 0.05, field.shape) + return TypedData(WIND_FIELD, + {"arch": model, "w": w, "result": (3, field)}) + + self._train = train + self._infer = infer + + async def _on_input(self, in_data: TypedData) -> None: + self.batch.append(in_data.data) + + async def main_loop(self, runtime: RuntimeAPI): + runtime.subscribe_to_topic(runtime.ON_INPUT, self._on_input) + runtime.set_inference_task(self._infer) + runtime.publish_new_model({"model": "na", "w": 0.0}) + + while True: + if len(self.batch) < WINDOW: + await asyncio.sleep(1.0) + continue + + points, self.batch = self.batch[:WINDOW], self.batch[WINDOW:] + # one simulation per window, shared and memoised across the + # three investigators -- the SIM_MASTER contract from twin.py + sim = await runtime.call_shared_subtask(SIM_MASTER, points[0]) + model = await self._train(self.arch, points, sim) + runtime.publish_new_model(model, {"quality": random.random()}) diff --git a/service/makefile b/service/makefile new file mode 100644 index 0000000..da3cac6 --- /dev/null +++ b/service/makefile @@ -0,0 +1,6 @@ + +PROJECT_TYPE = python +PROJECT_NAME = xgfabric.service + +include $(HOME)/.makefile + diff --git a/service/profiler.py b/service/profiler.py new file mode 100644 index 0000000..bb37a77 --- /dev/null +++ b/service/profiler.py @@ -0,0 +1,87 @@ +"""Profiling chain, service-safe: measure, then predict for the Pi. + +``ServiceProfiler`` really profiles: it unpickles the surrogate's +inference function and times a run with the example input -- inline as a +flow task on the endpoint, instead of twin.py's exported-pickle + +subprocess round trip (whose script path does not survive shipping by +value). Results are memoised per (function, model) exactly like the +original. + +``ServicePiPredictor`` stands in for the Pi endpoint's learned runtime +model: predicted Pi runtime = measured runtime x a Pi/HPC slowdown +factor with mild noise. The real EndpointInvestigator's training on +recorded Pi profiles is the phase-2 story. +""" + +import time + +import cloudpickle + +from digitaltwin.components import DataType, ModelInvestigator, TypedData +from digitaltwin.lru import LRUCache, freeze +from digitaltwin.runtime import RuntimeAPI + +TASK_DESCRIPTION_DTYPE = DataType("TASK_INFO") +PROFILE_RESULTS = DataType("PROFILE_RESULT") +PI_PREDICT_RUNTIME = DataType("pi_PREDICT_RUNTIME") + +PI_SLOWDOWN = 7.5 # a Raspberry Pi against one HPC core, rough + + +class ServiceProfiler(ModelInvestigator): + """Times one run of a shipped inference function; memoises per model.""" + + def __init__(self, flow): + super().__init__(flow) + self.flow = flow + + @flow.function_task + async def run_profiled(blob, example, model_kwargs): + fn = cloudpickle.loads(blob) + t0 = time.monotonic() + await fn(example, **(model_kwargs or {})) + return {"runtime": time.monotonic() - t0} + + cache = LRUCache(128) + + async def do_inference(in_data: TypedData): + blob, example, model_kwargs = in_data.data + key = freeze((blob, tuple(sorted((model_kwargs or {}).items())))) + + if await cache.exists(key): + profile = await cache.fetch_item(key) + else: + profile = await run_profiled(blob, example, model_kwargs) + await cache.put_item(key, profile) + + return TypedData(PROFILE_RESULTS, + {"profile": profile, "task": in_data.data}) + + self.inference_task = do_inference + + async def main_loop(self, runtime: RuntimeAPI): + runtime.set_inference_task(self.inference_task) + runtime.publish_new_model() + + +class ServicePiPredictor(ModelInvestigator): + """PROFILE_RESULT -> predicted runtime on the Pi endpoint.""" + + def __init__(self, flow): + super().__init__(flow) + self.flow = flow + + @flow.function_task + async def predict(in_data): + import random + measured = in_data.data["profile"]["runtime"] + return measured * PI_SLOWDOWN * random.uniform(0.9, 1.1) + + async def do_inference(in_data: TypedData): + return TypedData(PI_PREDICT_RUNTIME, await predict(in_data)) + + self.inference_task = do_inference + + async def main_loop(self, runtime: RuntimeAPI): + runtime.set_inference_task(self.inference_task) + runtime.publish_new_model() diff --git a/service/sensor_publisher.py b/service/sensor_publisher.py new file mode 100644 index 0000000..77cd69a --- /dev/null +++ b/service/sensor_publisher.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +"""Fake Davis wind sensor as an EXTERNAL channel publisher. + +In service mode the sensor is not a twin component: it publishes to a +shared channel and the twin binds it with ``add_input`` -- any number of +twins can listen. The records mirror what ``tasks.davis`` emits +(``strip_cols`` output), at the same 5 s cadence, fake for the same +reason twin.py fakes them. + +Environment: DT_STREAM_BACKEND / RADICAL_ORBIT_BROKER_URL(+_CERT) select +the data plane, exactly like every other client-side process. +""" + +import asyncio +import datetime +import random + +from digitaltwin.streaming import ChannelPublisher + +DAVIS_CHANNEL = "xgf/davis" + + +async def main() -> None: + publisher = await ChannelPublisher.open(DAVIS_CHANNEL) + n = 0 + try: + while True: + r = random.random() + record = { + "dt": datetime.datetime.now().isoformat(), + "wind_speed": round(r * 20, 1), + "wind_avg": round(r * 15, 1), + "wind_dir": round(r * 360), + } + await publisher.publish(record) + n += 1 + if n % 12 == 0: + print(f"davis: {n} records", flush=True) + await asyncio.sleep(5) + finally: + await publisher.close() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/service/sink.py b/service/sink.py new file mode 100644 index 0000000..e26e5dc --- /dev/null +++ b/service/sink.py @@ -0,0 +1,71 @@ +"""Terminal component: heatmap of the selected surrogate's wind field. + +The heatmap is rendered on the endpoint and returned inline (small PNG +bytes) so the runtime can surface it to the dashboard via +`record_output` -- the DT service has no file staging, and a downscaled +heatmap is well under the return-value cap. It is also written to +XGF_WORKSPACE on the endpoint for the record (resolved at runtime, since +this module ships by value). +""" + +import base64 +import os +from pathlib import Path + +from digitaltwin.components import TypedData, UtilityTask + + +def _workspace() -> Path: + base = Path(os.environ.get("XGF_WORKSPACE", "") + or Path.home() / "xgf_twin") + base.mkdir(parents=True, exist_ok=True) + return base + + +class ServiceSink(UtilityTask): + def __init__(self, flow): + super().__init__(flow) + self.flow = flow + self.count = 0 + + @flow.function_task + async def render_heatmap(data, arch, w, fname): + import io + + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig = plt.figure(figsize=(4, 3.2)) + im = plt.imshow(data, cmap="viridis", origin="lower", + vmin=0, vmax=2.5) + plt.colorbar(im) + plt.title(f"{arch} W={round(w, 3)}") + plt.xlabel("X") + plt.ylabel("Y") + + # to the endpoint filesystem, for the record + fig.savefig(fname, dpi=90, bbox_inches="tight") + # and to memory, small, for the inline return + buf = io.BytesIO() + fig.savefig(buf, format="png", dpi=72, bbox_inches="tight") + plt.close(fig) + return buf.getvalue() + + self._render = render_heatmap + + async def main_loop(self, runtime, in_data: TypedData): + arch = in_data.data["arch"] + print(f"[sink] field from {arch}", flush=True) + if arch == "na" or not in_data.data.get("result"): + return + + self.count += 1 + fname = str(_workspace() / f"field_{self.count:04d}_{arch}.png") + png = await self._render(in_data.data["result"][1], arch, + in_data.data["w"], fname) + + # surface it to the dashboard -- small enough to ride inline + b64 = base64.b64encode(png).decode("ascii") + runtime.record_output(f"field {self.count} ({arch})", + f"data:image/png;base64,{b64}") diff --git a/service/wind_agent.py b/service/wind_agent.py new file mode 100644 index 0000000..ad9894b --- /dev/null +++ b/service/wind_agent.py @@ -0,0 +1,97 @@ +"""The WindField agent, service-safe. + +Structure and selection logic mirror ``tasks.wind_agent.WindFieldAgent``: +three surrogate investigators under one agent, a shared (memoised) +simulation subtask, and a model selector that ranks published models by +their profiler-predicted Pi runtime -- fetched through the twin's own +``get_inference`` chain (task description -> profile -> predicted +runtime). Only the physics behind those steps is faked. +""" + +import asyncio +import random + +import cloudpickle + +from digitaltwin.components import SciAgent, TypedData, DataType +from digitaltwin.runtime import RuntimeAPI + +from tasks.common.dtypes import SIM_MASTER, DAVIS_WIND_SENSOR +from service.investigators import SurrogateInvestigator +from service.profiler import PI_PREDICT_RUNTIME, PROFILE_RESULTS, \ + TASK_DESCRIPTION_DTYPE + + +class ServiceWindFieldAgent(SciAgent): + def __init__(self, flow, learn_backend: str | None = None): + super().__init__(flow) + self.flow = flow + + self.fno = SurrogateInvestigator(flow, "fno", learn_backend) + self.pinn = SurrogateInvestigator(flow, "pinn", learn_backend) + self.pcr = SurrogateInvestigator(flow, "pcr", learn_backend) + + @flow.function_task + async def simulate(sensor_pt): + # twin.py's tk_do_simulation is already faked to a random + # quality plus a precomputed-sim path; same here, local path + await asyncio.sleep(1.0) + return (random.random(), f"precalc_sims/{random.randint(0, 71)}.csv") + + self.sim_master = simulate + + @flow.function_task + async def model_select(in_data: TypedData, models): + # the demo's crown jewel: pick the model with the shortest + # profiler-PREDICTED runtime on the Pi + if not models: + return 0 + best = min(models, key=lambda m: m["pi_runtime"]) + return best["investigator"] + + self.model_selector = model_select + + self.model_to_process: asyncio.Queue[dict] = asyncio.Queue() + self.models: list[dict] = [] + + async def model_publish_cb(self, inv, model_args: dict, acc: dict): + await self.model_to_process.put({ + "investigator": inv.get_id(), + "model_args": model_args, + "metrics": acc, + }) + + async def main_loop(self, runtime: RuntimeAPI): + runtime.start_investigator(self.fno) + runtime.start_investigator(self.pinn) + runtime.start_investigator(self.pcr) + + runtime.register_shared_subtask(SIM_MASTER, self.sim_master, 64) + runtime.subscribe_to_topic(runtime.ON_MODEL_PUBLISH, + self.model_publish_cb) + + runtime.set_model_selection_task(self.model_selector) + runtime.update_model_selector(models=[]) + + while True: + item = await self.model_to_process.get() + if item["model_args"]["model"] == "na": + continue + + raw = runtime.get_inference_tasks()[item["investigator"]] + raw = getattr(raw, "__wrapped__", raw) + + probe_pt = {"dt": "probe", "wind_speed": 10.0, + "wind_avg": 8.0, "wind_dir": 180.0} + description = (cloudpickle.dumps(raw), + TypedData(DAVIS_WIND_SENSOR, probe_pt), + item["model_args"]) + + profile = await runtime.get_inference( + TypedData(TASK_DESCRIPTION_DTYPE, description), + PROFILE_RESULTS) + pi = await runtime.get_inference(profile, PI_PREDICT_RUNTIME) + + item["pi_runtime"] = pi.data + self.models.append(item) + runtime.update_model_selector(models=self.models[-10:]) diff --git a/twin_service.py b/twin_service.py new file mode 100644 index 0000000..a8f9bae --- /dev/null +++ b/twin_service.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""The xGFabric twin as a service: twin.py's graph on a DTaaS broker. + +Same story as twin.py -- Davis wind sensor, a field agent whose three +surrogate architectures compete on profiler-predicted Pi runtime, a +heatmap sink -- but the twin lives on an ORBIT broker and its tasks run +on a rhapsody endpoint. The sensor is an external channel publisher +(service/sensor_publisher.py); components ship by value (service/*, +fake physics, see service/__init__.py). + +Environment (client side): + RADICAL_ORBIT_BROKER_URL(+_CERT) the broker + DT_STREAM_BACKEND=orbit the data plane + DT_SERVICE_HOST participant hosting `dt` (broker) + DT_INFERENCE_ENDPOINT / _BACKEND task placement (local endpoint / + concurrent by default) +""" + +import argparse +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from radical.orbit import EndpointRuntime + +from digitaltwin.components import NULL_DTYPE, TypedData +from digitaltwin.service import register_user_modules + +import service +import service.investigators +import service.profiler +import service.sink +import service.wind_agent +import tasks +import tasks.common +import tasks.common.dtypes + +from tasks.common.dtypes import DAVIS_WIND_SENSOR, WIND_FIELD +from service.profiler import ( + PI_PREDICT_RUNTIME, + PROFILE_RESULTS, + TASK_DESCRIPTION_DTYPE, + ServicePiPredictor, + ServiceProfiler, +) +from service.sensor_publisher import DAVIS_CHANNEL +from service.sink import ServiceSink +from service.wind_agent import ServiceWindFieldAgent + +register_user_modules([ + tasks, tasks.common, tasks.common.dtypes, + service, service.investigators, service.profiler, + service.sink, service.wind_agent, +]) + +DT_HOST = os.environ.get("DT_SERVICE_HOST", "broker") + +INFERENCE_EP = os.environ.get("DT_INFERENCE_ENDPOINT", "dt_inference_ep") +INFERENCE_BE = os.environ.get("DT_INFERENCE_BACKEND", "concurrent") +# learning defaults to the same endpoint on the concurrent executor: its +# own dashboard lane for the surrogate retraining, no second dragon +# runtime, and the training tasks stay clear of dragon's mp bridge +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(args) -> int: + runtime = EndpointRuntime() + runtime.start(wait=True) + + try: + dt = runtime.get_plugin(DT_HOST, "dt", config=ENGINES) + print(f"[client] session: {dt.sid} (reattach with this sid)") + + twin = dt.create_twin() + print(f"[client] twin: {twin}") + + agent = dt.package(ServiceWindFieldAgent, learn_backend="learning") + profiler = dt.package(ServiceProfiler) + pi = dt.package(ServicePiPredictor) + sink = dt.package(ServiceSink) + + # sensor is external: bind its channel to the input dtype + dt.add_input(twin, DAVIS_WIND_SENSOR, DAVIS_CHANNEL) + + dt.add_agent(twin, agent, DAVIS_WIND_SENSOR, WIND_FIELD) + dt.add_task(twin, sink, WIND_FIELD, NULL_DTYPE) + + # the profiling chain the agent's selector queries + dt.add_investigator(twin, profiler, TASK_DESCRIPTION_DTYPE, + PROFILE_RESULTS) + dt.add_investigator(twin, pi, PROFILE_RESULTS, PI_PREDICT_RUNTIME) + + dt.start(twin) + + # client-side feedback: the pipeline is stream driven, so poll the + # twin and probe the field inference -- a stuck twin shows in 10s + probe_pt = {"dt": "probe", "wind_speed": 12.0, "wind_avg": 9.0, + "wind_dir": 90.0} + deadline = time.time() + args.runtime + while time.time() < deadline: + time.sleep(10) + + info = dt.twin(twin) + print(f"[client] state={info['state']}" + f" calls={info.get('calls') or {}}", flush=True) + if info["state"] == "failed": + print(f"[client] twin failed: {info.get('last_error')}") + return 1 + + answer = dt.get_inference( + twin, TypedData(DAVIS_WIND_SENSOR, probe_pt), WIND_FIELD, + timeout=60) + arch = answer.data["arch"] + print(f"[client] field probe -> arch={arch}" + f" w={answer.data.get('w')}", flush=True) + + print("SHUTDOWN") + dt.twin_close(twin) + return 0 + + finally: + runtime.stop() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="xGFabric twin, service mode") + parser.add_argument("--runtime", type=int, default=240, + help="seconds to keep the twin running") + sys.exit(main(parser.parse_args()))