Braket hybrid

.github/workflows/test.yml

@@ -85,11 +85,11 @@ jobs:
           flux start flux submit --help 2>&1 | grep -E -- '--quantum-vendor' \
             || { echo "CLI plugin not discovered"; exit 1; }
 
-      - name: vendor backends contribute their own options
+      - name: the submit options are the same for every vendor
         run: |
           export FLUX_CLI_PLUGINPATH=$PWD/cli-plugins
           out=$(flux start flux submit --help 2>&1)
-          for opt in --quantum-ibm-backend --quantum-braket-device --quantum-mock-session; do
+          for opt in --quantum-device --quantum-hold --quantum-hold-max --quantum-wait --quantum-dry-run; do
             echo "$out" | grep -q -- "$opt" || { echo "missing $opt"; exit 1; }
             echo "  $opt OK"
           done
@@ -111,6 +111,20 @@ jobs:
           SH
           flux start bash /tmp/jt.sh
 
+      # the core budget over unfinished pairs, rejecting a pair whose classical
+      # half could not be placed rather than admitting it
+      - name: jobtap admission control
+        run: FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-admission.sh
+
+      # preemption is destructive, so it runs on its own. Skips itself when the
+      # runner has fewer than three cores.
+      - name: jobtap preemption
+        run: FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-preemption.sh
+
+      # unload with a grace timer pending, and reload with jobs running
+      - name: jobtap unload and reload
+        run: FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-reload.sh
+
       # one submit gives a held classical and a scout, and the qpu stays
       # allocated until the classical is done
       - name: end-to-end, single broker
@@ -124,14 +138,21 @@ jobs:
       - name: vendor discovery + --quantum-select
         run: flux start bash tests/integration/test-select-discovery.sh
 
+      - name: IonQ backend end-to-end, against the fake service
+        run: flux start bash tests/integration/test-ionq-e2e.sh
+
+      - name: IonQ backend end-to-end, account without sessions
+        run: FAKE_IONQ_NO_SESSIONS=1 flux start bash tests/integration/test-ionq-e2e.sh
+
       # this test expects fluxion already loaded, as rc1 does on a real cluster
       - name: system-instance test (fluxion preloaded, 2 ranks)
         run: |
           cat > /tmp/sys.sh <<'SH'
           set -e
           flux module remove -f sched-simple 2>/dev/null || true
           flux module load sched-fluxion-resource match-format=rv1
-          flux module load sched-fluxion-qmanager
+          flux module load sched-fluxion-qmanager queue-policy=coschedule \
+            || flux module load sched-fluxion-qmanager
           export FLUX_CLI_PLUGINPATH=$PWD/cli-plugins
           bash tests/integration/test-system-instance.sh
           SH

.gitignore

@@ -1,2 +1,15 @@
 dist
 *.egg-info
+
+# never commit filled in credentials
+.env
+*.env
+!env.example
+
+# build and python artifacts
+__pycache__/
+*.pyc
+*.so
+
+# a sibling checkout the experiments live in
+flux-quantum-experiments/

README.md

@@ -1,46 +1,139 @@
 # flux-quantum
-Quantum + classical coscheduling for Flux, including jobtap, validation, and client plugins.
-I am developing this against flux-core **flux-core 0.87.0** and my custom branch of flux
-sched with support for hold. The quantum backend is external (e,g., IBM/QRMI, AWS Braket)
-and informs the work for the flux-sched reservation being held for classical resources.
+
+Quantum and classical coscheduling for Flux. One `flux submit` produces a pair.
+The user's work is submitted held, and a small scout job that requests the vendor
+device opens the vendor session, hands the session id to the held job over its
+eventlog, and releases it. Developed against flux-core 0.87.0 and the `add-hold`
+branch of https://github.com/vsoch/flux-sched, and also loads against flux-core 0.89.0.
+
 ## Components
