FE-1162: Python optimisation script (Optuna) driving the Petrinaut CLI

libs/@hashintel/petrinaut-opt/.gitignore

@@ -0,0 +1,9 @@
+# Python
+__pycache__/
+*.pyc
+
+# Desktop store
+*.DS_Store
+
+# Virtual environments
+uv.lock

libs/@hashintel/petrinaut-opt/README.md

@@ -0,0 +1,170 @@
+# Petrinaut optimization
+
+Black-box optimization of a **Petri-net execution**. A request is submitted to a
+Petrinaut UNIX socket server that executes the Petri net for a set of parameters
+and initial states and returns a single metric value per run.
+[Optuna](https://optuna.org/) searches the input space to maximise (or minimise)
+that metric. Results stream out one evaluation at a time over Server-Sent Events,
+so a UI can watch the optimization live.
+
+The search space is **continuous and discrete**, and evaluations are treated as
+expensive, so a sample-efficient sampler (TPE by default) is used.
+
+## How it connects to Petrinaut
+
+This package does **not** run the Petri net itself — it drives the
+[`petrinaut-cli`](../petrinaut-cli) socket server. Follow the instructions
+on [there](../petrinaut-cli/README.md) first.
+
+## Components
+
+| File                                                 | Role                                                                                                                                                                                                                                   |
+| ---------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| [petrinaut_client.py](src/petrinaut_client.py)       | `PetrinautModel` — connects to the Petrinaut CLI socket, builds each `run` request, and returns the metric. `PetrinautModelSpec` configures the execution.                                                                             |
+| [petrinaut_optimizer.py](src/petrinaut_optimizer.py) | `PetrinautOptimizer` — drives the Optuna study: proposes inputs, runs the model, and streams evaluations (`stream_all` / `stream_best`) as Server-Sent Events. `OptimizationSpec` configures a run; `BOUNDS` defines the search space. |
+| [optimization_api.py](src/optimization_api.py)       | FastAPI service exposing `/init`, the two streaming endpoints, `/status`, and `/`.                                                                                                                                                     |
+
+## Setup
+
+This is a [uv](https://docs.astral.sh/uv/) project (Python ≥ 3.10.20):
+
+```bash
+uv sync
+```
+
+Imports are package-qualified (`from src...`), so run everything from the package
+root using module syntax.
+
+## Run the optimizer directly
+
+`main()` runs a short study against the socket and logs each evaluation:
+
+```bash
+# 1. petrinaut serve --model <...> --socket /tmp/petrinaut.sock   (in another shell)
+# 2. from libs/@hashintel/petrinaut-opt:
+uv run python -m src.petrinaut_optimizer
+```
+
+## Run the API
+
+```bash
+uv run uvicorn src.optimization_api:app --reload
+```
+
+### 1. Initialise a run — `POST /init`
+
+The body carries **two** objects: `opt_spec` (what to optimize) and `pn_spec`
+(the Petri-net execution model). It returns a `session_id`.
+
+```bash
+curl -s -X POST localhost:8000/init -H 'content-type: application/json' -d '{
+  "opt_spec": {
+    "parameters": {"infection_rate": 0.5},
+    "initial_states": {"Susceptible": 990},
+    "n_trials": 30,
+    "sampler": "tpe",
+    "direction": "minimize"
+  },
+  "pn_spec": {
+    "name": "SIR",
+    "metric": "Infected Fraction",
+    "steps": 100,
+    "dt": 0.1,
+    "seed": 1234
+  }
+}'
+# -> {"session_id": "...", "status": "initialised", "pn_model": "SIR", "opt_study": "input_opt_..."}
+```
+
+Here `infection_rate`, and `Susceptible` are **fixed** at the given
+values, and every other input in the search space (`recovery_rate`, `Recovered`, `Infected`)
+is **optimized** — see [Configuring a run](#configuring-a-run).
+
+### 2. Stream evaluations — `GET /optimize/{session_id}/stream/all`
+
+Opens a Server-Sent Events stream: one frame per finished trial, then a final
+`event: done`. Disconnecting the client stops the underlying study. Each frame
+reports the inputs that were **searched** this trial (fixed inputs are constant
+and omitted) plus the resulting metric.
+
+```bash
+curl -sN localhost:8000/optimize/<session_id>/stream/all
+# data: {"step": 0, "params": {"recovery_rate": 3.21}, "init_states": {"Recovered": 412.7}, "metric": 0.87, "state": "COMPLETE"}
+# data: {"step": 1, ...}
+# event: done
+# data: {}
+```
+
+`state` is the Optuna trial state (`COMPLETE`, `PRUNED`, `FAIL`); `metric` is
+`null` for a pruned trial. Pass `?n_trials=<n>` to override the run's `n_trials`.
+
+### 3. Stream the running best — `GET /optimize/{session_id}/stream/best`
+
+Same shape, but each frame reports the **best-so-far** inputs and metric rather
+than the latest trial. Frames are suppressed until at least one trial has
+completed.
+
+### Other endpoints
+
+- `GET /status` — lists the currently active sessions (`session_id`, `pn_model`,
+  `opt_study`).
+- `DELETE /optimize/{session_id}` — drops a session.
+- `GET /` — welcome message.
+
+## Configuring a run
+
+**Search space** — the universe of optimizable inputs is defined once in `BOUNDS`
+at the top of [petrinaut_optimizer.py](src/petrinaut_optimizer.py):
+
+```python
+BOUNDS = {
+    "parameters": {
+        "infection_rate": Bounds(0.01, 10.0),
+        "recovery_rate":  Bounds(0.01, 10.0),
+    },
+    "initial_states": {
+        "Susceptible": Bounds(0.0, 1000),
+        "Infected":    Bounds(0.0, 1000),
+        "Recovered":   Bounds(0.0, 1000),
+    },
+}
+```
+
+**`OptimizationSpec`** ([petrinaut_optimizer.py](src/petrinaut_optimizer.py))
+partitions those inputs per run:
+
+- `parameters` / `initial_states` — any input you give a value here is **held
+  fixed** at that value; any input you omit (leave `null`) is **optimized** over
+  its `BOUNDS` range. Provided values must fall within `BOUNDS` (validated on
+  `/init`).
+- `sampler` — `tpe` or `random`.
+- `direction` — `maximize` or `minimize`.
+- `n_trials` — number of evaluations (overridable per stream via the `n_trials`
+  query parameter).
+- `study_name` — optional label for the Optuna study.
+
+**`PetrinautModelSpec`** ([petrinaut_client.py](src/petrinaut_client.py))
+configures the execution sent to the CLI:
+
+- `name` — Petri-net class name (default `SIR`).
+- `metric` — metric name computed at the end of a run and used as the objective
+  (default `Infected Fraction`); must match a metric defined in the loaded model.
+- `steps` — number of steps per run (sent as `maxSteps`).
+- `dt` — timestep for the dynamics.
+- `seed` — RNG seed (fixed → deterministic runs).
+- `structure`, `outpath`, `store`, `eval_timeout` — accepted but currently unused
+  (the model structure is loaded server-side via the CLI's `--model`).
+
+## Notes
+
+- **Failures/timeouts**: any error returned by the CLI (or other exception during
+  a trial) marks that trial pruned; the study continues.
+- **Streaming model**: evaluations run in a background thread and are pushed to
+  the SSE client through an `asyncio.Queue`. Each optimizer instance holds a lock,
+  so one session can't be driven by two concurrent streams — a second concurrent
+  stream on the same session receives `event: error` (`already running`).
+- **Sessions are independent**: there is no global single-run guard; multiple
+  `/init` sessions can exist and run at once.
+- **Leave one input free per group**: each group (`parameters`, `initial_states`)
+  should leave at least one input unfixed so there is something to optimize.
+  </content>

libs/@hashintel/petrinaut-opt/pyproject.toml

@@ -0,0 +1,11 @@
+[project]
+name = "petrinaut-optimization"
+version = "0.1.0"
+description = "Petrinaut parameter and initial state optimisation using optuna"
+readme = "README.md"
+requires-python = ">=3.10.20"
+dependencies = [
+    "fastapi>=0.128.8",
+    "optuna>=3.6",
+    "uvicorn[standard]>=0.39.0",
+]

libs/@hashintel/petrinaut-opt/src/optimization_api.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""HTTP API for Petrinaut optimization.
+
+Endpoints
+---------
+POST /init             Initialise the Petri-net execution model (petrinet_model.initialize)
+                       and start an optimization run in the background.
+GET  /optimize         Server-Sent Events stream of objective evaluations,
+                       emitted as each trial completes, ending with a summary.       
+GET  /status           Current run model and state.
+
+Run with:  uv run uvicorn optimization_api:app --reload
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+import json
+
+from contextlib import asynccontextmanager
+from typing import Union, Generator
+
+import optuna
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+
+from src.petrinaut_client import PetrinautModelSpec, PetrinautModel
+from src.petrinaut_optimizer import OptimizationSpec, PetrinautOptimizer
+
+optuna.logging.set_verbosity(optuna.logging.WARNING)
+
+# Concurrency is scoped per session: each `PetrinautOptimizer` holds its own lock
+# so a single session can't be driven by two concurrent streams (see
+# `stream_all`/`stream_best`). Independent sessions run independently — there is
+# deliberately no global single-run guard.
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.sessions: dict[str, PetrinautOptimizer] = {}
+    yield
+    app.state.sessions.clear()
+
+app = FastAPI(title="Petrinaut optimization API", lifespan=lifespan)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Dummy functions
+# ─────────────────────────────────────────────────────────────────────────────
+
+def dummy_stream() -> Generator[dict[str,Union[float,int]]]:
+    """Dummy data stream to check that API endpoint works
+    """
+    from datetime import datetime
+    n = 0   
+    while n < 10:
+        time.sleep(2)
+        event = {"inputs":[1.2,3.4],"output":datetime.now().strftime('%H:%M:%S'),"step":n}
+        yield f"{json.dumps(event)}\n\n"
+        n +=1 
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Endpoints
+# ─────────────────────────────────────────────────────────────────────────────
+@app.post("/init")
+async def init(opt_spec: OptimizationSpec, pn_spec: PetrinautModelSpec) -> dict:
+    """Initialises the Petrinaut model and its optimization model from the Petrinaut UI
+
+    Args:
+        opt_spec (OptimizationSpec): Specification for Petri Net input optimization with respect to output
+        pn_spec (PetrinautModelSpec): Specification for Petri Net execution
+
+    Raises:
+        HTTPException: _description_
+
+    Returns:
+        dict: API status and Petrinaut model name
+    """
+    # Build the model + optimizer. Each session is independent, so a failure here
+    # only fails this request — it never wedges the service for future /init calls.
+    try:
+        # Build the Petri net from the client spec.
+        petrinet_model = PetrinautModel(pn_spec)
+        # Instantiate Petrinaut optimization class
+        optimizer = PetrinautOptimizer(
+            opt_spec = opt_spec,
+            pn_model = petrinet_model,
+        )
+    except Exception as exc:
+        raise HTTPException(500, f"failed to initialise optimization: {exc}")
+
+    session_id = str(uuid.uuid4())
+    app.state.sessions[session_id] = optimizer
+    return {"session_id": session_id, "status": "initialised", "pn_model": petrinet_model.name, "opt_study": optimizer.study_name}
+
+
+@app.get("/optimize/{session_id}/stream/all",response_class=StreamingResponse)
+async def optimize_stream_all(session_id: str, request: Request, n_trials:Union[int,None]=None) -> StreamingResponse:
+    """Streams optimization results per optimization step (trial) to json line
+
+    Args:
+        session_id (str): Optimization API curent session id
+        request (Request): Optimization API generic request
+
+    Raises:
+        HTTPException: Unknown session id
+
+    Returns:
+        StreamingResponse: 
+    """
+    optimizer = app.state.sessions.get(session_id)
+    if optimizer is None:
+        raise HTTPException(404, "unknown session_id — call /init first")
+
+    # Set default trials if no argument passed
+    n_trials = n_trials if (n_trials is not None) else optimizer.n_trials
+
+    # The optimiser's SSE generator acquires/releases the session lock itself, so
+    # ending the stream (completion, error, or client disconnect) never leaves the
+    # session wedged.
+    return StreamingResponse(
+        optimizer.stream_all(request, n_trials=n_trials),
+        media_type="text/event-stream",
+        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+    )
+
+@app.get("/optimize/{session_id}/stream/best",response_class=StreamingResponse)
+async def optimize_stream_best(session_id: str, request: Request, n_trials:Union[int,None]=None) -> StreamingResponse:
+    """Streams current best optimization results per optimization step (trial) to json line
+
+    Args:
+        session_id (str): Optimization API curent session id
+        request (Request): Optimization API generic request
+
+    Raises:
+        HTTPException: Unknown session id
+
+    Returns:
+        StreamingResponse: 
+    """
+    optimizer = app.state.sessions.get(session_id)
+    if optimizer is None:
+        raise HTTPException(404, "unknown session_id — call /init first")
+
+    # Set default trials if no argument passed
+    n_trials = n_trials if (n_trials is not None) else optimizer.n_trials
+
+    # The optimiser's SSE generator acquires/releases the session lock itself, so
+    # ending the stream (completion, error, or client disconnect) never leaves the
+    # session wedged.
+    return StreamingResponse(
+        optimizer.stream_best(request, n_trials=n_trials),
+        media_type="text/event-stream",
+        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+    )
+
+@app.delete("/optimize/{session_id}")
+async def teardown(session_id: str):
+    if app.state.sessions.pop(session_id, None) is None:
+        raise HTTPException(404, "unknown session_id")
+    return {"status": "deleted"}
+
+
+@app.get("/status")
+async def status() -> dict:
+    """Lists the currently active optimization sessions."""
+    return {
+        "sessions": [
+            {
+                "session_id": session_id,
+                "pn_model": optimizer.pn_model.name,
+                "opt_study": optimizer.study_name,
+            }
+            for session_id, optimizer in app.state.sessions.items()
+        ]
+    }
+
+@app.get("/")
+async def root() -> dict:
+    return {"message": "Welcome to Petrinaut optimization API"}