STDE — Process Scalability and Weibull Fault Injection

This repository contains the code and results for the 2026 SCSS paper:

Scalable Stochastic Differential Operator Estimation for PDE-Constrained Deep Learning under Parallel Execution and Fault Injection
Costin-Alexandru Deonise, Florin Pop
Student Scientific Communications Session (SCSS), 2026

The experiments are split across two hardware environments: a laptop (Apple M3 Max) for CPU process scalability, and a university HPC cluster (Intel Xeon Haswell) for Weibull fault injection.

This repo is a fork of the official STDE implementation by Shi et al. (NeurIPS 2024): github.com/sail-sg/stde. The stde/ package and sine_gordon.py come from that release and are used here as the benchmark workload, with a minor modification (fixed random seed for reproducibility across process counts and failure configurations).

---

Repository structure

.
├── sine_gordon.py              # benchmark workload (from Shi et al. 2024, seed fixed)
├── stde/                       # STDE package (from Shi et al. 2024, unmodified)
├── main.py                     # full STDE training entry point (Shi et al.)
├── requirements.txt
├── setup.py / setup.cfg
├── scripts/                    # HPC job submission scripts (Shi et al.)
│
├── # ── LAPTOP (Apple M3 Max, 14-core, macOS) ───────────────────────────────
├── find_max_procs.py           # scalability scan: sweeps p=1..N, records throughput
├── bench_cores.py              # alternative core-count benchmark
├── process_scaling_throughput.py  # generates throughput vs process-count plot
├── plot_scan.py                # plots max_procs_scan.jsonl
├── plot_cores.py               # plots bench_cores results
├── max_procs_scan.jsonl        # raw scalability results  (p = 1 .. 6)
├── pics/
│   ├── process_scaling_throughput.png
│   └── process_scaling_throughput.pdf
│
├── # ── HPC CLUSTER (Intel Xeon E5-2698 v3 Haswell, SLURM) ──────────────────
├── bench_faults_weibull.py     # live fault injection: spawns p workers,
│                               #   assigns Weibull lifetimes, sends SIGKILL
├── bench_relative_error_weibull.py  # statistical simulation of estimator
│                               #   error under sample loss (no live processes)
├── check_relative_error.py     # sanity checks on relative_error_agg.csv
├── plot_weibull.py             # plots all Weibull experiment results
├── plot_weibull+throughput.py  # combined throughput + Weibull overview plot
├── plot_throuput.py            # throughput-only plot
├── relative_error_runs.jsonl   # raw per-run results (90 runs = 9 configs × 10 reps)
├── relative_error_agg.csv      # aggregated results: mean/std per configuration
└── graphs/                     # all generated figures (PDF + PNG)

---

Laptop experiments — Apple M3 Max

Hardware: Apple MacBook Pro M3 Max, 14-core CPU (10 performance + 4 efficiency), 36 GB unified memory, macOS.

What this measures: How throughput (tasks/second) scales as the number of parallel worker processes increases. Each worker runs one full sine-Gordon training run end to end; the process manager just counts completions per unit time.

Key result: Throughput grows monotonically from 39.6 tasks/s at p=1 to 199.7 tasks/s at p=6. Diminishing-returns onset is at p=5 — the marginal gain at p=5 drops below 50% of the initial marginal gain. Parallel efficiency η ≥ 0.90 up to p=5; it falls to 0.84 at p=6.

Thread isolation

This is required for the results to be meaningful. By default, JAX's XLA backend and NumPy's BLAS each spawn their own thread pool, so a nominally single-process run can quietly occupy most of the machine. Set these before running anything:

export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export VECLIB_MAXIMUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export JAX_PLATFORM_NAME=cpu

Without this, the throughput curves measure a tangle of inter-process scaling and intra-process threading together, which is not what the experiment is about.

Running the scalability scan

python find_max_procs.py
# writes max_procs_scan.jsonl

python plot_scan.py
# writes pics/process_scaling_throughput.png

---