-- **flux_quantum/cli.py** the CLI plugin that runs in user space and (if needed) can use credentials to look at queue depth or cost.
-- **flux_quantum/backends/** the vendor libraries. E.g., handling braket will be different than IBM or qrmi, and how we validate will vary. These also run only in user-space.
-- **flux_quantum/selector.py** is generic discover and rank over backends.
-- **flux_quantum/validator.py** is going to check that we have the right envars, etc. per a backend choice.
-- **flux_quantum/jobtap/quantum.c** the owner-side policy gate (runs in the job manager, holds no credentials). At `job.validate` it enforces vendor policy, rejecting a vendor that isn't permitted/available. The scout and the vendor *selection* live in user space (the CLI plugin), since those need the user's credentials -- jobtap only mutates/enforces. (Selecting a vendor *type* in the graph is a future owned-hardware step; the external-vendor model has no qpu resource to select.)
-## Install / discovery (flux-core 0.87.0)
-The CLI plugin is found by flux via `{confdir}/cli/plugins`, `{libexecdir}/cli/plugins`, or any directory on `FLUX_CLI_PLUGINPATH`. To install:
+
+- **flux_quantum/cli.py** submit plugin. Runs as the user, picks a vendor, submits the held classical job and turns the submit itself into the scout.
+- **flux_quantum/backends/** one module or package per vendor, IBM through QRMI, AWS Braket, IonQ over its REST API, and a mock. The common submit options are mapped to each vendor's terms there, so a user changes nothing but the vendor and the device. They probe queue depth and cost and open and close the session. They only ever run as the user.
+- **flux_quantum/selector.py** discovers vendors from the graph and ranks them by policy.
+- **flux_quantum/scout.py** and **wrap.py** the scout, and the wrapper that exports `QUANTUM_SESSION_ID` before the user's command runs.
+- **flux_quantum/graph.py**, **qresource.py**, **populate.py** add `qdevice_<vendor> -> qpu` to the fluxion graph.
+- **flux_quantum/jobtap/quantum.c** owner side policy in the job manager. Allowed vendors, admission against reachable cores, protection of both halves, and optional preemption. Holds no credentials.
+
+## Install
+
+The CLI plugin is found through `{confdir}/cli/plugins`, `{libexecdir}/cli/plugins`, or a directory on `FLUX_CLI_PLUGINPATH`.
+
 ```bash
-flux python -m pip install -e .            # install into the python flux uses
-export FLUX_CLI_PLUGINPATH=$PWD/cli-plugins # dir holds ONLY the discovery shim
-flux submit --help | grep quantum          # options should appear
-```
-The plugin search dir (`cli-plugins/`) contains only `quantum.py`, a shim that
-imports `QuantumCLIPlugin` from the installed package -- do NOT point
-`FLUX_CLI_PLUGINPATH` at `flux_quantum/` itself, or flux will try to import the
-package's internal modules as standalone plugins and fail. For a system install,
-copy `cli-plugins/quantum.py` into `{confdir}/cli/plugins`.
-For the jobtap plugin:
+flux python -m pip install -e .              # into the python flux uses
+export FLUX_CLI_PLUGINPATH=$PWD/cli-plugins  # holds only the discovery shim
+flux submit --help | grep quantum
+```
+
+`cli-plugins/` contains only `quantum.py`, a shim that imports the installed package. Do not point `FLUX_CLI_PLUGINPATH` at `flux_quantum/` itself, or flux tries to import every module as a plugin. For a system install copy the shim into `{confdir}/cli/plugins`.
+
+Populate the graph once, before any job runs. Growing it while jobs hold resources corrupts fluxion.
+
 ```bash
-make -C flux_quantum/jobtap  &&  flux jobtap load .../quantum.so
-```
-Still a WIP I need to test all this with flux-sched! And we need a cute logo. Very important.
-## Testing (token-free)
-No real vendor token needed. `FLUX_QUANTUM_MOCK=1` registers `mock`/`mock_busy`
-backends (credentials always "present", no API calls) and the scout session is
-mocked, so the whole pipeline runs without credentials.
-- `testing/setup-mock-registry.sh` derives the real graph and injects `qvendor_*`
-  markers so the selector has vendors to discover (markers are plain graph vertices
-  added with `flux inject` -- no token).
-- `testing/test-mock-e2e.sh` runs it end to end: plugin discovery -> registry ->
-  `flux submit --quantum-vendor mock` (held) -> scout (mock session) -> handoff -> run.
+flux python -m flux_quantum.populate ibm braket ionq
+```
+
+Build and load the jobtap plugin. `flux jobtap query quantum.so` shows the settings in force and the capacity numbers.
+
+```bash
+make -C flux_quantum/jobtap
+flux jobtap load $PWD/flux_quantum/jobtap/quantum.so vendors=ibm,braket,ionq \
+    total_cores=128 reserve_cores=32 protect_types=qpu preempt=1
+```
+
+The plugin admits a pair only when its cores plus one for the scout are reachable, and only when a core is free right now for the scout, since nothing can be preempted for a pair until its scout runs and releases the classical half. Both halves run at `boost_priority` (default 1000, the scout one higher so it is placed before its classical reserves). A site with flux-accounting sets its own priorities and should load with `boost_priority=0`. Nothing then orders the scout ahead of its classical, which is an open gap for accounting sites and wants a factor inside flux-accounting rather than a second plugin.
+
+Preemption is off unless `preempt=1`. With it on, the plugin watches each classical half and cancels other work only when three things hold: the scout has released it, so the session is open and costing money; `settle` seconds later (default 2) it is still not running; and fluxion's annotation says it is reserved for later rather than placed. Then it cancels unprotected running jobs for the shortfall, the classical's cores less what is free, lowest priority first and youngest among equals, once per job. A cancelled job loses its work, which is why the switch is off by default.
+
+## Submitting
+
+The options are the same whichever vendor it is.
+
+```
+--quantum-vendor VENDOR   ibm, braket, ionq, or mock under FLUX_QUANTUM_MOCK
+--quantum-select POLICY   pick the vendor instead: any, queue or cost
+--quantum-device NAME     the vendor's name for the device: a Braket ARN, an IonQ
+                          backend, a QRMI resource id. Each vendor has a default
+--quantum-hold MODE       session takes the vendor's real hold, a hybrid job on
+                          Braket, a session on IonQ and IBM. probe submits a front
+                          of queue job and holds nothing. Default session
+--quantum-hold-max SECS   give the hold up after this long. Default the job's
+                          duration plus 300, or 900 without one
+--quantum-wait SECS       give up and cancel the held job if the hold is not ours
+                          after this long. Default 0, as long as the scout may run
+--quantum-dry-run         run on the vendor's simulator. FLUX_QUANTUM_MOCK implies it
+```
+
+```bash
+flux submit --quantum-vendor ionq --quantum-dry-run -n1 -- flux python examples/ionq/workload.py
+flux submit --quantum-vendor braket --quantum-device arn:aws:braket:us-east-1::device/qpu/ionq/Forte-Enterprise-1 -t 30m -n4 -- ./work
+flux submit --quantum-vendor ibm --quantum-device ibm_kingston --quantum-wait 600 -n4 -- flux python examples/qrmi/workload.py
+```
+
+Settings that tune an installation rather than describe a job are environment variables with the `FLUX_QUANTUM_` prefix, read at submit time. Each backend's docstring lists its own. The ones that matter most:
+
+```
+FLUX_QUANTUM_BRAKET_HOLD_INSTANCE   instance for the hybrid job, ml.m5.large
+FLUX_QUANTUM_BRAKET_SHOTS           shots for the probe task, 1
+FLUX_QUANTUM_IONQ_SHOTS             shots for the warm-up job, 100
+FLUX_QUANTUM_IBM_TYPE               QRMI resource type, qiskit-runtime-service
+FLUX_QUANTUM_IBM_SKIP_WARMUP        release as soon as the session opens
+FLUX_QUANTUM_IBM_SESSION_MODE       batch, billed for hardware time only, or
+                                    dedicated, billed for wall clock until closed
+FLUX_QUANTUM_MOCK_SESSION           force the mock's session id
+FLUX_QUANTUM_MOCK_QUEUE             the mock's simulated queue depth
+```
+
+## Testing
+
+`FLUX_QUANTUM_MOCK=1` registers the `mock` and `mock_busy` backends, so the whole pipeline runs without a vendor token. Export it before `flux start`, since the ingest validator is a separate process. The unit tests need no broker.
+
+```bash
+pytest tests/unit
+flux start bash tests/integration/test-mock-e2e.sh              # one submit, held classical and scout, handoff
+flux start -s2 bash tests/integration/test-mock-e2e.sh          # the two halves on different ranks
+FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-admission.sh    # jobtap admission
+FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-preemption.sh   # jobtap preemption
+FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-reload.sh       # jobtap unload and reload
+flux start bash tests/integration/test-select-discovery.sh      # vendors discovered from the graph
+flux start bash tests/integration/test-ionq-e2e.sh              # the IonQ backend against a fake api.ionq.co
+bash tests/integration/test-system-instance.sh                  # against a running instance
+```
+
+The e2e and discovery scripts load fluxion themselves. The jobtap and system instance scripts expect fluxion to be the scheduler already, as it is in the CI image. `tests/integration/setup-mock-registry.sh` adds `qdevice_*` vertices to a live graph so the selector has something to discover.
+
+## IonQ
+
+`flux_quantum/backends/ionq/` talks to the IonQ v0.4 REST API with the standard library and reads the key from `IONQ_API_KEY`. The session hold opens a session on the backend, warms it up with a one qubit job, and releases the classical job once IonQ reports the session active. The session id is what the classical job is handed, and jobs it submits with that id run inside the session. Sessions are in beta and not every account has them, and an account without them is told to use `--quantum-hold probe`, which submits the warm-up job alone and releases when it starts.
+
+A dry run sends every job to the simulator, which IonQ does not charge for, with the noise model of the backend asked for, so `--quantum-device qpu.forte-1 --quantum-dry-run` runs on the simulator with Forte's noise. The simulator has no queue, so a dry run exercises the plumbing and not the hold.
+
+```bash
+export IONQ_API_KEY=...
+flux submit --quantum-vendor ionq --quantum-dry-run -n1 -- flux python examples/ionq/workload.py
+flux submit --quantum-vendor ionq --quantum-device qpu.forte-1 -n1 -- flux python examples/ionq/workload.py
+```
+
+`python -m flux_quantum.backends.ionq.fake` runs a stand-in for the service, so the backend and the whole pipeline run with no key at all. `IONQ_API_URL` points the backend at it. It listens on the loopback, so on an instance of more than one node start it with `--bind 0.0.0.0` and use the URL it prints, which the other nodes can reach. `FAKE_IONQ_QUEUE=30-90` gives it a device queue, seconds a job waits before it is served, which a started session's jobs skip, and `POST /fake/config` changes it while running.
+
+## Braket priority probe
+
+`tests/probe_hold.py` measures the thing this branch depends on: whether a task submitted from outside the hybrid job, carrying its token, is served ahead of the device queue. It opens the hold, submits a control task without the token, then the token task after it, and watches both until the token task finishes. Braket only reports a task's queue, Normal or Priority, and its position when `GetQuantumTask` is asked for `QueueInfo`, so the probe asks, and it reads the finishing order from the service timestamps. Both come from the tasks themselves. The device-wide counters lag a submit by minutes and are not used for the verdict.
+
+```bash
+python3 tests/probe_hold.py --survey                    # which QPUs are open and have a queue
+python3 tests/probe_hold.py                             # SV1, shows only that the token associates
+AWS_DEFAULT_REGION=us-east-1 python3 tests/probe_hold.py \
+    --device arn:aws:braket:us-east-1::device/qpu/ionq/Forte-Enterprise-1 \
+    --shots 100 --max-seconds 1800
+```
+
+`--max-seconds` is the hold, and the token task has to finish inside it, so give a device where tasks take minutes half an hour. The hold is released as soon as the token task finishes, and the control is then followed to its end unless `--cancel-control` is passed. Each run appends a JSON line to `probe-runs.jsonl` with both tasks' queue, position, timestamps and a timeline, and everything printed, the backend's progress included, is also appended to `probe-hold.log` (`--log PATH` to change it, `--log ''` to disable). It costs two tasks at the device rate plus the hold instance, and prints the estimate before creating anything. QuEra Aquila takes an analog program rather than a circuit, and the probe sends one when the ARN is QuEra's. `--check-program` runs the program on the local simulator first, free, and the analog simulator applies Aquila's own limits.
 
 ## License
-DevTools is distributed under the terms of the MIT license.
+
+flux-quantum is distributed under the terms of the MIT license.
 All new contributions must be made under this license.
-See [LICENSE](https://github.com/converged-computing/cloud-select/blob/main/LICENSE),
-[COPYRIGHT](https://github.com/converged-computing/cloud-select/blob/main/COPYRIGHT), and
-[NOTICE](https://github.com/converged-computing/cloud-select/blob/main/NOTICE) for details.
+See [LICENSE](LICENSE), [COPYRIGHT](COPYRIGHT), and [NOTICE](NOTICE) for details.
 SPDX-License-Identifier: (MIT)
 LLNL-CODE- 842614

examples/ionq/workload.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+"""Example classical workload for IonQ.
+
+Reads the session the scout opened and the target it was given, submits a
+Bell circuit into the session, waits for it and prints the histogram. This
+is what the classical half of a pair does. The key comes from the user's
+environment, the rest from the scout.
+
+    flux submit --quantum-vendor ionq --quantum-dry-run -n1 \\
+        -- flux python examples/ionq/workload.py
+"""
+
+import json
+import os
+import sys
+import time
+import urllib.request
+
+BELL = {
+    "qubits": 2,
+    "gateset": "qis",
+    "circuit": [
+        {"gate": "h", "target": 0},
+        {"gate": "cnot", "control": 0, "target": 1},
+    ],
+}
+
+
+def call(method, path, body=None):
+    url = (os.environ.get("IONQ_API_URL") or "https://api.ionq.co/v0.4").rstrip("/")
+    req = urllib.request.Request(
+        url + path,
+        data=None if body is None else json.dumps(body).encode(),
+        method=method,
+    )
+    req.add_header("Authorization", "apiKey " + os.environ["IONQ_API_KEY"])
+    req.add_header("Content-Type", "application/json")
+    with urllib.request.urlopen(req, timeout=30) as r:
+        return json.loads(r.read() or b"{}")
+
+
+def main():
+    session = os.environ.get("QUANTUM_SESSION_ID")
+    if not session:
+        sys.exit("no QUANTUM_SESSION_ID, this job was not started by the scout")
+    if not os.environ.get("IONQ_API_KEY"):
+        sys.exit("no IONQ_API_KEY in the environment")
+    target = os.environ.get("IONQ_BACKEND", "simulator")
+    print("session:", session)
+    print("target: ", target)
+
+    body = {
+        "type": "ionq.circuit.v1",
+        "name": "flux-quantum-bell",
+        "shots": 100,
+        "backend": target,
+        "input": BELL,
+    }
+    # a probe hold hands over a job id, not a session, and a job id is not
+    # something another job can be submitted into
+    if not session.startswith("job:"):
+        body["session_id"] = session
+    noise = os.environ.get("IONQ_NOISE_MODEL")
+    if noise:
+        body["noise"] = {"model": noise}
+
+    job = call("POST", "/jobs", body)
+    print("job:", job["id"], job.get("status"))
+    while True:
+        state = call("GET", "/jobs/{}".format(job["id"]))
+        if state.get("status") in ("completed", "failed", "canceled"):
+            break
+        time.sleep(2)
+    print("status:", state.get("status"))
+    if state.get("status") == "completed":
+        print("results:", call("GET", "/jobs/{}/results".format(job["id"])))
+
+
+if __name__ == "__main__":
+    main()