diff --git a/demo/Sep_04_broker.sh b/demo/Sep_04_broker.sh new file mode 100755 index 0000000..479926e --- /dev/null +++ b/demo/Sep_04_broker.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# xGFabric demo -- DTaaS broker. Run on the broker host (radical.3). +# Sets up (install + fixes) then starts the broker. Idempotent. +# +# demo/Sep_04_broker.sh [broker-ip] +# +# Assumes broker_cert.pem / broker_key.pem / broker.token in ~/.radical/orbit/. +set -euo pipefail + +BROKER_IP="${1:-95.217.193.116}" + +# --- demo/site specific ---------------------------------------------------- +DT_DIR="${DT_DIR:-$HOME/digital_twins}" +RH="rhapsody-py[telemetry] @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" +# --------------------------------------------------------------------------- + +echo "--------------------------" +echo "xGFabric Demo September 04" +echo "Broker on $BROKER_IP ($(hostname -f))" +echo "--------------------------" + +# checkout + install the pinned stack, then pin the dragon-compatible +# rhapsody branch (extras backfilled) +[ -d "$DT_DIR" ] || git clone https://github.com/radical-cybertools/digital.twins.git "$DT_DIR" +( cd "$DT_DIR" && git checkout devel && git pull ) +( cd "$DT_DIR" + ./deploy/install.sh broker + # install.sh compares by version string and skips when it is unchanged, + # so it can leave stale digitaltwin code (e.g. missing record_output). + # Force the checked-out devel tip over whatever it left. + ./ve.demo/bin/pip install -q --force-reinstall --no-deps "$DT_DIR" + ./ve.demo/bin/pip install -q numpy pandas # the agent runs broker-side + ./ve.demo/bin/pip install -q --force-reinstall --no-deps \ + "rhapsody-py @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" + ./ve.demo/bin/pip install -q "$RH" ) + +echo "starting broker (wss://0.0.0.0:8000) ..." +exec "$DT_DIR/deploy/run-broker.sh" "$DT_DIR/ve.demo" diff --git a/demo/Sep_04_client.sh b/demo/Sep_04_client.sh new file mode 100755 index 0000000..56953f7 --- /dev/null +++ b/demo/Sep_04_client.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# xGFabric demo -- client driver. Run on the client. Ensures the light +# client-side deps, then drives the real twin against the HPC endpoint. +# +# demo/Sep_04_client.sh [broker-ip] +set -euo pipefail + +BROKER_IP="${1:-95.217.193.116}" + +# --- demo/site specific ---------------------------------------------------- +VENV="${DT_VENV:-$HOME/radical/digital_twins/ve.demo}" +DT_DIR="${DT_DIR:-$HOME/radical/digital_twins}" +HERE="$(cd "$(dirname "$0")/.." && pwd)" # xGFabric checkout root +RUNTIME="${RUNTIME:-600}" +# default to the fast fake-surrogate driver (CPU-friendly: sleeps + numpy, +# no TF/xgboost, heatmaps in seconds, learning lane built in). RUN_REAL=1 +# switches to the real TF workload (needs a GPU endpoint to be timely). +DRIVER="twin_service.py" +[ "${RUN_REAL:-}" = 1 ] && DRIVER="twin_service_real.py" +# --------------------------------------------------------------------------- + +echo "--------------------------" +echo "xGFabric Demo September 04" +echo "Client $DRIVER ($(hostname -f)) -> broker $BROKER_IP, endpoint HPC 'hpc'" +echo "--------------------------" + +cd "$HERE" + +# client-side deps the real components import at packaging time (TF and +# matplotlib stay lazy / endpoint-only); pyspot submodule for the sensor. +"$VENV/bin/pip" install -q python-dotenv numpy pandas +git submodule update --init --recursive +# keep the client digitaltwin current (record_output etc.) +( cd "$DT_DIR" && git checkout devel && git pull ) && "$VENV/bin/pip" install -q "$DT_DIR" + +# placement: run the twin's tasks on the HPC endpoint (dragon) +export RADICAL_ORBIT_BROKER_URL="wss://$BROKER_IP:8000" +export RADICAL_ORBIT_BROKER_CERT="$HOME/.radical/orbit/broker_cert.pem" +export DT_STREAM_BACKEND=orbit +export DT_INFERENCE_ENDPOINT=hpc +export DT_INFERENCE_BACKEND=dragon_v3 +export DT_LEARNING_ENDPOINT=hpc +export DT_LEARNING_BACKEND=concurrent + +echo "driving twin via $DRIVER (runtime ${RUNTIME}s) ..." +echo "dashboard: https://$BROKER_IP:8000/broker/dt/ui?live=1" +exec "$VENV/bin/python" "$DRIVER" --runtime "$RUNTIME" diff --git a/demo/Sep_04_endpoint.sh b/demo/Sep_04_endpoint.sh new file mode 100755 index 0000000..e7ed34b --- /dev/null +++ b/demo/Sep_04_endpoint.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# xGFabric demo -- Perlmutter endpoint RUN. Run INSIDE the compute +# allocation (no pip here; Sep_04_endpoint_setup.sh did that on the login +# node). Launches the rhapsody endpoint under the dragon launcher. +# +# demo/Sep_04_endpoint.sh [broker-ip] [endpoint-name] +set -euo pipefail + +BROKER_IP="${1:-95.217.193.116}" +EP="${2:-hpc}" + +# --- demo/site specific ---------------------------------------------------- +ENV_PREFIX="${DT_ENV:-$SCRATCH/dt-endpoint-env}" +XGF_DIR="${XGF_DIR:-$SCRATCH/xgfabric}" +PLAYGROUND_DIR="${PLAYGROUND_DIR:-$SCRATCH/xgf_playground}" +# --------------------------------------------------------------------------- + +echo "--------------------------" +echo "xGFabric Demo September 04" +echo "Endpoint HPC '$EP' on $(hostname -f) -> broker $BROKER_IP" +echo "--------------------------" + +# dragon resolves its helpers BY NAME via srun on the task side +export PATH="$ENV_PREFIX/bin:$PATH" +export PYTHONPATH="$XGF_DIR:${PYTHONPATH:-}" +export RADICAL_ORBIT_BROKER_URL="wss://$BROKER_IP: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 XGF_WORKSPACE="$PLAYGROUND_DIR" +# keep the orbit log off HOME (tiny NERSC quota) and modest +export RADICAL_ORBIT_LOG_LVL="${RADICAL_ORBIT_LOG_LVL:-WARNING}" +export RADICAL_ORBIT_LOG_FILE="$SCRATCH/orbit-logs/$EP.log" +mkdir -p "$SCRATCH/orbit-logs" "$PLAYGROUND_DIR" +rm -f "$HOME/.radical/orbit/logs/"*.log 2>/dev/null || true + +# --- GPU (opt-in: RUN_GPU=1) ----------------------------------------------- +# Only for a real GPU allocation. Prepending /usr/lib64 to LD_LIBRARY_PATH +# can shadow the conda env's own libs and break the dragon task launch, so +# it stays OFF by default -- the CPU/fake demo path must not touch the env. +if [ "${RUN_GPU:-}" = 1 ]; then + module load cudatoolkit 2>/dev/null || true + export LD_LIBRARY_PATH="/usr/lib64:$ENV_PREFIX/lib:${LD_LIBRARY_PATH:-}" + echo "GPU check (nvidia-smi -L):" + nvidia-smi -L 2>&1 | sed 's/^/ /' || echo " no GPU visible (need salloc -C gpu)" +fi + +echo "launching endpoint '$EP' under dragon ..." +exec "$ENV_PREFIX/bin/dragon" "$ENV_PREFIX/bin/radical-orbit-endpoint.py" -n "$EP" \ + 2>&1 | tee "$SCRATCH/orbit-logs/$EP.console.log" diff --git a/demo/Sep_04_endpoint_setup.sh b/demo/Sep_04_endpoint_setup.sh new file mode 100755 index 0000000..9a33d1d --- /dev/null +++ b/demo/Sep_04_endpoint_setup.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# xGFabric demo -- Perlmutter endpoint SETUP. Run on a LOGIN NODE (it +# pip-installs and clones; the allocation cannot). Idempotent. +# +# demo/Sep_04_endpoint_setup.sh [broker-ip] (broker-ip only for the banner) +# +# Assumes keys/token already staged in ~/.radical/orbit/. +set -euo pipefail + +BROKER_IP="${1:-95.217.193.116}" +EP=hpc + +# --- demo/site specific ---------------------------------------------------- +BRANCH=feature/dtaas-twin-real +DT_DIR="${DT_DIR:-$SCRATCH/digital_twins}" +XGF_DIR="${XGF_DIR:-$SCRATCH/xgfabric}" +ENV_PREFIX="${DT_ENV:-$SCRATCH/dt-endpoint-env}" +PLAYGROUND_DIR="${PLAYGROUND_DIR:-$SCRATCH/xgf_playground}" +BEN_ENVS="/global/common/software/m5290/bcarter/mconda/envs" +BASE_ENV=cfdaai +RH="rhapsody-py[telemetry,dragon] @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" +ROSE="rose @ git+https://github.com/radical-cybertools/ROSE@64330d9cb43c3e13ca67daf0d8ae84a2ae6c3f17" +# --------------------------------------------------------------------------- + +echo "--------------------------" +echo "xGFabric Demo September 04 -- ENDPOINT SETUP (login node)" +echo "Broker $BROKER_IP | Endpoint '$EP' | env $ENV_PREFIX | playground $PLAYGROUND_DIR" +echo "--------------------------" + +module load python/3.12 2>/dev/null || true +command -v conda >/dev/null || { echo "ERROR: conda not on PATH" >&2; exit 1; } + +# checkouts +[ -d "$DT_DIR" ] || git clone https://github.com/radical-cybertools/digital.twins.git "$DT_DIR" +( cd "$DT_DIR" && git checkout devel && git pull ) +[ -d "$XGF_DIR" ] || git clone https://github.com/radical-collaboration/xGFabric.git "$XGF_DIR" +( cd "$XGF_DIR" && git checkout "$BRANCH" && git pull && git submodule update --init --recursive ) + +# conda env: a clone of Ben's cfdaai (TensorFlow + CFD/ML stack), never his +conda config --add envs_dirs "$BEN_ENVS" 2>/dev/null || true +[ -d "$ENV_PREFIX" ] || conda create -y -p "$ENV_PREFIX" --clone "$BEN_ENVS/$BASE_ENV" +PY="$ENV_PREFIX/bin/python" + +ver="$("$PY" -c 'import sys;print("%d.%d"%sys.version_info[:2])')" +[ "$ver" = "3.12" ] || { echo "ERROR: $BASE_ENV is Python $ver, need 3.12" >&2; exit 1; } + +# our runtime into the clone. orbit force-reinstalled (a clone may carry a +# stale one that hides the endpoint script), then deps backfilled. +echo "==> installing runtime into $ENV_PREFIX" +"$PY" -m pip install -q "$ROSE" +"$PY" -m pip install -q --force-reinstall --no-deps "radical.orbit>=0.7" +"$PY" -m pip install -q "radical.orbit>=0.7" +"$PY" -m pip install -q "$DT_DIR" +# rhapsody LAST so the dragon-compatible branch wins over what digitaltwin +# pulled (main passes task_logs= to Batch(), which the pinned dragonhpc +# rejects); extras step backfills opentelemetry + dragonhpc. +"$PY" -m pip install -q --force-reinstall --no-deps \ + "rhapsody-py @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" +"$PY" -m pip install -q "$RH" +# endpoint_trainer.py / endpoint_eval.py need these; cfdaai lacks xgboost +"$PY" -m pip install -q xgboost scikit-learn + +# data dirs: seed the Pi-predictor dataset at the RUNTIME datastore (where +# endpoint_trainer.py reads it) +mkdir -p "$PLAYGROUND_DIR/profiler/pi_profiler" +cp -n "$XGF_DIR/tasks/profiler/pi_profiler/data.csv.sample" \ + "$PLAYGROUND_DIR/profiler/pi_profiler/data.csv" 2>/dev/null || true + +# sanity +echo "==> verify" +"$PY" -c "import radical.orbit, digitaltwin, rhapsody, opentelemetry; print(' imports ok')" +"$PY" -c "import rhapsody.backends.execution.dragon as d; n=open(d.__file__).read().count('task_logs'); print(' dragon task_logs (want 0):', n)" +[ -x "$ENV_PREFIX/bin/radical-orbit-endpoint.py" ] && echo " endpoint script ok" || echo " WARN: endpoint script missing" +[ -x "$ENV_PREFIX/bin/dragon" ] && echo " dragon ok" || echo " WARN: dragon missing" + +echo "--------------------------" +echo "done. get an allocation, then run:" +echo " $XGF_DIR/demo/Sep_04_endpoint.sh $BROKER_IP $EP" +echo "--------------------------" diff --git a/demo/Sep_04_sensor.sh b/demo/Sep_04_sensor.sh new file mode 100755 index 0000000..f144d96 --- /dev/null +++ b/demo/Sep_04_sensor.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# xGFabric demo -- external wind sensor (fake Davis). Run on the client. +# Publishes to the twin's input channel over the ORBIT data plane. +# +# demo/Sep_04_sensor.sh [broker-ip] +set -euo pipefail + +BROKER_IP="${1:-95.217.193.116}" + +# --- demo/site specific ---------------------------------------------------- +VENV="${DT_VENV:-$HOME/radical/digital_twins/ve.demo}" +HERE="$(cd "$(dirname "$0")/.." && pwd)" # xGFabric checkout root +# --------------------------------------------------------------------------- + +echo "--------------------------" +echo "xGFabric Demo September 04" +echo "Sensor davis-wind ($(hostname -f)) -> broker $BROKER_IP" +echo "--------------------------" + +export RADICAL_ORBIT_BROKER_URL="wss://$BROKER_IP:8000" +export RADICAL_ORBIT_BROKER_CERT="$HOME/.radical/orbit/broker_cert.pem" +export DT_STREAM_BACKEND=orbit + +cd "$HERE" +echo "publishing davis-wind to channel xgf/davis ..." +exec "$VENV/bin/python" service/sensor_publisher.py diff --git a/service/README.md b/service/README.md new file mode 100644 index 0000000..042ba68 --- /dev/null +++ b/service/README.md @@ -0,0 +1,134 @@ +# 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 + +Driver: `twin_service_real.py` wires the real components (done, but +experimental -- not yet validated end to end; it is the on-Perlmutter +starting point). The lazy-import refactor is in place, so the real +investigators import TF-free on the client/broker; TF loads only in the +task bodies on the endpoint. + +Client/broker still need the light deps the real components carry: + + git submodule update --init --recursive # pyspot (sensor) + /bin/pip install python-dotenv numpy pandas + +Two items remain open before a real run is trustworthy: + +- **Shared filesystem.** The real components write to + `config['PLAYGROUND_DIR']` from both main_loops (broker) and tasks + (endpoint); those line up only on a shared filesystem, so run the + **broker on Perlmutter too** for the real workload (PLAYGROUND_DIR on + `$SCRATCH`). A broker on radical.3 splits the playground across hosts. +- **cloudpickle parity.** The cloned conda env and `ve.demo` may carry + different cloudpickle versions; align them 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 | endpoint-side rhapsody telemetry + `collect_reports.py` | + +## Telemetry and reports + +Nothing to enable: the rhapsody plugin on the endpoint records every +task plus a resource poll on its own (the `[telemetry]` extra), into +`telemetry-output/session.*.telemetry.jsonl` under the endpoint's +working directory. After (or during) a run: + + /bin/python service/collect_reports.py [telemetry-dir] + +renders twin.py's report set — task waterfall, dependency wait, stage +timers, swimlane, concurrency/resources, and the gantt — next to the +jsonl. For a remote endpoint, scp the jsonl files first. + +Known gap: per-workflow grouping in the gantt needs +`asyncflow.workflow_id`, which only `workflow_scope()` stamps — the +service engine does not run one. That is an engine-side telemetry +feature for the DT service (digital.twins), tracked there; every other +report is complete without it. 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/collect_reports.py b/service/collect_reports.py new file mode 100644 index 0000000..bb73f7b --- /dev/null +++ b/service/collect_reports.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +"""Render twin.py's telemetry reports from a service-mode run. + +In service mode nobody calls ``flow.start_telemetry`` -- and nobody has +to: the rhapsody plugin on the ENDPOINT records every task and a +resource poll on its own (the ``[telemetry]`` extra), as +``telemetry-output/session.*.telemetry.jsonl`` in the endpoint's +working directory. This script points twin.py's existing report +generators at those files. + + python service/collect_reports.py [telemetry-dir] [--last N] + +``telemetry-dir`` defaults to ``telemetry-output/`` under the local +digital.twins checkout; for a remote endpoint, scp the jsonl files over +first. Reports land next to the jsonl. + +Known gap: the workflow-gantt renders but cannot group per workflow -- +service twins do not run inside ``workflow_scope()``, so no +``asyncflow.workflow_id`` is stamped. Grouping needs engine-side +telemetry in the DT service (a digital.twins feature, not a demo-side +one); the task waterfall / swimlane / resource dashboards are complete +without it. +""" + +import argparse +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from reports.plot_workflow_dashboard import plot_split +from reports.plot_workflow_gantt import plot as plot_gantt + +DEFAULT_DIR = os.path.expanduser("~/radical/digital_twins/telemetry-output") + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Render reports from endpoint telemetry") + parser.add_argument("telemetry_dir", nargs="?", default=DEFAULT_DIR) + parser.add_argument("--last", type=int, default=1, + help="how many of the newest sessions to render") + args = parser.parse_args() + + files = sorted(Path(args.telemetry_dir).glob("*.telemetry.jsonl"), + key=lambda p: p.stat().st_mtime, reverse=True) + if not files: + print(f"no *.telemetry.jsonl under {args.telemetry_dir}", + file=sys.stderr) + return 1 + + for f in files[:args.last]: + print(f"== {f.name}") + plot_split(str(f)) + plot_gantt(f) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) 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..bd8e5bd --- /dev/null +++ b/service/deploy/run-hpc-endpoint.sh @@ -0,0 +1,49 @@ +#!/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}" +# the orbit log defaults to ~/.radical/orbit/logs, which is HOME -- a tiny +# quota on NERSC. Put it on scratch and keep the level modest; also clear +# any old HOME log so a full quota there cannot block the run. +export RADICAL_ORBIT_LOG_LVL="${RADICAL_ORBIT_LOG_LVL:-WARNING}" +export RADICAL_ORBIT_LOG_FILE="${RADICAL_ORBIT_LOG_FILE:-${SCRATCH:-$HOME}/orbit-logs/hpc.log}" +mkdir -p "$(dirname "$RADICAL_ORBIT_LOG_FILE")" +rm -f "$HOME/.radical/orbit/logs/"*.log 2>/dev/null || true + +# 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..c10a5c7 --- /dev/null +++ b/service/deploy/setup-broker.sh @@ -0,0 +1,30 @@ +#!/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" +# backfill the telemetry extra's deps (opentelemetry) in case a fresh +# install did not carry them +./ve.demo/bin/pip install -q \ + "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-client.sh b/service/deploy/setup-client.sh new file mode 100755 index 0000000..e28f72f --- /dev/null +++ b/service/deploy/setup-client.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# Client-side deps for the real service twin, on top of digital.twins' +# `./deploy/install.sh client` (which provides the DT stack + rose). +# +# ./setup-client.sh [ve.demo path] +# +# The real components import these at packaging time -- TensorFlow and +# matplotlib stay lazy (endpoint-only), so the client needs only the +# light ones plus the pyspot submodule the sensor uses. +set -euo pipefail +VENV="${1:-$HOME/radical/digital_twins/ve.demo}" + +"$VENV/bin/pip" install -q python-dotenv numpy pandas + +# pyspot (tasks/common/pyspot) is a git submodule -- davis.py imports it +git submodule update --init --recursive + +echo "done. run from this checkout:" +echo " source service/deploy/client-env.sh remote" +echo " $VENV/bin/python service/sensor_publisher.py # terminal 1" +echo " $VENV/bin/python twin_service_real.py # terminal 2" diff --git a/service/deploy/setup-hpc-endpoint-real.sh b/service/deploy/setup-hpc-endpoint-real.sh new file mode 100755 index 0000000..3af7ee4 --- /dev/null +++ b/service/deploy/setup-hpc-endpoint-real.sh @@ -0,0 +1,100 @@ +#!/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}" +# demo-specific: must match PLAYGROUND_DIR in tasks/common/config.sh (the +# client ships that value to the components, which use it here on the +# endpoint). The Pi-predictor's datastore is seeded under it below. +PLAYGROUND_DIR="${PLAYGROUND_DIR:-${SCRATCH:-$HOME}/xgf_playground}" + +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 ) +RH="rhapsody-py[telemetry,dragon] @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" + +"$PY" -m pip install -q \ + "rose @ git+https://github.com/radical-cybertools/ROSE@64330d9cb43c3e13ca67daf0d8ae84a2ae6c3f17" +# force-reinstall orbit: a cloned env may already carry a stale radical.orbit +# that pip would treat as satisfied, leaving the endpoint script uninstalled. +# --no-deps places the script/version; the plain install backfills orbit's +# own deps (fastapi, uvicorn, websockets) without disturbing the pins above. +"$PY" -m pip install -q --force-reinstall --no-deps "radical.orbit>=0.7" +"$PY" -m pip install -q "radical.orbit>=0.7" +"$PY" -m pip install -q "$DT_DIR" +# rhapsody LAST, so the dragon-compatible branch wins over whatever +# digitaltwin's install pulled (main's dragon backend passes task_logs= to +# Batch(), which the pinned dragonhpc rejects). Two steps: --no-deps pins +# the branch build, then the extras install backfills opentelemetry +# (telemetry) and dragonhpc (dragon) without re-resolving the rest. +"$PY" -m pip install -q --force-reinstall --no-deps \ + "rhapsody-py @ git+https://github.com/radical-cybertools/rhapsody@fix/dragon-cancel-idempotent" +"$PY" -m pip install -q "$RH" + +# 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-real && git pull \ + && git submodule update --init --recursive ) # pyspot (davis/sensor) +# the Pi predictor's clean-slate dataset (Ben: the sample IS the real one), +# seeded at the RUNTIME datastore under PLAYGROUND_DIR -- that is where +# endpoint_trainer.py reads it (the tasks-tree copy is not that path) +mkdir -p "$PLAYGROUND_DIR/profiler/pi_profiler" +cp -n "$XGF_DIR/tasks/profiler/pi_profiler/data.csv.sample" \ + "$PLAYGROUND_DIR/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..4567695 --- /dev/null +++ b/service/sink.py @@ -0,0 +1,74 @@ +"""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. A copy is also written to +XGF_WORKSPACE, but that path is resolved and created *inside* the task +(endpoint-side); the sink's main_loop runs on the broker, a different +host, so it must not build or create that path itself. +""" + +import base64 + +from digitaltwin.components import TypedData, UtilityTask + + +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, count): + import io + import os + from pathlib import Path + + 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") + + # a copy to the endpoint filesystem, for the record -- resolved + # here (endpoint XGF_WORKSPACE) and best-effort, so a missing or + # read-only path never fails the field + try: + base = Path(os.environ.get("XGF_WORKSPACE", "") + or Path.home() / "xgf_twin") + base.mkdir(parents=True, exist_ok=True) + fig.savefig(str(base / f"field_{count:04d}_{arch}.png"), + dpi=90, bbox_inches="tight") + except OSError: + pass + + # 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 + png = await self._render(in_data.data["result"][1], arch, + in_data.data["w"], self.count) + + # 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/tasks/common/config.sh b/tasks/common/config.sh index e274c09..9cd236d 100644 --- a/tasks/common/config.sh +++ b/tasks/common/config.sh @@ -12,7 +12,7 @@ NODE_COUNT=4 # ROSE / Python config -PLAYGROUND_DIR=/pscratch/sd/b/bcarter/playground_dt +PLAYGROUND_DIR=/pscratch/sd/m/merzky/xgf_playground COMMON_DIR="../common" diff --git a/tasks/do_fno/__init__.py b/tasks/do_fno/__init__.py index 15b6a64..0791656 100644 --- a/tasks/do_fno/__init__.py +++ b/tasks/do_fno/__init__.py @@ -1 +1,9 @@ -from .main import * +# Lazy re-export: `from .main import *` pulled TensorFlow/scikit-learn at +# package import, which forced those onto the client and broker just to +# import a submodule (e.g. the investigator wrapper). PEP 562 defers it +# to attribute access -- `tasks..tk_*` still resolves (on the +# endpoint, where the stack lives; used by wrapper.py), but importing +# `tasks..` no longer drags the heavy deps in. +def __getattr__(name): + from . import main + return getattr(main, name) diff --git a/tasks/do_fno/fno_investigator.py b/tasks/do_fno/fno_investigator.py index 80e7e1c..f82fbfa 100644 --- a/tasks/do_fno/fno_investigator.py +++ b/tasks/do_fno/fno_investigator.py @@ -4,7 +4,11 @@ from radical.asyncflow import WorkflowEngine from digitaltwin.components import Any, ModelInvestigator, TypedData from digitaltwin.runtime import DTRuntime, RuntimeAPI -from .main import tk_do_fno, tk_fno_eval + +# NOTE: `.main` pulls TensorFlow at import. It is imported lazily inside +# the task bodies below (which run on the endpoint, where TF lives), so +# packaging/instantiating this class on the client and broker stays +# TF-free. See service/README.md "Real workload". from rose.al.active_learner import Learner @@ -52,6 +56,8 @@ async def do_inference(in_data: TypedData, config=None, model="na"): # pass through, no model yet return TypedData(WIND_FIELD, {"arch": "na"}) + from .main import tk_fno_eval + # model available! # logger.info(f"Do FNO inference: {model}") wind = in_data.data["wind_speed"] @@ -76,6 +82,7 @@ async def do_train(self, runtime: RuntimeAPI, batch): @self.learner.training_task(as_executable=False) async def do_fno(config, sims): + from .main import tk_do_fno return tk_do_fno(config, sims) self.config["TASK_COUNTER"] = self.task_counter diff --git a/tasks/do_pcr/__init__.py b/tasks/do_pcr/__init__.py index 15b6a64..0791656 100644 --- a/tasks/do_pcr/__init__.py +++ b/tasks/do_pcr/__init__.py @@ -1 +1,9 @@ -from .main import * +# Lazy re-export: `from .main import *` pulled TensorFlow/scikit-learn at +# package import, which forced those onto the client and broker just to +# import a submodule (e.g. the investigator wrapper). PEP 562 defers it +# to attribute access -- `tasks..tk_*` still resolves (on the +# endpoint, where the stack lives; used by wrapper.py), but importing +# `tasks..` no longer drags the heavy deps in. +def __getattr__(name): + from . import main + return getattr(main, name) diff --git a/tasks/do_pcr/pcr_investigator.py b/tasks/do_pcr/pcr_investigator.py index d8adf20..6c3595d 100644 --- a/tasks/do_pcr/pcr_investigator.py +++ b/tasks/do_pcr/pcr_investigator.py @@ -6,9 +6,9 @@ from digitaltwin.components import Any, ModelInvestigator, TypedData from digitaltwin.runtime import RuntimeAPI -from .main import tk_pcr_eval, tk_pcr_partition -from .main import tk_do_pcr -from .main import tk_do_pcr_pack +# NOTE: `.main` pulls scikit-learn (and the PCR train stack) at import -- +# the tk_* helpers are imported lazily in the task bodies below (which run +# on the endpoint) so packaging/instantiating on client/broker stays light. from rose.al.active_learner import Learner @@ -32,6 +32,8 @@ async def do_inference(in_data: TypedData, config=None, model="na"): # PCR works on a range rather than a single data point. # for now, assume the wind is stable. + from .main import tk_pcr_eval + wind_single = in_data.data["wind_speed"] wind = np.ones(13) * wind_single result = tk_pcr_eval(config, model, wind) @@ -87,14 +89,17 @@ async def do_train(self, runtime: RuntimeAPI, batch): @self.flow.function_task async def partition(config, sims, sensor_vals): + from .main import tk_pcr_partition return tk_pcr_partition(config, sims, sensor_vals) @self.flow.function_task async def do_pcr(config, machine_data_output): + from .main import tk_do_pcr return tk_do_pcr(config, machine_data_output) @self.flow.function_task async def do_pack(config, *pcr_output_dirs): + from .main import tk_do_pcr_pack return tk_do_pcr_pack(config, *pcr_output_dirs) self.config["TASK_NAME"] = "partition" diff --git a/tasks/do_pinn/__init__.py b/tasks/do_pinn/__init__.py index 15b6a64..0791656 100644 --- a/tasks/do_pinn/__init__.py +++ b/tasks/do_pinn/__init__.py @@ -1 +1,9 @@ -from .main import * +# Lazy re-export: `from .main import *` pulled TensorFlow/scikit-learn at +# package import, which forced those onto the client and broker just to +# import a submodule (e.g. the investigator wrapper). PEP 562 defers it +# to attribute access -- `tasks..tk_*` still resolves (on the +# endpoint, where the stack lives; used by wrapper.py), but importing +# `tasks..` no longer drags the heavy deps in. +def __getattr__(name): + from . import main + return getattr(main, name) diff --git a/tasks/do_pinn/pinn_investigator.py b/tasks/do_pinn/pinn_investigator.py index fe82984..59aad7f 100644 --- a/tasks/do_pinn/pinn_investigator.py +++ b/tasks/do_pinn/pinn_investigator.py @@ -4,7 +4,9 @@ from radical.asyncflow import WorkflowEngine from digitaltwin.components import Any, ModelInvestigator, TypedData from digitaltwin.runtime import RuntimeAPI -from .main import tk_do_pinn, tk_pinn_eval + +# NOTE: `.main` pulls TensorFlow at import -- imported lazily in the task +# bodies below (which run on the endpoint) so client/broker stay TF-free. from rose.al.active_learner import Learner @@ -54,6 +56,8 @@ async def do_inference(in_data: TypedData, config=None, model="na"): # pass through, no model yet return TypedData(WIND_FIELD, {"arch": "na"}) + from .main import tk_pinn_eval + # model available! wind = in_data.data["wind_speed"] result = tk_pinn_eval(config, model, wind) @@ -76,6 +80,7 @@ async def do_train(self, runtime: RuntimeAPI, batch): @self.learner.training_task(as_executable=False) async def do_pinn(config, sims): + from .main import tk_do_pinn return tk_do_pinn(config, sims) self.config["TASK_COUNTER"] = self.task_counter diff --git a/tasks/profiler/PLAN-telemetry-profiling.md b/tasks/profiler/PLAN-telemetry-profiling.md new file mode 100644 index 0000000..567d428 --- /dev/null +++ b/tasks/profiler/PLAN-telemetry-profiling.md @@ -0,0 +1,80 @@ +# Plan: profile from endpoint telemetry (approach 2), for later + +Status: **not started.** Approach (1), in-process measurement, is live +(`inproc.py` + `ProfilerInvestigator`/`EndpointInvestigator`). This is +the follow-on that turns "profile" into "run once, read telemetry." + +## Why + +The surrogate selector needs a resource fingerprint of each candidate's +inference to predict its edge (Pi) runtime. Approach (1) measures the +call in-process with `resource.getrusage` + `psutil.io_counters`. That +works and is staging-free, but: + +- it re-runs the inference *only to measure it* (a second execution on + top of the one the twin already does for real), and +- in-process counters include a little of the task-runner's own + activity and give peak RSS, not the isolated child's PSS. + +The rhapsody endpoint **already records** per-task telemetry -- task +lifecycle plus `ResourceUpdate` (cpu %, memory %, disk read/write bytes) +polled while the task runs -- into +`telemetry-output/session.*.telemetry.jsonl` (the `[telemetry]` extra; +this is what `service/collect_reports.py` renders). So the cost of a +candidate's inference is captured for free when it runs as a normal +task. Approach (2): read that instead of re-running under a profiler. + +## Shape + +1. **Correlate a task to its telemetry.** The DT already ties a task uid + to its twin/component (`DTRuntime.note_task`, carried in `twin_list` + as `tasks`/`task_components`). The telemetry events carry the same + `task_id`. A helper reads the session jsonl and, for a given uid, + returns the `{total_seconds, cpu_seconds, disk_read_bytes, + disk_write_bytes, memory_bytes}` aggregated from that task's + `TaskStarted`/`TaskCompleted` span and its `ResourceUpdate` samples. + +2. **Profiler investigator becomes a reader, not a runner.** When the + agent asks to profile a candidate, the profiler no longer executes + the inference: it (a) ensures the candidate has run at least once + (the twin's normal inference already does this), (b) looks up that + task's uid, (c) reads the telemetry record for it, (d) returns the + same six-key profile dict. No second execution, no subprocess. + +3. **Keep the schema.** Emit the exact keys `inproc.py` and the + subprocess profiler use, so `data.csv`, `endpoint_trainer.py`, and + the Pi predictor are unchanged; approach (2) is a drop-in source + swap. + +## Open questions to resolve when building it + +- **Telemetry access from the profiler.** The jsonl lives on the + endpoint; the profiler component's `main_loop` runs on the broker. + Reading it needs either (i) a small function-task on the endpoint that + greps the jsonl for the uid and returns the record (endpoint-side, + staging-free -- preferred), or (ii) an engine-side telemetry API on + the DT service (see digital.twins#36, engine-side telemetry) that + surfaces per-task metrics to the runtime. +- **Timing / flush.** `ResourceUpdate` is polled (default ~0.5 s); a + very fast inference may produce few samples. Fall back to the task + span's wall/cpu when samples are sparse, or lower the poll interval + for the profiling window. +- **Attribution granularity.** Confirm the telemetry `task_id` matches + the uid the DT records for the *inference* task specifically (not a + wrapping flow task), mirroring the exact-attribution work the + dashboard relies on. +- **Isolation.** Telemetry measures the task as it actually ran + (shared with whatever else the endpoint was doing). For a + comparative fingerprint that is acceptable; note it if absolute + numbers ever matter. + +## Relationship to the remaining executable-task hand-offs + +The Pi-predictor's `train_model` / `call_inference` (`endpoint_trainer.py` +/ `endpoint_eval.py`) are still executable tasks whose *command-builder* +bodies prepare input files (`model.json`, `inf.json`). `data.csv` is now +endpoint-local (the append moved to a function-task); `inf.json` in +`call_inference` is still written broker-side. Those are the prediction +*model*, not profiling, and are out of scope here -- but the same fix +applies: write their inputs in a preceding function-task (endpoint-side) +or pass them as command arguments. Track that alongside this. diff --git a/tasks/profiler/components.py b/tasks/profiler/components.py index 542d90f..f6c50a0 100644 --- a/tasks/profiler/components.py +++ b/tasks/profiler/components.py @@ -25,7 +25,6 @@ TASK_DESCRIPTION_DTYPE = DataType("TASK_INFO") PROFILE_RESULTS = DataType("PROFILE_RESULT") -script_path = os.path.dirname(os.path.realpath(__file__)) class ProfilerInvestigator(ModelInvestigator): @@ -33,50 +32,36 @@ def __init__(self, flow: WorkflowEngine, workdir: str = "."): super().__init__(flow) self.flow = flow self.workdir = workdir - os.makedirs(self.workdir, exist_ok=True) + # no makedirs here: the profiling runs in-process on the endpoint + # (below), so there is no file to prepare on this (broker) side. - @self.flow.executable_task - async def exec_profiler(task, example_data: TypedData, model_kwargs: dict): - if model_kwargs is None: - model_kwargs = {} - print("Exec profiler request") - # fix to use a unique file name. - export_inference_function( - f"{self.workdir}/meta-profiler.pkl", task, example_data, **model_kwargs - ) + @self.flow.function_task + async def run_profiled(blob, example_data, model_kwargs): + # runs on the endpoint: reconstruct the candidate's inference + # callable and measure one call in-process -- no cloudpickle + # file, no subprocess, nothing written on the broker. + from .inproc import load_callable, profile_call - # call profiler - return shlex.join( - [ - "python3", - f"{script_path}/profiler.py", - f"{self.workdir}/meta-profiler.pkl", - ] - ) + func = load_callable(blob) + return await profile_call(func, example_data, model_kwargs or {}) sim_lock = asyncio.Lock() sim_lru = LRUCache(128) # store 128 different sims async def do_inference(in_data: TypedData): - # # for inference, just run the simulation. async with sim_lock: - # ignore example_data task, example_data, model_kwargs = in_data.data key = freeze((task, model_kwargs)) if await sim_lru.exists(key): profile = await sim_lru.fetch_item(key) - task_data = in_data.data - out = {"profile": profile, "task": task_data} - return TypedData(PROFILE_RESULTS, out) + return TypedData(PROFILE_RESULTS, + {"profile": profile, "task": in_data.data}) - result = await exec_profiler(task, example_data, model_kwargs) - r = json.loads(result) + r = await run_profiled(task, example_data, model_kwargs) await sim_lru.put_item(key, r) - - task_data = in_data.data - out = {"profile": r, "task": task_data} - return TypedData(PROFILE_RESULTS, out) + return TypedData(PROFILE_RESULTS, + {"profile": r, "task": in_data.data}) self.inference_task = do_inference @@ -111,32 +96,29 @@ def __init__( else: self.datastore = datastore_path - os.makedirs(self.datastore, exist_ok=True) + # NOT created here: __init__ runs on the broker, but the datastore + # is an endpoint path (where the tasks read/write it). The + # endpoint-side tasks makedirs it (stage_inf / append_row). + try: + os.makedirs(self.datastore, exist_ok=True) + except OSError: + pass self.callback_jobs: asyncio.Queue = asyncio.Queue() self.done_jobs: set = set() self.learner = Learner(flow) - @self.flow.executable_task + @self.flow.function_task async def exec_profiler(task_data): - task, example_data, model_kwargs = task_data - - # fix to use a unique file name. - export_inference_function( - f"{self.datastore}/profile.pkl", task, example_data, **model_kwargs - ) + # the endpoint ("Pi") measurement, in-process on this endpoint: + # run the candidate's inference and read the counters, no + # cloudpickle file and no subprocess (approach 1). + from .inproc import load_callable, profile_call - # call profiler - print("endpoint profiler request") - return shlex.join( - [ - "python3", - f"{script_path}/profiler.py", - "--csv", - f"{self.datastore}/profile.pkl", - ] - ) + task, example_data, model_kwargs = task_data + func = load_callable(task) + return await profile_call(func, example_data, model_kwargs or {}) self.exec_profiler = exec_profiler @@ -144,8 +126,7 @@ async def exec_profiler(task_data): async def train_model(): return shlex.join( [ - "python3", - f"{script_path}/endpoint_trainer.py", + "python3", "-m", "tasks.profiler.endpoint_trainer", f"{self.datastore}/data.csv", f"{self.datastore}/model.json", ] @@ -153,16 +134,44 @@ async def train_model(): self.train_task = train_model + @self.flow.function_task + async def append_row(datastore, row): + # runs on the endpoint (a function_task body IS the task), so the + # row lands in the same data.csv that train_model (executable, + # also endpoint) reads -- see main_loop. Doing this append in + # main_loop instead would write it on the broker, on a different + # filesystem from the training task. + import os + + os.makedirs(datastore, exist_ok=True) + with open(f"{datastore}/data.csv", "a") as fh: + fh.write(row + "\n") + + self._append_row = append_row + + @self.flow.function_task + async def stage_inf(datastore, pf): + # write endpoint_eval's input on the endpoint (a function_task + # body runs there), so the executable command below reads it on + # the same host -- not the broker, where the executable-task + # command-builder runs. + import json as _json + import os + + os.makedirs(datastore, exist_ok=True) + with open(f"{datastore}/inf.json", "w") as fh: + _json.dump(pf, fh) + + self._stage_inf = stage_inf + @self.flow.executable_task async def call_inference(in_data: TypedData, model=None, name=""): - # for inference, just run the simulation. pf = in_data.data["profile"] - with open(f"{self.datastore}/inf.json", "w") as f: - json.dump(pf, f) + # stage the input endpoint-side before the command runs + await self._stage_inf(self.datastore, pf) return shlex.join( [ - "python3", - f"{script_path}/endpoint_eval.py", + "python3", "-m", "tasks.profiler.endpoint_eval", model, f"{self.datastore}/inf.json", ] @@ -215,17 +224,19 @@ async def main_loop(self, runtime: RuntimeAPI): while True: item = await self.callback_jobs.get() - pi_out = await self.exec_profiler(item["task"]) - # I only want the first column - pi_time = pi_out.split(",", 1)[0] + # exec_profiler now returns the profile dict (in-process); the + # Pi runtime is its wall time + pi_profile = await self.exec_profiler(item["task"]) + pi_time = str(pi_profile["total_seconds"]) # label the endpoint_time as "pi_seconds" nersc_profile = item["profile"] out = ",".join([str(f) for f in nersc_profile.values()]) + "," out += pi_time - with open(f"{self.datastore}/data.csv", "a") as f: - f.write(out + "\n") + # append on the endpoint (where train_model reads data.csv), not + # here on the broker -- see append_row + await self._append_row(self.datastore, out) out = json.loads(await self.train_task()) model = out["model"] diff --git a/tasks/profiler/inproc.py b/tasks/profiler/inproc.py new file mode 100644 index 0000000..7264c80 --- /dev/null +++ b/tasks/profiler/inproc.py @@ -0,0 +1,82 @@ +"""In-process resource profiling of an inference call. + +Approach (1) from the profiler discussion: instead of cloudpickling the +inference function to a file and running the standalone ``profiler.py`` +under a fresh process group, run the call in-process and read the +process's own resource counters around it. This runs entirely where the +task runs (the endpoint) -- no file hand-off between the executable-task +command-builder (broker) and the command (endpoint), which is what the +subprocess profiler needed a shared filesystem for. + +Trade-off vs the subprocess profiler: no process-group isolation, so the +numbers include a little of the task runner's own activity during the +call, and ``memory_bytes`` is the process peak RSS (a high-water mark) +rather than the child's PSS. For a *comparative* fingerprint across the +candidate surrogates -- which is all the selector needs -- that is fine. + +The returned dict keeps the exact keys and order the subprocess profiler +emitted, so ``data.csv`` and ``endpoint_trainer.py`` are unchanged. +""" + +import os +import resource +import time + +import cloudpickle +import psutil + + +def load_callable(blob): + """Reconstruct the inference callable the agent cloudpickled.""" + obj = cloudpickle.loads(blob) + # the agent ships either the raw callable or the {"func": ...} payload + if isinstance(obj, dict) and "func" in obj: + obj = obj["func"] + return obj + + +async def profile_call(func, example_data, kwargs=None) -> dict: + """Run ``func(example_data, **kwargs)`` once and measure its cost. + + Returns the profiler's schema: + total_seconds, cpu_seconds, disk_read_bytes, disk_write_bytes, + sys_read_bytes, memory_bytes + """ + + kwargs = kwargs or {} + proc = psutil.Process(os.getpid()) + + try: + io0 = proc.io_counters() + except Exception: + io0 = None + ru0 = resource.getrusage(resource.RUSAGE_SELF) + t0 = time.perf_counter() + + result = func(example_data, **kwargs) + if hasattr(result, "__await__"): + await result + + wall = time.perf_counter() - t0 + ru1 = resource.getrusage(resource.RUSAGE_SELF) + try: + io1 = proc.io_counters() + except Exception: + io1 = None + + def io_delta(attr): + if io0 is None or io1 is None: + return 0 + return max(0, getattr(io1, attr, 0) - getattr(io0, attr, 0)) + + return { + "total_seconds": wall, + "cpu_seconds": (ru1.ru_utime - ru0.ru_utime) + + (ru1.ru_stime - ru0.ru_stime), + "disk_read_bytes": io_delta("read_bytes"), + "disk_write_bytes": io_delta("write_bytes"), + "sys_read_bytes": io_delta("read_chars"), + # ru_maxrss is peak RSS in KiB on Linux (bytes on macOS); the peak + # high-water mark, matching the subprocess profiler's peak metric. + "memory_bytes": int(ru1.ru_maxrss) * 1024, + } diff --git a/tasks/sink.py b/tasks/sink.py index e6af586..87fa98c 100644 --- a/tasks/sink.py +++ b/tasks/sink.py @@ -1,62 +1,77 @@ -import asyncio +"""Terminal component: heatmap of the selected surrogate's wind field. + +The heatmap renders on the endpoint and is 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 +PLAYGROUND_DIR on the endpoint for the record. + +matplotlib is imported lazily inside the task body (which runs on the +endpoint) so packaging/instantiating this class on the client and broker +does not need it. +""" + +import base64 +import os from radical.asyncflow import WorkflowEngine from digitaltwin.components import UtilityTask, TypedData from .common.dtypes import * import logging -import matplotlib.pyplot as plt logger = logging.getLogger(__name__) -def graph(data, fname, model_name, z, w): - - plt.figure(figsize=(6, 5)) - - # Plot heatmap - im = plt.imshow(data, cmap="viridis", origin="lower", vmin=0, vmax=2.5) - - # Add colorbar - plt.colorbar(im) - - # Optional labels - plt.title(f"Heatmap of {model_name} at Z={z}, W={round(w,3)}") - plt.xlabel("X") - plt.ylabel("Y") - - # Save image - plt.savefig(fname, dpi=300, bbox_inches="tight") - - plt.close() - - class CUPS_Sink(UtilityTask): def __init__(self, flow: WorkflowEngine, config): super().__init__(flow) self.flow = flow self.config = config + self.count = 0 @self.flow.function_task - async def save_photo(in_data, config): - fname = config["PLAYGROUND_DIR"] + "/out.png" - graph( - in_data.data["result"][1], - fname, - in_data.data["arch"], - 3, - w=in_data.data["w"], - ) - print("\n\n" + "=" * 30 + "\n" + str(fname) + "\n" + "=" * 30) - - self.save_photo = save_photo + async def render(field, 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(field, 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") + + if fname: + os.makedirs(os.path.dirname(fname), exist_ok=True) + fig.savefig(fname, dpi=150, bbox_inches="tight") + buf = io.BytesIO() + fig.savefig(buf, format="png", dpi=72, bbox_inches="tight") + plt.close(fig) + return buf.getvalue() + + self._render = render async def main_loop(self, runtime, in_data: TypedData): - print(f"Received Inference: {in_data.data['arch']}") + arch = in_data.data["arch"] + print(f"Received Inference: {arch}") - if in_data.data["arch"] == "na": + if arch == "na" or not in_data.data.get("result"): print("No surrogate ready yet") return - # create graph - await self.save_photo(in_data, self.config) + self.count += 1 + fname = os.path.join(self.config.get("PLAYGROUND_DIR", "."), + f"out_{self.count:04d}.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}") + print("\n" + "=" * 30 + f"\n{fname}\n" + "=" * 30) diff --git a/tasks/wind_agent.py b/tasks/wind_agent.py index ae32fa8..fb5cb54 100644 --- a/tasks/wind_agent.py +++ b/tasks/wind_agent.py @@ -25,8 +25,11 @@ class WindFieldAgent(SciAgent): - def __init__(self, flow: WorkflowEngine, config: dict, logger): + def __init__(self, flow: WorkflowEngine, config: dict, logger=None): super().__init__(flow) + # the DTaaS service instantiates with flow + config only; fall + # back to the module logger when no logger is injected + logger = logger or logging.getLogger(__name__) self.flow = flow self.config = config.copy() self.config["AGENT_NAME"] = "WindField" diff --git a/tasks/wind_learner.py b/tasks/wind_learner.py new file mode 100644 index 0000000..c9ec639 --- /dev/null +++ b/tasks/wind_learner.py @@ -0,0 +1,66 @@ +"""A visible learning lane for the real twin. + +The real workload's own learners (FNO/PINN/PCR and the Pi predictor) run +on the twin's default ('inference') engine, so the dashboard shows no +distinct learning lane. This component adds one: it consumes the same +wind-sensor stream, and every window of points runs a training task +routed to the 'learning' backend -- exactly the labeling trick from the +service demo's SurrogateInvestigator (@function_task(backend=...)). It +publishes to its own WIND_TREND dtype, so it never competes with the +real WIND_FIELD path feeding the sink. +""" + +import asyncio +import random + +from digitaltwin.components import DataType, ModelInvestigator, TypedData +from digitaltwin.runtime import RuntimeAPI + +WIND_TREND = DataType("WindTrend_dt") + +WINDOW = 4 # sensor points per training window +TRAIN_SECONDS = 4.0 # endpoint-side "training" cost, visible in the lane + + +class WindTrendLearner(ModelInvestigator): + """Online mean-wind learner; its training runs in the learning lane.""" + + def __init__(self, flow, learn_backend: str | None = None): + super().__init__(flow) + self.flow = flow + self.batch: list = [] + + # the label rides here: with learn_backend='learning' the training + # task routes to the session's learning engine and shows in its own + # lane; inference stays on the default engine. backend=None is the + # plain default (asyncflow has no aliasing -- a label is set only + # when the caller declares the matching engine). + @flow.function_task(backend=learn_backend) + async def train(points): + await asyncio.sleep(TRAIN_SECONDS) + speeds = [p["wind_speed"] for p in points] + return {"model": "wind_trend", "mean": sum(speeds) / len(speeds)} + + @flow.function_task + async def infer(in_data, model="na", mean=0.0, **_): + return TypedData(WIND_TREND, {"model": model, "mean": mean}) + + 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", "mean": 0.0}) + + while True: + if len(self.batch) < WINDOW: + await asyncio.sleep(1.0) + continue + + points, self.batch = self.batch[:WINDOW], self.batch[WINDOW:] + model = await self._train(points) + runtime.publish_new_model(model, {"quality": random.random()}) diff --git a/twin_service.py b/twin_service.py new file mode 100644 index 0000000..ba5c323 --- /dev/null +++ b/twin_service.py @@ -0,0 +1,139 @@ +#!/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) + 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())) diff --git a/twin_service_real.py b/twin_service_real.py new file mode 100644 index 0000000..a6ebdf0 --- /dev/null +++ b/twin_service_real.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +"""The xGFabric twin as a service, with the REAL workload. + +twin.py's graph, run on a DTaaS broker like twin_service.py -- but the +field agent, its FNO/PINN/PCR surrogates, the profiler and the Pi +predictor are the real components from tasks/, not the service/* fakes. +Physics stays where twin.py already fakes it (the sensor, and +tk_do_simulation's precalc-CSV shortcut); the trainings, the profiler +timing, and the selection are real. + +STATUS: runs end to end on the DTaaS stack -- broker on radical.3, the +ORBIT data plane, and a rhapsody/dragon endpoint on Perlmutter. The real +TF surrogates are GPU-bound, though: on a CPU allocation each inference is +tens of seconds, too slow for a live demo, so the September demo defaults +to the fake driver (twin_service.py) and selects this one only with +RUN_REAL=1 (see demo/Sep_04_client.sh) on a GPU endpoint. + +Notes: + + * Endpoint env: demo/Sep_04_endpoint_setup.sh clones Ben's cfdaai + (TF + CFD/ML) and adds xgboost. The lazy-import refactor keeps the + client/broker TF-free -- they need only dotenv + numpy + pandas and + the pyspot submodule (`git submodule update --init`). + * No shared-filesystem requirement: the compute chain runs entirely on + the endpoint, and the cross-host hand-offs were moved endpoint-side + (data.csv append and inf.json staging are function_tasks, the profiler + measures in-process), so a broker on radical.3 with the endpoint on + Perlmutter works. + * Backends: inference on dragon_v3, learning on concurrent, set via the + DT_* env knobs (see demo/Sep_04_client.sh). + +Environment: config.sh (tasks/common/config.sh) supplies PLAYGROUND_DIR, +CSPOT_LIMIT, endpoint/model paths etc., loaded via dotenv as in twin.py. +Client wire env as usual: RADICAL_ORBIT_BROKER_URL(+_CERT), +DT_STREAM_BACKEND=orbit, DT_SERVICE_HOST. +""" + +import argparse +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import dotenv + +from radical.orbit import EndpointRuntime + +from digitaltwin.components import NULL_DTYPE, DataType, TypedData +from digitaltwin.service import register_user_modules + +# real components (TF is deferred to the task bodies -- import is light) +import tasks +import tasks.common +import tasks.common.dtypes +import tasks.wind_agent +import tasks.wind_learner +import tasks.sink +import tasks.profiler.components +import tasks.do_fno.fno_investigator +import tasks.do_pinn.pinn_investigator +import tasks.do_pcr.pcr_investigator + +from tasks.common.dtypes import DAVIS_WIND_SENSOR, WIND_FIELD +from tasks.wind_agent import WindFieldAgent +from tasks.wind_learner import WindTrendLearner, WIND_TREND +from tasks.sink import CUPS_Sink +from tasks.profiler.components import ( + ProfilerInvestigator, + EndpointInvestigator, + TASK_DESCRIPTION_DTYPE, + PROFILE_RESULTS, +) +from service.sensor_publisher import DAVIS_CHANNEL + +PI_PREDICT_RUNTIME = DataType("pi_PREDICT_RUNTIME") + +dotenv.load_dotenv("tasks/common/config.sh") + +# Ship every user module in the agent's reference graph by value: the +# broker has no xGFabric checkout, so anything pickled by reference +# (tasks.davis -> utils_architecture, tasks.do_simulation, the common +# helpers, ...) fails to import there. Sweep the whole tasks.* subtree +# plus the repo-root helpers rather than curate a list that drifts. +register_user_modules([ + m for name, m in list(sys.modules.items()) + if m is not None and (name == "tasks" or name.startswith("tasks.") + or name == "utils_architecture") +]) + +DT_HOST = os.environ.get("DT_SERVICE_HOST", "broker") + +INFERENCE_EP = os.environ.get("DT_INFERENCE_ENDPOINT", "hpc") +INFERENCE_BE = os.environ.get("DT_INFERENCE_BACKEND", "concurrent") +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: + config = dict(os.environ) + playground = config.get("PLAYGROUND_DIR", os.path.expanduser("~/xgf_twin")) + + 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}") + + field = dt.package(WindFieldAgent, config) + sink = dt.package(CUPS_Sink, config) + # a visible learning lane on the wind stream (see wind_learner.py); + # its training routes to the 'learning' engine declared in ENGINES + learner = dt.package(WindTrendLearner, learn_backend="learning") + base_profiler = dt.package( + ProfilerInvestigator, playground + "/profiler/nersc_profiler") + pi_profiler = dt.package( + EndpointInvestigator, "pi", playground + "/profiler/pi_profiler") + + # sensor is external (fake Davis, as twin.py fakes it); bind its + # channel to the input dtype + dt.add_input(twin, DAVIS_WIND_SENSOR, DAVIS_CHANNEL) + + dt.add_agent(twin, field, DAVIS_WIND_SENSOR, WIND_FIELD) + dt.add_task(twin, sink, WIND_FIELD, NULL_DTYPE) + + dt.add_investigator(twin, base_profiler, + TASK_DESCRIPTION_DTYPE, PROFILE_RESULTS) + dt.add_investigator(twin, pi_profiler, + PROFILE_RESULTS, PI_PREDICT_RUNTIME) + + # visible learning lane, fed by the same sensor stream + dt.add_investigator(twin, learner, DAVIS_WIND_SENSOR, WIND_TREND) + + dt.start(twin) + + 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 + + print("SHUTDOWN") + dt.twin_close(twin) + return 0 + finally: + runtime.stop() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="xGFabric twin, service mode, REAL workload") + parser.add_argument("--runtime", type=int, default=600, + help="seconds to keep the twin running") + sys.exit(main(parser.parse_args()))