HPC experiments — HASWELL cluster

Hardware: Intel Xeon E5-2698 v3 (Haswell), 32 logical cores, 127 GB RAM, 16 GB per-job memory limit, Linux, SLURM.

What this measures: Estimator robustness when workers are killed mid-run. Each worker is assigned a Weibull-distributed lifetime at the start; when its elapsed time exceeds that lifetime, it receives SIGKILL. The experiment records makespan, killed count, and relative estimation error across 9 parameter configurations × 10 repetitions.

Weibull configurations

Shape k	Scale λ (s)	Failure regime
0.7	90, 180, 360	Early-life (infant mortality)
1.0	90, 180, 360	Constant rate (memoryless)
1.3	90, 180, 360	Wear-out failures

Key result

Under the most aggressive configuration (k=0.7, λ=90 s, mean 3.3 out of 4 workers killed), mean relative estimation error stays below 0.54%. The estimator degrades gracefully because the independent-sample structure of STDE means that losing a worker reduces the effective sample count but does not invalidate the contributions of surviving workers.

Important: two separate experiments

bench_faults_weibull.py and bench_relative_error_weibull.py are not the same experiment:

• bench_faults_weibull.py — live process injection. Actually spawns subprocesses with subprocess.Popen and sends SIGKILL. Measures makespan and killed count. This is the wall-clock fault injection experiment.

• bench_relative_error_weibull.py — statistical simulation. Simulates what fraction of samples would complete given a Weibull-distributed runtime, without terminating any real processes. Produces relative_error_agg.csv.

Running fault injection

# Live process injection (run on HPC with sufficient cores)
python bench_faults_weibull.py \
    --procs 4 \
    --repeats 10 \
    --shape_k 0.7 \
    --scale_lambda_s 90 \
    --dim 100000 \
    --rand_batch_size 16 \
    --out bench_faults_k07_l90.jsonl

# Relative error simulation
python bench_relative_error_weibull.py
# writes relative_error_runs.jsonl and relative_error_agg.csv

# Generate all figures
python plot_weibull.py
# writes to graphs/

---

Installation

pip install -e .

Or manually:

pip install jax jaxlib dm-haiku optax jaxtyping tqdm numpy pandas matplotlib

---

Results summary

Experiment	Hardware	Configuration	Result
Process scalability	M3 Max (laptop)	p=1..6	Throughput 39.6 → 199.7 tasks/s
Process scalability	M3 Max (laptop)	Diminishing-returns onset	p=5 (η drops below 0.90)
Fault injection	HASWELL (HPC)	k=0.7, λ=90 s	Mean relative error 0.54%
Fault injection	HASWELL (HPC)	k=0.7, λ=90 s	Mean workers killed 3.3 / 4

---

Attribution

The stde/ package, main.py, sine_gordon.py, setup.py, requirements.txt, and scripts/ are from the original STDE repository by Shi et al. (NeurIPS 2024):

@inproceedings{shi2024stochastic,
  author    = {Zekun Shi and Zheyuan Hu and Min Lin and Kenji Kawaguchi},
  title     = {Stochastic {Taylor} Derivative Estimator: Efficient Amortization
               for Arbitrary Differential Operators},
  booktitle = {Advances in Neural Information Processing Systems},
  volume    = {37},
  year      = {2024},
  note      = {Best Paper Award at NeurIPS 2024}
}

The scalability and fault-injection experiments (find_max_procs.py, bench_cores.py, bench_faults_weibull.py, bench_relative_error_weibull.py, and all plot_* and process_scaling_* files) are original contributions from this work.

---

Citation

@misc{deonise2026stde,
  author = {Costin-Alexandru Deonise and Florin Pop},
  title  = {Scalable Stochastic Differential Operator Estimation for
            {PDE}-Constrained Deep Learning under Parallel Execution
            and Fault Injection},
  year   = {2026},
  note   = {Student Scientific Communications Session (SCSS),
            National University of Science and Technology Politehnica Bucharest}
}