From 201493a9d9b73e40af5f0bdf5cd1021a1f029f07 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 12 Aug 2026 09:09:48 -0700 Subject: [PATCH 01/26] wip: qrmi backend Signed-off-by: vsoch --- .github/workflows/test.yml | 2 +- .gitignore | 5 + examples/qrmi/README.md | 135 +++++++++++ examples/qrmi/env.example | 21 ++ examples/qrmi/workload.py | 35 +++ flux_quantum/backends/base.py | 5 + flux_quantum/backends/ibm.py | 64 +----- flux_quantum/backends/qrmi.py | 413 ++++++++++++++++++++++++++++++++++ flux_quantum/cli.py | 8 +- flux_quantum/wrap.py | 20 +- pyproject.toml | 4 + tests/unit/test_backends.py | 75 +++++- tests/unit/test_cli_plugin.py | 24 +- tests/unit/test_qrmi.py | 321 ++++++++++++++++++++++++++ tests/unit/test_selector.py | 33 ++- 15 files changed, 1091 insertions(+), 74 deletions(-) create mode 100644 examples/qrmi/README.md create mode 100644 examples/qrmi/env.example create mode 100755 examples/qrmi/workload.py create mode 100644 flux_quantum/backends/qrmi.py create mode 100644 tests/unit/test_qrmi.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c043988..769ab71 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -89,7 +89,7 @@ jobs: 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-ibm-resource --quantum-braket-device --quantum-mock-session; do echo "$out" | grep -q -- "$opt" || { echo "missing $opt"; exit 1; } echo " $opt OK" done diff --git a/.gitignore b/.gitignore index 8a09f29..dbb9afa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,7 @@ dist *.egg-info + +# never commit filled in credentials +.env +*.env +!env.example diff --git a/examples/qrmi/README.md b/examples/qrmi/README.md new file mode 100644 index 0000000..3ee35be --- /dev/null +++ b/examples/qrmi/README.md @@ -0,0 +1,135 @@ +# Running against IBM through QRMI + +This is the demo for using qrmi with flux-quantum. + +## Install + +QRMI needs python 3.11 or newer. It is not a dependency of flux-quantum, so + + pip install "qrmi[ibm]" + +Install it into whatever python `flux python` uses, on every node that could run +the scout. qiskit comes along with it and we need that for the warmup task. + +## Credentials + +QRMI reads its config out of the environment. Every variable is prefixed with the +resource id, so for ibm_kingston the api key lives in +`ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY`. Four of them are required. + + cp env.example .env + chmod 600 .env + $EDITOR .env + set -a && . ./.env && set +a + +Check before you submit anything. + + flux python -c " + from flux_quantum.backends import get_backend + print(get_backend('ibm').credentials_present())" + +That names any variable you are missing. It never prints a value. + +## Submit + + flux submit -t 30m --quantum-vendor ibm --quantum-ibm-resource ibm_kingston \ + -n4 flux python examples/qrmi/workload.py + +stdout is the scout id. The classical job id comes back on stderr. If only one +resource has credentials in the environment you can drop +`--quantum-ibm-resource` and we work it out. + + flux jobs -a + flux job eventlog | grep memo + flux job attach + +## You need an account that can open sessions + +This is the thing that will bite you first. The qiskit-runtime-service type opens +a session, and IBM only allows that on plans that support sessions, like Premium. +On Open or Pay As You Go you can submit tasks but you cannot hold a session, so +there is nothing for the scout to co-allocate. + +You get told that in those words rather than a 403. The held classical job is +cancelled, so nothing is left sitting in the queue. + +If you have direct access instead, use `--quantum-ibm-type ibm-quantum-system`. + +## A session is not the same as having the QPU + +An IBM session goes active when its first task reaches the head of the queue. +After that, tasks in the session keep that priority. So opening a session tells +you very little on its own. + +After acquiring, the scout submits a warmup task, one qubit measured once, and +waits for it to leave the queue. Once that task is running we are at the head, +the session is live, and only then do we release the classical job. + +That is the whole point. The classical job sits held with its nodes reserved +while the scout waits in the QPU queue. It starts the moment the QPU is really +ours, so neither side pays to wait for the other. + +If the warmup never runs we release the session and fail the submit. The +classical job never starts against a QPU we do not have. + + --quantum-ibm-warmup-timeout SECONDS give up after this long. Default 0, + which means wait as long as the + scout job is allowed to live + --quantum-ibm-skip-warmup release as soon as the session opens, + no priority guarantee + --quantum-ibm-ready-timeout SECONDS default 120, how long to wait for the + backend to report itself up + +The warmup costs one shot. + +## The scout holds the session for as long as the classical runs + +That is deliberate. The fluxion allocation and the vendor session start and end +together. It does mean you are billed for the length of the classical job and not +the length of the quantum work, so set `-t`. + +A killed scout still releases, the release is in a finally and SIGTERM is +handled. If the node dies you are down to the vendor timeout. + +## Credentials, and where they end up + +The plugin runs in your process at submit time and reads the variables out of +your environment. The scout uses them, and it runs as you. + +Flux copies the submitting environment into the jobspec, so they are in the job +record and the instance owner can read them. Same as any other job environment. +What does not happen is a credential owned by the system user, or one shared +between users. If that is not good enough for you, do not export the key in the +shell you submit from. + +Only the session id travels to the classical job, and it goes over the job +eventlog. + +## What the classical job sees + + QUANTUM_SESSION_ID the open session + _QRMI_JOB_ACQUISITION_TOKEN same value, where QRMI looks + QRMI_JOB_QPU_RESOURCES the resource id + QRMI_JOB_QPU_TYPES the resource type + +All but the first are what the Slurm and LSF QRMI plugins set, so +`get_job_qpu_resources_and_types()` works and a workload written for Slurm runs +here untouched. See workload.py. + +## Other QRMI settings + +QRMI reads these itself, prefixed with the resource id. Put them in `.env` if you +need them. + + _QRMI_IBM_QRS_SESSION_MODE batch or dedicated + _QRMI_IBM_QRS_SESSION_MAX_TTL session lifetime + _QRMI_IBM_QRS_TIMEOUT_SECONDS request timeout + _QRMI_IBM_QRS_SESSION_ID join a session that already exists + +## When it breaks + + flux job attach + +The scout says why it gave up. Credential problems name the variable. Set +`RUST_LOG=debug` before submitting if you want QRMI's own logs, including the +HTTP calls. diff --git a/examples/qrmi/env.example b/examples/qrmi/env.example new file mode 100644 index 0000000..bb0d669 --- /dev/null +++ b/examples/qrmi/env.example @@ -0,0 +1,21 @@ +# QRMI credentials for one IBM resource. +# Copy to .env, fill in, then chmod 600 .env +# +# These comes from the IBM Quantum Platform console. SERVICE_CRN is the CRN of +# your instance and IAM_APIKEY is an IBM Cloud API key. + +ibm_kingston_QRMI_IBM_QRS_ENDPOINT=https://quantum.cloud.ibm.com/api/v1 +ibm_kingston_QRMI_IBM_QRS_IAM_ENDPOINT=https://iam.cloud.ibm.com +ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY=REPLACE_ME +ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN=REPLACE_ME + +# Optional. batch or dedicated. QRMI defaults this when unset. +# ibm_kingston_QRMI_IBM_QRS_SESSION_MODE=batch + +# For a direct access system use the QS prefix and +# --quantum-ibm-type ibm-quantum-system instead: +# +# ibm_kingston_QRMI_IBM_QS_ENDPOINT= +# ibm_kingston_QRMI_IBM_QS_IAM_ENDPOINT=https://iam.cloud.ibm.com +# ibm_kingston_QRMI_IBM_QS_IAM_APIKEY= +# ibm_kingston_QRMI_IBM_QS_SERVICE_CRN= diff --git a/examples/qrmi/workload.py b/examples/qrmi/workload.py new file mode 100755 index 0000000..6e5569a --- /dev/null +++ b/examples/qrmi/workload.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +"""Example classical workload. + +Reads the session the scout opened and the QPU it was given, then would run the +quantum part. Nothing here is flux specific. The same script works under the +Slurm and LSF QRMI plugins, because they set the same two variables. +""" + +import os +import sys + + +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") + print("session:", session) + + try: + from qrmi import get_job_qpu_resources_and_types + except ImportError: + sys.exit("qrmi is not installed, pip install 'qrmi[ibm]'") + + resources, types = get_job_qpu_resources_and_types() + print("qpu resources:", resources) + print("qpu types:", types) + + # From here the real work would build a payload and call task_start on the + # resource. The session is already open and the classical side already has + # its nodes, so nothing waits on anything. + print("both halves are allocated, ready to submit quantum tasks") + + +if __name__ == "__main__": + main() diff --git a/flux_quantum/backends/base.py b/flux_quantum/backends/base.py index 82220cb..54f701a 100644 --- a/flux_quantum/backends/base.py +++ b/flux_quantum/backends/base.py @@ -57,6 +57,11 @@ def open_session(self, options): "{}: open_session is not implemented for this vendor".format(self.name) ) + def job_environment(self, options): + """Env vars to add to the classical job, for example to tell it which + QPU it has. Defaults to none.""" + return {} + def close_session(self, session=None): """Release the session opened by open_session. diff --git a/flux_quantum/backends/ibm.py b/flux_quantum/backends/ibm.py index 8a18081..4b53a1f 100644 --- a/flux_quantum/backends/ibm.py +++ b/flux_quantum/backends/ibm.py @@ -1,63 +1,15 @@ -"""IBM Quantum backend, via QRMI or the IBM Runtime API. +"""IBM Quantum backend, through QRMI. -The probe is a stub. The real one reads queue depth and cost with the user -token. +Defaults to the qiskit-runtime-service resource type. That type opens a real +session, so the account has to be one that supports sessions. Use +--quantum-ibm-type ibm-quantum-system for a direct access system. """ -import os -from .base import Backend, Signals, register +from .base import register +from .qrmi import QRMIBackend @register -class IBMBackend(Backend): +class IBMBackend(QRMIBackend): name = "ibm" - # env vars the user must have exported for IBM - required_env = ("QISKIT_IBM_TOKEN",) - - @classmethod - def add_options(cls, add_option): - add_option( - "--ibm-backend", - metavar="NAME", - default=None, - help="IBM: target backend/device name (e.g. ibm_brisbane)", - ) - add_option( - "--ibm-instance", - metavar="HUB/GROUP/PROJECT", - default=None, - help="IBM: instance (hub/group/project)", - ) - add_option( - "--ibm-shots", - metavar="N", - default=None, - help="IBM: number of shots for the session workload", - ) - - def scout_options(self, args): - return { - "backend": getattr(args, "ibm_backend", None), - "instance": getattr(args, "ibm_instance", None), - "shots": getattr(args, "ibm_shots", None), - } - - def open_session(self, options): - # TODO open a QRMI or Runtime session on the requested backend with - # the user token and return the session id - raise NotImplementedError( - "IBM open_session not yet implemented; use --quantum-vendor mock " - "with FLUX_QUANTUM_MOCK for now (options captured: {})".format(options) - ) - - def credentials_present(self): - missing = [v for v in self.required_env if not os.environ.get(v)] - if missing: - return False, "IBM: missing env var(s): " + ", ".join(missing) - return True, "IBM: credentials present" - - def probe(self): - # TODO query IBM or QRMI for live queue depth and cost - return Signals( - available=True, queue_depth=None, cost=None, detail={"note": "mock probe"} - ) + default_type = "qiskit-runtime-service" diff --git a/flux_quantum/backends/qrmi.py b/flux_quantum/backends/qrmi.py new file mode 100644 index 0000000..d4f0735 --- /dev/null +++ b/flux_quantum/backends/qrmi.py @@ -0,0 +1,413 @@ +"""QRMI backed vendor sessions. + +QRMI is itself vendor agnostic, so one implementation covers every resource +type it supports. acquire maps onto open_session and release onto +close_session, which is the scout lifecycle exactly. + +QRMI reads its own configuration from the environment, and every variable is +prefixed with the resource id, so the api key for ibm_kingston is + + ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY + +Credentials stay in the user environment. We only ever check that a variable is +set and report its name, never its value. + +qrmi is an optional dependency and needs python 3.11 or newer, so it is +imported inside the methods that use it. +""" + +import json +import os +import time + +from .base import Backend, Signals + +# type strings QRMI itself uses in qrmi_config.json and in QRMI_JOB_QPU_TYPES, +# mapped to the ResourceType attribute +RESOURCE_TYPES = { + "qiskit-runtime-service": "IBMQiskitRuntimeService", + "ibm-quantum-system": "IBMQuantumSystem", + "pasqal-cloud": "PasqalCloud", + "pasqal-local": "PasqalLocal", + "iqm-server": "IQMServer", + "alice-bob-felis": "AliceBobFelis", +} + +# suffixes QRMI requires per type, discovered by constructing a QuantumResource +# with an empty environment and reading back what it asked for. Each one is +# looked up as . +REQUIRED_ENV = { + "qiskit-runtime-service": ( + "_QRMI_IBM_QRS_ENDPOINT", + "_QRMI_IBM_QRS_IAM_ENDPOINT", + "_QRMI_IBM_QRS_IAM_APIKEY", + "_QRMI_IBM_QRS_SERVICE_CRN", + ), + "ibm-quantum-system": ( + "_QRMI_IBM_QS_ENDPOINT", + "_QRMI_IBM_QS_IAM_ENDPOINT", + "_QRMI_IBM_QS_IAM_APIKEY", + "_QRMI_IBM_QS_SERVICE_CRN", + ), + "iqm-server": ( + "_QRMI_IQM_ISA_ENDPOINT", + "_QRMI_IQM_ISA_TOKEN", + ), + "pasqal-cloud": (), + "pasqal-local": (), + "alice-bob-felis": (), +} + + +def resource_type(name): + """Return the QRMI ResourceType for a type string like ibm-quantum-system.""" + from qrmi import ResourceType + + attr = RESOURCE_TYPES.get(name, name) + if not hasattr(ResourceType, attr): + known = ", ".join(sorted(RESOURCE_TYPES)) + raise ValueError( + "unknown QRMI resource type {}, expected one of {}".format(name, known) + ) + return getattr(ResourceType, attr) + + +def resources_in_environment(rtype): + """Return the resource ids that have credentials set for this type. + + QRMI prefixes every variable with the resource id, so the ids can be read + back out of the environment rather than being configured twice. + """ + suffixes = REQUIRED_ENV.get(rtype, ()) + if not suffixes: + return [] + suffix = suffixes[0] + return sorted(k[: -len(suffix)] for k in os.environ if k.endswith(suffix)) + + +def missing_env(resource, rtype): + """Return the variables this resource still needs.""" + return [s for s in REQUIRED_ENV.get(rtype, ()) if not os.environ.get(resource + s)] + + +# QRMI passes the acquisition token to the job in this variable, prefixed with +# the resource id. The Slurm and LSF plugins set it, so we set it too and a +# workload written for either runs here unchanged. +ACQUISITION_TOKEN = "_QRMI_JOB_ACQUISITION_TOKEN" + + +def warmup_payload(shots=1): + """Smallest valid sampler input, one qubit measured once. + + Shape copied from the QRMI sampler in qrmi/primitives/base_sampler.py, and + qiskit generates the qasm3 so we are not hand writing a circuit. A measure + only circuit needs no basis gate translation, so it is ISA valid anywhere. + """ + try: + from qiskit import QuantumCircuit, qasm3 + except ImportError: + raise RuntimeError( + "qiskit is needed to build the warmup task. Install it with " + "pip install 'qrmi[ibm]', or skip the priority check" + ) + + qc = QuantumCircuit(1, 1) + qc.measure(0, 0) + qasm = qasm3.dumps(qc, disable_constants=True, allow_aliasing=True) + return { + "pubs": [[qasm, None, shots]], + "shots": shots, + "options": {}, + "version": 2, + "support_qiskit": True, + } + + +def explain_acquire_failure(error, resource, rtype): + """Turn a QRMI acquire error into something a user can act on. + + QRMI raises with the HTTP status in the message, so the status is what we + match on. + """ + text = str(error) + if "403" in text or "Forbidden" in text: + return ( + "{r}: not allowed to open a session on {t}. IBM only permits " + "sessions on plans that support them, such as Premium. On an Open " + "or Pay As You Go plan you can submit tasks but not hold a " + "session, so there is nothing for the scout to co-allocate.\n" + " use --quantum-ibm-type ibm-quantum-system if you have direct " + "access, or --quantum-vendor mock to test the pipeline.\n" + " QRMI said: {e}".format(r=resource, t=rtype, e=text) + ) + if "401" in text or "Unauthorized" in text or "Token renewal failed" in text: + return ( + "{r}: credentials rejected. Check {r}_QRMI_IBM_QRS_IAM_APIKEY and " + "{r}_QRMI_IBM_QRS_SERVICE_CRN.\n QRMI said: {e}".format(r=resource, e=text) + ) + if "Invalid session mode" in text: + return "{r}: bad {r}_QRMI_IBM_QRS_SESSION_MODE. QRMI said: {e}".format( + r=resource, e=text + ) + if "404" in text: + return "{r}: no such resource on {t}. QRMI said: {e}".format( + r=resource, t=rtype, e=text + ) + return "{r}: could not acquire on {t}: {e}".format(r=resource, t=rtype, e=text) + + +class QRMIBackend(Backend): + """Shared implementation for any vendor reachable through QRMI. + + A vendor subclass sets name and default_type. Everything else, including + the option names, comes from name so the options stay namespaced. + """ + + # QRMI type string used when the user does not pass one + default_type = None + + def __init__(self): + self._resource = None + self._lock = None + + @classmethod + def add_options(cls, add_option): + add_option( + "--{}-resource".format(cls.name), + metavar="ID", + default=None, + help="{}: QRMI resource id, for example ibm_kingston".format(cls.name), + ) + add_option( + "--{}-ready-timeout".format(cls.name), + metavar="SECONDS", + default=None, + help="{}: wait this long after acquiring for the resource to report " + "itself accessible, 0 to skip. Default 120".format(cls.name), + ) + add_option( + "--{}-skip-warmup".format(cls.name), + action="store_true", + help="{}: release the classical job as soon as the session opens, " + "without waiting for a warmup task to reach the head of the " + "queue. Faster, but priority is not confirmed".format(cls.name), + ) + add_option( + "--{}-warmup-timeout".format(cls.name), + metavar="SECONDS", + default=None, + help="{}: give up if the warmup task is still queued after this " + "long. Default 0, meaning wait for as long as the scout job " + "is allowed to run".format(cls.name), + ) + add_option( + "--{}-type".format(cls.name), + metavar="TYPE", + default=None, + help="{}: QRMI resource type, default {}".format( + cls.name, cls.default_type + ), + ) + + def scout_options(self, args): + rtype = getattr(args, "{}_type".format(self.name), None) or self.default_type + resource = getattr(args, "{}_resource".format(self.name), None) + if not resource: + # only one resource configured in the environment is unambiguous + found = resources_in_environment(rtype) + if len(found) == 1: + resource = found[0] + ready = getattr(args, "{}_ready_timeout".format(self.name), None) + warmup = getattr(args, "{}_warmup_timeout".format(self.name), None) + return { + "resource": resource, + "type": rtype, + "ready_timeout": 120 if ready is None else float(ready), + "skip_warmup": bool( + getattr(args, "{}_skip_warmup".format(self.name), False) + ), + "warmup_timeout": 0 if warmup is None else float(warmup), + } + + def job_environment(self, options): + """Tell the classical job which QPU it has, the same way the Slurm and + LSF plugins do, so user code can call get_job_qpu_resources_and_types + without knowing which workload manager it is running under.""" + if not options.get("resource"): + return {} + return { + "QRMI_JOB_QPU_RESOURCES": options["resource"], + "QRMI_JOB_QPU_TYPES": options.get("type") or self.default_type, + } + + @staticmethod + def acquisition_token_env(resource, session): + """The variable QRMI expects the acquisition token in.""" + return {resource + ACQUISITION_TOKEN: session} + + def open_session(self, options): + from qrmi import QuantumResource + + resource = options.get("resource") + rtype = options.get("type") or self.default_type + if not resource: + raise ValueError( + "{}: no QRMI resource id. Pass --quantum-{}-resource or export " + "the credentials for exactly one resource".format(self.name, self.name) + ) + missing = missing_env(resource, rtype) + if missing: + raise ValueError( + "{}: {} is missing {}".format( + self.name, resource, ", ".join(resource + s for s in missing) + ) + ) + self._resource = QuantumResource(resource, resource_type(rtype)) + try: + self._lock = self._resource.acquire() + except Exception as e: + self._resource = None + raise RuntimeError(explain_acquire_failure(e, resource, rtype)) + + # Hold nothing we cannot use. If the resource never becomes usable, + # give the session back rather than letting the classical job start. + ready, why = self.wait_until_ready(float(options.get("ready_timeout") or 0)) + if not ready: + self.close_session() + raise RuntimeError( + "{}: acquired a session but {} did not become usable: {}".format( + self.name, resource, why + ) + ) + + # Reaching the head of the queue is what gives us priority, so wait for + # it before letting the classical job start. + if not options.get("skip_warmup"): + try: + ok, task, reason = self.wait_for_priority( + timeout=float(options.get("warmup_timeout") or 0) + ) + except Exception as e: + self.close_session() + raise RuntimeError( + "{}: session opened on {} but the warmup task could not " + "run, so priority is unconfirmed: {}\n pass " + "--quantum-{}-skip-warmup to release without this " + "check".format(self.name, resource, e, self.name) + ) + if not ok: + self.close_session() + raise RuntimeError( + "{}: session opened on {} but the warmup task never " + "reached the head of the queue ({}), so we do not have " + "priority".format(self.name, resource, reason) + ) + return str(self._lock) + + def wait_until_ready(self, timeout=0.0, interval=5.0, sleep=time.sleep): + """Poll until the resource reports itself accessible. + + Returns (ready, reason). timeout of 0 skips the wait. + + This is as far as QRMI lets us check. is_accessible reads the backend + status, so it tells us the backend is online and taking work. QRMI does + not expose session state or queue position to python, so it is not proof + that our session has been activated. + """ + if timeout <= 0 or self._resource is None: + return True, "not checked" + deadline = time.time() + timeout + last = "no response" + while True: + try: + if self._resource.is_accessible(): + return True, "accessible" + last = "not accepting work" + except Exception as e: + last = str(e) + if time.time() + interval >= deadline: + return False, last + sleep(interval) + + def wait_for_priority(self, timeout=0.0, interval=5.0, sleep=time.sleep, shots=1): + """Submit a one shot task and wait for it to leave the queue. + + An IBM session activates when its first task reaches the head of the + queue, and later tasks in the session inherit that priority, so a task + that has started running is the signal that we actually hold the QPU. + This is the only check available that means priority rather than just + reachability. + + Returns (ok, task_id, reason). timeout of 0 waits as long as the scout + job is allowed to live. + """ + from qrmi import Payload, TaskStatus + + payload = Payload.QiskitPrimitive( + input=json.dumps(warmup_payload(shots)), program_id="sampler" + ) + task = self._resource.task_start(payload) + deadline = time.time() + timeout if timeout > 0 else None + while True: + status = self._resource.task_status(task) + if status in (TaskStatus.Running, TaskStatus.Completed): + return True, task, str(status) + if status in (TaskStatus.Failed, TaskStatus.Cancelled): + return False, task, str(status) + if deadline is not None and time.time() + interval >= deadline: + try: + self._resource.task_stop(task) + except Exception: + pass + return False, task, "still queued after {}s".format(timeout) + sleep(interval) + + def close_session(self, session=None): + if self._resource is not None and self._lock is not None: + self._resource.release(self._lock) + self._resource = None + self._lock = None + + def credentials_present(self): + found = resources_in_environment(self.default_type) + if not found: + wanted = ", ".join( + "" + s for s in REQUIRED_ENV.get(self.default_type, ()) + ) + return False, "{}: no QRMI credentials in the environment, set {}".format( + self.name, wanted + ) + incomplete = {r: missing_env(r, self.default_type) for r in found} + usable = [r for r, m in incomplete.items() if not m] + if not usable: + r = found[0] + return False, "{}: {} is missing {}".format( + self.name, r, ", ".join(r + s for s in incomplete[r]) + ) + return True, "{}: credentials present for {}".format( + self.name, " ".join(usable) + ) + + def probe(self): + """Ask QRMI whether the resource is reachable. + + Only possible when exactly one resource is configured, since the + selector runs before the submit options are parsed. Otherwise report + available and let the scout find out for real. + """ + found = [ + r + for r in resources_in_environment(self.default_type) + if not missing_env(r, self.default_type) + ] + if len(found) != 1: + return Signals(available=bool(found), detail={"resources": found}) + try: + from qrmi import QuantumResource + + qr = QuantumResource(found[0], resource_type(self.default_type)) + return Signals( + available=bool(qr.is_accessible()), detail={"resource": found[0]} + ) + except Exception as e: + return Signals(available=False, detail={"error": str(e)}) diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index d796f90..632d272 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -54,6 +54,7 @@ def prepare_pair( vendor, scout_cores=1, options=None, + job_env=None, submit_fn=None, populate_fn=None, get_graph_fn=None, @@ -90,6 +91,8 @@ def prepare_pair( sysattr["hold"] = 1 quantum = sysattr.setdefault("quantum", {}) quantum["vendor"] = vendor + if job_env: + sysattr.setdefault("environment", {}).update(job_env) # if feasibility validation is on, an unsatisfiable request is rejected here # and we abort before spending any quantum quota @@ -237,8 +240,11 @@ def modify_jobspec(self, args, jobspec): backend = get_backend(vendor) options = backend.scout_options(args) if backend else {} + job_env = backend.job_environment(options) if backend else {} handle = flux.Flux() - main_id = prepare_pair(handle, jobspec, vendor, scout_cores=1, options=options) + main_id = prepare_pair( + handle, jobspec, vendor, scout_cores=1, options=options, job_env=job_env + ) print( "flux quantum: held classical job {} (vendor={}); this submit " "launches its scout".format(main_id, vendor), diff --git a/flux_quantum/wrap.py b/flux_quantum/wrap.py index c9ec0ac..7077bb0 100644 --- a/flux_quantum/wrap.py +++ b/flux_quantum/wrap.py @@ -42,6 +42,24 @@ def _alarm(signum, frame): raise RuntimeError("job eventlog ended without a session memo") +def session_environment(session, resources=None): + """Return the env the wrapped program should see. + + QRMI reads the acquisition token from _QRMI_JOB_ACQUISITION_TOKEN, + the same variable the Slurm and LSF plugins set, so a workload written for + either runs here unchanged. The resource names arrive at submit time, the + token only exists once the scout has acquired. + """ + if resources is None: + resources = os.environ.get("QRMI_JOB_QPU_RESOURCES", "") + env = {"QUANTUM_SESSION_ID": session} + for resource in resources.split(","): + resource = resource.strip() + if resource: + env[resource + "_QRMI_JOB_ACQUISITION_TOKEN"] = session + return env + + def main(): ap = argparse.ArgumentParser(prog="quantum-wrap") ap.add_argument("--timeout", type=float, default=60.0) @@ -60,7 +78,7 @@ def main(): except Exception as e: sys.exit("quantum-wrap: {}".format(e)) - os.environ["QUANTUM_SESSION_ID"] = session + os.environ.update(session_environment(session)) cmd = args.command if cmd and cmd[0] == "--": diff --git a/pyproject.toml b/pyproject.toml index 266232f..a06f055 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,10 @@ requires-python = ">=3.8" # flux-core releases. Developed/tested against flux-core 0.87.0. Pin to a # contemporaneous flux-core, per Flux packaging guidance. +[project.optional-dependencies] +# qrmi needs python 3.11 or newer, so it stays optional +qrmi = ["qrmi[ibm]"] + [project.urls] Homepage = "https://github.com/flux-framework" diff --git a/tests/unit/test_backends.py b/tests/unit/test_backends.py index 3c597e4..cf2923a 100644 --- a/tests/unit/test_backends.py +++ b/tests/unit/test_backends.py @@ -1,3 +1,11 @@ +IBM_CREDS = { + "ibm_kingston_QRMI_IBM_QRS_ENDPOINT": "supersecret", + "ibm_kingston_QRMI_IBM_QRS_IAM_ENDPOINT": "supersecret", + "ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY": "supersecret", + "ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN": "supersecret", +} + + def test_real_backends_registered(backends_real): assert backends_real.known_vendors() == {"ibm", "braket"} @@ -17,13 +25,25 @@ def test_mock_present_with_env(backends_mock): def test_missing_creds_reported(backends_real): ok, msg = backends_real.get_backend("ibm").credentials_present() - assert not ok and "QISKIT_IBM_TOKEN" in msg + # names the variables it wants, never a value + assert not ok and "_QRMI_IBM_QRS_IAM_APIKEY" in msg def test_creds_present_when_env_set(fresh): - b = fresh(mock=False, env={"QISKIT_IBM_TOKEN": "x"}) + b = fresh(mock=False, env=IBM_CREDS) ok, msg = b.get_backend("ibm").credentials_present() - assert ok + assert ok and "ibm_kingston" in msg + + +def test_partial_creds_name_only_the_missing_variable(fresh): + partial = dict(IBM_CREDS) + del partial["ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN"] + b = fresh(mock=False, env=partial) + ok, msg = b.get_backend("ibm").credentials_present() + assert not ok + assert "ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN" in msg + # the api key value must never be echoed back + assert "supersecret" not in msg def test_mock_needs_no_creds(backends_mock): @@ -39,15 +59,37 @@ def test_backend_declares_options_and_scout_options_roundtrip(): ibm = get_backend("ibm") declared = [] ibm.add_options(lambda name, **kw: declared.append(name)) - assert "--ibm-backend" in declared and "--ibm-shots" in declared + assert "--ibm-resource" in declared and "--ibm-type" in declared class _Args: - ibm_backend = "ibm_brisbane" - ibm_instance = None - ibm_shots = "1024" + ibm_resource = "ibm_kingston" + ibm_type = None opts = ibm.scout_options(_Args()) - assert opts["backend"] == "ibm_brisbane" and opts["shots"] == "1024" + assert opts["resource"] == "ibm_kingston" + assert opts["type"] == "qiskit-runtime-service" + + +def test_scout_options_infer_the_resource_from_the_environment(fresh): + """One resource configured is unambiguous, so the id need not be repeated.""" + b = fresh(mock=False, env=IBM_CREDS) + + class _Args: + ibm_resource = None + ibm_type = None + + assert b.get_backend("ibm").scout_options(_Args())["resource"] == "ibm_kingston" + + +def test_job_environment_follows_the_qrmi_convention(fresh): + """The classical job gets the same variables Slurm and LSF set, so user + code can call get_job_qpu_resources_and_types unchanged.""" + b = fresh(mock=False, env=IBM_CREDS) + env = b.get_backend("ibm").job_environment( + {"resource": "ibm_kingston", "type": "qiskit-runtime-service"} + ) + assert env["QRMI_JOB_QPU_RESOURCES"] == "ibm_kingston" + assert env["QRMI_JOB_QPU_TYPES"] == "qiskit-runtime-service" def test_mock_open_session_uses_options(monkeypatch): @@ -63,9 +105,20 @@ def test_mock_open_session_uses_options(monkeypatch): assert b.open_session({}).startswith("mock-session-") -def test_ibm_open_session_deferred_is_explicit(): +def test_open_session_without_a_resource_is_explicit(): from flux_quantum.backends import get_backend import flux_quantum.backends.ibm # noqa - with __import__("pytest").raises(NotImplementedError): - get_backend("ibm").open_session({"backend": "x"}) + with __import__("pytest").raises(ValueError, match="no QRMI resource id"): + get_backend("ibm").open_session({}) + + +def test_open_session_names_the_missing_variable(fresh): + partial = dict(IBM_CREDS) + del partial["ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY"] + b = fresh(mock=False, env=partial) + with __import__("pytest").raises(ValueError) as e: + b.get_backend("ibm").open_session( + {"resource": "ibm_kingston", "type": "qiskit-runtime-service"} + ) + assert "ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY" in str(e.value) diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index 0cb7e0e..4db7d87 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -66,7 +66,7 @@ def test_plugin_registers_options_with_quantum_prefix(stub_flux_cli): # core options always present assert {"--vendor", "--select"} <= names # each registered vendor backend contributes its own namespaced options - assert "--ibm-backend" in names and "--braket-device" in names + assert "--ibm-resource" in names and "--braket-device" in names def test_plugin_inactive_for_other_progs(stub_flux_cli): @@ -296,3 +296,25 @@ def test_unlimited_classical_keeps_scout_unlimited(stub_flux_cli): cancel_fn=lambda *a: None, ) assert js.jobspec["attributes"]["system"]["duration"] == 0 + + +def test_job_env_reaches_the_classical(stub_flux_cli): + """A backend can add env to the classical job, which is how the QPU + assignment gets there.""" + from flux_quantum import cli + + js = _fake_jobspec(["myprog"]) + submitted = {} + + cli.prepare_pair( + handle=None, + jobspec=js, + vendor="mock", + job_env={"QRMI_JOB_QPU_RESOURCES": "ibm_kingston"}, + submit_fn=lambda h, j: submitted.setdefault("js", json.loads(j)) and 0 or 7, + populate_fn=lambda h, v: None, + get_graph_fn=lambda h: _live_graph(), + cancel_fn=lambda *a: None, + ) + env = submitted["js"]["attributes"]["system"]["environment"] + assert env["QRMI_JOB_QPU_RESOURCES"] == "ibm_kingston" diff --git a/tests/unit/test_qrmi.py b/tests/unit/test_qrmi.py new file mode 100644 index 0000000..230c67d --- /dev/null +++ b/tests/unit/test_qrmi.py @@ -0,0 +1,321 @@ +"""QRMI backed sessions, driven against a fake qrmi module. + +qrmi needs python 3.11 or newer and real credentials, so the tests stub it and +check the calls we make rather than talking to a vendor. +""" + +import sys +import types + +import pytest + +RESOURCE = "ibm_kingston" +SUFFIXES = ( + "_QRMI_IBM_QRS_ENDPOINT", + "_QRMI_IBM_QRS_IAM_ENDPOINT", + "_QRMI_IBM_QRS_IAM_APIKEY", + "_QRMI_IBM_QRS_SERVICE_CRN", +) +CREDS = {RESOURCE + s: "supersecret" for s in SUFFIXES} + + +class FakeStatus: + Queued = "Queued" + Running = "Running" + Completed = "Completed" + Failed = "Failed" + Cancelled = "Cancelled" + + +class FakeResource: + calls = [] + # statuses handed back by task_status, in order + statuses = ["Running"] + + def __init__(self, rid, rtype): + self.rid = rid + FakeResource.calls.append(("ctor", rid, str(rtype))) + + def acquire(self): + FakeResource.calls.append(("acquire", self.rid)) + return "LOCK-123" + + def release(self, id): + FakeResource.calls.append(("release", self.rid, id)) + + def is_accessible(self): + return True + + def task_start(self, payload): + FakeResource.calls.append(("task_start", self.rid)) + return "TASK-1" + + def task_status(self, task): + return FakeResource.statuses.pop(0) + + def task_stop(self, task): + FakeResource.calls.append(("task_stop", task)) + + +@pytest.fixture +def fake_qrmi(monkeypatch): + FakeResource.calls = [] + FakeResource.statuses = ["Running"] + mod = types.ModuleType("qrmi") + mod.QuantumResource = FakeResource + mod.ResourceType = type( + "ResourceType", (), {"IBMQiskitRuntimeService": "IBMQiskitRuntimeService"} + ) + mod.TaskStatus = FakeStatus + mod.Payload = type( + "Payload", (), {"QiskitPrimitive": staticmethod(lambda input, program_id: {})} + ) + monkeypatch.setitem(sys.modules, "qrmi", mod) + for k, v in CREDS.items(): + monkeypatch.setenv(k, v) + return FakeResource + + +def _opts(**kw): + base = {"resource": RESOURCE, "type": "qiskit-runtime-service"} + base.update(kw) + return base + + +def _ibm(): + from flux_quantum.backends import get_backend + import flux_quantum.backends.ibm # noqa + + return get_backend("ibm") + + +def test_acquire_and_release_round_trip(fake_qrmi): + """open_session acquires and close_session releases the same lock.""" + b = _ibm() + session = b.open_session(_opts(skip_warmup=True)) + assert session == "LOCK-123" + b.close_session(session) + assert ("acquire", RESOURCE) in fake_qrmi.calls + assert ("release", RESOURCE, "LOCK-123") in fake_qrmi.calls + + +def test_close_without_open_is_harmless(fake_qrmi): + """The scout always calls close in a finally, even if open never ran.""" + _ibm().close_session("anything") + assert fake_qrmi.calls == [] + + +def test_resource_is_inferred_when_only_one_is_configured(fake_qrmi): + class Args: + ibm_resource = None + ibm_type = None + + assert _ibm().scout_options(Args())["resource"] == RESOURCE + + +def test_explicit_resource_wins(fake_qrmi): + class Args: + ibm_resource = "ibm_fez" + ibm_type = None + + assert _ibm().scout_options(Args())["resource"] == "ibm_fez" + + +def test_job_environment_matches_the_qrmi_convention(fake_qrmi): + """Slurm and LSF set these two, so a workload moves between them unchanged.""" + env = _ibm().job_environment( + {"resource": RESOURCE, "type": "qiskit-runtime-service"} + ) + assert env == { + "QRMI_JOB_QPU_RESOURCES": RESOURCE, + "QRMI_JOB_QPU_TYPES": "qiskit-runtime-service", + } + + +def test_unknown_resource_type_is_rejected(fake_qrmi): + from flux_quantum.backends.qrmi import resource_type + + with pytest.raises(ValueError, match="unknown QRMI resource type"): + resource_type("not-a-real-type") + + +def test_secrets_never_appear_in_messages(fake_qrmi, monkeypatch): + """Credential checks report variable names and never values.""" + monkeypatch.delenv(RESOURCE + "_QRMI_IBM_QRS_IAM_APIKEY") + ok, msg = _ibm().credentials_present() + assert not ok + assert RESOURCE + "_QRMI_IBM_QRS_IAM_APIKEY" in msg + assert "supersecret" not in msg + + +def test_premium_message_when_sessions_are_not_allowed(): + """A 403 on acquire means the plan cannot hold a session. Say so plainly + instead of surfacing an HTTP status.""" + from flux_quantum.backends.qrmi import explain_acquire_failure + + msg = explain_acquire_failure( + RuntimeError("status code error in create_session: 403"), + "ibm_kingston", + "qiskit-runtime-service", + ) + assert "Premium" in msg + assert "submit tasks but not hold a session" in msg + assert "--quantum-vendor mock" in msg + assert "403" in msg # the original is still there for debugging + + +def test_bad_credentials_point_at_the_variables(): + from flux_quantum.backends.qrmi import explain_acquire_failure + + msg = explain_acquire_failure( + RuntimeError("Token renewal failed: 401"), "ibm_fez", "qiskit-runtime-service" + ) + assert "ibm_fez_QRMI_IBM_QRS_IAM_APIKEY" in msg + + +def test_acquire_failure_is_translated(fake_qrmi, monkeypatch): + def boom(self): + raise RuntimeError("status code error in create_session: 403 Forbidden") + + monkeypatch.setattr(fake_qrmi, "acquire", boom) + b = _ibm() + with pytest.raises(RuntimeError, match="Premium"): + b.open_session(_opts(skip_warmup=True)) + + +def test_session_is_released_if_the_resource_never_becomes_usable( + fake_qrmi, monkeypatch +): + """Never hand the classical job a session we could not confirm.""" + monkeypatch.setattr(fake_qrmi, "is_accessible", lambda self: False) + b = _ibm() + with pytest.raises(RuntimeError, match="did not become usable"): + b.open_session(_opts(ready_timeout=1, skip_warmup=True)) + # the acquired session must not be left dangling + assert ("release", RESOURCE, "LOCK-123") in fake_qrmi.calls + + +def test_ready_wait_polls_until_accessible(fake_qrmi): + b = _ibm() + b.open_session(_opts(skip_warmup=True)) + states = [False, False, True] + fake_qrmi.is_accessible = lambda self: states.pop(0) + slept = [] + ready, why = b.wait_until_ready(timeout=60, interval=1, sleep=slept.append) + assert ready and why == "accessible" + assert len(slept) == 2 # two failures, so two waits + + +def test_wrap_exports_the_qrmi_acquisition_token(): + """The token goes into the variable QRMI reads, not only ours.""" + from flux_quantum.wrap import session_environment + + env = session_environment("SESS-1", resources="ibm_kingston,ibm_fez") + assert env["QUANTUM_SESSION_ID"] == "SESS-1" + assert env["ibm_kingston_QRMI_JOB_ACQUISITION_TOKEN"] == "SESS-1" + assert env["ibm_fez_QRMI_JOB_ACQUISITION_TOKEN"] == "SESS-1" + + +def test_wrap_without_qrmi_resources_sets_only_our_variable(): + from flux_quantum.wrap import session_environment + + assert session_environment("SESS-2", resources="") == { + "QUANTUM_SESSION_ID": "SESS-2" + } + + +# --------------------------------------------------------------------------- +# priority. An IBM session activates when its first task reaches the head of +# the queue, so a task that is running is the signal that we hold the QPU. +# --------------------------------------------------------------------------- + + +def _acquired(fake, statuses): + """A backend with a session already open, so wait_for_priority can run.""" + b = _ibm() + b.open_session(_opts(ready_timeout=0, skip_warmup=True)) + fake.statuses = list(statuses) + fake.calls = [] + return b + + +def test_warmup_waits_until_the_task_leaves_the_queue(fake_qrmi, monkeypatch): + monkeypatch.setattr( + "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} + ) + b = _acquired(fake_qrmi, ["Queued", "Queued", "Running"]) + slept = [] + ok, task, reason = b.wait_for_priority(interval=1, sleep=slept.append) + assert ok and reason == "Running" + assert len(slept) == 2 + + +def test_warmup_failure_means_no_priority(fake_qrmi, monkeypatch): + monkeypatch.setattr( + "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} + ) + b = _acquired(fake_qrmi, ["Queued", "Failed"]) + ok, task, reason = b.wait_for_priority(interval=0, sleep=lambda s: None) + assert not ok and reason == "Failed" + + +def test_warmup_timeout_stops_the_task(fake_qrmi, monkeypatch): + monkeypatch.setattr( + "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} + ) + b = _acquired(fake_qrmi, ["Queued"] * 10) + ok, task, reason = b.wait_for_priority(timeout=1, interval=1, sleep=lambda s: None) + assert not ok and "still queued" in reason + assert ("task_stop", "TASK-1") in fake_qrmi.calls + + +def test_session_released_when_priority_is_never_reached(fake_qrmi, monkeypatch): + """Never release the classical job against a session we do not own.""" + monkeypatch.setattr( + "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} + ) + fake_qrmi.statuses = ["Failed"] + b = _ibm() + with pytest.raises(RuntimeError, match="do not have priority"): + b.open_session(_opts(ready_timeout=0)) + assert ("release", RESOURCE, "LOCK-123") in fake_qrmi.calls + + +def test_skip_warmup_releases_immediately(fake_qrmi, monkeypatch): + monkeypatch.setattr( + "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} + ) + fake_qrmi.statuses = [] + b = _ibm() + assert b.open_session(_opts(ready_timeout=0, skip_warmup=True)) == "LOCK-123" + assert not any(c[0] == "task_start" for c in fake_qrmi.calls) + + +def test_warmup_payload_is_a_single_shot_measure(fake_qrmi): + """Cheapest thing that still occupies the QPU, so priority costs one shot.""" + pytest.importorskip("qiskit") + from flux_quantum.backends.qrmi import warmup_payload + + p = warmup_payload() + assert p["shots"] == 1 and p["version"] == 2 and p["support_qiskit"] is True + qasm = p["pubs"][0][0] + assert "measure" in qasm + + +def test_missing_qiskit_says_what_to_install(fake_qrmi, monkeypatch): + """The warmup needs qiskit. Say so, and offer the way out.""" + import builtins + + real = builtins.__import__ + + def no_qiskit(name, *a, **k): + if name == "qiskit": + raise ImportError("no qiskit") + return real(name, *a, **k) + + monkeypatch.setattr(builtins, "__import__", no_qiskit) + b = _ibm() + with pytest.raises(RuntimeError) as e: + b.open_session(_opts(ready_timeout=0)) + assert "qrmi[ibm]" in str(e.value) + assert "skip-warmup" in str(e.value) diff --git a/tests/unit/test_selector.py b/tests/unit/test_selector.py index 0d3ecbd..0c99e4f 100644 --- a/tests/unit/test_selector.py +++ b/tests/unit/test_selector.py @@ -1,3 +1,11 @@ +IBM_CREDS = { + "ibm_kingston_QRMI_IBM_QRS_ENDPOINT": "x", + "ibm_kingston_QRMI_IBM_QRS_IAM_ENDPOINT": "x", + "ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY": "x", + "ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN": "x", +} + + import importlib import pytest @@ -11,16 +19,35 @@ def test_no_usable_vendor_raises(backends_real): with pytest.raises(sel.SelectionError) as e: sel.select_vendor() # error lists the missing creds for each vendor - assert "QISKIT_IBM_TOKEN" in str(e.value) + assert "_QRMI_IBM_QRS_IAM_APIKEY" in str(e.value) + +def test_select_ibm_when_reachable(fresh, monkeypatch): + """Credentials alone are not enough. The probe asks QRMI whether the + resource is reachable, so stub that rather than calling out to IBM.""" + b = fresh(mock=False, env=IBM_CREDS) + from flux_quantum.backends.base import Signals -def test_select_ibm_with_token(fresh): - fresh(mock=False, env={"QISKIT_IBM_TOKEN": "x"}) + monkeypatch.setattr( + b.get_backend("ibm").__class__, "probe", lambda self: Signals(available=True) + ) sel = _selector() vendor, sig, log = sel.select_vendor(policy="any") assert vendor == "ibm" +def test_unreachable_ibm_is_not_a_candidate(fresh, monkeypatch): + b = fresh(mock=False, env=IBM_CREDS) + from flux_quantum.backends.base import Signals + + monkeypatch.setattr( + b.get_backend("ibm").__class__, "probe", lambda self: Signals(available=False) + ) + sel = _selector() + with pytest.raises(sel.SelectionError): + sel.select_vendor(policy="any") + + def test_explicit_braket_without_creds_raises(backends_real): sel = _selector() with pytest.raises(sel.SelectionError): From 8cf6199c590ab09c8d2f4e8479fa51f1100685a3 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 12 Aug 2026 14:27:28 -0700 Subject: [PATCH 02/26] test: fix smoke test Signed-off-by: vsoch --- tests/integration/setup-mock-registry.sh | 3 +++ tests/integration/test-mock-e2e.sh | 8 ++++++++ tests/integration/test-select-discovery.sh | 1 + tests/integration/test-system-instance.sh | 12 +++++++++++- 4 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/integration/setup-mock-registry.sh b/tests/integration/setup-mock-registry.sh index 8b4393e..08de5d9 100755 --- a/tests/integration/setup-mock-registry.sh +++ b/tests/integration/setup-mock-registry.sh @@ -15,10 +15,13 @@ SUB="${SUB:-/tmp/quantum-subgraph.json}" LIVE="${LIVE:-/tmp/live-graph.json}" echo "-- release preloaded scheduler; load fluxion + qmanager NORMALLY --" +# reverse dependency order, feasibility sits between resource and qmanager flux module remove -f sched-fluxion-qmanager 2>/dev/null || true +flux module remove -f sched-fluxion-feasibility 2>/dev/null || true flux module remove -f sched-fluxion-resource 2>/dev/null || true flux module remove -f sched-simple 2>/dev/null || true flux module load sched-fluxion-resource +flux module load sched-fluxion-feasibility 2>/dev/null || true flux module load sched-fluxion-qmanager echo "-- populate qdevice_ -> qpu into the live graph (find + add_subgraph RPC) --" diff --git a/tests/integration/test-mock-e2e.sh b/tests/integration/test-mock-e2e.sh index 737c65e..6bb7918 100755 --- a/tests/integration/test-mock-e2e.sh +++ b/tests/integration/test-mock-e2e.sh @@ -20,10 +20,17 @@ flux submit --help 2>&1 | grep -E 'quantum' || { echo "FAIL: CLI plugin not disc echo "" echo "=== 1. load fluxion + qmanager (normal acquire, no load-file) ===" +# rc1 loads three fluxion modules and feasibility sits between the other two, +# so remove in reverse dependency order. Leaving feasibility running while +# resource goes away kills it with a notify failure. flux module remove -f sched-fluxion-qmanager 2>/dev/null || true +flux module remove -f sched-fluxion-feasibility 2>/dev/null || true flux module remove -f sched-fluxion-resource 2>/dev/null || true flux module remove -f sched-simple 2>/dev/null || true flux module load sched-fluxion-resource +# feasibility is what rejects an unsatisfiable request at ingest, which is the +# gate the plugin relies on, so put it back rather than running without it. +flux module load sched-fluxion-feasibility 2>/dev/null || true flux module load sched-fluxion-qmanager echo "fluxion + qmanager loaded" @@ -89,6 +96,7 @@ fi echo "=== 5. graceful teardown (remove fluxion so shutdown is clean) ===" flux module remove -f sched-fluxion-qmanager 2>/dev/null || true +flux module remove -f sched-fluxion-feasibility 2>/dev/null || true flux module remove -f sched-fluxion-resource 2>/dev/null || true rm -f "$ERR" diff --git a/tests/integration/test-select-discovery.sh b/tests/integration/test-select-discovery.sh index 8dc7ef9..f3aedcf 100755 --- a/tests/integration/test-select-discovery.sh +++ b/tests/integration/test-select-discovery.sh @@ -75,6 +75,7 @@ fi # graceful teardown so shutdown does not log an acquire failure flux module remove -f sched-fluxion-qmanager 2>/dev/null || true +flux module remove -f sched-fluxion-feasibility 2>/dev/null || true flux module remove -f sched-fluxion-resource 2>/dev/null || true rm -f "$ERR" diff --git a/tests/integration/test-system-instance.sh b/tests/integration/test-system-instance.sh index 996e80d..038a24b 100755 --- a/tests/integration/test-system-instance.sh +++ b/tests/integration/test-system-instance.sh @@ -28,6 +28,12 @@ flux submit --help 2>&1 | grep -q -- --quantum-vendor || { echo "FAIL: no --quantum options. Is FLUX_CLI_PLUGINPATH set?"; exit 1; } flux submit --help 2>&1 | grep -- --quantum- | sed 's/^/ /' +# Remember how much log there is before we start. This is a long lived +# instance, so everything since boot is still in the ring buffer and an +# unrelated error from an hour ago is not our problem. Counting lines rather +# than parsing timestamps keeps this working whatever the log format is. +DMESG_MARK=$(flux dmesg 2>/dev/null | wc -l) + echo "" echo "=== 4. one quantum submit -> held classical + scout ===" SID="systest-$(date +%s)" @@ -91,7 +97,11 @@ fi echo "" echo "=== 8. instance log errors during the run ===" -errs=$(flux dmesg 2>&1 | grep -iE "\.err\[[0-9]+\]|: error:|fatal" || true) +# only what was logged after the mark, and rexec complaints are not ours. They +# come from flux exec, which the scout and the classical never use. +errs=$(flux dmesg 2>/dev/null | tail -n +$((DMESG_MARK + 1)) \ + | grep -iE "\.err\[[0-9]+\]|: error:|fatal" \ + | grep -v "rexec" || true) if [ -n "$errs" ]; then echo "FAIL: errors logged:"; echo "$errs" | sed 's/^/ /'; rc=1 fi From 17438700676add59d863864b489d7919e9735d25 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 12 Aug 2026 18:49:25 -0700 Subject: [PATCH 03/26] test: qrmi ibm run example Signed-off-by: vsoch --- examples/qrmi/README.md | 65 +++++++-------- examples/qrmi/ibm-run.sh | 169 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+), 38 deletions(-) create mode 100755 examples/qrmi/ibm-run.sh diff --git a/examples/qrmi/README.md b/examples/qrmi/README.md index 3ee35be..0a3d988 100644 --- a/examples/qrmi/README.md +++ b/examples/qrmi/README.md @@ -15,12 +15,16 @@ the scout. qiskit comes along with it and we need that for the warmup task. QRMI reads its config out of the environment. Every variable is prefixed with the resource id, so for ibm_kingston the api key lives in -`ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY`. Four of them are required. +`ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY`. Four of them are required. Note that this used to work for me, and it stopped. I now login with the web interface copy paste, which looks like this: - cp env.example .env - chmod 600 .env - $EDITOR .env - set -a && . ./.env && set +a +```bash +ibmcloud login -a https://cloud.ibm.com -u passcode -p +``` +```bash +export IBM_CLOUD_TOKEN= +ibmcloud login --apikey $IBM_CLOUD_TOKEN +export IBM_CLOUD_CRN=$(ibmcloud resource service-instances --service-name quantum-computing --output json | jq -r '.[] | {name: .name, crn: .crn}' | jq -r .crn) +``` Check before you submit anything. @@ -30,6 +34,17 @@ Check before you submit anything. That names any variable you are missing. It never prints a value. + +unset FLUX_QUANTUM_MOCK +flux submit -t 30m \ + --quantum-vendor ibm --quantum-ibm-resource \ + --quantum-ibm-warmup-timeout 600 \ + -n4 flux python /opt/flux-quantum/examples/qrmi/workload.py + +flux jobs -a +flux job attach # this is where a Premium refusal shows up +flux job eventlog | grep memo + ## Submit flux submit -t 30m --quantum-vendor ibm --quantum-ibm-resource ibm_kingston \ @@ -45,32 +60,15 @@ resource has credentials in the environment you can drop ## You need an account that can open sessions -This is the thing that will bite you first. The qiskit-runtime-service type opens -a session, and IBM only allows that on plans that support sessions, like Premium. -On Open or Pay As You Go you can submit tasks but you cannot hold a session, so -there is nothing for the scout to co-allocate. - -You get told that in those words rather than a 403. The held classical job is -cancelled, so nothing is left sitting in the queue. - -If you have direct access instead, use `--quantum-ibm-type ibm-quantum-system`. +The qiskit-runtime-service type opens a session, and IBM only allows that on plans that support sessions (Premium). On Open or Pay As You Go you can submit tasks but you cannot hold a session. You will get an error message (403). Our library will cancel the held classical job. If you have direct access instead, use `--quantum-ibm-type ibm-quantum-system`. ## A session is not the same as having the QPU An IBM session goes active when its first task reaches the head of the queue. After that, tasks in the session keep that priority. So opening a session tells -you very little on its own. - -After acquiring, the scout submits a warmup task, one qubit measured once, and -waits for it to leave the queue. Once that task is running we are at the head, -the session is live, and only then do we release the classical job. - -That is the whole point. The classical job sits held with its nodes reserved -while the scout waits in the QPU queue. It starts the moment the QPU is really -ours, so neither side pays to wait for the other. - -If the warmup never runs we release the session and fail the submit. The -classical job never starts against a QPU we do not have. +you very little on its own. After acquire we have to submit a warmup task and monitor it, +and it is small (one qubit and one measurement) and wait for it to run. Then we can at least +say that we reached the top and have priority and we release classical. If the warmup never runs we release the session and fail the submit. The classical job never starts against a QPU we do not have. --quantum-ibm-warmup-timeout SECONDS give up after this long. Default 0, which means wait as long as the @@ -84,26 +82,17 @@ The warmup costs one shot. ## The scout holds the session for as long as the classical runs -That is deliberate. The fluxion allocation and the vendor session start and end +The fluxion allocation and the vendor session start and end together. It does mean you are billed for the length of the classical job and not the length of the quantum work, so set `-t`. A killed scout still releases, the release is in a finally and SIGTERM is handled. If the node dies you are down to the vendor timeout. -## Credentials, and where they end up +## Credentials The plugin runs in your process at submit time and reads the variables out of -your environment. The scout uses them, and it runs as you. - -Flux copies the submitting environment into the jobspec, so they are in the job -record and the instance owner can read them. Same as any other job environment. -What does not happen is a credential owned by the system user, or one shared -between users. If that is not good enough for you, do not export the key in the -shell you submit from. - -Only the session id travels to the classical job, and it goes over the job -eventlog. +your environment. The scout uses them, and it runs as you. Only the session id travels to the classical job, and it goes over the job eventlog. ## What the classical job sees diff --git a/examples/qrmi/ibm-run.sh b/examples/qrmi/ibm-run.sh new file mode 100755 index 0000000..b3d92c8 --- /dev/null +++ b/examples/qrmi/ibm-run.sh @@ -0,0 +1,169 @@ +#!/bin/bash +# One cheap end to end run against IBM, and a report of what it cost. +# +# export IBM_CLOUD_TOKEN= +# ibmcloud login --apikey "$IBM_CLOUD_TOKEN" +# ibm-run +# +# Override anything with an env var, for example +# RESOURCE=ibm_kingston NTASKS=4 ibm-run +# SKIP_WARMUP=1 ibm-run free, opens and closes a session, no task +# SLEEP=5 ibm-run let the warmup task finish instead of being +# cancelled when the session closes + +set -u + +NTASKS="${NTASKS:-1}" # flux -n is tasks, not nodes. One is enough here. +# How long the classical job stays alive once it has the session. Zero means +# print and exit, which is cheapest, but the session then closes while the +# warmup task is still running so the warmup ends up CANCELLED. Around 5 +# seconds lets it finish instead. Every second of this is billed. +SLEEP="${SLEEP:-0}" +WALLTIME="${WALLTIME:-1m}" + +# waiting is free, so be generous +WARMUP_TIMEOUT="${WARMUP_TIMEOUT:-1800}" +SKIP_WARMUP="${SKIP_WARMUP:-0}" +OUT="${OUT:-$HOME/ibm-run-$(date +%Y%m%d-%H%M%S)}" + +say() { echo "== $* =="; } +die() { echo "$*" >&2; exit 1; } + +# 1. credentials and CRN +[ -n "${IBM_CLOUD_TOKEN:-}" ] || die "set IBM_CLOUD_TOKEN to your ibm cloud api key" + +if [ -z "${IBM_CLOUD_CRN:-}" ]; then + command -v ibmcloud >/dev/null || die "need the ibmcloud cli, or set IBM_CLOUD_CRN" + command -v jq >/dev/null || die "need jq" + say "looking up the quantum instance" + crns=$(ibmcloud resource service-instances --service-name quantum-computing \ + --output json 2>/dev/null | jq -r '.[].crn') + n=$(printf '%s\n' "$crns" | grep -c .) + [ "$n" -eq 0 ] && die "no quantum instances. logged in? ibmcloud login --apikey \"\$IBM_CLOUD_TOKEN\"" + if [ "$n" -gt 1 ]; then + # one per line, so picking blindly would export a multiline value + echo "found $n instances, set IBM_CLOUD_CRN to the one you want:" >&2 + ibmcloud resource service-instances --service-name quantum-computing \ + --output json | jq -r '.[] | " \(.name) \(.crn)"' >&2 + exit 1 + fi + IBM_CLOUD_CRN="$crns" +fi + +export QISKIT_IBM_TOKEN="$IBM_CLOUD_TOKEN" +export QISKIT_IBM_INSTANCE="$IBM_CLOUD_CRN" +unset FLUX_QUANTUM_MOCK +# mock would win --quantum-select on queue depth + +# 2. pick the least busy QPU, unless told which +if [ -z "${RESOURCE:-}" ]; then + say "picking the least busy QPU" + RESOURCE=$(flux python -c " +from qiskit_ibm_runtime import QiskitRuntimeService +rows = [] +for b in QiskitRuntimeService().backends(operational=True, simulator=False): + try: + rows.append((b.status().pending_jobs, b.name)) + except Exception: + pass +for q, name in sorted(rows): + print(name, q) +" 2>/dev/null | tee /dev/stderr | head -1 | awk '{print $1}') +fi +[ -n "$RESOURCE" ] || die "could not find a QPU. is the plan active?" + +# QRMI prefixes every variable with the resource id +export ${RESOURCE}_QRMI_IBM_QRS_ENDPOINT="https://quantum.cloud.ibm.com/api/v1" +export ${RESOURCE}_QRMI_IBM_QRS_IAM_ENDPOINT="https://iam.cloud.ibm.com" +export ${RESOURCE}_QRMI_IBM_QRS_IAM_APIKEY="$IBM_CLOUD_TOKEN" +export ${RESOURCE}_QRMI_IBM_QRS_SERVICE_CRN="$IBM_CLOUD_CRN" +# session, not batch. Batch is cheaper but gives no exclusive access, and +# exclusive access is the thing being tested. +export ${RESOURCE}_QRMI_IBM_QRS_SESSION_MODE="dedicated" + +mkdir -p "$OUT" +echo "resource $RESOURCE" | tee "$OUT/run.txt" +echo "crn $IBM_CLOUD_CRN" >> "$OUT/run.txt" + +say "credentials" +flux python -c " +from flux_quantum.backends import get_backend +ok, msg = get_backend('ibm').credentials_present() +print(msg) +raise SystemExit(0 if ok else 1)" | tee -a "$OUT/run.txt" || die "credentials incomplete" + +# 3. submit. The classical job prints and exits, so the billed window is warmup plus a few seconds. +warmup_args="--quantum-ibm-warmup-timeout $WARMUP_TIMEOUT" +[ "$SKIP_WARMUP" = 1 ] && warmup_args="--quantum-ibm-skip-warmup" + +say "submitting" +t0=$(date +%s) +err=$(mktemp) +scout=$(flux submit -t "$WALLTIME" --quantum-vendor ibm \ + --quantum-ibm-resource "$RESOURCE" $warmup_args -n"$NTASKS" \ + -- sh -c "echo classical got \$QUANTUM_SESSION_ID; sleep $SLEEP" 2>"$err") +tee -a "$OUT/run.txt" < "$err" +main=$(grep -oE 'held classical job [0-9]+' "$err" | awk '{print $NF}') +rm -f "$err" +[ -n "$main" ] || die "no held classical job, see above" +echo "scout=$scout classical=$main" | tee -a "$OUT/run.txt" + +say "waiting. the queue wait is free, so this can take a while" +flux job wait-event -t "$((WARMUP_TIMEOUT + 600))" "$main" clean >/dev/null 2>&1 +flux job wait-event -t 300 "$scout" clean >/dev/null 2>&1 +t1=$(date +%s) + +# 4. collect +say "results" +{ + echo "--- wall clock ---" + echo "submit to both jobs done: $((t1 - t0))s" + echo "--- scout ---" + flux job attach "$scout" 2>&1 + echo "--- classical ---" + flux job attach "$main" 2>&1 + echo "--- classical eventlog ---" + flux job eventlog "$main" 2>&1 + echo "--- scout eventlog ---" + flux job eventlog "$scout" 2>&1 +} | tee -a "$OUT/run.txt" + +say "what IBM billed" +# qpu_charge_time_seconds is what you are charged. usage_estimation is the +# finer grained figure. A CANCELLED warmup is normal when SLEEP is 0, it means +# the classical finished and we closed the session while the one shot task was +# still in flight. It had already done its job by reaching the head of the queue. +flux python -c " +import os +from qiskit_ibm_runtime import QiskitRuntimeService +svc = QiskitRuntimeService(channel='ibm_quantum_platform', + token=os.environ['QISKIT_IBM_TOKEN'], + instance=os.environ['QISKIT_IBM_INSTANCE']) +jobs = list(svc.jobs(limit=5)) +if not jobs: + print('no jobs, which is expected with SKIP_WARMUP=1') +total = 0 +for j in jobs: + m = {} + try: + m = j.metrics() or {} + except Exception as e: + print('metrics unavailable', e) + charged = m.get('usage', {}).get('qpu_charge_time_seconds') + est = None + try: + est = (j.usage_estimation or {}).get('quantum_seconds') + except Exception: + pass + if charged: + total += charged + print(j.job_id(), j.status(), 'backend=%s' % j.backend().name, + 'session=%s' % getattr(j, 'session_id', None), + 'charged_s=%s' % charged, 'estimate_s=%s' % est) +if total: + print() + print('charged %ss, about \$%.2f at 1.60 a second' % (total, total * 1.60)) +" 2>&1 | tee -a "$OUT/run.txt" + +echo +echo "saved to $OUT/run.txt" From 2b33fd8eeb7d150de0ed42239ca8d53e630576e8 Mon Sep 17 00:00:00 2001 From: vsoch Date: Tue, 18 Aug 2026 20:37:24 -0700 Subject: [PATCH 04/26] backend: add braket qrmi with ibm quantum is too expensive for experiments. Signed-off-by: vsoch --- flux_quantum/backends/braket.py | 187 ++++++++++++++++++++++++++++---- flux_quantum/scout.py | 35 ++++-- tests/unit/test_braket.py | 180 ++++++++++++++++++++++++++++++ tests/unit/test_handoff.py | 27 +++++ 4 files changed, 397 insertions(+), 32 deletions(-) create mode 100644 tests/unit/test_braket.py diff --git a/flux_quantum/backends/braket.py b/flux_quantum/backends/braket.py index 499d4a1..a10343d 100644 --- a/flux_quantum/backends/braket.py +++ b/flux_quantum/backends/braket.py @@ -1,17 +1,54 @@ """AWS Braket backend. -The probe is a stub. The real one uses boto3 with the user AWS credentials. -Braket has no acquire and release like QRMI, so it needs its own backend. +Braket has no session to acquire. You submit a task, it waits in the device +queue, and it runs. So the scout submits a trivial no-op task and waits for it +to reach the front of the queue, which is as close to priority as Braket +offers. + +Note the limitation, because it is the vendor's and not ours. Reaching the +front of the queue does not reserve anything, so another user can take the +device between our task running and the classical job starting. IBM sessions do +hold the device, at the price of billing wall clock for it. Braket exposes the +queue position but no way to hold it. + +The SDK is an optional dependency, so it is imported inside the methods. """ import os +import time + from .base import Backend, Signals, register +# state simulator, no queue to speak of and cents per task +SV1 = "arn:aws:braket:::device/quantum-simulator/amazon/sv1" + +# a task that has left the queue reports no position at all, so these are how +# we notice we missed the window rather than waiting for a value that can no +# longer appear +RAN = ("RUNNING", "COMPLETED") +FAILED = ("FAILED", "CANCELLED") + + +def region_for(device_arn, override=None): + """Region to talk to. Devices are not all in one, and a QPU searched in the + wrong region simply never appears.""" + parts = device_arn.split(":") + from_arn = parts[3] if len(parts) > 3 else "" + return ( + override + or from_arn + or os.environ.get("AWS_DEFAULT_REGION") + or os.environ.get("AWS_REGION") + or "us-east-1" + ) + @register class BraketBackend(Backend): name = "braket" - required_env = ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY") + + def __init__(self): + self._task = None @classmethod def add_options(cls, add_option): @@ -19,43 +56,151 @@ def add_options(cls, add_option): "--braket-device", metavar="ARN", default=None, - help="Braket: device ARN to target", + help="Braket: device ARN, defaults to the SV1 simulator", ) add_option( "--braket-region", metavar="REGION", default=None, - help="Braket: AWS region for the device", + help="Braket: AWS region, defaults to the one in the device ARN", ) add_option( "--braket-shots", metavar="N", default=None, - help="Braket: number of shots for the session workload", + help="Braket: shots for the queue probe task, default 1", + ) + add_option( + "--braket-queue-timeout", + metavar="SECONDS", + default=None, + help="Braket: give up if the probe task is still queued after this " + "long. Default 0, meaning wait as long as the scout may run", ) def scout_options(self, args): + device = getattr(args, "braket_device", None) or SV1 + shots = getattr(args, "braket_shots", None) + timeout = getattr(args, "braket_queue_timeout", None) return { - "device": getattr(args, "braket_device", None), - "region": getattr(args, "braket_region", None), - "shots": getattr(args, "braket_shots", None), + "device": device, + "region": region_for(device, getattr(args, "braket_region", None)), + "shots": 1 if shots is None else int(shots), + "queue_timeout": 0 if timeout is None else float(timeout), + } + + def job_environment(self, options): + """The classical job may want to fetch the result, and it needs the same + region to find the task.""" + return { + "BRAKET_DEVICE": options.get("device") or SV1, + "AWS_DEFAULT_REGION": options.get("region") or "us-east-1", } def open_session(self, options): - # TODO validate the reservation ARN for the requested device - raise NotImplementedError( - "Braket open_session not yet implemented; use --quantum-vendor mock " - "with FLUX_QUANTUM_MOCK for now (options captured: {})".format(options) + """Submit the probe task and wait for it to reach the front of the queue. + + Returns the task ARN, which is what the classical job is handed. + """ + from braket.aws import AwsDevice + from braket.circuits import Circuit + + device_arn = options.get("device") or SV1 + region = options.get("region") or region_for(device_arn) + os.environ.setdefault("AWS_DEFAULT_REGION", region) + + device = AwsDevice(device_arn) + # identity on one qubit. Braket measures every qubit, so this is the + # smallest thing that still occupies the device. + self._task = device.run(Circuit().i(0), shots=int(options.get("shots") or 1)) + arn = self._task.id + + ok, reason = self.wait_for_priority( + timeout=float(options.get("queue_timeout") or 0) ) + if not ok: + raise RuntimeError( + "braket: probe task {} never reached the front of the queue " + "({})".format(arn, reason) + ) + return arn + + def wait_for_priority(self, timeout=0.0, interval=5.0, sleep=time.sleep): + """Poll until the probe task is next in line. + + Ready at position 1, and also once the task is RUNNING or COMPLETED, + because a task that has left the queue reports no position and the + value we are waiting for can never arrive. FAILED and CANCELLED are not + ready, they are a failure, so the caller can cancel the classical job + rather than start it against nothing. + + Returns (ok, reason). + """ + deadline = time.time() + timeout if timeout > 0 else None + while True: + state = self._task.state() + if state in FAILED: + return False, state + if state in RAN: + return True, state + pos = self.queue_position() + # positions over 2000 come back as the string >2000, so compare as + # a string rather than an int + if str(pos) == "1": + return True, "queue position 1" + print("braket: queued at position {}".format(pos), flush=True) + if deadline is not None and time.time() + interval >= deadline: + return False, "still queued at position {} after {}s".format( + pos, timeout + ) + sleep(interval) + + def queue_position(self): + """Position in the device queue, or None once the task has left it.""" + try: + return self._task.queue_position().queue_position + except Exception: + return None + + def close_session(self, session=None): + """Nothing to release. The probe task finishes or is already terminal.""" + return def credentials_present(self): - missing = [v for v in self.required_env if not os.environ.get(v)] - if missing: - return False, "Braket: missing env var(s): " + ", ".join(missing) - return True, "Braket: credentials present" + """Ask boto3 to resolve credentials the normal way, and say where they + came from. An EC2 instance role works but belongs to the node rather + than to you, so it is called out.""" + try: + import boto3 + except ImportError: + return ( + False, + "braket: boto3 is not installed, pip install amazon-braket-sdk", + ) + creds = boto3.Session().get_credentials() + if creds is None: + return False, ( + "braket: no AWS credentials. Set AWS_ACCESS_KEY_ID and " + "AWS_SECRET_ACCESS_KEY, or configure ~/.aws/credentials" + ) + method = getattr(creds, "method", "unknown") + if method in ("iam-role", "instance-metadata"): + return True, ( + "braket: using the EC2 instance role, which belongs to the node " + "and not to you. Export your own keys to keep the credential in " + "user space" + ) + return True, "braket: credentials present via {}".format(method) def probe(self): - # TODO braket device status, queue and pricing with the user creds - return Signals( - available=True, queue_depth=None, cost=None, detail={"note": "mock probe"} - ) + """Report the device status. Cheap, it is a metadata call.""" + ok, _ = self.credentials_present() + if not ok: + return Signals(available=False) + try: + from braket.aws import AwsDevice + + device = AwsDevice(SV1) + return Signals(available=device.status == "ONLINE", detail={"device": SV1}) + except Exception as e: + return Signals(available=False, detail={"error": str(e)}) diff --git a/flux_quantum/scout.py b/flux_quantum/scout.py index f26a3a0..1340013 100644 --- a/flux_quantum/scout.py +++ b/flux_quantum/scout.py @@ -38,6 +38,26 @@ def wait_for_job(handle, jobid, waiter=None): waiter(handle, jobid, "clean", raiseJobException=False) +def abort_held(handle, jobid, why, cancel=None): + """Cancel the held classical and exit. + + Anything that fails before the release leaves the job held with its + reservation and nothing on the way to free it, so cancel it rather than + leave nodes parked on work that will never start. + """ + if cancel is None: + from flux.job import cancel + try: + cancel(handle, jobid, why) + print("quantum-scout: cancelled held job {}".format(jobid), file=sys.stderr) + except Exception as e: + print( + "quantum-scout: WARNING could not cancel held job {}: {}".format(jobid, e), + file=sys.stderr, + ) + sys.exit("quantum-scout: {}".format(why)) + + def _install_signal_handlers(): """Unwind on SIGTERM/SIGINT so the session still gets closed.""" @@ -79,6 +99,7 @@ def main(): from flux.job import JobID jobid = int(JobID(args.job)) + h = flux.Flux() # --session bypasses the backend for testing backend = None @@ -93,23 +114,15 @@ def main(): ) try: session = backend.open_session(json.loads(args.options)) - except NotImplementedError as e: - sys.exit("quantum-scout: {}".format(e)) except Exception as e: - sys.exit( - "quantum-scout: opening {} session failed: {}".format(args.vendor, e) - ) - - h = flux.Flux() + abort_held(h, jobid, "opening {} session failed: {}".format(args.vendor, e)) # must land before the release, or the job could start with no session try: post_session(h, jobid, session) except Exception as e: - sys.exit( - "quantum-scout: could not post the session to job {}: {}".format( - args.job, e - ) + abort_held( + h, jobid, "could not post the session to job {}: {}".format(args.job, e) ) _install_signal_handlers() diff --git a/tests/unit/test_braket.py b/tests/unit/test_braket.py new file mode 100644 index 0000000..fe14c90 --- /dev/null +++ b/tests/unit/test_braket.py @@ -0,0 +1,180 @@ +"""Braket backend, driven against a fake SDK. + +Braket has no session to hold, so the scout submits a no-op task and waits for +it to reach the front of the queue. The interesting behaviour is what counts as +ready and what counts as a failure. +""" + +import sys +import types + +import pytest + +ARN = "arn:aws:braket:us-west-1:123:quantum-task/abc" + + +class FakeTask: + def __init__(self, states, positions): + self.id = ARN + self._states = list(states) + self._positions = list(positions) + self.state_calls = 0 + + def state(self): + self.state_calls += 1 + return self._states.pop(0) if self._states else "QUEUED" + + def queue_position(self): + pos = self._positions.pop(0) if self._positions else None + return types.SimpleNamespace(queue_position=pos) + + +@pytest.fixture +def fake_braket(monkeypatch): + """Minimal braket.aws and braket.circuits.""" + submitted = {} + + class FakeDevice: + status = "ONLINE" + + def __init__(self, arn): + submitted["device"] = arn + + def run(self, circuit, shots=1): + submitted["shots"] = shots + return submitted["task"] + + aws = types.ModuleType("braket.aws") + aws.AwsDevice = FakeDevice + circuits = types.ModuleType("braket.circuits") + + class FakeCircuit: + def i(self, q): + submitted["qubit"] = q + return self + + circuits.Circuit = FakeCircuit + pkg = types.ModuleType("braket") + monkeypatch.setitem(sys.modules, "braket", pkg) + monkeypatch.setitem(sys.modules, "braket.aws", aws) + monkeypatch.setitem(sys.modules, "braket.circuits", circuits) + return submitted + + +def _backend(): + from flux_quantum.backends import get_backend + import flux_quantum.backends.braket # noqa + + return get_backend("braket") + + +def test_defaults_to_the_sv1_simulator(): + """SV1 is cents per task, so it is the sane default for a probe.""" + from flux_quantum.backends.braket import SV1 + + class Args: + braket_device = None + braket_region = None + braket_shots = None + braket_queue_timeout = None + + opts = _backend().scout_options(Args()) + assert opts["device"] == SV1 + assert opts["shots"] == 1 + + +def test_region_comes_from_the_device_arn(monkeypatch): + """A QPU searched in the wrong region simply never appears.""" + from flux_quantum.backends.braket import region_for + + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") + assert region_for("arn:aws:braket:eu-north-1::device/qpu/iqm/x") == "eu-north-1" + # SV1 has an empty region field, so fall back + assert ( + region_for("arn:aws:braket:::device/quantum-simulator/amazon/sv1") + == "us-east-1" + ) + # an explicit override always wins + assert region_for("arn:aws:braket:eu-north-1::device/x", "us-west-2") == "us-west-2" + + +def test_ready_at_queue_position_one(fake_braket): + b = _backend() + b._task = FakeTask(states=["QUEUED", "QUEUED"], positions=["3", "1"]) + ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + assert ok and "position 1" in why + + +def test_ready_when_the_task_already_left_the_queue(fake_braket): + """A task that is RUNNING reports no position, so waiting for 1 would hang.""" + b = _backend() + b._task = FakeTask(states=["RUNNING"], positions=[None]) + ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + assert ok and why == "RUNNING" + + +def test_completed_is_also_ready(fake_braket): + """On SV1 a trivial task can finish before we ever observe position 1.""" + b = _backend() + b._task = FakeTask(states=["COMPLETED"], positions=[None]) + ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + assert ok and why == "COMPLETED" + + +def test_failed_is_not_ready(fake_braket): + """The classical must be cancelled, not started against a task with no result.""" + b = _backend() + b._task = FakeTask(states=["FAILED"], positions=[None]) + ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + assert not ok and why == "FAILED" + + +def test_cancelled_is_not_ready(fake_braket): + b = _backend() + b._task = FakeTask(states=["CANCELLED"], positions=[None]) + ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + assert not ok + + +def test_deep_queue_position_is_compared_as_a_string(fake_braket): + """Braket reports anything over 2000 as the string >2000.""" + b = _backend() + b._task = FakeTask(states=["QUEUED", "QUEUED"], positions=[">2000", "1"]) + ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + assert ok + + +def test_timeout_reports_the_position(fake_braket): + b = _backend() + b._task = FakeTask(states=["QUEUED"] * 5, positions=["7"] * 5) + ok, why = b.wait_for_priority(timeout=1, interval=1, sleep=lambda s: None) + assert not ok and "still queued at position 7" in why + + +def test_open_session_submits_a_one_qubit_no_op(fake_braket): + from flux_quantum.backends.braket import SV1 + + fake_braket["task"] = FakeTask(states=["QUEUED"], positions=["1"]) + b = _backend() + arn = b.open_session({"device": SV1, "shots": 1, "queue_timeout": 0}) + assert arn == ARN + assert fake_braket["device"] == SV1 + assert fake_braket["qubit"] == 0 # identity on qubit 0, a real no-op + assert fake_braket["shots"] == 1 + + +def test_open_session_raises_when_the_task_fails(fake_braket): + """So the scout cancels the held classical instead of releasing it.""" + fake_braket["task"] = FakeTask(states=["FAILED"], positions=[None]) + with pytest.raises(RuntimeError, match="never reached the front of the queue"): + _backend().open_session({"queue_timeout": 0}) + + +def test_close_session_is_a_no_op(fake_braket): + assert _backend().close_session(ARN) is None + + +def test_job_environment_carries_device_and_region(fake_braket): + env = _backend().job_environment({"device": "arn:x", "region": "eu-north-1"}) + assert env["BRAKET_DEVICE"] == "arn:x" + assert env["AWS_DEFAULT_REGION"] == "eu-north-1" diff --git a/tests/unit/test_handoff.py b/tests/unit/test_handoff.py index 5fda012..c15fb1f 100644 --- a/tests/unit/test_handoff.py +++ b/tests/unit/test_handoff.py @@ -145,3 +145,30 @@ def test_signal_handlers_unwind_so_the_session_is_closed(): assert closed == ["closed"] finally: signal.signal(signal.SIGTERM, previous) + + +def test_abort_cancels_the_held_job(): + """A failure before the release must not leave the classical held forever.""" + from flux_quantum.scout import abort_held + + cancelled = {} + + def fake_cancel(handle, jobid, why): + cancelled["jobid"] = jobid + cancelled["why"] = why + + with pytest.raises(SystemExit): + abort_held(None, 4021041664, "opening ibm session failed", cancel=fake_cancel) + assert cancelled["jobid"] == 4021041664 + assert "opening ibm session failed" in cancelled["why"] + + +def test_abort_still_exits_when_the_cancel_fails(): + """A broker we cannot reach must not turn into a hang.""" + from flux_quantum.scout import abort_held + + def boom(handle, jobid, why): + raise RuntimeError("no broker") + + with pytest.raises(SystemExit): + abort_held(None, 1, "whatever", cancel=boom) From b7707f3e824588014b03b065f041f9f60055bd82 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 2 Sep 2026 19:59:02 -0700 Subject: [PATCH 05/26] test mock queue Signed-off-by: vsoch --- flux_quantum/backends/base.py | 13 ++++ flux_quantum/backends/braket.py | 61 +++++++++++------- flux_quantum/backends/mock.py | 84 ++++++++++++++++++++++++ flux_quantum/backends/qrmi.py | 41 ++++-------- flux_quantum/scout.py | 17 ++++- tests/unit/test_braket.py | 26 ++++---- tests/unit/test_mock_queue.py | 111 ++++++++++++++++++++++++++++++++ tests/unit/test_qrmi.py | 39 ++++++----- 8 files changed, 307 insertions(+), 85 deletions(-) create mode 100644 tests/unit/test_mock_queue.py diff --git a/flux_quantum/backends/base.py b/flux_quantum/backends/base.py index 54f701a..59fd6c6 100644 --- a/flux_quantum/backends/base.py +++ b/flux_quantum/backends/base.py @@ -62,6 +62,19 @@ def job_environment(self, options): QPU it has. Defaults to none.""" return {} + def wait_for_priority(self, options=None): + """Block until the vendor is actually ours, then return (ok, reason). + + Opening a session is not the same as having the device. IBM activates a + session when its first task reaches the head of the queue, and Braket + reports a queue position and never holds anything. So this is a separate + step from open_session, and the scout does not release the classical job + until it returns ok. + + Defaults to ready, for vendors with no queue to wait on. + """ + return True, "not applicable" + def close_session(self, session=None): """Release the session opened by open_session. diff --git a/flux_quantum/backends/braket.py b/flux_quantum/backends/braket.py index a10343d..da536e2 100644 --- a/flux_quantum/backends/braket.py +++ b/flux_quantum/backends/braket.py @@ -29,6 +29,18 @@ FAILED = ("FAILED", "CANCELLED") +def position_at_most(pos, threshold): + """True when the queue position is at or under threshold. + + Braket reports anything over 2000 as the string >2000, so a position that + is not a number counts as far away. + """ + try: + return int(pos) <= int(threshold) + except (TypeError, ValueError): + return False + + def region_for(device_arn, override=None): """Region to talk to. Devices are not all in one, and a QPU searched in the wrong region simply never appears.""" @@ -70,6 +82,14 @@ def add_options(cls, add_option): default=None, help="Braket: shots for the queue probe task, default 1", ) + add_option( + "--braket-ungate-position", + metavar="N", + default=None, + help="Braket: release the classical job once the probe task is at " + "this queue position or closer. Default 1, next in line. Raise it " + "when the classical job is slow to start", + ) add_option( "--braket-queue-timeout", metavar="SECONDS", @@ -87,6 +107,10 @@ def scout_options(self, args): "region": region_for(device, getattr(args, "braket_region", None)), "shots": 1 if shots is None else int(shots), "queue_timeout": 0 if timeout is None else float(timeout), + "ungate_position": int( + getattr(args, "braket_ungate_position", None) + or os.environ.get("QUANTUM_BRAKET_UNGATE_POSITION", 1) + ), } def job_environment(self, options): @@ -98,10 +122,8 @@ def job_environment(self, options): } def open_session(self, options): - """Submit the probe task and wait for it to reach the front of the queue. - - Returns the task ARN, which is what the classical job is handed. - """ + """Submit the probe task and return its ARN, which is what the + classical job is handed. The scout calls wait_for_priority next.""" from braket.aws import AwsDevice from braket.circuits import Circuit @@ -113,29 +135,22 @@ def open_session(self, options): # identity on one qubit. Braket measures every qubit, so this is the # smallest thing that still occupies the device. self._task = device.run(Circuit().i(0), shots=int(options.get("shots") or 1)) - arn = self._task.id - - ok, reason = self.wait_for_priority( - timeout=float(options.get("queue_timeout") or 0) - ) - if not ok: - raise RuntimeError( - "braket: probe task {} never reached the front of the queue " - "({})".format(arn, reason) - ) - return arn + return self._task.id - def wait_for_priority(self, timeout=0.0, interval=5.0, sleep=time.sleep): + def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep): """Poll until the probe task is next in line. - Ready at position 1, and also once the task is RUNNING or COMPLETED, - because a task that has left the queue reports no position and the - value we are waiting for can never arrive. FAILED and CANCELLED are not - ready, they are a failure, so the caller can cancel the classical job - rather than start it against nothing. + Ready at the requested position or closer, and also once the task is + RUNNING or COMPLETED, because a task that left the queue reports no + position and the value we wait for can never arrive. FAILED and + CANCELLED are a failure and not readiness, so the caller can cancel the + classical job rather than start it against nothing. Returns (ok, reason). """ + options = options or {} + timeout = float(options.get("queue_timeout") or 0) + position = int(options.get("ungate_position") or 1) deadline = time.time() + timeout if timeout > 0 else None while True: state = self._task.state() @@ -146,8 +161,8 @@ def wait_for_priority(self, timeout=0.0, interval=5.0, sleep=time.sleep): pos = self.queue_position() # positions over 2000 come back as the string >2000, so compare as # a string rather than an int - if str(pos) == "1": - return True, "queue position 1" + if position_at_most(pos, position): + return True, "queue position {}".format(pos) print("braket: queued at position {}".format(pos), flush=True) if deadline is not None and time.time() + interval >= deadline: return False, "still queued at position {} after {}s".format( diff --git a/flux_quantum/backends/mock.py b/flux_quantum/backends/mock.py index c4d6de6..e89cf10 100644 --- a/flux_quantum/backends/mock.py +++ b/flux_quantum/backends/mock.py @@ -7,9 +7,19 @@ mock_busy queue_depth 9, so --select queue prefers mock FLUX_QUANTUM_MOCK_QUEUE and FLUX_QUANTUM_MOCK_COST override the mock vendor. + +The mock also simulates a vendor queue, which is the only way to control queue +depth as an experiment variable. No real vendor lets you do that. The wait is + + base_overhead + depth * service_time + +where base_overhead is the fixed cost of getting a task onto the device even +with nothing ahead of it. On ibm_marrakesh at depth 0 that measured 10 to 12 +seconds, hence the default. """ import os +import random import time from .base import Backend, Signals, register @@ -33,13 +43,87 @@ def add_options(cls, add_option): default=None, help="mock: delay this many seconds before opening a session", ) + add_option( + "--mock-queue-depth", + metavar="N", + default=None, + help="mock: tasks ahead of us in the simulated vendor queue", + ) + add_option( + "--mock-service-time", + metavar="SECONDS", + default=None, + help="mock: seconds per task ahead of us, default 0.1 so a deep " + "queue still finishes in a runnable time", + ) + add_option( + "--mock-base-overhead", + metavar="SECONDS", + default=None, + help="mock: fixed wait even at depth 0, default 10 which is what " + "ibm_marrakesh measured", + ) + add_option( + "--mock-jitter", + metavar="FRACTION", + default=None, + help="mock: multiplicative noise on the wait, 0 for a " + "deterministic run which is what an experiment wants", + ) + add_option( + "--mock-seed", + metavar="N", + default=None, + help="mock: seed for the jitter, so a run can be repeated", + ) def scout_options(self, args): + def num(name, default): + v = getattr(args, name, None) + return default if v is None else float(v) + return { "session": getattr(args, "mock_session", None), "latency": getattr(args, "mock_latency", None), + "queue_depth": int(num("mock_queue_depth", 0)), + "service_time": num("mock_service_time", 0.1), + "base_overhead": num("mock_base_overhead", 10.0), + "jitter": num("mock_jitter", 0.0), + "seed": getattr(args, "mock_seed", None), } + def queue_wait(self, options): + """Seconds the simulated queue makes us wait.""" + depth = int(options.get("queue_depth") or 0) + wait = float(options.get("base_overhead") or 0) + depth * float( + options.get("service_time") or 0 + ) + jitter = float(options.get("jitter") or 0) + if jitter: + rng = random.Random(options.get("seed")) + wait *= 1.0 + rng.uniform(-jitter, jitter) + return max(0.0, wait) + + def wait_for_priority(self, options=None, sleep=time.sleep, now=time.time): + """Drain the simulated queue, reporting position as it goes. + + Reports like a real vendor so the experiment can read the same series + out of the scout output whichever backend it ran against. + """ + options = options or {} + depth = int(options.get("queue_depth") or 0) + wait = self.queue_wait(options) + started = now() + per = wait / depth if depth else 0.0 + for ahead in range(depth, 0, -1): + print("mock: queued at position {}".format(ahead), flush=True) + sleep(per) + remaining = wait - (now() - started) + if remaining > 0: + sleep(remaining) + self._waited = wait + return True, "queue drained after {:.1f}s at depth {}".format(wait, depth) + def open_session(self, options): latency = options.get("latency") if latency: diff --git a/flux_quantum/backends/qrmi.py b/flux_quantum/backends/qrmi.py index d4f0735..160c2b1 100644 --- a/flux_quantum/backends/qrmi.py +++ b/flux_quantum/backends/qrmi.py @@ -108,7 +108,8 @@ def warmup_payload(shots=1): except ImportError: raise RuntimeError( "qiskit is needed to build the warmup task. Install it with " - "pip install 'qrmi[ibm]', or skip the priority check" + "pip install 'qrmi[ibm]', or pass --quantum-ibm-skip-warmup to " + "release without confirming priority" ) qc = QuantumCircuit(1, 1) @@ -280,28 +281,6 @@ def open_session(self, options): ) ) - # Reaching the head of the queue is what gives us priority, so wait for - # it before letting the classical job start. - if not options.get("skip_warmup"): - try: - ok, task, reason = self.wait_for_priority( - timeout=float(options.get("warmup_timeout") or 0) - ) - except Exception as e: - self.close_session() - raise RuntimeError( - "{}: session opened on {} but the warmup task could not " - "run, so priority is unconfirmed: {}\n pass " - "--quantum-{}-skip-warmup to release without this " - "check".format(self.name, resource, e, self.name) - ) - if not ok: - self.close_session() - raise RuntimeError( - "{}: session opened on {} but the warmup task never " - "reached the head of the queue ({}), so we do not have " - "priority".format(self.name, resource, reason) - ) return str(self._lock) def wait_until_ready(self, timeout=0.0, interval=5.0, sleep=time.sleep): @@ -329,7 +308,7 @@ def wait_until_ready(self, timeout=0.0, interval=5.0, sleep=time.sleep): return False, last sleep(interval) - def wait_for_priority(self, timeout=0.0, interval=5.0, sleep=time.sleep, shots=1): + def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep, shots=1): """Submit a one shot task and wait for it to leave the queue. An IBM session activates when its first task reaches the head of the @@ -338,11 +317,15 @@ def wait_for_priority(self, timeout=0.0, interval=5.0, sleep=time.sleep, shots=1 This is the only check available that means priority rather than just reachability. - Returns (ok, task_id, reason). timeout of 0 waits as long as the scout - job is allowed to live. + Returns (ok, reason). A timeout of 0 waits as long as the scout job is + allowed to live. """ from qrmi import Payload, TaskStatus + options = options or {} + if options.get("skip_warmup"): + return True, "warmup skipped, priority unconfirmed" + timeout = float(options.get("warmup_timeout") or 0) payload = Payload.QiskitPrimitive( input=json.dumps(warmup_payload(shots)), program_id="sampler" ) @@ -351,15 +334,15 @@ def wait_for_priority(self, timeout=0.0, interval=5.0, sleep=time.sleep, shots=1 while True: status = self._resource.task_status(task) if status in (TaskStatus.Running, TaskStatus.Completed): - return True, task, str(status) + return True, "warmup task {} is {}".format(task, status) if status in (TaskStatus.Failed, TaskStatus.Cancelled): - return False, task, str(status) + return False, "warmup task {} is {}".format(task, status) if deadline is not None and time.time() + interval >= deadline: try: self._resource.task_stop(task) except Exception: pass - return False, task, "still queued after {}s".format(timeout) + return False, "still queued after {}s".format(timeout) sleep(interval) def close_session(self, session=None): diff --git a/flux_quantum/scout.py b/flux_quantum/scout.py index 1340013..7368473 100644 --- a/flux_quantum/scout.py +++ b/flux_quantum/scout.py @@ -112,11 +112,26 @@ def main(): "quantum-scout: no backend registered for vendor {}, " "set FLUX_QUANTUM_MOCK for the mock vendor".format(args.vendor) ) + opts = json.loads(args.options) try: - session = backend.open_session(json.loads(args.options)) + session = backend.open_session(opts) except Exception as e: abort_held(h, jobid, "opening {} session failed: {}".format(args.vendor, e)) + # opening is not the same as having the device, so wait until it is ours + # before letting the classical job start + try: + ok, why = backend.wait_for_priority(opts) + except Exception as e: + backend.close_session(session) + abort_held( + h, jobid, "waiting for {} priority failed: {}".format(args.vendor, e) + ) + if not ok: + backend.close_session(session) + abort_held(h, jobid, "{} never became ours: {}".format(args.vendor, why)) + print("quantum-scout: {}".format(why)) + # must land before the release, or the job could start with no session try: post_session(h, jobid, session) diff --git a/tests/unit/test_braket.py b/tests/unit/test_braket.py index fe14c90..68d9b92 100644 --- a/tests/unit/test_braket.py +++ b/tests/unit/test_braket.py @@ -101,7 +101,7 @@ def test_region_comes_from_the_device_arn(monkeypatch): def test_ready_at_queue_position_one(fake_braket): b = _backend() b._task = FakeTask(states=["QUEUED", "QUEUED"], positions=["3", "1"]) - ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda s: None) assert ok and "position 1" in why @@ -109,7 +109,7 @@ def test_ready_when_the_task_already_left_the_queue(fake_braket): """A task that is RUNNING reports no position, so waiting for 1 would hang.""" b = _backend() b._task = FakeTask(states=["RUNNING"], positions=[None]) - ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda s: None) assert ok and why == "RUNNING" @@ -117,7 +117,7 @@ def test_completed_is_also_ready(fake_braket): """On SV1 a trivial task can finish before we ever observe position 1.""" b = _backend() b._task = FakeTask(states=["COMPLETED"], positions=[None]) - ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda s: None) assert ok and why == "COMPLETED" @@ -125,14 +125,14 @@ def test_failed_is_not_ready(fake_braket): """The classical must be cancelled, not started against a task with no result.""" b = _backend() b._task = FakeTask(states=["FAILED"], positions=[None]) - ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda s: None) assert not ok and why == "FAILED" def test_cancelled_is_not_ready(fake_braket): b = _backend() b._task = FakeTask(states=["CANCELLED"], positions=[None]) - ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda s: None) assert not ok @@ -140,14 +140,16 @@ def test_deep_queue_position_is_compared_as_a_string(fake_braket): """Braket reports anything over 2000 as the string >2000.""" b = _backend() b._task = FakeTask(states=["QUEUED", "QUEUED"], positions=[">2000", "1"]) - ok, why = b.wait_for_priority(interval=0, sleep=lambda s: None) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda s: None) assert ok def test_timeout_reports_the_position(fake_braket): b = _backend() b._task = FakeTask(states=["QUEUED"] * 5, positions=["7"] * 5) - ok, why = b.wait_for_priority(timeout=1, interval=1, sleep=lambda s: None) + ok, why = b.wait_for_priority( + {"queue_timeout": 1}, interval=1, sleep=lambda s: None + ) assert not ok and "still queued at position 7" in why @@ -163,11 +165,13 @@ def test_open_session_submits_a_one_qubit_no_op(fake_braket): assert fake_braket["shots"] == 1 -def test_open_session_raises_when_the_task_fails(fake_braket): - """So the scout cancels the held classical instead of releasing it.""" +def test_failed_task_is_reported_not_raised(fake_braket): + """The scout cancels the held classical on a False from here.""" fake_braket["task"] = FakeTask(states=["FAILED"], positions=[None]) - with pytest.raises(RuntimeError, match="never reached the front of the queue"): - _backend().open_session({"queue_timeout": 0}) + b = _backend() + b.open_session({"queue_timeout": 0}) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda s: None) + assert not ok and why == "FAILED" def test_close_session_is_a_no_op(fake_braket): diff --git a/tests/unit/test_mock_queue.py b/tests/unit/test_mock_queue.py new file mode 100644 index 0000000..f818d4e --- /dev/null +++ b/tests/unit/test_mock_queue.py @@ -0,0 +1,111 @@ +"""The simulated vendor queue. + +This is the experiment instrument, so the wait has to be a controlled function +of depth rather than something approximate. No real vendor lets you set queue +depth, which is why the measurements come from here. +""" + +import pytest + + +@pytest.fixture +def mock(monkeypatch): + monkeypatch.setenv("FLUX_QUANTUM_MOCK", "1") + import importlib + from flux_quantum.backends import mock as mockmod + + importlib.reload(mockmod) + return mockmod.MockBackend() + + +class Args: + mock_session = None + mock_latency = None + mock_queue_depth = None + mock_service_time = None + mock_base_overhead = None + mock_jitter = None + mock_seed = None + + +def test_defaults_match_what_ibm_measured(mock): + """ibm_marrakesh at depth 0 took 10 to 12 seconds to dequeue, so the fixed + overhead is not invented.""" + opts = mock.scout_options(Args()) + assert opts["base_overhead"] == 10.0 + assert opts["queue_depth"] == 0 + assert opts["service_time"] == 0.1 + assert opts["jitter"] == 0.0 + + +def test_wait_is_linear_in_depth(mock): + base = {"base_overhead": 10.0, "service_time": 0.5, "jitter": 0.0} + assert mock.queue_wait(dict(base, queue_depth=0)) == 10.0 + assert mock.queue_wait(dict(base, queue_depth=10)) == 15.0 + assert mock.queue_wait(dict(base, queue_depth=100)) == 60.0 + + +def test_deterministic_without_jitter(mock): + opts = {"queue_depth": 7, "service_time": 1.0, "base_overhead": 2.0} + assert len({mock.queue_wait(opts) for _ in range(5)}) == 1 + + +def test_jitter_is_repeatable_with_a_seed(mock): + a = { + "queue_depth": 10, + "service_time": 1.0, + "base_overhead": 0.0, + "jitter": 0.5, + "seed": 7, + } + b = dict(a, seed=8) + assert mock.queue_wait(a) == mock.queue_wait(a) + assert mock.queue_wait(a) != mock.queue_wait(b) + + +def test_priority_wait_sleeps_the_computed_time(mock): + """The scout blocks here while the classical sits held, which is the whole + point, so the total slept has to be the queue wait.""" + slept = [] + clock = [0.0] + + def fake_sleep(s): + slept.append(s) + clock[0] += s + + ok, why = mock.wait_for_priority( + {"queue_depth": 5, "service_time": 2.0, "base_overhead": 0.0}, + sleep=fake_sleep, + now=lambda: clock[0], + ) + assert ok + assert sum(slept) == pytest.approx(10.0) + assert "depth 5" in why + + +def test_depth_zero_still_waits_the_overhead(mock): + slept = [] + clock = [0.0] + + def fake_sleep(s): + slept.append(s) + clock[0] += s + + ok, why = mock.wait_for_priority( + {"queue_depth": 0, "base_overhead": 10.0}, + sleep=fake_sleep, + now=lambda: clock[0], + ) + assert ok and sum(slept) == pytest.approx(10.0) + + +def test_position_is_reported_while_draining(mock, capsys): + """The scout output carries the same position series a real vendor gives, + so the experiment reads one format either way.""" + mock.wait_for_priority( + {"queue_depth": 3, "service_time": 0.0, "base_overhead": 0.0}, + sleep=lambda s: None, + now=lambda: 0.0, + ) + out = capsys.readouterr().out + assert "position 3" in out and "position 1" in out diff --git a/tests/unit/test_qrmi.py b/tests/unit/test_qrmi.py index 230c67d..b59e527 100644 --- a/tests/unit/test_qrmi.py +++ b/tests/unit/test_qrmi.py @@ -245,8 +245,8 @@ def test_warmup_waits_until_the_task_leaves_the_queue(fake_qrmi, monkeypatch): ) b = _acquired(fake_qrmi, ["Queued", "Queued", "Running"]) slept = [] - ok, task, reason = b.wait_for_priority(interval=1, sleep=slept.append) - assert ok and reason == "Running" + ok, reason = b.wait_for_priority({}, interval=1, sleep=slept.append) + assert ok and "Running" in reason assert len(slept) == 2 @@ -255,8 +255,8 @@ def test_warmup_failure_means_no_priority(fake_qrmi, monkeypatch): "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} ) b = _acquired(fake_qrmi, ["Queued", "Failed"]) - ok, task, reason = b.wait_for_priority(interval=0, sleep=lambda s: None) - assert not ok and reason == "Failed" + ok, reason = b.wait_for_priority({}, interval=0, sleep=lambda s: None) + assert not ok and "Failed" in reason def test_warmup_timeout_stops_the_task(fake_qrmi, monkeypatch): @@ -264,30 +264,27 @@ def test_warmup_timeout_stops_the_task(fake_qrmi, monkeypatch): "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} ) b = _acquired(fake_qrmi, ["Queued"] * 10) - ok, task, reason = b.wait_for_priority(timeout=1, interval=1, sleep=lambda s: None) + ok, reason = b.wait_for_priority( + {"warmup_timeout": 1}, interval=1, sleep=lambda s: None + ) assert not ok and "still queued" in reason assert ("task_stop", "TASK-1") in fake_qrmi.calls -def test_session_released_when_priority_is_never_reached(fake_qrmi, monkeypatch): - """Never release the classical job against a session we do not own.""" +def test_priority_not_reached_is_reported(fake_qrmi, monkeypatch): + """The scout cancels the held job on a False here, so never say ok.""" monkeypatch.setattr( "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} ) - fake_qrmi.statuses = ["Failed"] - b = _ibm() - with pytest.raises(RuntimeError, match="do not have priority"): - b.open_session(_opts(ready_timeout=0)) - assert ("release", RESOURCE, "LOCK-123") in fake_qrmi.calls + b = _acquired(fake_qrmi, ["Failed"]) + ok, reason = b.wait_for_priority({}, interval=0, sleep=lambda s: None) + assert not ok and "Failed" in reason -def test_skip_warmup_releases_immediately(fake_qrmi, monkeypatch): - monkeypatch.setattr( - "flux_quantum.backends.qrmi.warmup_payload", lambda shots=1: {"pubs": []} - ) - fake_qrmi.statuses = [] - b = _ibm() - assert b.open_session(_opts(ready_timeout=0, skip_warmup=True)) == "LOCK-123" +def test_skip_warmup_submits_nothing(fake_qrmi): + b = _acquired(fake_qrmi, []) + ok, reason = b.wait_for_priority({"skip_warmup": True}) + assert ok and "skipped" in reason assert not any(c[0] == "task_start" for c in fake_qrmi.calls) @@ -313,9 +310,9 @@ def no_qiskit(name, *a, **k): raise ImportError("no qiskit") return real(name, *a, **k) + b = _acquired(fake_qrmi, ["Queued"]) monkeypatch.setattr(builtins, "__import__", no_qiskit) - b = _ibm() with pytest.raises(RuntimeError) as e: - b.open_session(_opts(ready_timeout=0)) + b.wait_for_priority({}) assert "qrmi[ibm]" in str(e.value) assert "skip-warmup" in str(e.value) From 33c82e8abee4e22e25dcd67d00b0b888d607de3d Mon Sep 17 00:00:00 2001 From: vsoch Date: Thu, 3 Sep 2026 11:41:06 -0700 Subject: [PATCH 06/26] feat: jobtap plugin Signed-off-by: vsoch --- .github/workflows/test.yml | 10 + flux_quantum/cli.py | 6 +- flux_quantum/jobtap/quantum.c | 562 +++++++++++++++++++++++++-- flux_quantum/qresource.py | 23 ++ tests/integration/test-admission.sh | 149 +++++++ tests/integration/test-preemption.sh | 94 +++++ tests/unit/test_cli_plugin.py | 29 ++ tests/unit/test_qresource.py | 75 ++++ 8 files changed, 924 insertions(+), 24 deletions(-) create mode 100755 tests/integration/test-admission.sh create mode 100755 tests/integration/test-preemption.sh diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 769ab71..5896107 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -111,6 +111,16 @@ 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 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 start bash tests/integration/test-preemption.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 diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index 632d272..56186d3 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -21,7 +21,7 @@ import flux from flux.cli.plugin import CLIPlugin -from . import graph +from . import graph, qresource from .selector import select_vendor, SelectionError, discover_registry_vendors from .backends import get_backend, backend_classes from .launch import build_scout_jobspec @@ -91,6 +91,10 @@ def prepare_pair( sysattr["hold"] = 1 quantum = sysattr.setdefault("quantum", {}) quantum["vendor"] = vendor + # the jobtap plugin keeps a core budget over every unfinished pair, and it + # needs the classical size to do that. Counting here rather than walking the + # jobspec in C. + quantum["cores"] = qresource.count_cores(classical) if job_env: sysattr.setdefault("environment", {}).update(job_env) diff --git a/flux_quantum/jobtap/quantum.c b/flux_quantum/jobtap/quantum.c index 4d21c86..380e804 100644 --- a/flux_quantum/jobtap/quantum.c +++ b/flux_quantum/jobtap/quantum.c @@ -1,26 +1,69 @@ /* - * quantum-jobtap: a Flux job-manager jobtap plugin (policy layer) for - * quantum + classical coscheduling. + * quantum-jobtap: site policy for quantum and classical coscheduling. * - * External-quota model: a quantum job is identified by the attribute + * A quantum job is the held classical half of a pair, identified by * attributes.system.quantum.vendor = "" - * (set by `flux quantum-submit`). This plugin enforces SITE POLICY on such - * jobs at job.validate: it permits only configured vendors and rejects others - * before submission completes. It consumes no resources and submits nothing. + * attributes.system.quantum.cores = + * both stamped by the submit plugin. The scout carries neither, so it is not + * counted twice. * - * Load: flux jobtap load /path/quantum-jobtap.so [vendors="ibm,braket"] - * List: flux jobtap list + * Three jobs enforced here. * - * Extension points (deliberately NOT enabled here; see NOTES): - * - central hold enforcement (set attributes.system.hold from the plugin) - * - owned-hardware resource scoping (rewrite qpu -> qvendor_ -> qpu) - * Both depend on which jobspec-update keys the job-manager permits, which must - * be confirmed on the target flux-core before relying on them. The reliable - * home for the hold today is `flux quantum-submit` (--setattr=system.hold=1), - * which is already proven. + * Vendors. Only configured vendors are permitted. + * + * Protection. Both halves of a pair are marked + * attributes.system.protected = "quantum" + * which exempts them from preemption. The key is generic on purpose, so a + * plugin protecting any other scarce external resource can reuse it. It is set + * here and never by the submitter, and a submission carrying it is rejected. + * + * Admission. Launching a scout commits real money the moment it acquires a + * vendor session, so the classical half has to be placeable. Every unfinished + * pair, held or running, reserves cores + 1 against a budget, and a new pair + * that would exceed the budget is rejected at submit rather than admitted and + * left to wait with a meter running. Held pairs count because a held pair is a + * promise. + * + * Accounting lives in job.new and not job.validate, because the job manager + * replays job.new for every active job on restart or plugin reload, so the + * budget rebuilds itself. job.destroy gives the cores back. + * + * Load: + * flux jobtap load quantum.so vendors="ibm,braket" \ + * total_cores=128 reserve_cores=32 + * + * total_cores cores available to quantum pairs. 0 disables admission + * control and leaves only the vendor check. + * reserve_cores held back for ordinary classical work. + * protect_types resource types that mark a job as needing external access, + * default qpu. Used to protect the scout, which carries no + * quantum attributes. + * preempt_after seconds to wait after the scout releases the classical + * before cancelling unprotected jobs to make room. 0, the + * default, never preempts and the pair just waits. + * + * Preemption. Admission promises that the classical half can be placed, but + * ordinary work may be sitting on the cores when the scout releases it. So the + * plugin watches for the session memo, which the scout posts immediately before + * releasing, waits preempt_after seconds, and if the job still has not started + * it cancels unprotected jobs until enough cores are freed. Youngest first, + * because flux cancels rather than requeues and the youngest job loses the + * least work. + * + * NOTES + * Cancelling is destructive. The victim loses whatever it had done, so + * preemption is off unless preempt_after is set. Protected jobs are never + * victims, and both halves of a pair are protected, so preemption cannot eat + * another pair. + * + * Freeing cores is not instant. Cancel, then epilog, then the scheduler + * notices. The classical job starts after that, not at the moment of the + * cancel, and the vendor session covers the gap. */ #include #include +#include +#include #include #include #include @@ -28,6 +71,199 @@ /* default vendor allowlist; over/replace via load-time config "vendors=" */ static char *g_vendors = NULL; /* comma-separated, e.g. "ibm,braket" */ +/* Jobs the plugin knows about, keyed by jobid as a string. + * id -> cores, protected, run + * A running unprotected job is a preemption candidate, and the run time orders + * them so the youngest goes first. + */ +static json_t *g_jobs = NULL; + +static double g_preempt_after = 0.0; /* 0 disables preemption */ + +static int g_total_cores = 0; /* 0 means admission control is off */ +static int g_reserve_cores = 0; +static int g_used_cores = 0; /* cores promised to unfinished pairs */ + +/* marks a job as counted, so job.destroy only gives back what was taken */ +static const char *AUX_KEY = "quantum::cores"; + +/* Marks a job as not preemptible. Deliberately not under quantum, so any + * policy plugin protecting a scarce external resource can use the same key and + * the preemption logic stays generic. The value is a reason, and preemption + * cares only that it is present. */ +static const char *PROTECT_KEY = "attributes.system.protected"; + +/* resource types that identify a job as needing external quantum access. The + * scout carries no quantum attributes, only this request, and a request cannot + * be faked into existence because fluxion has to match it against the graph. */ +static char *g_protect_types = NULL; /* comma separated, e.g. "qpu" */ + +static int type_protected (const char *type) +{ + const char *list = g_protect_types ? g_protect_types : "qpu"; + size_t tlen = strlen (type); + const char *p = list; + while (*p) { + const char *comma = strchr (p, ','); + size_t seg = comma ? (size_t)(comma - p) : strlen (p); + if (seg == tlen && strncmp (p, type, tlen) == 0) + return 1; + if (!comma) + break; + p = comma + 1; + } + return 0; +} + +/* walk a v1 resource tree looking for a protected type */ +static int tree_wants_protection (json_t *resources) +{ + size_t i; + json_t *entry; + + if (!json_is_array (resources)) + return 0; + json_array_foreach (resources, i, entry) { + const char *type = NULL; + json_t *with = NULL; + if (json_unpack (entry, "{s:s}", "type", &type) == 0 + && type + && type_protected (type)) + return 1; + if (json_unpack (entry, "{s:o}", "with", &with) == 0 + && tree_wants_protection (with)) + return 1; + } + return 0; +} + +/* True when the job is one half of a pair. The classical carries the vendor + * attribute, the scout carries the resource request. */ +static int needs_protection (json_t *jobspec, const char *vendor) +{ + json_t *resources = NULL; + + if (vendor) + return 1; + if (json_unpack (jobspec, "{s:o}", "resources", &resources) == 0) + return tree_wants_protection (resources); + return 0; +} + +/* Set by this plugin, never by the submitter. A user who sets it would make + * their own job unpreemptible, so it is rejected rather than stripped, because + * silently ignoring an attempted privilege grab hides it. */ +static int user_set_protection (json_t *jobspec) +{ + json_t *val = NULL; + (void) json_unpack (jobspec, "{s:{s:{s:o}}}", + "attributes", "system", "protected", &val); + return val ? 1 : 0; +} + +static const char *idkey (flux_jobid_t id, char *buf, size_t len) +{ + snprintf (buf, len, "%ju", (uintmax_t) id); + return buf; +} + +/* remember a job so it can be considered as a victim later */ +static void track (flux_jobid_t id, int cores, int protected_) +{ + char key[32]; + json_t *entry; + + if (!g_jobs) + return; + entry = json_pack ("{s:i s:b s:f}", + "cores", cores, "protected", protected_, "run", 0.0); + if (entry) + (void) json_object_set_new (g_jobs, idkey (id, key, sizeof (key)), entry); +} + +static void forget (flux_jobid_t id) +{ + char key[32]; + if (g_jobs) + (void) json_object_del (g_jobs, idkey (id, key, sizeof (key))); +} + +static json_t *lookup (flux_jobid_t id) +{ + char key[32]; + if (!g_jobs) + return NULL; + return json_object_get (g_jobs, idkey (id, key, sizeof (key))); +} + +static int budget (void) +{ + int b = g_total_cores - g_reserve_cores; + return b > 0 ? b : 0; +} + +/* cores this pair reserves. The scout needs one of its own, so a classical of + * n cores makes the pair n + 1. */ +static int pair_cores (json_t *jobspec) +{ + json_int_t cores = 0; + if (json_unpack (jobspec, "{s:{s:{s:{s:I}}}}", + "attributes", "system", "quantum", "cores", &cores) < 0) + return 0; + if (cores < 0) + return 0; + return (int)cores + 1; +} + +/* cores an ordinary job asks for. The submit plugin stamps the count for our + * own pairs, but a victim candidate is any job, so count its tree. */ +static int count_tree_cores (json_t *resources, int factor); + +static int count_tree_cores_top (json_t *jobspec) +{ + json_t *resources = NULL; + if (json_unpack (jobspec, "{s:o}", "resources", &resources) < 0) + return 0; + return count_tree_cores (resources, 1); +} + +static int count_tree_cores (json_t *resources, int factor) +{ + size_t i; + json_t *entry; + int total = 0; + + if (!json_is_array (resources)) + return 0; + json_array_foreach (resources, i, entry) { + const char *type = NULL; + json_t *with = NULL; + json_int_t count = 1; + int n; + + (void) json_unpack (entry, "{s?I}", "count", &count); + if (count < 1) + count = 1; + n = factor * (int) count; + if (json_unpack (entry, "{s:s}", "type", &type) == 0 + && type + && strcmp (type, "core") == 0) + total += n; + if (json_unpack (entry, "{s:o}", "with", &with) == 0) + total += count_tree_cores (with, n); + } + return total; +} + +/* the vendor attribute is what makes a job one of ours */ +static const char *job_vendor (json_t *jobspec) +{ + const char *vendor = NULL; + (void) json_unpack (jobspec, "{s:{s:{s:{s:s}}}}", + "attributes", "system", "quantum", "vendor", &vendor); + return vendor; +} + static int vendor_allowed (const char *vendor) { const char *list = g_vendors ? g_vendors : "ibm,braket"; @@ -57,9 +293,26 @@ static int validate_cb (flux_plugin_t *p, "{s:o}", "jobspec", &jobspec) < 0) return 0; /* can't read jobspec; do not block the job */ + if (user_set_protection (jobspec)) + return flux_jobtap_reject_job (p, args, + "attributes.system.protected is set by the scheduler and " + "must not be set on a submission. Remove it. It marks a job " + "as not preemptible, which is not yours to decide"); + /* attributes.system.quantum.vendor -- absent => not a quantum job */ (void) json_unpack (jobspec, "{s:{s:{s:{s:s}}}}", "attributes", "system", "quantum", "vendor", &vendor); + + /* the scout has no quantum attributes, only the resource request, and it + * must be protected too or preempting it would leak a vendor session */ + if (needs_protection (jobspec, vendor)) { + if (flux_jobtap_jobspec_update_pack (p, "{s:s}", + PROTECT_KEY, "quantum") < 0) + flux_log (flux_jobtap_get_flux (p), LOG_ERR, + "quantum: could not mark the job protected, it would be " + "preemptible"); + } + if (!vendor) return 0; @@ -67,24 +320,287 @@ static int validate_cb (flux_plugin_t *p, return flux_jobtap_reject_job (p, args, "quantum: vendor '%s' is not permitted at this site", vendor); - /* policy satisfied; the job proceeds. The hold is set at submission by - * flux quantum-submit, and the scout unholds it once the session is live. */ + if (g_total_cores > 0) { + int want = pair_cores (jobspec); + if (want <= 0) + return flux_jobtap_reject_job (p, args, + "quantum: attributes.system.quantum.cores is missing, " + "so this pair cannot be accounted for"); + if (g_used_cores + want > budget ()) + return flux_jobtap_reject_job (p, args, + "quantum: no room for another pair. %d of %d cores are " + "promised to unfinished pairs and this one needs %d. " + "Wait for one to finish or ask for fewer cores", + g_used_cores, budget (), want); + } + + /* policy satisfied; the job proceeds. The hold is set at submission by the + * submit plugin, and the scout releases it once the session is live. */ return 0; } +/* Cancel unprotected running jobs, youngest first, until need cores are freed. + * Returns the cores freed, which may be less than asked for if there is not + * enough unprotected work to take. */ +static int preempt_for (flux_plugin_t *p, flux_jobid_t sparing, int need) +{ + flux_t *h = flux_jobtap_get_flux (p); + int freed = 0; + + while (freed < need) { + const char *key, *chosen = NULL; + json_t *entry; + double newest = -1.0; + int cores = 0; + + json_object_foreach (g_jobs, key, entry) { + int prot = 0, c = 0; + double run = 0.0; + char spare[32]; + + if (json_unpack (entry, "{s:i s:b s:F}", + "cores", &c, "protected", &prot, "run", &run) < 0) + continue; + if (prot || run <= 0.0 || c <= 0) + continue; /* protected, or not running, so not a candidate */ + if (strcmp (key, idkey (sparing, spare, sizeof (spare))) == 0) + continue; + if (run > newest) { + newest = run; + chosen = key; + cores = c; + } + } + if (!chosen) + break; /* nothing left to take */ + + if (flux_jobtap_raise_exception (p, + (flux_jobid_t) strtoull (chosen, NULL, 10), + "preempt", + 0, + "preempted to free %d cores for a " + "quantum job whose session is already " + "open", need) < 0) { + flux_log_error (h, "quantum: could not preempt job %s", chosen); + break; + } + flux_log (h, LOG_WARNING, + "quantum: preempted job %s for %d cores", chosen, cores); + freed += cores; + /* it is going away, so do not pick it again on the next pass */ + (void) json_object_del (g_jobs, chosen); + } + return freed; +} + +struct grace { + flux_plugin_t *p; + flux_jobid_t id; + int cores; + flux_watcher_t *w; +}; + +static void grace_cb (flux_reactor_t *r, + flux_watcher_t *w, + int revents, + void *arg) +{ + struct grace *g = arg; + json_t *entry = lookup (g->id); + double run = 0.0; + + /* started on its own, so there is nothing to do */ + if (entry + && json_unpack (entry, "{s:F}", "run", &run) == 0 + && run > 0.0) + goto done; + if (!entry) + goto done; /* finished or vanished */ + + flux_log (flux_jobtap_get_flux (g->p), LOG_WARNING, + "quantum: job %ju still not started after %.0fs, preempting for " + "%d cores", (uintmax_t) g->id, g_preempt_after, g->cores); + (void) preempt_for (g->p, g->id, g->cores); +done: + flux_watcher_destroy (w); + free (g); +} + +/* The scout posts the session memo immediately before releasing the classical + * job, so the memo is the signal that the job should now be able to start. */ +static int event_cb (flux_plugin_t *p, + const char *topic, + flux_plugin_arg_t *args, + void *arg) +{ + flux_jobid_t id; + const char *name = NULL; + json_t *entry; + struct grace *g; + int cores = 0; + flux_t *h; + + if (g_preempt_after <= 0.0) + return 0; + if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, + "{s:I s:{s:s}}", + "id", &id, + "entry", "name", &name) < 0) + return 0; + if (!name || strcmp (name, "memo") != 0) + return 0; + if (!(entry = lookup (id)) + || json_unpack (entry, "{s:i}", "cores", &cores) < 0 + || cores <= 0) + return 0; + + h = flux_jobtap_get_flux (p); + if (!(g = calloc (1, sizeof (*g)))) + return 0; + g->p = p; + g->id = id; + /* the pair reserves cores + 1, but the scout already holds its own core */ + g->cores = cores - 1 > 0 ? cores - 1 : cores; + if (!(g->w = flux_timer_watcher_create (flux_get_reactor (h), + g_preempt_after, 0., + grace_cb, g))) { + free (g); + return 0; + } + flux_watcher_start (g->w); + return 0; +} + +/* Accounting. Called for new jobs and replayed for active ones on restart or + * plugin reload, which is how the budget survives both. */ +static int new_cb (flux_plugin_t *p, + const char *topic, + flux_plugin_arg_t *args, + void *arg) +{ + json_t *jobspec = NULL; + flux_jobid_t id; + const char *vendor; + int want; + + if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, + "{s:I s:o}", "id", &id, + "jobspec", &jobspec) < 0) + return 0; + vendor = job_vendor (jobspec); + + /* Track every job, not only ours. An unprotected running job is what a + * quantum pair preempts, so the plugin has to know they exist. Called + * again on replay, and json_object_set overwrites, so this is idempotent + * apart from losing the run time, which the state callback restores. */ + track (id, + vendor ? pair_cores (jobspec) : count_tree_cores_top (jobspec), + needs_protection (jobspec, vendor)); + if (g_preempt_after > 0.0) + (void) flux_jobtap_job_subscribe (p, id); + + if (g_total_cores <= 0) + return 0; + if (!vendor) + return 0; + if ((want = pair_cores (jobspec)) <= 0) + return 0; + + /* record what was taken, so destroy gives back the same amount even if the + * jobspec is no longer readable by then */ + if (flux_jobtap_job_aux_set (p, id, AUX_KEY, + (void *)(intptr_t)want, NULL) < 0) + return 0; + g_used_cores += want; + return 0; +} + +static int destroy_cb (flux_plugin_t *p, + const char *topic, + flux_plugin_arg_t *args, + void *arg) +{ + flux_jobid_t id; + intptr_t want; + + if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, + "{s:I}", "id", &id) < 0) + return 0; + /* a job rejected at validate was never counted, so there is no aux and + * nothing to give back */ + forget (id); + want = (intptr_t) flux_jobtap_job_aux_get (p, id, AUX_KEY); + if (want <= 0) + return 0; + g_used_cores -= (int)want; + if (g_used_cores < 0) + g_used_cores = 0; + return 0; +} + +/* Note when a job starts, so preemption can take the youngest first and so the + * grace timer can tell whether the classical job got going on its own. */ +static int run_cb (flux_plugin_t *p, + const char *topic, + flux_plugin_arg_t *args, + void *arg) +{ + flux_jobid_t id; + json_t *entry; + + if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, + "{s:I}", "id", &id) < 0) + return 0; + if ((entry = lookup (id))) + (void) json_object_set_new (entry, "run", + json_real (flux_reactor_now ( + flux_get_reactor ( + flux_jobtap_get_flux (p))))); + return 0; +} + +static const struct flux_plugin_handler handlers[] = { + { "job.validate", validate_cb, NULL }, + { "job.new", new_cb, NULL }, + { "job.state.run", run_cb, NULL }, + { "job.event.memo", event_cb, NULL }, + { "job.destroy", destroy_cb, NULL }, + { 0 }, +}; + int flux_plugin_init (flux_plugin_t *p) { const char *vendors = NULL; + const char *protect_types = NULL; + int total = 0, reserve = 0; + double preempt_after = 0.0; - if (flux_plugin_set_name (p, "quantum") < 0) + + if (!(g_jobs = json_object ())) return -1; - /* optional load-time config: vendors="ibm,braket" */ - if (flux_plugin_conf_unpack (p, "{s?s}", "vendors", &vendors) == 0 - && vendors) - g_vendors = strdup (vendors); + if (flux_plugin_conf_unpack (p, "{s?s s?s s?i s?i s?F}", + "vendors", &vendors, + "protect_types", &protect_types, + "total_cores", &total, + "reserve_cores", &reserve, + "preempt_after", &preempt_after) == 0) { + if (vendors) + g_vendors = strdup (vendors); + if (protect_types) + g_protect_types = strdup (protect_types); + g_total_cores = total; + g_reserve_cores = reserve; + g_preempt_after = preempt_after; + if (g_preempt_after > 0.0) + flux_log (flux_jobtap_get_flux (p), LOG_WARNING, + "quantum: preemption is on. An unprotected job may be " + "cancelled to make room for a pair whose session is " + "already open, and a cancelled job loses its work"); + } - return flux_plugin_add_handler (p, "job.validate", validate_cb, NULL); + /* registers the table and sets the plugin name in one call */ + return flux_plugin_register (p, "quantum", handlers); } /* vi: ts=4 sw=4 expandtab diff --git a/flux_quantum/qresource.py b/flux_quantum/qresource.py index 5fd6e15..8a9a48c 100644 --- a/flux_quantum/qresource.py +++ b/flux_quantum/qresource.py @@ -144,3 +144,26 @@ def classical_resource(live_graph, ncores=1, label="scout"): for t in reversed(intermediates): inner = {"type": t, "count": 1, "with": [inner]} return {"type": "node", "count": 1, "with": [inner]} + + +def count_cores(jobspec): + """Total cores a v1 jobspec asks for. + + Multiplies counts down the resource tree and sums the core leaves. The + jobtap plugin needs this to keep a core budget, and doing it here means it + is testable without a broker rather than parsed in C. + """ + + def walk(entries, factor): + total = 0 + for entry in entries or []: + count = entry.get("count", 1) + if isinstance(count, dict): # a range, take the minimum we must get + count = count.get("min", 1) + n = factor * int(count) + if entry.get("type") == "core": + total += n + total += walk(entry.get("with"), n) + return total + + return walk(jobspec.get("resources"), 1) diff --git a/tests/integration/test-admission.sh b/tests/integration/test-admission.sh new file mode 100755 index 0000000..ee75655 --- /dev/null +++ b/tests/integration/test-admission.sh @@ -0,0 +1,149 @@ +#!/bin/bash +# The jobtap plugin keeps a core budget over every unfinished pair, held or +# running, and rejects a pair that would exceed it. Launching a scout commits +# money, so the classical half has to be placeable. +# +# flux start -s1 bash tests/integration/test-admission.sh + +set -u +export FLUX_QUANTUM_MOCK=1 +HERE=$(cd "$(dirname "$0")/../.." && pwd) +export FLUX_CLI_PLUGINPATH="$HERE/cli-plugins" +rc=0 + +echo "=== build and load with a small budget ===" +make -s -C "$HERE/flux_quantum/jobtap" || { echo "FAIL build"; exit 1; } +# 20 cores for quantum, none held back, so a 4 core pair takes 5 and four fit +flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ + vendors="mock,ibm,braket" total_cores=20 reserve_cores=0 \ + protect_types="qpu" || { + echo "FAIL load"; exit 1; } +flux jobtap list + +echo "" +echo "=== a pair is accounted for as its cores plus one for the scout ===" +# submitted held, so it stays in the budget without running +ids="" +for i in 1 2 3 4; do + if id=$(flux submit --quantum-vendor mock -n4 sleep 300 2>&1 | tail -1); then + ids="$ids $id" + echo " pair $i admitted" + else + echo "FAIL pair $i should have been admitted, budget is 20 and 4 pairs need 20" + rc=1 + fi +done + +echo "" +echo "=== the fifth is rejected rather than admitted and left to wait ===" +if err=$(flux submit --quantum-vendor mock -n4 sleep 300 2>&1); then + echo "FAIL the fifth pair was admitted, the budget is not being enforced" + rc=1 +else + if echo "$err" | grep -q "no room for another pair"; then + echo " rejected, and the message says why" + echo "$err" | grep -o "no room for another pair.*" | head -1 | sed 's/^/ /' + else + echo "FAIL rejected for the wrong reason: $err" + rc=1 + fi +fi + +echo "" +echo "=== a pair that does not fit at all is rejected on its own ===" +if flux submit --quantum-vendor mock -n64 sleep 300 >/dev/null 2>&1; then + echo "FAIL a 64 core pair should not fit a 20 core budget"; rc=1 +else + echo " rejected" +fi + +echo "" +echo "=== the budget comes back as pairs finish ===" +# shellcheck disable=SC2086 +for id in $ids; do flux cancel "$id" 2>/dev/null; done +for _ in $(seq 1 30); do + flux jobs -no "{id}" 2>/dev/null | grep -q . || break + sleep 1 +done +if flux submit --quantum-vendor mock -n4 sleep 1 >/dev/null 2>&1; then + echo " a new pair is admitted again" +else + echo "FAIL the budget did not recover, cores are leaking on job.destroy" + rc=1 +fi + +echo "" +echo "=== the total rebuilds itself when the plugin is reloaded ===" +# job.new is replayed for active jobs, which is how the budget survives a +# restart. Hold two pairs, reload, and check the budget is still spent. +held="" +for i in 1 2; do + id=$(flux submit --quantum-vendor mock -n8 sleep 300 2>&1 | tail -1) && held="$held $id" +done +flux jobtap remove quantum 2>/dev/null +flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ + vendors="mock" total_cores=20 reserve_cores=0 +if flux submit --quantum-vendor mock -n8 sleep 300 >/dev/null 2>&1; then + echo "FAIL after reload the budget was forgotten, so a third pair got in" + rc=1 +else + echo " still full after reload, so job.new replay rebuilt the total" +fi +# shellcheck disable=SC2086 +for id in $held; do flux cancel "$id" 2>/dev/null; done + +echo "" +echo "=== both halves of a pair are marked protected, by the plugin ===" +err=$(mktemp) +scout=$(flux submit --quantum-vendor mock -n2 sleep 60 2>"$err") +main=$(grep -oE 'held classical job [0-9]+' "$err" | awk '{print $NF}') +rm -f "$err" +for half in "classical:$main" "scout:$scout"; do + name=${half%%:*}; id=${half#*:} + [ -z "$id" ] && { echo "FAIL no $name id"; rc=1; continue; } + if flux job info "$id" jobspec 2>/dev/null | grep -q '"protected"'; then + echo " $name is protected" + else + echo "FAIL the $name half was not marked protected, so it is preemptible" + rc=1 + fi +done +flux cancel "$main" "$scout" 2>/dev/null + +echo "" +echo "=== a submitter cannot mark their own job protected ===" +if err=$(flux submit --setattr=system.protected=quantum -n1 true 2>&1); then + echo "FAIL a user set the protection flag and was allowed to" + rc=1 +else + if echo "$err" | grep -q "set by the scheduler"; then + echo " rejected, and the message says why" + else + echo "FAIL rejected for the wrong reason: $err" + rc=1 + fi +fi + +echo "" +echo "=== an ordinary job is left preemptible ===" +id=$(flux submit -n1 sleep 30) +if flux job info "$id" jobspec 2>/dev/null | grep -q '"protected"'; then + echo "FAIL an ordinary job was marked protected" + rc=1 +else + echo " not protected, so it can be preempted" +fi +flux cancel "$id" 2>/dev/null + +echo "" +echo "=== vendor policy still applies ===" +if flux submit --quantum-vendor rigetti -n1 true >/dev/null 2>&1; then + echo "FAIL an unconfigured vendor was accepted"; rc=1 +else + echo " unconfigured vendor rejected" +fi + +flux jobtap remove quantum 2>/dev/null +echo "" +echo "=== admission test rc=$rc ===" +exit "$rc" diff --git a/tests/integration/test-preemption.sh b/tests/integration/test-preemption.sh new file mode 100755 index 0000000..d0dfda4 --- /dev/null +++ b/tests/integration/test-preemption.sh @@ -0,0 +1,94 @@ +#!/bin/bash +# Admission promises the classical half can be placed, but ordinary work may be +# sitting on the cores when the scout releases it. With preempt_after set, the +# plugin waits and then cancels unprotected jobs to make room. +# +# flux start bash tests/integration/test-preemption.sh +# +# Needs at least 3 cores, so the pair has somewhere to go once room is made. + +set -u +export FLUX_QUANTUM_MOCK=1 +HERE=$(cd "$(dirname "$0")/../.." && pwd) +export FLUX_CLI_PLUGINPATH="$HERE/cli-plugins" +rc=0 + +CORES=$(flux resource list -no "{ncores}" 2>/dev/null | head -1) +CORES=${CORES:-0} +echo "=== cluster has $CORES cores ===" +if [ "$CORES" -lt 3 ]; then + echo "SKIP need at least 3 cores" + exit 0 +fi + +make -s -C "$HERE/flux_quantum/jobtap" || { echo "FAIL build"; exit 1; } +flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ + vendors="mock" total_cores="$CORES" reserve_cores=0 \ + protect_types="qpu" preempt_after=3 || { echo "FAIL load"; exit 1; } + +echo "" +echo "=== fill every core with unprotected work ===" +fill="" +for _ in $(seq 1 "$CORES"); do + id=$(flux submit --quiet -n1 sleep 600) && fill="$fill $id" +done +sleep 3 +echo " submitted $CORES one core jobs" +flux jobs -no "{id} {state} {ntasks}" | head -5 | sed 's/^/ /' + +echo "" +echo "=== submit a pair, its classical half cannot be placed ===" +err=$(mktemp) +scout=$(flux submit --quantum-vendor mock -n1 -- sh -c 'echo classical ran' 2>"$err") +main=$(grep -oE 'held classical job [0-9]+' "$err" | awk '{print $NF}') +sed 's/^/ /' "$err"; rm -f "$err" +if [ -z "$main" ]; then + echo "FAIL no pair was created" + exit 1 +fi +echo " scout=$scout classical=$main" + +echo "" +echo "=== after the grace period something unprotected is cancelled ===" +if flux job wait-event -t 120 "$main" clean /dev/null 2>&1; then + echo " the classical half ran, so room was made" +else + echo "FAIL the classical half never started, preemption did not free cores" + flux jobs -a | sed 's/^/ /' + rc=1 +fi + +victims=0 +# shellcheck disable=SC2086 +for id in $fill; do + if flux job eventlog "$id" 2>/dev/null | grep -q 'type="preempt"'; then + victims=$((victims + 1)) + fi +done +if [ "$victims" -gt 0 ]; then + echo " $victims unprotected job(s) preempted" +else + echo "FAIL nothing was preempted, so the cores came free some other way" + rc=1 +fi + +echo "" +echo "=== neither half of the pair was a victim ===" +for half in "classical:$main" "scout:$scout"; do + name=${half%%:*}; id=${half#*:} + if flux job eventlog "$id" 2>/dev/null | grep -q 'type="preempt"'; then + echo "FAIL the $name half was preempted, protection is not working" + rc=1 + else + echo " $name was not preempted" + fi +done + +# shellcheck disable=SC2086 +for id in $fill; do flux cancel "$id" 2>/dev/null; done +flux cancel "$scout" 2>/dev/null +flux jobtap remove quantum 2>/dev/null + +echo "" +echo "=== preemption test rc=$rc ===" +exit "$rc" diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index 4db7d87..bacb747 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -318,3 +318,32 @@ def test_job_env_reaches_the_classical(stub_flux_cli): ) env = submitted["js"]["attributes"]["system"]["environment"] assert env["QRMI_JOB_QPU_RESOURCES"] == "ibm_kingston" + + +def test_classical_carries_its_core_count(stub_flux_cli): + """The jobtap plugin budgets on this, and only the classical carries it so + the scout is not counted twice.""" + from flux_quantum import cli + + js = _fake_jobspec(["myprog"]) + js.jobspec["resources"] = [ + {"type": "slot", "count": 4, "with": [{"type": "core", "count": 1}]} + ] + submitted = {} + + cli.prepare_pair( + handle=None, + jobspec=js, + vendor="mock", + submit_fn=lambda h, j: submitted.setdefault("js", json.loads(j)) and 0 or 7, + populate_fn=lambda h, v: None, + get_graph_fn=lambda h: _live_graph(), + cancel_fn=lambda *a: None, + ) + quantum = submitted["js"]["attributes"]["system"]["quantum"] + assert quantum["cores"] == 4 + assert quantum["vendor"] == "mock" + + # what flux submits is the scout, and it carries no quantum attributes + scout_sys = js.jobspec["attributes"]["system"] + assert "quantum" not in scout_sys or "cores" not in scout_sys.get("quantum", {}) diff --git a/tests/unit/test_qresource.py b/tests/unit/test_qresource.py index 76fafd0..1db6b94 100644 --- a/tests/unit/test_qresource.py +++ b/tests/unit/test_qresource.py @@ -165,3 +165,78 @@ def test_classical_resource_fallback_no_core(): } r = qr.classical_resource(empty) assert r["type"] == "node" and r["with"][0]["type"] == "slot" + + +def test_count_cores_walks_the_resource_tree(): + """The jobtap plugin budgets on this number, so it has to be right for the + shapes the plugin will actually see.""" + from flux_quantum.qresource import count_cores + + node_slot_core = { + "resources": [ + { + "type": "node", + "count": 1, + "with": [ + {"type": "slot", "count": 8, "with": [{"type": "core", "count": 1}]} + ], + } + ] + } + assert count_cores(node_slot_core) == 8 + + # counts multiply down the tree + socketed = { + "resources": [ + { + "type": "node", + "count": 2, + "with": [ + { + "type": "socket", + "count": 2, + "with": [ + { + "type": "slot", + "count": 4, + "with": [{"type": "core", "count": 1}], + } + ], + } + ], + } + ] + } + assert count_cores(socketed) == 16 + + # the qpu is not a core, so a scout jobspec counts as one + scout = { + "resources": [ + {"type": "slot", "count": 1, "with": [{"type": "core", "count": 1}]}, + { + "type": "qdevice_ibm", + "count": 1, + "with": [{"type": "qpu", "count": 1, "exclusive": True}], + }, + ] + } + assert count_cores(scout) == 1 + + assert count_cores({"resources": []}) == 0 + assert count_cores({}) == 0 + + +def test_count_cores_takes_the_minimum_of_a_range(): + """A range means the job may get more, but only the minimum is promised.""" + from flux_quantum.qresource import count_cores + + ranged = { + "resources": [ + { + "type": "slot", + "count": {"min": 2, "max": 8}, + "with": [{"type": "core", "count": 1}], + } + ] + } + assert count_cores(ranged) == 2 From 1de122206e08af22eee40b904c8dec5ed2e4d6d8 Mon Sep 17 00:00:00 2001 From: vsoch Date: Sat, 5 Sep 2026 16:03:09 -0700 Subject: [PATCH 07/26] populate Signed-off-by: vsoch --- flux_quantum/graph.py | 33 ++++- flux_quantum/jobtap/quantum.c | 44 ++++--- flux_quantum/populate.py | 65 ++++++++++ pyproject.toml | 5 + tests/integration/test-admission.sh | 178 +++++++++++++-------------- tests/integration/test-preemption.sh | 16 ++- 6 files changed, 227 insertions(+), 114 deletions(-) create mode 100644 flux_quantum/populate.py diff --git a/flux_quantum/graph.py b/flux_quantum/graph.py index ac3d44d..fb750ad 100644 --- a/flux_quantum/graph.py +++ b/flux_quantum/graph.py @@ -52,10 +52,27 @@ def vendors_present(graph): return out -def populate(handle, vendors, qpus=1, graph=None): +def allocated_cores(handle, lister=None): + """Cores the scheduler currently has allocated.""" + if lister is None: + from flux.resource import resource_list as lister + return int(lister(handle).get().allocated.ncores) + + +def populate(handle, vendors, qpus=1, graph=None, lister=None, force=False): """Ensure qdevice_ -> qpu exists in the live graph. Idempotent, only missing vendors are added. Returns what was added. + + Growing the graph while jobs hold resources corrupts fluxion. The spans of + the running jobs no longer match the graph, so every later free fails with + planner_multi_rem_span returned -1, those resources are never released, and + the instance stops scheduling anything at all. It looks like free cores that + nothing will use. + + So refuse to grow a busy graph. Populate at startup instead, before any + jobs run, which is what the vendors are there for anyway. force=True is for + a caller that knows the instance is idle. """ if isinstance(vendors, str): vendors = [vendors] @@ -64,6 +81,20 @@ def populate(handle, vendors, qpus=1, graph=None): missing = [v for v in vendors if v not in vendors_present(graph)] if not missing: return set() + + if not force: + try: + busy = allocated_cores(handle, lister) + except Exception: + busy = 0 # cannot tell, so do not block the caller + if busy: + raise RuntimeError( + "quantum: {} would have to be added to the fluxion graph, but " + "{} cores are allocated. Growing the graph now would break " + "resource release for every running job and wedge the " + "scheduler. Populate at startup instead, with " + "flux python -m flux_quantum.populate".format(", ".join(missing), busy) + ) subgraph = qresource.graph_subgraph(graph, missing, qpus=qpus) try: handle.rpc( diff --git a/flux_quantum/jobtap/quantum.c b/flux_quantum/jobtap/quantum.c index 380e804..b5021ee 100644 --- a/flux_quantum/jobtap/quantum.c +++ b/flux_quantum/jobtap/quantum.c @@ -168,15 +168,16 @@ static const char *idkey (flux_jobid_t id, char *buf, size_t len) } /* remember a job so it can be considered as a victim later */ -static void track (flux_jobid_t id, int cores, int protected_) +static void track (flux_jobid_t id, int cores, int protected_, int pair) { char key[32]; json_t *entry; if (!g_jobs) return; - entry = json_pack ("{s:i s:b s:f}", - "cores", cores, "protected", protected_, "run", 0.0); + entry = json_pack ("{s:i s:b s:b s:f}", + "cores", cores, "protected", protected_, + "pair", pair, "run", 0.0); if (entry) (void) json_object_set_new (g_jobs, idkey (id, key, sizeof (key)), entry); } @@ -426,41 +427,43 @@ static void grace_cb (flux_reactor_t *r, free (g); } -/* The scout posts the session memo immediately before releasing the classical - * job, so the memo is the signal that the job should now be able to start. */ -static int event_cb (flux_plugin_t *p, +/* Start the clock when the classical half reaches SCHED. + * + * The memo would be the natural signal, since the scout posts it immediately + * before releasing. But on a full machine the scout cannot get a core either, + * so there is no memo and the pair waits forever with nothing to trigger on. + * Admission already promised this pair room, so the promise is what the clock + * hangs off, not the handover. + */ +static int sched_cb (flux_plugin_t *p, const char *topic, flux_plugin_arg_t *args, void *arg) { flux_jobid_t id; - const char *name = NULL; json_t *entry; struct grace *g; - int cores = 0; + int cores = 0, pair = 0; flux_t *h; if (g_preempt_after <= 0.0) return 0; - if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, - "{s:I s:{s:s}}", - "id", &id, - "entry", "name", &name) < 0) - return 0; - if (!name || strcmp (name, "memo") != 0) + if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, "{s:I}", "id", &id) < 0) return 0; if (!(entry = lookup (id)) - || json_unpack (entry, "{s:i}", "cores", &cores) < 0 + || json_unpack (entry, "{s:i s:b}", "cores", &cores, "pair", &pair) < 0 + || !pair || cores <= 0) - return 0; + return 0; /* not the classical half of a pair */ h = flux_jobtap_get_flux (p); if (!(g = calloc (1, sizeof (*g)))) return 0; g->p = p; g->id = id; - /* the pair reserves cores + 1, but the scout already holds its own core */ - g->cores = cores - 1 > 0 ? cores - 1 : cores; + /* cores is the whole pair, classical plus one for the scout. Neither may + * be running, so make room for both. */ + g->cores = cores; if (!(g->w = flux_timer_watcher_create (flux_get_reactor (h), g_preempt_after, 0., grace_cb, g))) { @@ -495,7 +498,8 @@ static int new_cb (flux_plugin_t *p, * apart from losing the run time, which the state callback restores. */ track (id, vendor ? pair_cores (jobspec) : count_tree_cores_top (jobspec), - needs_protection (jobspec, vendor)); + needs_protection (jobspec, vendor), + vendor ? 1 : 0); if (g_preempt_after > 0.0) (void) flux_jobtap_job_subscribe (p, id); @@ -563,7 +567,7 @@ static const struct flux_plugin_handler handlers[] = { { "job.validate", validate_cb, NULL }, { "job.new", new_cb, NULL }, { "job.state.run", run_cb, NULL }, - { "job.event.memo", event_cb, NULL }, + { "job.state.sched", sched_cb, NULL }, { "job.destroy", destroy_cb, NULL }, { 0 }, }; diff --git a/flux_quantum/populate.py b/flux_quantum/populate.py new file mode 100644 index 0000000..00602cf --- /dev/null +++ b/flux_quantum/populate.py @@ -0,0 +1,65 @@ +############################################################## +# Copyright 2024 Lawrence Livermore National Security, LLC +# (c.f. AUTHORS, NOTICE.LLNS, COPYING) +# +# SPDX-License-Identifier: LGPL-3.0 +############################################################## + +"""Add vendor devices to the fluxion graph once, at startup. + +This has to happen before any job holds resources. Growing the graph while jobs +are allocated corrupts fluxion, every later free fails with +planner_multi_rem_span returned -1, and the instance stops scheduling. Doing it +at startup means the first quantum submit never has to grow anything. + + flux python -m flux_quantum.populate ibm braket mock + +Run it with flux python, not the system python. The flux bindings are installed +for whichever interpreter flux was built against, so a console script on the +default python cannot import them. +""" + +import argparse +import sys + +import flux + +from . import graph + + +def main(argv=None): + ap = argparse.ArgumentParser( + prog="flux-quantum-populate", + description="add vendor devices to the fluxion graph, once, at startup", + ) + ap.add_argument("vendors", nargs="+", help="vendor names, e.g. ibm braket mock") + ap.add_argument( + "--qpus", + type=int, + default=1, + help="qpu vertices per vendor. This is the concurrency limit, the qpu " + "is matched exclusively so it caps how many pairs run at once", + ) + ap.add_argument( + "--force", + action="store_true", + help="grow the graph even with jobs allocated. This corrupts fluxion, " + "only use it if you know the instance is idle", + ) + args = ap.parse_args(argv) + + handle = flux.Flux() + try: + added = graph.populate(handle, args.vendors, qpus=args.qpus, force=args.force) + except Exception as exc: + sys.exit("flux-quantum-populate: {}".format(exc)) + + if added: + print("added {}".format(", ".join(sorted(added)))) + else: + print("nothing to add, every vendor is already in the graph") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/pyproject.toml b/pyproject.toml index a06f055..f0a3162 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,11 @@ requires-python = ">=3.8" # flux-core releases. Developed/tested against flux-core 0.87.0. Pin to a # contemporaneous flux-core, per Flux packaging guidance. +[project.scripts] +# populate the graph once at startup. Growing it while jobs hold resources +# corrupts fluxion, so this must run before any job does. +flux-quantum-populate = "flux_quantum.populate:main" + [project.optional-dependencies] # qrmi needs python 3.11 or newer, so it stays optional qrmi = ["qrmi[ibm]"] diff --git a/tests/integration/test-admission.sh b/tests/integration/test-admission.sh index ee75655..438022b 100755 --- a/tests/integration/test-admission.sh +++ b/tests/integration/test-admission.sh @@ -3,137 +3,134 @@ # running, and rejects a pair that would exceed it. Launching a scout commits # money, so the classical half has to be placeable. # -# flux start -s1 bash tests/integration/test-admission.sh +# FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-admission.sh +# +# The variable has to be set BEFORE flux start. The CLIPlugin validate hook is +# also called by the job ingest validator, a separate process started with the +# broker, so a variable exported inside the instance never reaches it and every +# submit fails with no backend for vendor mock. set -u -export FLUX_QUANTUM_MOCK=1 HERE=$(cd "$(dirname "$0")/../.." && pwd) export FLUX_CLI_PLUGINPATH="$HERE/cli-plugins" rc=0 -echo "=== build and load with a small budget ===" +if [ -z "${FLUX_QUANTUM_MOCK:-}" ]; then + echo "FAIL FLUX_QUANTUM_MOCK must be exported before flux start" + exit 1 +fi + +# submit a pair, print both job ids, or fail. Checking the exit status is not +# enough, the plugin prints its errors and still exits 0. +submit_pair () { + local n="$1" err out + err=$(mktemp) + out=$(flux submit --quantum-vendor mock -n"$n" sleep 300 2>"$err") + if grep -q "held classical job" "$err"; then + printf '%s %s\n' "$out" \ + "$(grep -oE 'held classical job [0-9]+' "$err" | awk '{print $NF}')" + rm -f "$err" + return 0 + fi + cat "$err" >&2 + rm -f "$err" + return 1 +} + +echo "=== build and load ===" make -s -C "$HERE/flux_quantum/jobtap" || { echo "FAIL build"; exit 1; } -# 20 cores for quantum, none held back, so a 4 core pair takes 5 and four fit + +# add the vendor before anything is allocated. Growing the fluxion graph while +# jobs hold resources breaks every later free and wedges the scheduler. +flux python -m flux_quantum.populate mock >/dev/null 2>&1 \ + || echo "WARNING could not populate the graph up front" + +# a one core pair reserves two, so a budget of four takes exactly two pairs flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ - vendors="mock,ibm,braket" total_cores=20 reserve_cores=0 \ - protect_types="qpu" || { - echo "FAIL load"; exit 1; } -flux jobtap list + vendors="mock,ibm,braket" total_cores=4 reserve_cores=0 \ + protect_types="qpu" || { echo "FAIL load"; exit 1; } echo "" -echo "=== a pair is accounted for as its cores plus one for the scout ===" -# submitted held, so it stays in the budget without running -ids="" -for i in 1 2 3 4; do - if id=$(flux submit --quantum-vendor mock -n4 sleep 300 2>&1 | tail -1); then - ids="$ids $id" +echo "=== a pair reserves its cores plus one for the scout ===" +pairs="" +for i in 1 2; do + if ids=$(submit_pair 1); then + pairs="$pairs $ids" echo " pair $i admitted" else - echo "FAIL pair $i should have been admitted, budget is 20 and 4 pairs need 20" + echo "FAIL pair $i should have been admitted, two one core pairs fit a budget of four" rc=1 fi done echo "" -echo "=== the fifth is rejected rather than admitted and left to wait ===" -if err=$(flux submit --quantum-vendor mock -n4 sleep 300 2>&1); then - echo "FAIL the fifth pair was admitted, the budget is not being enforced" +echo "=== the third is rejected rather than left waiting ===" +err=$(mktemp) +if submit_pair 1 >/dev/null 2>"$err"; then + echo "FAIL a third pair was admitted, the budget is not being enforced" rc=1 +elif grep -q "no room for another pair" "$err"; then + echo " rejected, and the message says why" + grep -o "no room for another pair.*" "$err" | head -1 | cut -c1-100 | sed 's/^/ /' else - if echo "$err" | grep -q "no room for another pair"; then - echo " rejected, and the message says why" - echo "$err" | grep -o "no room for another pair.*" | head -1 | sed 's/^/ /' - else - echo "FAIL rejected for the wrong reason: $err" - rc=1 - fi + echo "FAIL rejected for the wrong reason"; sed 's/^/ /' "$err"; rc=1 fi +rm -f "$err" echo "" -echo "=== a pair that does not fit at all is rejected on its own ===" -if flux submit --quantum-vendor mock -n64 sleep 300 >/dev/null 2>&1; then - echo "FAIL a 64 core pair should not fit a 20 core budget"; rc=1 +echo "=== a pair too large for the budget is rejected on its own ===" +if submit_pair 64 >/dev/null 2>&1; then + echo "FAIL a 64 core pair should not fit a budget of four"; rc=1 else echo " rejected" fi echo "" -echo "=== the budget comes back as pairs finish ===" -# shellcheck disable=SC2086 -for id in $ids; do flux cancel "$id" 2>/dev/null; done -for _ in $(seq 1 30); do - flux jobs -no "{id}" 2>/dev/null | grep -q . || break - sleep 1 -done -if flux submit --quantum-vendor mock -n4 sleep 1 >/dev/null 2>&1; then - echo " a new pair is admitted again" -else - echo "FAIL the budget did not recover, cores are leaking on job.destroy" - rc=1 -fi - -echo "" -echo "=== the total rebuilds itself when the plugin is reloaded ===" -# job.new is replayed for active jobs, which is how the budget survives a -# restart. Hold two pairs, reload, and check the budget is still spent. -held="" -for i in 1 2; do - id=$(flux submit --quantum-vendor mock -n8 sleep 300 2>&1 | tail -1) && held="$held $id" -done -flux jobtap remove quantum 2>/dev/null -flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ - vendors="mock" total_cores=20 reserve_cores=0 -if flux submit --quantum-vendor mock -n8 sleep 300 >/dev/null 2>&1; then - echo "FAIL after reload the budget was forgotten, so a third pair got in" - rc=1 -else - echo " still full after reload, so job.new replay rebuilt the total" -fi -# shellcheck disable=SC2086 -for id in $held; do flux cancel "$id" 2>/dev/null; done - -echo "" -echo "=== both halves of a pair are marked protected, by the plugin ===" -err=$(mktemp) -scout=$(flux submit --quantum-vendor mock -n2 sleep 60 2>"$err") -main=$(grep -oE 'held classical job [0-9]+' "$err" | awk '{print $NF}') -rm -f "$err" -for half in "classical:$main" "scout:$scout"; do - name=${half%%:*}; id=${half#*:} - [ -z "$id" ] && { echo "FAIL no $name id"; rc=1; continue; } +echo "=== both halves are marked protected, by the plugin and not the user ===" +for id in $pairs; do if flux job info "$id" jobspec 2>/dev/null | grep -q '"protected"'; then - echo " $name is protected" + echo " $id is protected" else - echo "FAIL the $name half was not marked protected, so it is preemptible" - rc=1 + echo "FAIL $id was not marked protected, so it is preemptible"; rc=1 fi done -flux cancel "$main" "$scout" 2>/dev/null echo "" echo "=== a submitter cannot mark their own job protected ===" -if err=$(flux submit --setattr=system.protected=quantum -n1 true 2>&1); then - echo "FAIL a user set the protection flag and was allowed to" - rc=1 +err=$(mktemp) +if flux submit --setattr=system.protected=quantum -n1 true >/dev/null 2>"$err"; then + echo "FAIL a user set the protection flag and was allowed to"; rc=1 +elif grep -q "set by the scheduler" "$err"; then + echo " rejected, and the message says why" else - if echo "$err" | grep -q "set by the scheduler"; then - echo " rejected, and the message says why" - else - echo "FAIL rejected for the wrong reason: $err" - rc=1 - fi + echo "FAIL rejected for the wrong reason"; sed 's/^/ /' "$err"; rc=1 fi +rm -f "$err" echo "" echo "=== an ordinary job is left preemptible ===" -id=$(flux submit -n1 sleep 30) -if flux job info "$id" jobspec 2>/dev/null | grep -q '"protected"'; then - echo "FAIL an ordinary job was marked protected" - rc=1 +plain=$(flux submit -n1 sleep 300) +if flux job info "$plain" jobspec 2>/dev/null | grep -q '"protected"'; then + echo "FAIL an ordinary job was marked protected"; rc=1 else echo " not protected, so it can be preempted" fi -flux cancel "$id" 2>/dev/null + +echo "" +echo "=== the budget is given back as pairs finish ===" +# shellcheck disable=SC2086 +flux cancel $pairs "$plain" >/dev/null 2>&1 +for _ in $(seq 1 30); do + flux jobs -no "{id}" 2>/dev/null | grep -q . || break + sleep 1 +done +if ids=$(submit_pair 1); then + echo " a new pair is admitted again" + # shellcheck disable=SC2086 + flux cancel $ids >/dev/null 2>&1 +else + echo "FAIL the budget did not recover, cores are leaking on job.destroy"; rc=1 +fi echo "" echo "=== vendor policy still applies ===" @@ -143,7 +140,8 @@ else echo " unconfigured vendor rejected" fi -flux jobtap remove quantum 2>/dev/null +# jobtap plugins are listed and removed by file name, not by registered name +flux jobtap remove quantum.so >/dev/null 2>&1 echo "" echo "=== admission test rc=$rc ===" exit "$rc" diff --git a/tests/integration/test-preemption.sh b/tests/integration/test-preemption.sh index d0dfda4..a3e5f05 100755 --- a/tests/integration/test-preemption.sh +++ b/tests/integration/test-preemption.sh @@ -8,11 +8,15 @@ # Needs at least 3 cores, so the pair has somewhere to go once room is made. set -u -export FLUX_QUANTUM_MOCK=1 HERE=$(cd "$(dirname "$0")/../.." && pwd) export FLUX_CLI_PLUGINPATH="$HERE/cli-plugins" rc=0 +if [ -z "${FLUX_QUANTUM_MOCK:-}" ]; then + echo "FAIL FLUX_QUANTUM_MOCK must be exported before flux start" + exit 1 +fi + CORES=$(flux resource list -no "{ncores}" 2>/dev/null | head -1) CORES=${CORES:-0} echo "=== cluster has $CORES cores ===" @@ -22,6 +26,12 @@ if [ "$CORES" -lt 3 ]; then fi make -s -C "$HERE/flux_quantum/jobtap" || { echo "FAIL build"; exit 1; } + +# Add the vendor to the graph before anything is allocated. Growing the graph +# while jobs hold resources corrupts fluxion, every later free fails, and the +# instance stops scheduling entirely. +flux python -m flux_quantum.populate mock >/dev/null 2>&1 \ + || echo "WARNING could not populate the graph up front" flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ vendors="mock" total_cores="$CORES" reserve_cores=0 \ protect_types="qpu" preempt_after=3 || { echo "FAIL load"; exit 1; } @@ -30,7 +40,7 @@ echo "" echo "=== fill every core with unprotected work ===" fill="" for _ in $(seq 1 "$CORES"); do - id=$(flux submit --quiet -n1 sleep 600) && fill="$fill $id" + id=$(flux submit -n1 sleep 600) && fill="$fill $id" done sleep 3 echo " submitted $CORES one core jobs" @@ -87,7 +97,7 @@ done # shellcheck disable=SC2086 for id in $fill; do flux cancel "$id" 2>/dev/null; done flux cancel "$scout" 2>/dev/null -flux jobtap remove quantum 2>/dev/null +flux jobtap remove quantum.so 2>/dev/null echo "" echo "=== preemption test rc=$rc ===" From 34661b092c773a3784a2e20c17ab10b8dd37024b Mon Sep 17 00:00:00 2001 From: vsoch Date: Sat, 5 Sep 2026 18:53:26 -0700 Subject: [PATCH 08/26] update scout Signed-off-by: vsoch --- .github/workflows/test.yml | 4 +- flux_quantum/cli.py | 40 ++++++ flux_quantum/graph.py | 16 ++- flux_quantum/jobtap/quantum.c | 190 +++++++++++++++++++++++++-- flux_quantum/scout.py | 9 +- tests/integration/test-preemption.sh | 5 +- tests/unit/test_cli_plugin.py | 58 ++++++++ tests/unit/test_graph.py | 32 ++++- tests/unit/test_handoff.py | 23 ++++ 9 files changed, 356 insertions(+), 21 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 5896107..515c156 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -114,12 +114,12 @@ jobs: # 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 start bash tests/integration/test-admission.sh + 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 start bash tests/integration/test-preemption.sh + run: FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-preemption.sh # one submit gives a held classical and a scout, and the qpu stays # allocated until the classical is done diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index 56186d3..8570bd0 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -37,6 +37,31 @@ def _wrap_commands(jobspec_dict, wrap_path): task["command"] = ["flux", "python", wrap_path, "--"] + cmd +def _check_pair_fits(handle, cores, rpc_fn=None): + """Refuse a pair that cannot be placed, before either half is created.""" + try: + if rpc_fn is None: + resp = handle.rpc( + "job-manager.quantum.check", {"pair_cores": int(cores)} + ).get() + else: + resp = rpc_fn(handle, int(cores)) + except Exception: + # The plugin may be absent, or too old to serve the method. Admission is + # then whatever the plugin does at validate, which is how this worked + # before the check existed, so do not block the submit on it. + return + if not resp.get("ok", True): + raise SystemExit( + "flux quantum: no room for this pair. It needs {} cores plus one " + "for the scout, and the cluster has {} free and {} preemptible. " + "Submitting anyway would open a vendor session for a pair that " + "cannot run.".format( + cores, resp.get("free", "?"), resp.get("preemptible", "?") + ) + ) + + def _safe_cancel(cancel_fn, handle, jobid, reason): """Cancel the held job so a failed setup does not park it forever.""" try: @@ -98,6 +123,21 @@ def prepare_pair( if job_env: sysattr.setdefault("environment", {}).update(job_env) + # Ask before creating anything. The plugin's own admission check only fires + # once a job exists, and it compares outstanding pairs against total + # capacity, which is a quota across pairs rather than a feasibility test: a + # 9 core pair was admitted onto a cluster with 7 cores free, could not be + # placed, and left preemption to clean up. + # + # job-manager.quantum.check answers "does free + preemptible leave room for + # this classical and a scout", from occupancy the plugin already tracks. + # Asking here means a pair that cannot be placed never opens a session. + # + # There is a race between this answer and the submits below. That is + # accepted: the plugin's job.validate check is still the backstop, and + # either half is cancelled if the other fails. + _check_pair_fits(handle, quantum["cores"]) + # if feasibility validation is on, an unsatisfiable request is rejected here # and we abort before spending any quantum quota try: diff --git a/flux_quantum/graph.py b/flux_quantum/graph.py index fb750ad..b0dd5ee 100644 --- a/flux_quantum/graph.py +++ b/flux_quantum/graph.py @@ -83,10 +83,22 @@ def populate(handle, vendors, qpus=1, graph=None, lister=None, force=False): return set() if not force: + # Being unable to read the allocation count is not evidence that the + # instance is idle. Growing the graph while jobs hold resources + # corrupts fluxion irrecoverably, so an unknown answer has to be + # treated the same as a busy one. try: busy = allocated_cores(handle, lister) - except Exception: - busy = 0 # cannot tell, so do not block the caller + except Exception as exc: + raise RuntimeError( + "quantum: {} would have to be added to the fluxion graph, but " + "the allocated core count could not be read ({}), so there is " + "no way to tell whether the instance is idle. Growing a busy " + "graph breaks resource release for every running job. Populate " + "at startup with flux python -m flux_quantum.populate, or pass " + "force=True if you know the instance is " + "idle".format(", ".join(missing), exc) + ) if busy: raise RuntimeError( "quantum: {} would have to be added to the fluxion graph, but " diff --git a/flux_quantum/jobtap/quantum.c b/flux_quantum/jobtap/quantum.c index b5021ee..0181cc9 100644 --- a/flux_quantum/jobtap/quantum.c +++ b/flux_quantum/jobtap/quantum.c @@ -63,6 +63,7 @@ #include #include #include +#include #include #include #include @@ -563,6 +564,108 @@ static int run_cb (flux_plugin_t *p, return 0; } +/* Sum cores over running jobs, and separately over running jobs that a pair is + * allowed to preempt. track() records every job, not only pairs, so this is + * derivable from state the plugin already holds and needs no query to fluxion. + * + * Note fluxion cannot answer this question. It knows free and allocated, but + * "allocated and preemptible" depends on attributes.system.protected, which is + * a concept of this plugin that fluxion has never seen. + */ +static void occupancy (int *runningp, int *preemptiblep) +{ + const char *key; + json_t *entry; + int running = 0, preemptible = 0; + + if (g_jobs) { + json_object_foreach (g_jobs, key, entry) { + int prot = 0, c = 0; + double run = 0.0; + + if (json_unpack (entry, "{s:i s:b s:F}", + "cores", &c, "protected", &prot, "run", &run) < 0) + continue; + if (run <= 0.0 || c <= 0) + continue; /* not running, so occupying nothing */ + running += c; + if (!prot) + preemptible += c; + } + } + if (runningp) + *runningp = running; + if (preemptiblep) + *preemptiblep = preemptible; +} + +/* Can a pair of pair_cores classical cores plus one scout core be placed, now + * or by preempting? Returns 1 for yes, 0 for no, and fills the counts. + * + * The scout's core comes out of the budget before the comparison, which is why + * this is "- 1" rather than comparing against pair_cores + 1. Same condition, + * written the way the requirement was specified. + */ +static int pair_fits (int pair_cores, int *freep, int *preemptiblep) +{ + int running = 0, preemptible = 0, budget_free; + + occupancy (&running, &preemptible); + budget_free = budget () - running; + if (budget_free < 0) + budget_free = 0; + if (freep) + *freep = budget_free; + if (preemptiblep) + *preemptiblep = preemptible; + if (pair_cores <= 0) + return 0; + return (budget_free + preemptible - 1) >= pair_cores; +} + +/* job-manager.quantum.check + * + * Asked by the submitting client BEFORE either half of a pair is created, so + * that a pair which cannot be placed never opens a vendor session. The old + * check compared outstanding pairs against total capacity, which is a quota + * across pairs rather than a feasibility test, so a 9 core pair was admitted + * onto a cluster with 7 cores free. + * + * There is a race between this answer and the submits that follow it. That is + * accepted: the caller cancels whichever half survives if the other fails, and + * the job.validate check below remains as a backstop. + */ +static void check_cb (flux_t *h, flux_msg_handler_t *mh, + const flux_msg_t *msg, void *arg) +{ + int pair_cores = 0, free_cores = 0, preemptible = 0, ok; + + if (flux_request_unpack (msg, NULL, "{s:i}", "pair_cores", &pair_cores) < 0) { + if (flux_respond_error (h, msg, EPROTO, NULL) < 0) + flux_log_error (h, "quantum.check: error responding"); + return; + } + if (pair_cores <= 0) { + if (flux_respond_error (h, msg, EINVAL, + "pair_cores must be greater than zero") < 0) + flux_log_error (h, "quantum.check: error responding"); + return; + } + if (g_total_cores <= 0) { + /* admission control is off, so there is nothing to promise */ + if (flux_respond_pack (h, msg, "{s:b s:i s:i s:b}", + "ok", 1, "free", 0, "preemptible", 0, + "enforced", 0) < 0) + flux_log_error (h, "quantum.check: error responding"); + return; + } + ok = pair_fits (pair_cores, &free_cores, &preemptible); + if (flux_respond_pack (h, msg, "{s:b s:i s:i s:b}", + "ok", ok, "free", free_cores, + "preemptible", preemptible, "enforced", 1) < 0) + flux_log_error (h, "quantum.check: error responding"); +} + static const struct flux_plugin_handler handlers[] = { { "job.validate", validate_cb, NULL }, { "job.new", new_cb, NULL }, @@ -578,17 +681,36 @@ int flux_plugin_init (flux_plugin_t *p) const char *protect_types = NULL; int total = 0, reserve = 0; double preempt_after = 0.0; - + flux_t *h = flux_jobtap_get_flux (p); + int rc; if (!(g_jobs = json_object ())) return -1; - if (flux_plugin_conf_unpack (p, "{s?s s?s s?i s?i s?F}", - "vendors", &vendors, - "protect_types", &protect_types, - "total_cores", &total, - "reserve_cores", &reserve, - "preempt_after", &preempt_after) == 0) { + /* One unpack sets everything, so any single key of an unexpected type + * returns -1 and none of it is applied. The defaults that would then be + * left in place are total_cores 0, which turns admission control off, and + * preempt_after 0, which turns preemption off. Loading unenforced and + * saying nothing is worse than not loading, so refuse. + * + * ENOENT is different. It means no conf was given at all, which is a + * legitimate way to ask for the defaults. + */ + rc = flux_plugin_conf_unpack (p, "{s?s s?s s?i s?i s?F}", + "vendors", &vendors, + "protect_types", &protect_types, + "total_cores", &total, + "reserve_cores", &reserve, + "preempt_after", &preempt_after); + if (rc < 0 && errno != ENOENT) { + flux_log (h, LOG_ERR, + "quantum: cannot read plugin config: %s. Not loading, " + "because loading with defaults would leave admission control " + "and preemption off without saying so", + flux_plugin_strerror (p)); + return -1; + } + if (rc == 0) { if (vendors) g_vendors = strdup (vendors); if (protect_types) @@ -596,15 +718,55 @@ int flux_plugin_init (flux_plugin_t *p) g_total_cores = total; g_reserve_cores = reserve; g_preempt_after = preempt_after; - if (g_preempt_after > 0.0) - flux_log (flux_jobtap_get_flux (p), LOG_WARNING, - "quantum: preemption is on. An unprotected job may be " - "cancelled to make room for a pair whose session is " - "already open, and a cancelled job loses its work"); } - /* registers the table and sets the plugin name in one call */ - return flux_plugin_register (p, "quantum", handlers); + /* State what is in force. A campaign was once run with preemption off + * because the key was absent from the config, and nothing in the logs or + * the results said so. + */ + flux_log (h, LOG_INFO, + "quantum: vendors=%s total_cores=%d reserve_cores=%d " + "protect_types=%s preempt_after=%.1f", + g_vendors ? g_vendors : "(default)", + g_total_cores, + g_reserve_cores, + g_protect_types ? g_protect_types : "(default)", + g_preempt_after); + if (g_total_cores <= 0) + flux_log (h, LOG_WARNING, + "quantum: total_cores is not set, so admission control is " + "off and a pair may be admitted that cannot be placed"); + if (g_preempt_after > 0.0) + flux_log (h, LOG_WARNING, + "quantum: preemption is on. An unprotected job may be " + "cancelled to make room for a pair whose session is " + "already open, and a cancelled job loses its work"); + else + flux_log (h, LOG_WARNING, + "quantum: preemption is off. A pair that cannot be placed " + "will hold an open vendor session until the machine frees " + "up on its own"); + + /* The plugin name has to be set first: the service is named + * job-manager.., so registering it before + * flux_plugin_register fails with EINVAL and takes the job manager down + * with it. + */ + if (flux_plugin_register (p, "quantum", handlers) < 0) { + flux_log_error (h, "quantum: could not register handlers"); + return -1; + } + + /* Expose job-manager.quantum.check. The _ex form with FLUX_ROLE_USER is + * required because the system instance sets allow-guest-user, so ordinary + * users must be able to ask before they submit. + */ + if (flux_jobtap_service_register_ex (p, "check", FLUX_ROLE_USER, + check_cb, NULL) < 0) { + flux_log_error (h, "quantum: could not register the check service"); + return -1; + } + return 0; } /* vi: ts=4 sw=4 expandtab diff --git a/flux_quantum/scout.py b/flux_quantum/scout.py index 7368473..921803b 100644 --- a/flux_quantum/scout.py +++ b/flux_quantum/scout.py @@ -24,10 +24,17 @@ def post_session(handle, jobid, session, rpc=None): The memo RPC is FLUX_ROLE_USER and authorized against the job owner, so this needs no shared filesystem and no instance owner privilege. + + The released key is the durable record that this job was let go. A release + otherwise lives only in the scheduler memory, so a qmanager restart would + read the hold attribute again and park the job a second time, with nobody + left to release it and a vendor session already open and being paid for. + The eventlog survives a restart, so fluxion reads the marker from there. """ if rpc is None: rpc = handle.rpc - rpc("job-manager.memo", {"id": int(jobid), "memo": {SESSION_KEY: session}}).get() + memo = {SESSION_KEY: session, "released": 1} + rpc("job-manager.memo", {"id": int(jobid), "memo": memo}).get() def wait_for_job(handle, jobid, waiter=None): diff --git a/tests/integration/test-preemption.sh b/tests/integration/test-preemption.sh index a3e5f05..e00a657 100755 --- a/tests/integration/test-preemption.sh +++ b/tests/integration/test-preemption.sh @@ -17,7 +17,10 @@ if [ -z "${FLUX_QUANTUM_MOCK:-}" ]; then exit 1 fi -CORES=$(flux resource list -no "{ncores}" 2>/dev/null | head -1) +# -s all, not a bare -no "{ncores}". Without a state field in the format, +# flux resource list merges free, allocated and down onto one line, so the bare +# form returns the total while reading like the free count. +CORES=$(flux resource list -s all -no "{ncores}" 2>/dev/null) CORES=${CORES:-0} echo "=== cluster has $CORES cores ===" if [ "$CORES" -lt 3 ]; then diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index bacb747..91ff8e5 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -347,3 +347,61 @@ def test_classical_carries_its_core_count(stub_flux_cli): # what flux submits is the scout, and it carries no quantum attributes scout_sys = js.jobspec["attributes"]["system"] assert "quantum" not in scout_sys or "cores" not in scout_sys.get("quantum", {}) + + +def test_check_pair_fits_rejects_when_there_is_no_room(stub_flux_cli): + """A pair that cannot be placed is refused before either half is created.""" + from flux_quantum import cli + + def rpc(handle, cores): + return {"ok": False, "free": 7, "preemptible": 0, "enforced": True} + + try: + cli._check_pair_fits(None, 8, rpc_fn=rpc) + except SystemExit as exc: + assert "no room for this pair" in str(exc) + assert "7 free" in str(exc) + else: + raise AssertionError("a pair with no room was allowed through") + + +def test_check_pair_fits_allows_when_there_is_room(stub_flux_cli): + from flux_quantum import cli + + def rpc(handle, cores): + return {"ok": True, "free": 8, "preemptible": 120, "enforced": True} + + cli._check_pair_fits(None, 8, rpc_fn=rpc) + + +def test_check_pair_fits_is_permissive_when_the_service_is_absent(stub_flux_cli): + """An older plugin serves no check method. Admission then falls back to the + job.validate hook, which is how this worked before, so do not block.""" + from flux_quantum import cli + + def rpc(handle, cores): + raise OSError("Function not implemented") + + cli._check_pair_fits(None, 8, rpc_fn=rpc) + + +def test_prepare_pair_asks_before_submitting_anything(stub_flux_cli, monkeypatch): + """The check must happen before the classical is created, or a rejected pair + still leaves a held job behind.""" + from flux_quantum import cli + + order = [] + + def fake_check(handle, cores, rpc_fn=None): + order.append(("checked", cores)) + + monkeypatch.setattr(cli, "_check_pair_fits", fake_check) + + def submit(handle, jobspec_json): + order.append(("submitted", None)) + return 12345 + + _run_prepare(stub_flux_cli, submit=submit) + assert order[0][0] == "checked", order + assert order[0][1] > 0 + assert any(x[0] == "submitted" for x in order) diff --git a/tests/unit/test_graph.py b/tests/unit/test_graph.py index 770b892..b10504f 100644 --- a/tests/unit/test_graph.py +++ b/tests/unit/test_graph.py @@ -44,6 +44,20 @@ def _base_graph(): } +def _idle_lister(ncores=0): + """A resource_list stand-in reporting ncores allocated.""" + + class _Alloc: + def __init__(self, n): + self.ncores = n + + class _List: + def __init__(self, n): + self.allocated = _Alloc(n) + + return lambda handle: _FakeRPC(_List(ncores)) + + def test_get_live_graph_parses_find(): h = _FakeHandle(_base_graph()) g = qgraph.get_live_graph(h) @@ -52,13 +66,29 @@ def test_get_live_graph_parses_find(): def test_populate_adds_missing_vendor(): h = _FakeHandle(_base_graph()) - added = qgraph.populate(h, ["ibm"]) + added = qgraph.populate(h, ["ibm"], lister=_idle_lister(0)) assert added == {"ibm"} # one add_subgraph call carrying a qdevice_ibm vertex types = [n["metadata"]["type"] for n in h.added[0]["graph"]["nodes"]] assert "qdevice_ibm" in types and "qpu" in types +def test_populate_refuses_when_the_allocation_is_unreadable(): + """An unknown allocation count is not evidence the instance is idle.""" + + def _broken(handle): + raise RuntimeError("no scheduler") + + h = _FakeHandle(_base_graph()) + try: + qgraph.populate(h, ["ibm"], lister=_broken) + except RuntimeError as exc: + assert "could not be read" in str(exc) + else: + raise AssertionError("populate grew the graph without checking") + assert h.added == [] + + def test_populate_is_idempotent(): g = _base_graph() g["nodes"].append( diff --git a/tests/unit/test_handoff.py b/tests/unit/test_handoff.py index c15fb1f..6af1d87 100644 --- a/tests/unit/test_handoff.py +++ b/tests/unit/test_handoff.py @@ -172,3 +172,26 @@ def boom(handle, jobid, why): with pytest.raises(SystemExit): abort_held(None, 1, "whatever", cancel=boom) + + +def test_memo_carries_a_durable_release_marker(): + """A release lives only in scheduler memory, so a qmanager restart would + park the job again. The eventlog survives, and fluxion reads it there.""" + from flux_quantum.scout import post_session, SESSION_KEY + + sent = {} + + class Fut: + def get(self): + return None + + def rpc(topic, payload): + sent["topic"] = topic + sent["payload"] = payload + return Fut() + + post_session(None, 4021041664, "sess-1", rpc=rpc) + memo = sent["payload"]["memo"] + assert memo[SESSION_KEY] == "sess-1" + assert memo["released"] == 1 + assert sent["payload"]["id"] == 4021041664 From 755862c276a9356d6b42e7e790454317557708b9 Mon Sep 17 00:00:00 2001 From: vsoch Date: Sun, 6 Sep 2026 21:34:21 -0700 Subject: [PATCH 09/26] add match_coschedule Signed-off-by: vsoch --- flux_quantum/cli.py | 108 +++++++++----- flux_quantum/jobtap/quantum.c | 210 +++++++++++++--------------- flux_quantum/qresource.py | 2 + tests/integration/test-admission.sh | 32 +++++ tests/unit/test_cli_plugin.py | 93 ++++++++---- 5 files changed, 268 insertions(+), 177 deletions(-) diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index 8570bd0..1c8a064 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -14,6 +14,7 @@ """ import copy +import errno import json import os import sys @@ -37,29 +38,45 @@ def _wrap_commands(jobspec_dict, wrap_path): task["command"] = ["flux", "python", wrap_path, "--"] + cmd -def _check_pair_fits(handle, cores, rpc_fn=None): - """Refuse a pair that cannot be placed, before either half is created.""" +def _check_pair_fits(handle, classical, scout, rpc_fn=None): + """Refuse a pair that cannot be scheduled, before either half is created. + + Asks fluxion whether both halves could be placed together. It answers by + matching the pair for real and undoing it, so the answer is against the + live graph rather than a core count, and nothing is left placed. Neither + job exists yet, so no jobids are sent. + + A pair that is admitted and then cannot start is the expensive case: the + scout opens a metered vendor session and the classical half waits. + """ + payload = { + "check": True, + "jobs": [ + {"jobspec": scout, "op": "allocate"}, + {"jobspec": classical, "op": "reserve"}, + ], + } try: if rpc_fn is None: - resp = handle.rpc( - "job-manager.quantum.check", {"pair_cores": int(cores)} - ).get() + handle.rpc("sched-fluxion-resource.match_coschedule", payload).get() else: - resp = rpc_fn(handle, int(cores)) - except Exception: - # The plugin may be absent, or too old to serve the method. Admission is - # then whatever the plugin does at validate, which is how this worked - # before the check existed, so do not block the submit on it. - return - if not resp.get("ok", True): + rpc_fn(handle, payload) + except OSError as exc: + # EBUSY is fluxion saying the pair does not fit, ENODEV that it never + # could. Anything else means the question could not be asked at all, + # an older fluxion or none loaded, and that must not block a submit + # that would have worked before the check existed. + if exc.errno not in (errno.EBUSY, errno.ENODEV, errno.EINVAL): + return raise SystemExit( - "flux quantum: no room for this pair. It needs {} cores plus one " - "for the scout, and the cluster has {} free and {} preemptible. " - "Submitting anyway would open a vendor session for a pair that " - "cannot run.".format( - cores, resp.get("free", "?"), resp.get("preemptible", "?") - ) + "flux quantum: this pair cannot be scheduled together. The " + "classical half needs its cores and the scout needs one more, and " + "the cluster cannot place both. Submitting anyway would open a " + "vendor session for a pair that cannot run. ({})".format(exc) ) + except Exception: + # no usable handle, so there is nothing to ask + return def _safe_cancel(cancel_fn, handle, jobid, reason): @@ -120,6 +137,7 @@ def prepare_pair( # needs the classical size to do that. Counting here rather than walking the # jobspec in C. quantum["cores"] = qresource.count_cores(classical) + if job_env: sysattr.setdefault("environment", {}).update(job_env) @@ -129,42 +147,56 @@ def prepare_pair( # 9 core pair was admitted onto a cluster with 7 cores free, could not be # placed, and left preemption to clean up. # - # job-manager.quantum.check answers "does free + preemptible leave room for - # this classical and a scout", from occupancy the plugin already tracks. - # Asking here means a pair that cannot be placed never opens a session. + # So ask fluxion instead, which matches the pair against the live graph and + # undoes it. The scout jobspec built here is only for that question, so the + # job it would release does not exist yet and the id is a placeholder. Only + # the resources matter to the answer. # # There is a race between this answer and the submits below. That is # accepted: the plugin's job.validate check is still the backstop, and # either half is cancelled if the other fails. - _check_pair_fits(handle, quantum["cores"]) - - # if feasibility validation is on, an unsatisfiable request is rejected here - # and we abort before spending any quantum quota + # the foothold comes from the live graph so node->...->core matches reality. + # Done before either half is created, both because the check below needs the + # vendor device present to answer honestly, and because a graph failure then + # leaves no held job to clean up. try: - main_id = int(submit_fn(handle, json.dumps(classical))) + populate_fn(handle, [vendor]) + live = get_graph_fn(handle) except SystemExit: raise except Exception as exc: raise SystemExit( - "flux quantum: classical request not accepted " - "(unsatisfiable, or ingest error): {}".format(exc) + "flux quantum: could not prepare the quantum graph: {}".format(exc) ) - # the foothold comes from the live graph so node->...->core matches reality + # The scout jobspec built here is only for the question. The job it would + # release does not exist yet, so its id is a placeholder and only the + # resources matter to the answer. try: - populate_fn(handle, [vendor]) - live = get_graph_fn(handle) + probe = build_scout_jobspec( + vendor, + 0, + ncores=scout_cores, + duration=scout_duration, + live_graph=live, + scout_path=scout_path, + options=options, + ) + except Exception: + probe = None + if probe is not None: + _check_pair_fits(handle, classical, probe) + + # if feasibility validation is on, an unsatisfiable request is rejected here + # and we abort before spending any quantum quota + try: + main_id = int(submit_fn(handle, json.dumps(classical))) except SystemExit: raise except Exception as exc: - _safe_cancel( - cancel_fn, - handle, - main_id, - "flux quantum: aborting held classical (graph setup failed)", - ) raise SystemExit( - "flux quantum: could not prepare the quantum graph: {}".format(exc) + "flux quantum: classical request not accepted " + "(unsatisfiable, or ingest error): {}".format(exc) ) try: diff --git a/flux_quantum/jobtap/quantum.c b/flux_quantum/jobtap/quantum.c index 0181cc9..e663471 100644 --- a/flux_quantum/jobtap/quantum.c +++ b/flux_quantum/jobtap/quantum.c @@ -204,6 +204,55 @@ static int budget (void) return b > 0 ? b : 0; } +/* What a new pair could actually have. + * + * Everything here comes from the job table the plugin already keeps, so no + * query is needed. Three quantities, and the distinctions matter: + * + * free capacity not occupied by anything running + * preemptible cores held by running work that is not protected, so a pair + * can take them back + * promised cores already promised to pairs that were admitted but have + * not started. Their cores are not allocated, so they count as + * free, and would otherwise be handed out twice + * + * A running pair is deliberately not in promised: its cores are already out of + * free, and being protected they are not in preemptible either. + */ +static void capacity (int *freep, int *preemptiblep, int *promisedp) +{ + const char *key; + json_t *entry; + int running = 0, preemptible = 0, promised = 0; + + if (g_jobs) { + json_object_foreach (g_jobs, key, entry) { + int cores = 0, prot = 0, pair = 0; + double run = 0.0; + + if (json_unpack (entry, + "{s:i s:b s:b s:F}", + "cores", &cores, + "protected", &prot, + "pair", &pair, + "run", &run) < 0) + continue; + if (cores <= 0) + continue; + if (run > 0.0) { + running += cores; + if (!prot) + preemptible += cores; + } else if (pair) { + promised += cores; + } + } + } + *freep = budget () - running > 0 ? budget () - running : 0; + *preemptiblep = preemptible; + *promisedp = promised; +} + /* cores this pair reserves. The scout needs one of its own, so a classical of * n cores makes the pair n + 1. */ static int pair_cores (json_t *jobspec) @@ -328,12 +377,24 @@ static int validate_cb (flux_plugin_t *p, return flux_jobtap_reject_job (p, args, "quantum: attributes.system.quantum.cores is missing, " "so this pair cannot be accounted for"); - if (g_used_cores + want > budget ()) + /* Admission has to mean this pair will be able to run, not just that + * the site has not sold too many. So ask against capacity that is + * actually reachable: what is free, plus what can be taken back from + * unprotected work, less what is already promised to pairs that have + * not started. Comparing against total capacity instead admitted a + * pair onto a machine with no room for it, and left preemption to + * discover that afterwards with a metered session already open. + */ + int freec = 0, preemptible = 0, promised = 0; + capacity (&freec, &preemptible, &promised); + if (want + promised > freec + preemptible) return flux_jobtap_reject_job (p, args, - "quantum: no room for another pair. %d of %d cores are " - "promised to unfinished pairs and this one needs %d. " - "Wait for one to finish or ask for fewer cores", - g_used_cores, budget (), want); + "quantum: no room for another pair. This one needs %d " + "cores, %d are already promised to pairs that have not " + "started, and only %d are reachable (%d free, %d that " + "could be preempted). Wait for one to finish or ask for " + "fewer cores", + want, promised, freec + preemptible, freec, preemptible); } /* policy satisfied; the job proceeds. The hold is set at submission by the @@ -462,9 +523,11 @@ static int sched_cb (flux_plugin_t *p, return 0; g->p = p; g->id = id; - /* cores is the whole pair, classical plus one for the scout. Neither may - * be running, so make room for both. */ - g->cores = cores; + /* The entry records what this job itself asks for, so add the scout's core + * back to get what the pair needs. Neither half may be running yet, so + * make room for both. + */ + g->cores = cores + 1; if (!(g->w = flux_timer_watcher_create (flux_get_reactor (h), g_preempt_after, 0., grace_cb, g))) { @@ -497,8 +560,12 @@ static int new_cb (flux_plugin_t *p, * quantum pair preempts, so the plugin has to know they exist. Called * again on replay, and json_object_set overwrites, so this is idempotent * apart from losing the run time, which the state callback restores. */ + /* Record what this job itself asks for, not what its pair asks for. The + * scout has its own entry, so charging the classical for n + 1 would count + * the scout's core twice once it starts running. + */ track (id, - vendor ? pair_cores (jobspec) : count_tree_cores_top (jobspec), + count_tree_cores_top (jobspec), needs_protection (jobspec, vendor), vendor ? 1 : 0); if (g_preempt_after > 0.0) @@ -564,109 +631,37 @@ static int run_cb (flux_plugin_t *p, return 0; } -/* Sum cores over running jobs, and separately over running jobs that a pair is - * allowed to preempt. track() records every job, not only pairs, so this is - * derivable from state the plugin already holds and needs no query to fluxion. - * - * Note fluxion cannot answer this question. It knows free and allocated, but - * "allocated and preemptible" depends on attributes.system.protected, which is - * a concept of this plugin that fluxion has never seen. - */ -static void occupancy (int *runningp, int *preemptiblep) -{ - const char *key; - json_t *entry; - int running = 0, preemptible = 0; - - if (g_jobs) { - json_object_foreach (g_jobs, key, entry) { - int prot = 0, c = 0; - double run = 0.0; - if (json_unpack (entry, "{s:i s:b s:F}", - "cores", &c, "protected", &prot, "run", &run) < 0) - continue; - if (run <= 0.0 || c <= 0) - continue; /* not running, so occupying nothing */ - running += c; - if (!prot) - preemptible += c; - } - } - if (runningp) - *runningp = running; - if (preemptiblep) - *preemptiblep = preemptible; -} -/* Can a pair of pair_cores classical cores plus one scout core be placed, now - * or by preempting? Returns 1 for yes, 0 for no, and fills the counts. - * - * The scout's core comes out of the budget before the comparison, which is why - * this is "- 1" rather than comparing against pair_cores + 1. Same condition, - * written the way the requirement was specified. - */ -static int pair_fits (int pair_cores, int *freep, int *preemptiblep) -{ - int running = 0, preemptible = 0, budget_free; - - occupancy (&running, &preemptible); - budget_free = budget () - running; - if (budget_free < 0) - budget_free = 0; - if (freep) - *freep = budget_free; - if (preemptiblep) - *preemptiblep = preemptible; - if (pair_cores <= 0) - return 0; - return (budget_free + preemptible - 1) >= pair_cores; -} -/* job-manager.quantum.check - * - * Asked by the submitting client BEFORE either half of a pair is created, so - * that a pair which cannot be placed never opens a vendor session. The old - * check compared outstanding pairs against total capacity, which is a quota - * across pairs rather than a feasibility test, so a 9 core pair was admitted - * onto a cluster with 7 cores free. +/* Report the configuration this plugin was actually given. * - * There is a race between this answer and the submits that follow it. That is - * accepted: the caller cancels whichever half survives if the other fails, and - * the job.validate check below remains as a backstop. + * flux jobtap query otherwise reports only a name and a path, so there is no + * way afterwards to tell what a run was configured with. A whole campaign was + * once collected with preemption off and nobody could tell from the data, so + * this exists to make the settings recoverable. */ -static void check_cb (flux_t *h, flux_msg_handler_t *mh, - const flux_msg_t *msg, void *arg) +static int query_cb (flux_plugin_t *p, + const char *topic, + flux_plugin_arg_t *args, + void *arg) { - int pair_cores = 0, free_cores = 0, preemptible = 0, ok; - - if (flux_request_unpack (msg, NULL, "{s:i}", "pair_cores", &pair_cores) < 0) { - if (flux_respond_error (h, msg, EPROTO, NULL) < 0) - flux_log_error (h, "quantum.check: error responding"); - return; - } - if (pair_cores <= 0) { - if (flux_respond_error (h, msg, EINVAL, - "pair_cores must be greater than zero") < 0) - flux_log_error (h, "quantum.check: error responding"); - return; - } - if (g_total_cores <= 0) { - /* admission control is off, so there is nothing to promise */ - if (flux_respond_pack (h, msg, "{s:b s:i s:i s:b}", - "ok", 1, "free", 0, "preemptible", 0, - "enforced", 0) < 0) - flux_log_error (h, "quantum.check: error responding"); - return; - } - ok = pair_fits (pair_cores, &free_cores, &preemptible); - if (flux_respond_pack (h, msg, "{s:b s:i s:i s:b}", - "ok", ok, "free", free_cores, - "preemptible", preemptible, "enforced", 1) < 0) - flux_log_error (h, "quantum.check: error responding"); + if (flux_plugin_arg_pack (args, + FLUX_PLUGIN_ARG_OUT, + "{s:i s:i s:f s:s s:s s:i}", + "total_cores", g_total_cores, + "reserve_cores", g_reserve_cores, + "preempt_after", g_preempt_after, + "vendors", g_vendors ? g_vendors : "", + "protect_types", g_protect_types ? g_protect_types : "", + "used_cores", g_used_cores) + < 0) + return -1; + return 0; } static const struct flux_plugin_handler handlers[] = { + { "plugin.query", query_cb, NULL }, { "job.validate", validate_cb, NULL }, { "job.new", new_cb, NULL }, { "job.state.run", run_cb, NULL }, @@ -757,15 +752,6 @@ int flux_plugin_init (flux_plugin_t *p) return -1; } - /* Expose job-manager.quantum.check. The _ex form with FLUX_ROLE_USER is - * required because the system instance sets allow-guest-user, so ordinary - * users must be able to ask before they submit. - */ - if (flux_jobtap_service_register_ex (p, "check", FLUX_ROLE_USER, - check_cb, NULL) < 0) { - flux_log_error (h, "quantum: could not register the check service"); - return -1; - } return 0; } diff --git a/flux_quantum/qresource.py b/flux_quantum/qresource.py index 8a9a48c..152c959 100644 --- a/flux_quantum/qresource.py +++ b/flux_quantum/qresource.py @@ -15,6 +15,8 @@ exclusivity. No flux imports, just dicts. """ +import uuid + QDEVICE_PREFIX = "qdevice_" QPU = "qpu" diff --git a/tests/integration/test-admission.sh b/tests/integration/test-admission.sh index 438022b..53670ba 100755 --- a/tests/integration/test-admission.sh +++ b/tests/integration/test-admission.sh @@ -132,6 +132,38 @@ else echo "FAIL the budget did not recover, cores are leaking on job.destroy"; rc=1 fi +echo "" +echo "=== a pair is refused when the machine cannot reach its cores ===" +# Admission has to mean the pair will be able to run. A quota against total +# capacity would let this through: the numbers fit the budget, but the cores +# are not reachable because protected work holds them. +flux jobtap remove quantum.so >/dev/null 2>&1 +flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ + vendors="mock" total_cores=16 reserve_cores=0 protect_types="qpu" +if big=$(submit_pair 12); then + echo " a 12 core pair is admitted: 12 promised, 1 running scout, of 16" + err=$(mktemp) + if submit_pair 4 >/dev/null 2>"$err"; then + echo "FAIL a second pair was admitted, it needs 5 and only 3 are reachable" + rc=1 + elif grep -q "no room for another pair" "$err"; then + echo " the second pair is refused, and the message says what is reachable" + grep -o "needs .*reachable" "$err" | head -1 | cut -c1-92 | sed 's/^/ /' + else + echo "FAIL refused for the wrong reason"; sed 's/^/ /' "$err"; rc=1 + fi + rm -f "$err" + # shellcheck disable=SC2086 + flux cancel $big >/dev/null 2>&1 +else + echo "FAIL the first pair should have fit a budget of 16" + rc=1 +fi +for _ in $(seq 1 20); do + flux jobs -no "{id}" 2>/dev/null | grep -q . || break + sleep 1 +done + echo "" echo "=== vendor policy still applies ===" if flux submit --quantum-vendor rigetti -n1 true >/dev/null 2>&1; then diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index 91ff8e5..c59cae2 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -245,19 +245,34 @@ def rejecting_submit(handle, jobspec_json): _run_prepare(stub_flux_cli, submit=rejecting_submit) -def test_prepare_pair_cancels_held_classical_on_populate_failure(stub_flux_cli): - calls_seen = {} +def test_populate_failure_leaves_no_held_classical_behind(stub_flux_cli): + """The graph is set up before either half is created, so a graph failure + has nothing to clean up. It used to happen after the classical was + submitted, which meant cancelling a job that had just been held.""" + seen = {} def boom_populate(handle, vendors): raise RuntimeError("add_subgraph failed") def rec_cancel(handle, jobid, reason): - calls_seen["cancelled"] = jobid + seen["cancelled"] = jobid + + def rec_submit(handle, jobspec_json): + seen["submitted"] = True + return 12345 with pytest.raises(SystemExit): - _run_prepare(stub_flux_cli, populate=boom_populate, cancel=rec_cancel) + _run_prepare( + stub_flux_cli, + populate=boom_populate, + cancel=rec_cancel, + submit=rec_submit, + ) - assert calls_seen.get("cancelled") == 12345 + assert ( + "submitted" not in seen + ), "the classical was created before the graph was ready" + assert "cancelled" not in seen, "nothing should need cancelling" def test_scout_duration_covers_the_classical(stub_flux_cli): @@ -349,40 +364,63 @@ def test_classical_carries_its_core_count(stub_flux_cli): assert "quantum" not in scout_sys or "cores" not in scout_sys.get("quantum", {}) -def test_check_pair_fits_rejects_when_there_is_no_room(stub_flux_cli): - """A pair that cannot be placed is refused before either half is created.""" +def test_check_pair_fits_rejects_when_the_pair_cannot_be_placed(stub_flux_cli): + """A pair that cannot be scheduled together is refused before either half + is created, because the scout would open a metered session for a classical + job that cannot start.""" + import errno + from flux_quantum import cli - def rpc(handle, cores): - return {"ok": False, "free": 7, "preemptible": 0, "enforced": True} + def rpc(handle, payload): + raise OSError(errno.EBUSY, "the group does not fit right now") try: - cli._check_pair_fits(None, 8, rpc_fn=rpc) + cli._check_pair_fits(None, {"a": 1}, {"b": 2}, rpc_fn=rpc) except SystemExit as exc: - assert "no room for this pair" in str(exc) - assert "7 free" in str(exc) + assert "cannot be scheduled together" in str(exc) else: - raise AssertionError("a pair with no room was allowed through") + raise AssertionError("a pair that cannot be placed was allowed through") -def test_check_pair_fits_allows_when_there_is_room(stub_flux_cli): +def test_check_pair_fits_sends_both_halves_and_their_ops(stub_flux_cli): + """The question is about the pair, so both jobspecs go, the scout to be + allocated and the classical to be reserved. No jobids, since neither job + exists yet.""" from flux_quantum import cli - def rpc(handle, cores): - return {"ok": True, "free": 8, "preemptible": 120, "enforced": True} + seen = {} + + def rpc(handle, payload): + seen.update(payload) + return {"fits": True} - cli._check_pair_fits(None, 8, rpc_fn=rpc) + cli._check_pair_fits(None, {"classical": 1}, {"scout": 1}, rpc_fn=rpc) + assert seen["check"] is True + ops = [j["op"] for j in seen["jobs"]] + assert sorted(ops) == ["allocate", "reserve"] + assert {"scout": 1} in [j["jobspec"] for j in seen["jobs"]] + assert {"classical": 1} in [j["jobspec"] for j in seen["jobs"]] + assert all("jobid" not in j for j in seen["jobs"]) -def test_check_pair_fits_is_permissive_when_the_service_is_absent(stub_flux_cli): - """An older plugin serves no check method. Admission then falls back to the - job.validate hook, which is how this worked before, so do not block.""" +def test_check_pair_fits_is_permissive_when_it_cannot_ask(stub_flux_cli): + """An older fluxion serves no such method, and none may be loaded at all. + Admission then falls back to the plugin's job.validate hook, which is how + this worked before the check existed, so do not block a submit.""" + import errno + from flux_quantum import cli - def rpc(handle, cores): - raise OSError("Function not implemented") + def missing(handle, payload): + raise OSError(errno.ENOSYS, "Function not implemented") + + cli._check_pair_fits(None, {"a": 1}, {"b": 2}, rpc_fn=missing) + + def broken(handle, payload): + raise RuntimeError("no usable handle") - cli._check_pair_fits(None, 8, rpc_fn=rpc) + cli._check_pair_fits(None, {"a": 1}, {"b": 2}, rpc_fn=broken) def test_prepare_pair_asks_before_submitting_anything(stub_flux_cli, monkeypatch): @@ -392,16 +430,17 @@ def test_prepare_pair_asks_before_submitting_anything(stub_flux_cli, monkeypatch order = [] - def fake_check(handle, cores, rpc_fn=None): - order.append(("checked", cores)) + def fake_check(handle, classical, scout, rpc_fn=None): + order.append(("checked", classical, scout)) monkeypatch.setattr(cli, "_check_pair_fits", fake_check) def submit(handle, jobspec_json): - order.append(("submitted", None)) + order.append(("submitted", None, None)) return 12345 _run_prepare(stub_flux_cli, submit=submit) assert order[0][0] == "checked", order - assert order[0][1] > 0 + # both halves were described to the check + assert order[0][1] is not None and order[0][2] is not None assert any(x[0] == "submitted" for x in order) From d33432c1faa51e693ccc61bffa5935e1b29eac56 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 9 Sep 2026 06:14:34 +0000 Subject: [PATCH 10/26] refactor: drop unused code and defer the qrmi import validator.py, submit_scout and acquisition_token_env had no callers. wrap.py imports the token suffix instead of repeating the literal. open_session now checks its arguments before importing qrmi, so a missing resource is reported even where qrmi is not installed, which is what two of the unit tests already expected. --- flux_quantum/__init__.py | 4 ++-- flux_quantum/backends/qrmi.py | 11 ++++------- flux_quantum/launch.py | 8 -------- flux_quantum/validator.py | 32 -------------------------------- flux_quantum/wrap.py | 3 ++- 5 files changed, 8 insertions(+), 50 deletions(-) delete mode 100644 flux_quantum/validator.py diff --git a/flux_quantum/__init__.py b/flux_quantum/__init__.py index 02d1d92..66e14e1 100644 --- a/flux_quantum/__init__.py +++ b/flux_quantum/__init__.py @@ -1,8 +1,8 @@ """Quantum and classical coscheduling components for Flux. Bundled and versioned together, following Flux packaging guidance. A submit -time CLI plugin, a job validator, an owner side jobtap gate, vendor backends, +time CLI plugin, an owner side jobtap gate, vendor backends, and the scout and wrapper pair that hand off the session. """ -__version__ = "0.0.1" +__version__ = "0.0.11" diff --git a/flux_quantum/backends/qrmi.py b/flux_quantum/backends/qrmi.py index 160c2b1..e392ae3 100644 --- a/flux_quantum/backends/qrmi.py +++ b/flux_quantum/backends/qrmi.py @@ -241,14 +241,7 @@ def job_environment(self, options): "QRMI_JOB_QPU_TYPES": options.get("type") or self.default_type, } - @staticmethod - def acquisition_token_env(resource, session): - """The variable QRMI expects the acquisition token in.""" - return {resource + ACQUISITION_TOKEN: session} - def open_session(self, options): - from qrmi import QuantumResource - resource = options.get("resource") rtype = options.get("type") or self.default_type if not resource: @@ -263,6 +256,10 @@ def open_session(self, options): self.name, resource, ", ".join(resource + s for s in missing) ) ) + # deferred so a missing resource or variable is reported as a + # ValueError above even when qrmi itself is not installed + from qrmi import QuantumResource + self._resource = QuantumResource(resource, resource_type(rtype)) try: self._lock = self._resource.acquire() diff --git a/flux_quantum/launch.py b/flux_quantum/launch.py index fd7504d..c8b4f22 100644 --- a/flux_quantum/launch.py +++ b/flux_quantum/launch.py @@ -74,14 +74,6 @@ def build_scout_jobspec( } -def submit_scout(handle, vendor, hold_job, ncores=1, duration=0, session=None): - """Submit the scout job and return the flux JobID.""" - from flux.job import submit - - jobspec = build_scout_jobspec(vendor, hold_job, ncores, duration, session=session) - return submit(handle, json.dumps(jobspec)) - - def main(): ap = argparse.ArgumentParser( prog="quantum-scout-launch", diff --git a/flux_quantum/validator.py b/flux_quantum/validator.py deleted file mode 100644 index 2bcd09d..0000000 --- a/flux_quantum/validator.py +++ /dev/null @@ -1,32 +0,0 @@ -"""Ingest side validator that checks the vendor is set and known. - -The credential check runs in the CLI plugin, where the user environment is -visible. This is the server side backstop for attributes that are not secret, -rejecting before the job takes resources. - -Add this module to the ingest validator plugins to wire it up. -""" - -from .backends import known_vendors - - -def _vendor_of(jobspec): - try: - return ( - jobspec.get("attributes", {}) - .get("system", {}) - .get("quantum", {}) - .get("vendor") - ) - except AttributeError: - return None - - -def validate(jobspec): - """Return None when ok, otherwise an error string. No secret checks.""" - vendor = _vendor_of(jobspec) - if vendor is None: - return None # not a quantum job - if vendor not in known_vendors(): - return "quantum: unknown/unpermitted vendor '{}'".format(vendor) - return None diff --git a/flux_quantum/wrap.py b/flux_quantum/wrap.py index 7077bb0..78c6678 100644 --- a/flux_quantum/wrap.py +++ b/flux_quantum/wrap.py @@ -11,6 +11,7 @@ import signal import sys +from flux_quantum.backends.qrmi import ACQUISITION_TOKEN from flux_quantum.scout import SESSION_KEY @@ -56,7 +57,7 @@ def session_environment(session, resources=None): for resource in resources.split(","): resource = resource.strip() if resource: - env[resource + "_QRMI_JOB_ACQUISITION_TOKEN"] = session + env[resource + ACQUISITION_TOKEN] = session return env From 0fc4d4c8b4febe16c29fa74d43f42a5878399525 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 9 Sep 2026 06:14:34 +0000 Subject: [PATCH 11/26] fix: jobtap survives reload and unload Only job.new is replayed on reload, so running jobs looked idle afterwards and their cores could be handed out twice. The plugin now uses the job state it is given. Grace timers are tracked and destroyed on unload, since one firing later would crash the job manager. The used_cores budget fed no decision and is gone. query reports the capacity numbers instead. Rejected jobs are no longer marked protected, which left a pending update for job.destroy to complain about. A new integration test covers unload and reload. --- .github/workflows/test.yml | 4 + flux_quantum/jobtap/quantum.c | 352 ++++++++++++------------------- tests/integration/test-reload.sh | 124 +++++++++++ 3 files changed, 266 insertions(+), 214 deletions(-) create mode 100644 tests/integration/test-reload.sh diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 515c156..83d36d3 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -121,6 +121,10 @@ jobs: - 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 diff --git a/flux_quantum/jobtap/quantum.c b/flux_quantum/jobtap/quantum.c index e663471..9d47370 100644 --- a/flux_quantum/jobtap/quantum.c +++ b/flux_quantum/jobtap/quantum.c @@ -1,64 +1,39 @@ /* - * quantum-jobtap: site policy for quantum and classical coscheduling. + * Site policy for quantum and classical coscheduling. * - * A quantum job is the held classical half of a pair, identified by - * attributes.system.quantum.vendor = "" - * attributes.system.quantum.cores = - * both stamped by the submit plugin. The scout carries neither, so it is not - * counted twice. - * - * Three jobs enforced here. + * The submit plugin stamps the held classical half of a pair with + * attributes.system.quantum.vendor and attributes.system.quantum.cores. + * The scout carries neither, only its qpu request. * * Vendors. Only configured vendors are permitted. * - * Protection. Both halves of a pair are marked - * attributes.system.protected = "quantum" - * which exempts them from preemption. The key is generic on purpose, so a - * plugin protecting any other scarce external resource can reuse it. It is set - * here and never by the submitter, and a submission carrying it is rejected. - * - * Admission. Launching a scout commits real money the moment it acquires a - * vendor session, so the classical half has to be placeable. Every unfinished - * pair, held or running, reserves cores + 1 against a budget, and a new pair - * that would exceed the budget is rejected at submit rather than admitted and - * left to wait with a meter running. Held pairs count because a held pair is a - * promise. - * - * Accounting lives in job.new and not job.validate, because the job manager - * replays job.new for every active job on restart or plugin reload, so the - * budget rebuilds itself. job.destroy gives the cores back. + * Protection. Both halves of a pair get attributes.system.protected, which + * exempts them from preemption. The plugin sets it and a submission that + * carries it is rejected. * - * Load: - * flux jobtap load quantum.so vendors="ibm,braket" \ - * total_cores=128 reserve_cores=32 + * Admission. A scout costs money once it holds a vendor session, so a pair + * is admitted only if its cores plus one for the scout fit in what is + * reachable. Reachable means cores nothing is running on, plus cores held by + * unprotected work that could be preempted, less cores promised to pairs + * that were admitted but have not started. The job table behind this is + * built in job.new, which the job manager replays on restart and reload. * - * total_cores cores available to quantum pairs. 0 disables admission - * control and leaves only the vendor check. - * reserve_cores held back for ordinary classical work. - * protect_types resource types that mark a job as needing external access, - * default qpu. Used to protect the scout, which carries no - * quantum attributes. - * preempt_after seconds to wait after the scout releases the classical - * before cancelling unprotected jobs to make room. 0, the - * default, never preempts and the pair just waits. + * Preemption. When the classical half reaches SCHED a timer starts, and if + * it has not started by preempt_after seconds the youngest unprotected jobs + * are cancelled until enough cores are free. The clock starts at submission + * rather than at the scout's release, because on a full machine the scout + * cannot get a core and there would be no release to wait for. A vendor + * queue longer than the timer therefore frees cores before the session + * exists. Cancelling is destructive, so preemption is off unless + * preempt_after is set. * - * Preemption. Admission promises that the classical half can be placed, but - * ordinary work may be sitting on the cores when the scout releases it. So the - * plugin watches for the session memo, which the scout posts immediately before - * releasing, waits preempt_after seconds, and if the job still has not started - * it cancels unprotected jobs until enough cores are freed. Youngest first, - * because flux cancels rather than requeues and the youngest job loses the - * least work. + * Load with + * flux jobtap load quantum.so vendors=ibm,braket total_cores=128 \ + * reserve_cores=32 protect_types=qpu preempt_after=30 * - * NOTES - * Cancelling is destructive. The victim loses whatever it had done, so - * preemption is off unless preempt_after is set. Protected jobs are never - * victims, and both halves of a pair are protected, so preemption cannot eat - * another pair. - * - * Freeing cores is not instant. Cancel, then epilog, then the scheduler - * notices. The classical job starts after that, not at the moment of the - * cancel, and the vendor session covers the gap. + * total_cores 0 disables admission control. protect_types names the resource + * types that mark a job as one of ours, default qpu. flux jobtap query + * reports the configuration in force and the capacity numbers. */ #include #include @@ -69,36 +44,22 @@ #include #include -/* default vendor allowlist; over/replace via load-time config "vendors=" */ -static char *g_vendors = NULL; /* comma-separated, e.g. "ibm,braket" */ - -/* Jobs the plugin knows about, keyed by jobid as a string. - * id -> cores, protected, run - * A running unprotected job is a preemption candidate, and the run time orders - * them so the youngest goes first. - */ -static json_t *g_jobs = NULL; - +static char *g_vendors = NULL; /* comma separated, default ibm,braket */ +static char *g_protect_types = NULL; /* comma separated, default qpu */ static double g_preempt_after = 0.0; /* 0 disables preemption */ - -static int g_total_cores = 0; /* 0 means admission control is off */ +static int g_total_cores = 0; /* 0 disables admission control */ static int g_reserve_cores = 0; -static int g_used_cores = 0; /* cores promised to unfinished pairs */ -/* marks a job as counted, so job.destroy only gives back what was taken */ -static const char *AUX_KEY = "quantum::cores"; +/* every job the plugin has seen, keyed by jobid, with its cores, whether it + * is protected, whether it is the classical half of a pair, and the time it + * started or 0 */ +static json_t *g_jobs = NULL; -/* Marks a job as not preemptible. Deliberately not under quantum, so any - * policy plugin protecting a scarce external resource can use the same key and - * the preemption logic stays generic. The value is a reason, and preemption - * cares only that it is present. */ +/* Not under quantum on purpose, so a plugin protecting some other scarce + * resource can use the same key. The value is a reason and only presence + * matters. */ static const char *PROTECT_KEY = "attributes.system.protected"; -/* resource types that identify a job as needing external quantum access. The - * scout carries no quantum attributes, only this request, and a request cannot - * be faked into existence because fluxion has to match it against the graph. */ -static char *g_protect_types = NULL; /* comma separated, e.g. "qpu" */ - static int type_protected (const char *type) { const char *list = g_protect_types ? g_protect_types : "qpu"; @@ -151,9 +112,8 @@ static int needs_protection (json_t *jobspec, const char *vendor) return 0; } -/* Set by this plugin, never by the submitter. A user who sets it would make - * their own job unpreemptible, so it is rejected rather than stripped, because - * silently ignoring an attempted privilege grab hides it. */ +/* A user who sets the key would make their own job unpreemptible, so the job + * is rejected rather than quietly stripped. */ static int user_set_protection (json_t *jobspec) { json_t *val = NULL; @@ -168,8 +128,11 @@ static const char *idkey (flux_jobid_t id, char *buf, size_t len) return buf; } -/* remember a job so it can be considered as a victim later */ -static void track (flux_jobid_t id, int cores, int protected_, int pair) +static void track (flux_jobid_t id, + int cores, + int protected_, + int pair, + double run) { char key[32]; json_t *entry; @@ -178,7 +141,7 @@ static void track (flux_jobid_t id, int cores, int protected_, int pair) return; entry = json_pack ("{s:i s:b s:b s:f}", "cores", cores, "protected", protected_, - "pair", pair, "run", 0.0); + "pair", pair, "run", run); if (entry) (void) json_object_set_new (g_jobs, idkey (id, key, sizeof (key)), entry); } @@ -204,21 +167,10 @@ static int budget (void) return b > 0 ? b : 0; } -/* What a new pair could actually have. - * - * Everything here comes from the job table the plugin already keeps, so no - * query is needed. Three quantities, and the distinctions matter: - * - * free capacity not occupied by anything running - * preemptible cores held by running work that is not protected, so a pair - * can take them back - * promised cores already promised to pairs that were admitted but have - * not started. Their cores are not allocated, so they count as - * free, and would otherwise be handed out twice - * - * A running pair is deliberately not in promised: its cores are already out of - * free, and being protected they are not in preemptible either. - */ +/* Free is the budget less everything running. Preemptible is what running + * unprotected work holds. Promised is what admitted pairs that have not + * started will need, which still looks free and must not be handed out + * twice. A running pair is in neither of the last two. */ static void capacity (int *freep, int *preemptiblep, int *promisedp) { const char *key; @@ -253,8 +205,7 @@ static void capacity (int *freep, int *preemptiblep, int *promisedp) *promisedp = promised; } -/* cores this pair reserves. The scout needs one of its own, so a classical of - * n cores makes the pair n + 1. */ +/* the classical's cores plus one for the scout */ static int pair_cores (json_t *jobspec) { json_int_t cores = 0; @@ -266,8 +217,7 @@ static int pair_cores (json_t *jobspec) return (int)cores + 1; } -/* cores an ordinary job asks for. The submit plugin stamps the count for our - * own pairs, but a victim candidate is any job, so count its tree. */ +/* cores any job asks for, from its resource tree */ static int count_tree_cores (json_t *resources, int factor); static int count_tree_cores_top (json_t *jobspec) @@ -351,40 +301,18 @@ static int validate_cb (flux_plugin_t *p, "as not preemptible, which is not yours to decide"); /* attributes.system.quantum.vendor -- absent => not a quantum job */ - (void) json_unpack (jobspec, "{s:{s:{s:{s:s}}}}", - "attributes", "system", "quantum", "vendor", &vendor); - - /* the scout has no quantum attributes, only the resource request, and it - * must be protected too or preempting it would leak a vendor session */ - if (needs_protection (jobspec, vendor)) { - if (flux_jobtap_jobspec_update_pack (p, "{s:s}", - PROTECT_KEY, "quantum") < 0) - flux_log (flux_jobtap_get_flux (p), LOG_ERR, - "quantum: could not mark the job protected, it would be " - "preemptible"); - } - - if (!vendor) - return 0; + vendor = job_vendor (jobspec); - if (!vendor_allowed (vendor)) + if (vendor && !vendor_allowed (vendor)) return flux_jobtap_reject_job (p, args, "quantum: vendor '%s' is not permitted at this site", vendor); - if (g_total_cores > 0) { + if (vendor && g_total_cores > 0) { int want = pair_cores (jobspec); if (want <= 0) return flux_jobtap_reject_job (p, args, "quantum: attributes.system.quantum.cores is missing, " "so this pair cannot be accounted for"); - /* Admission has to mean this pair will be able to run, not just that - * the site has not sold too many. So ask against capacity that is - * actually reachable: what is free, plus what can be taken back from - * unprotected work, less what is already promised to pairs that have - * not started. Comparing against total capacity instead admitted a - * pair onto a machine with no room for it, and left preemption to - * discover that afterwards with a metered session already open. - */ int freec = 0, preemptible = 0, promised = 0; capacity (&freec, &preemptible, &promised); if (want + promised > freec + preemptible) @@ -397,8 +325,16 @@ static int validate_cb (flux_plugin_t *p, want, promised, freec + preemptible, freec, preemptible); } - /* policy satisfied; the job proceeds. The hold is set at submission by the - * submit plugin, and the scout releases it once the session is live. */ + /* Mark it protected last. A jobspec update left pending on a rejected + * job makes job.destroy complain. The scout is protected too, or + * preempting it would leak a vendor session. */ + if (needs_protection (jobspec, vendor)) { + if (flux_jobtap_jobspec_update_pack (p, "{s:s}", + PROTECT_KEY, "quantum") < 0) + flux_log (flux_jobtap_get_flux (p), LOG_ERR, + "quantum: could not mark the job protected, it would be " + "preemptible"); + } return 0; } @@ -456,13 +392,32 @@ static int preempt_for (flux_plugin_t *p, flux_jobid_t sparing, int need) return freed; } +/* Pending grace timers are kept on a list so unloading the plugin can destroy + * them. One firing after unload would take the job manager down. */ struct grace { flux_plugin_t *p; flux_jobid_t id; int cores; flux_watcher_t *w; + struct grace *next; }; +static struct grace *g_grace = NULL; + +static void grace_destroy (struct grace *g) +{ + struct grace **pp; + + for (pp = &g_grace; *pp; pp = &(*pp)->next) { + if (*pp == g) { + *pp = g->next; + break; + } + } + flux_watcher_destroy (g->w); + free (g); +} + static void grace_cb (flux_reactor_t *r, flux_watcher_t *w, int revents, @@ -485,18 +440,10 @@ static void grace_cb (flux_reactor_t *r, "%d cores", (uintmax_t) g->id, g_preempt_after, g->cores); (void) preempt_for (g->p, g->id, g->cores); done: - flux_watcher_destroy (w); - free (g); + grace_destroy (g); } -/* Start the clock when the classical half reaches SCHED. - * - * The memo would be the natural signal, since the scout posts it immediately - * before releasing. But on a full machine the scout cannot get a core either, - * so there is no memo and the pair waits forever with nothing to trigger on. - * Admission already promised this pair room, so the promise is what the clock - * hangs off, not the handover. - */ +/* start the grace timer when the classical half of a pair reaches SCHED */ static int sched_cb (flux_plugin_t *p, const char *topic, flux_plugin_arg_t *args, @@ -523,23 +470,20 @@ static int sched_cb (flux_plugin_t *p, return 0; g->p = p; g->id = id; - /* The entry records what this job itself asks for, so add the scout's core - * back to get what the pair needs. Neither half may be running yet, so - * make room for both. - */ - g->cores = cores + 1; + g->cores = cores + 1; /* neither half may be running, make room for both */ if (!(g->w = flux_timer_watcher_create (flux_get_reactor (h), g_preempt_after, 0., grace_cb, g))) { free (g); return 0; } + g->next = g_grace; + g_grace = g; flux_watcher_start (g->w); return 0; } -/* Accounting. Called for new jobs and replayed for active ones on restart or - * plugin reload, which is how the budget survives both. */ +/* called for new jobs, and replayed for active ones on restart and reload */ static int new_cb (flux_plugin_t *p, const char *topic, flux_plugin_arg_t *args, @@ -548,42 +492,33 @@ static int new_cb (flux_plugin_t *p, json_t *jobspec = NULL; flux_jobid_t id; const char *vendor; - int want; + int state = 0; + double t_submit = 0.0; + double run = 0.0; if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, - "{s:I s:o}", "id", &id, - "jobspec", &jobspec) < 0) + "{s:I s:o s?i s?F}", + "id", &id, + "jobspec", &jobspec, + "state", &state, + "t_submit", &t_submit) < 0) return 0; vendor = job_vendor (jobspec); - /* Track every job, not only ours. An unprotected running job is what a - * quantum pair preempts, so the plugin has to know they exist. Called - * again on replay, and json_object_set overwrites, so this is idempotent - * apart from losing the run time, which the state callback restores. */ - /* Record what this job itself asks for, not what its pair asks for. The - * scout has its own entry, so charging the classical for n + 1 would count - * the scout's core twice once it starts running. - */ + /* job.state.run is not replayed, so a job already running on replay + * would otherwise look idle. The start time is not in the args either, so + * the submit time stands in and keeps the youngest first order. */ + if (state == FLUX_JOB_STATE_RUN || state == FLUX_JOB_STATE_CLEANUP) + run = t_submit > 0.0 ? t_submit : 1.0; + + /* Every job is tracked, since any unprotected one is a preemption + * candidate. Each entry holds what that job alone asks for. The scout has + * its own entry, so the classical is not charged for it. */ track (id, count_tree_cores_top (jobspec), needs_protection (jobspec, vendor), - vendor ? 1 : 0); - if (g_preempt_after > 0.0) - (void) flux_jobtap_job_subscribe (p, id); - - if (g_total_cores <= 0) - return 0; - if (!vendor) - return 0; - if ((want = pair_cores (jobspec)) <= 0) - return 0; - - /* record what was taken, so destroy gives back the same amount even if the - * jobspec is no longer readable by then */ - if (flux_jobtap_job_aux_set (p, id, AUX_KEY, - (void *)(intptr_t)want, NULL) < 0) - return 0; - g_used_cores += want; + vendor ? 1 : 0, + run); return 0; } @@ -593,25 +528,14 @@ static int destroy_cb (flux_plugin_t *p, void *arg) { flux_jobid_t id; - intptr_t want; if (flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN, "{s:I}", "id", &id) < 0) return 0; - /* a job rejected at validate was never counted, so there is no aux and - * nothing to give back */ forget (id); - want = (intptr_t) flux_jobtap_job_aux_get (p, id, AUX_KEY); - if (want <= 0) - return 0; - g_used_cores -= (int)want; - if (g_used_cores < 0) - g_used_cores = 0; return 0; } -/* Note when a job starts, so preemption can take the youngest first and so the - * grace timer can tell whether the classical job got going on its own. */ static int run_cb (flux_plugin_t *p, const char *topic, flux_plugin_arg_t *args, @@ -631,35 +555,45 @@ static int run_cb (flux_plugin_t *p, return 0; } - - - -/* Report the configuration this plugin was actually given. - * - * flux jobtap query otherwise reports only a name and a path, so there is no - * way afterwards to tell what a run was configured with. A whole campaign was - * once collected with preemption off and nobody could tell from the data, so - * this exists to make the settings recoverable. - */ +/* flux jobtap query. Reports the configuration in force, so an experiment can + * record what it actually ran with, and the current capacity numbers. */ static int query_cb (flux_plugin_t *p, const char *topic, flux_plugin_arg_t *args, void *arg) { + int freec = 0, preemptible = 0, promised = 0; + + capacity (&freec, &preemptible, &promised); if (flux_plugin_arg_pack (args, FLUX_PLUGIN_ARG_OUT, - "{s:i s:i s:f s:s s:s s:i}", + "{s:i s:i s:f s:s s:s s:i s:i s:i s:i}", "total_cores", g_total_cores, "reserve_cores", g_reserve_cores, "preempt_after", g_preempt_after, "vendors", g_vendors ? g_vendors : "", "protect_types", g_protect_types ? g_protect_types : "", - "used_cores", g_used_cores) + "free_cores", freec, + "preemptible_cores", preemptible, + "promised_cores", promised, + "tracked_jobs", g_jobs ? (int) json_object_size (g_jobs) : 0) < 0) return -1; return 0; } +static void plugin_destroy (void *arg) +{ + while (g_grace) + grace_destroy (g_grace); + json_decref (g_jobs); + g_jobs = NULL; + free (g_vendors); + g_vendors = NULL; + free (g_protect_types); + g_protect_types = NULL; +} + static const struct flux_plugin_handler handlers[] = { { "plugin.query", query_cb, NULL }, { "job.validate", validate_cb, NULL }, @@ -681,16 +615,15 @@ int flux_plugin_init (flux_plugin_t *p) if (!(g_jobs = json_object ())) return -1; + if (flux_plugin_aux_set (p, NULL, g_jobs, plugin_destroy) < 0) { + json_decref (g_jobs); + g_jobs = NULL; + return -1; + } - /* One unpack sets everything, so any single key of an unexpected type - * returns -1 and none of it is applied. The defaults that would then be - * left in place are total_cores 0, which turns admission control off, and - * preempt_after 0, which turns preemption off. Loading unenforced and - * saying nothing is worse than not loading, so refuse. - * - * ENOENT is different. It means no conf was given at all, which is a - * legitimate way to ask for the defaults. - */ + /* A config that cannot be read would leave admission control and + * preemption off without saying so, so refuse to load. ENOENT just means + * no config was given. */ rc = flux_plugin_conf_unpack (p, "{s?s s?s s?i s?i s?F}", "vendors", &vendors, "protect_types", &protect_types, @@ -715,10 +648,6 @@ int flux_plugin_init (flux_plugin_t *p) g_preempt_after = preempt_after; } - /* State what is in force. A campaign was once run with preemption off - * because the key was absent from the config, and nothing in the logs or - * the results said so. - */ flux_log (h, LOG_INFO, "quantum: vendors=%s total_cores=%d reserve_cores=%d " "protect_types=%s preempt_after=%.1f", @@ -742,11 +671,6 @@ int flux_plugin_init (flux_plugin_t *p) "will hold an open vendor session until the machine frees " "up on its own"); - /* The plugin name has to be set first: the service is named - * job-manager.., so registering it before - * flux_plugin_register fails with EINVAL and takes the job manager down - * with it. - */ if (flux_plugin_register (p, "quantum", handlers) < 0) { flux_log_error (h, "quantum: could not register handlers"); return -1; diff --git a/tests/integration/test-reload.sh b/tests/integration/test-reload.sh new file mode 100644 index 0000000..724c9a2 --- /dev/null +++ b/tests/integration/test-reload.sh @@ -0,0 +1,124 @@ +#!/bin/bash +# Unload the plugin with a grace timer pending, then reload it with jobs +# running. The timer must die with the plugin, and the reloaded plugin must +# still see the running jobs, since job.new is replayed but job.state.run is +# not. +# +# FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-reload.sh +# +# Needs at least 3 cores. + +set -u +HERE=$(cd "$(dirname "$0")/../.." && pwd) +export FLUX_CLI_PLUGINPATH="$HERE/cli-plugins" +PLUGIN="$HERE/flux_quantum/jobtap/quantum.so" +rc=0 + +if [ -z "${FLUX_QUANTUM_MOCK:-}" ]; then + echo "FAIL FLUX_QUANTUM_MOCK must be exported before flux start" + exit 1 +fi + +CORES=$(flux resource list -s all -no "{ncores}" 2>/dev/null) +CORES=${CORES:-0} +echo "=== cluster has $CORES cores ===" +if [ "$CORES" -lt 3 ]; then + echo "SKIP need at least 3 cores" + exit 0 +fi + +make -s -C "$HERE/flux_quantum/jobtap" || { echo "FAIL build"; exit 1; } +flux python -m flux_quantum.populate mock >/dev/null 2>&1 \ + || echo "WARNING could not populate the graph up front" + +# query is by the registered name; fall back to the file name for older cores +query () { + flux jobtap query quantum 2>/dev/null || flux jobtap query quantum.so +} + +load () { + flux jobtap load "$PLUGIN" vendors="mock" total_cores="$CORES" \ + reserve_cores=0 protect_types="qpu" preempt_after=4 \ + || { echo "FAIL load"; exit 1; } +} + +echo "" +echo "=== fill every core with unprotected work ===" +load +fill="" +for _ in $(seq 1 "$CORES"); do + id=$(flux submit -n1 sleep 600) && fill="$fill $id" +done +for _ in $(seq 1 20); do + [ "$(flux jobs -no '{state}' | grep -c RUN)" -eq "$CORES" ] && break + sleep 0.5 +done +echo " $(flux jobs -no '{state}' | grep -c RUN) running" + +echo "" +echo "=== the plugin reports what it sees ===" +before=$(query) +echo "$before" | sed 's/^/ /' +if ! echo "$before" | grep -q '"preemptible_cores": *'"$CORES"; then + echo "FAIL expected $CORES preemptible cores with the machine full" + rc=1 +fi + +echo "" +echo "=== remove the plugin while a grace timer is pending ===" +# the classical half reaching SCHED arms the timer, and cancelling the pair +# afterwards leaves it armed with nothing promised +err=$(mktemp) +scout=$(flux submit --quantum-vendor mock -n1 -- true 2>"$err") +main=$(grep -oE 'held classical job [0-9]+' "$err" | awk '{print $NF}') +rm -f "$err" +if [ -z "$main" ]; then + echo "FAIL no pair was created"; exit 1 +fi +flux cancel "$main" "$scout" 2>/dev/null +flux jobtap remove quantum.so || { echo "FAIL remove"; rc=1; } +sleep 6 +if flux jobs -a >/dev/null 2>&1 && flux ping -c1 job-manager >/dev/null 2>&1; then + echo " job manager is still alive after the timer would have fired" +else + echo "FAIL the job manager is gone, the timer fired into the unloaded plugin" + rc=1 +fi +if flux dmesg 2>/dev/null | grep -qiE "segfault|job-manager.*(exited|crash)"; then + echo "FAIL the log shows a crash"; rc=1 +fi + +echo "" +echo "=== reload with jobs running: they must still count as running ===" +load +after=$(query) +echo "$after" | sed 's/^/ /' +if ! echo "$after" | grep -q '"preemptible_cores": *'"$CORES"; then + echo "FAIL after reload the running jobs are not counted, cores would be handed out twice" + rc=1 +fi +if ! echo "$after" | grep -q '"free_cores": *0'; then + echo "FAIL after reload the plugin thinks cores are free on a full machine" + rc=1 +fi + +echo "" +echo "=== and a pair too big for what is reachable is still refused ===" +err=$(mktemp) +if flux submit --quantum-vendor mock -n"$CORES" -- true >/dev/null 2>"$err"; then + echo "FAIL a pair needing $((CORES + 1)) was admitted onto $CORES reachable cores" + rc=1 +elif grep -q "no room for another pair" "$err"; then + echo " refused, and the message says why" +else + echo "FAIL refused for the wrong reason"; sed 's/^/ /' "$err"; rc=1 +fi +rm -f "$err" + +# shellcheck disable=SC2086 +flux cancel $fill 2>/dev/null +flux jobtap remove quantum.so 2>/dev/null + +echo "" +echo "=== reload test rc=$rc ===" +exit "$rc" From 6591c70f9b3ddd8d078fc977a0ec5211a1b9b103 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 9 Sep 2026 06:14:34 +0000 Subject: [PATCH 12/26] test: size the admission test from the machine The last section assumed 16 cores. Fluxion refuses a pair bigger than the graph before the budget is asked, so on a smaller box the test failed for the wrong reason. --- tests/integration/test-admission.sh | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/integration/test-admission.sh b/tests/integration/test-admission.sh index 53670ba..f2ae75c 100755 --- a/tests/integration/test-admission.sh +++ b/tests/integration/test-admission.sh @@ -134,14 +134,21 @@ fi echo "" echo "=== a pair is refused when the machine cannot reach its cores ===" -# Admission has to mean the pair will be able to run. A quota against total -# capacity would let this through: the numbers fit the budget, but the cores -# are not reachable because protected work holds them. +# The numbers fit the budget but the cores are held by protected work. Sized +# from the machine because fluxion refuses a pair bigger than the graph before +# the budget is asked. A pair of C-4 leaves C-4 promised and one scout +# running, so 3 are reachable, and a 4 core pair needs 5. +CORES=$(flux resource list -s all -no "{ncores}" 2>/dev/null) +CORES=${CORES:-0} +if [ "$CORES" -lt 6 ]; then + echo "SKIP need at least 6 cores, have $CORES" +else +FIRST=$((CORES - 4)) flux jobtap remove quantum.so >/dev/null 2>&1 flux jobtap load "$HERE/flux_quantum/jobtap/quantum.so" \ - vendors="mock" total_cores=16 reserve_cores=0 protect_types="qpu" -if big=$(submit_pair 12); then - echo " a 12 core pair is admitted: 12 promised, 1 running scout, of 16" + vendors="mock" total_cores="$CORES" reserve_cores=0 protect_types="qpu" +if big=$(submit_pair "$FIRST"); then + echo " a $FIRST core pair is admitted: $FIRST promised, 1 running scout, of $CORES" err=$(mktemp) if submit_pair 4 >/dev/null 2>"$err"; then echo "FAIL a second pair was admitted, it needs 5 and only 3 are reachable" @@ -156,9 +163,10 @@ if big=$(submit_pair 12); then # shellcheck disable=SC2086 flux cancel $big >/dev/null 2>&1 else - echo "FAIL the first pair should have fit a budget of 16" + echo "FAIL the first pair should have fit a budget of $CORES" rc=1 fi +fi for _ in $(seq 1 20); do flux jobs -no "{id}" 2>/dev/null | grep -q . || break sleep 1 From f32e3bdf905c37a7e969322eb67c7bda58d2f1e6 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 9 Sep 2026 06:14:34 +0000 Subject: [PATCH 13/26] fix: describe what the pair check actually tests The match_coschedule check in add-hold is a satisfiability test against the empty graph, not a live fit, and the comments claimed otherwise. A malformed request now warns and continues instead of refusing the submit as if the pair did not fit. --- flux_quantum/cli.py | 64 ++++++++++++++++++--------------------------- 1 file changed, 26 insertions(+), 38 deletions(-) diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index 1c8a064..51a6ac6 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -39,15 +39,13 @@ def _wrap_commands(jobspec_dict, wrap_path): def _check_pair_fits(handle, classical, scout, rpc_fn=None): - """Refuse a pair that cannot be scheduled, before either half is created. + """Refuse a pair the graph can never hold, before either half exists. - Asks fluxion whether both halves could be placed together. It answers by - matching the pair for real and undoing it, so the answer is against the - live graph rather than a core count, and nothing is left placed. Neither - job exists yet, so no jobids are sent. - - A pair that is admitted and then cannot start is the expensive case: the - scout opens a metered vendor session and the classical half waits. + The match_coschedule check in the add-hold branch is a satisfiability + test. Each half is checked alone against the whole graph with nothing + else on it, so this catches a missing qdevice or a half bigger than the + machine, not a busy machine. Fit against what is reachable right now is + the jobtap plugin's admission check. """ payload = { "check": True, @@ -62,17 +60,24 @@ def _check_pair_fits(handle, classical, scout, rpc_fn=None): else: rpc_fn(handle, payload) except OSError as exc: - # EBUSY is fluxion saying the pair does not fit, ENODEV that it never - # could. Anything else means the question could not be asked at all, - # an older fluxion or none loaded, and that must not block a submit - # that would have worked before the check existed. - if exc.errno not in (errno.EBUSY, errno.ENODEV, errno.EINVAL): + # ENODEV means the graph can never hold the pair. EBUSY would mean a + # fluxion that checks the live graph found no room now. Anything else + # means the question could not be asked, and that must not block a + # submit that worked before the check existed. + if exc.errno not in (errno.EBUSY, errno.ENODEV): + if exc.errno in (errno.EINVAL, errno.EPROTO): + print( + "flux quantum: WARNING fluxion did not understand the pair " + "check, skipping it: {}".format(exc), + file=sys.stderr, + ) return raise SystemExit( - "flux quantum: this pair cannot be scheduled together. The " - "classical half needs its cores and the scout needs one more, and " - "the cluster cannot place both. Submitting anyway would open a " - "vendor session for a pair that cannot run. ({})".format(exc) + "flux quantum: this pair cannot be scheduled together on this " + "cluster. The classical half needs its cores and the scout needs " + "one more plus the vendor device, and the graph cannot hold both. " + "Submitting anyway would open a vendor session for a pair that " + "cannot run. ({})".format(exc) ) except Exception: # no usable handle, so there is nothing to ask @@ -141,24 +146,9 @@ def prepare_pair( if job_env: sysattr.setdefault("environment", {}).update(job_env) - # Ask before creating anything. The plugin's own admission check only fires - # once a job exists, and it compares outstanding pairs against total - # capacity, which is a quota across pairs rather than a feasibility test: a - # 9 core pair was admitted onto a cluster with 7 cores free, could not be - # placed, and left preemption to clean up. - # - # So ask fluxion instead, which matches the pair against the live graph and - # undoes it. The scout jobspec built here is only for that question, so the - # job it would release does not exist yet and the id is a placeholder. Only - # the resources matter to the answer. - # - # There is a race between this answer and the submits below. That is - # accepted: the plugin's job.validate check is still the backstop, and - # either half is cancelled if the other fails. - # the foothold comes from the live graph so node->...->core matches reality. - # Done before either half is created, both because the check below needs the - # vendor device present to answer honestly, and because a graph failure then - # leaves no held job to clean up. + # Populate and ask fluxion before creating anything, so a pair the graph + # can never hold is refused with no held job to clean up. The scout shape + # comes from the live graph so the node to core path matches reality. try: populate_fn(handle, [vendor]) live = get_graph_fn(handle) @@ -169,9 +159,7 @@ def prepare_pair( "flux quantum: could not prepare the quantum graph: {}".format(exc) ) - # The scout jobspec built here is only for the question. The job it would - # release does not exist yet, so its id is a placeholder and only the - # resources matter to the answer. + # only the resources matter to the question, so the job id is a placeholder try: probe = build_scout_jobspec( vendor, From 841194129e9a5b6cb15f0123a303946d65116527 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 9 Sep 2026 06:14:34 +0000 Subject: [PATCH 14/26] docs: fix stale paths, names and the license section The integration scripts live in tests/integration, the graph type is qdevice, the CLI shim is cli-plugins/quantum.py, and the license text was copied from another project. Ignore build artifacts. --- .gitignore | 5 +++++ README.md | 26 ++++++++++++++++---------- pyproject.toml | 10 ++++++---- 3 files changed, 27 insertions(+), 14 deletions(-) diff --git a/.gitignore b/.gitignore index dbb9afa..a9b62e7 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,8 @@ dist .env *.env !env.example + +# build and python artifacts +__pycache__/ +*.pyc +*.so diff --git a/README.md b/README.md index a52b694..d4d032c 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,16 @@ # 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. +Developed against flux-core 0.87.0 and the `add-hold` branch of +https://github.com/vsoch/flux-sched (also loads against flux-core 0.89.0). 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. + ## 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: ```bash @@ -26,21 +28,25 @@ For the jobtap plugin: 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_*` +- `tests/integration/setup-mock-registry.sh` derives the real graph and injects `qdevice_*` 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 -> +- `tests/integration/test-mock-e2e.sh` runs it end to end: plugin discovery -> registry -> `flux submit --quantum-vendor mock` (held) -> scout (mock session) -> handoff -> run. +- `tests/integration/test-admission.sh` covers the jobtap admission gate: `FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-admission.sh`. +- `tests/integration/test-preemption.sh` covers preemption of a held job: `flux start bash tests/integration/test-preemption.sh`. +- `tests/integration/test-reload.sh` unloads the jobtap plugin with a grace timer pending and reloads it with jobs running: `FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-reload.sh`. +- `tests/integration/test-select-discovery.sh` covers auto vendor discovery via `--quantum-select any`: `flux start bash tests/integration/test-select-discovery.sh`. +- `tests/integration/test-system-instance.sh` runs the same pipeline as test-mock-e2e.sh against a running system instance with fluxion already loaded: `bash tests/integration/test-system-instance.sh`. ## 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 diff --git a/pyproject.toml b/pyproject.toml index f0a3162..e48a711 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "flux-quantum" version = "0.0.11" -description = "Quantum + classical coscheduling components for Flux (CLI plugin, validator, jobtap, backends, scout)" +description = "Quantum + classical coscheduling components for Flux (CLI plugin, jobtap, backends, scout)" requires-python = ">=3.8" # NOTE: interfaces (esp. flux.cli.plugin.CLIPlugin) are not yet stable across # flux-core releases. Developed/tested against flux-core 0.87.0. Pin to a @@ -27,6 +27,8 @@ Homepage = "https://github.com/flux-framework" include = ["flux_quantum*"] # The CLI plugin is discovered by flux from a *.py in the CLI plugin search -# path, NOT by normal import. Install/symlink flux_quantum/cli_plugin.py into -# one of: {confdir}/cli/plugins, {libexecdir}/cli/plugins, or a dir on -# FLUX_CLI_PLUGINPATH. See README (Install). +# path, NOT by normal import. Install/symlink cli-plugins/quantum.py (the shim +# that imports QuantumCLIPlugin from the installed package) into one of: +# {confdir}/cli/plugins, {libexecdir}/cli/plugins, or a dir on +# FLUX_CLI_PLUGINPATH -- that dir must contain only the shim. See README +# (Install). From 96fcb035c89e1d6f14a9c8a80f296e8cd48cc085 Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 9 Sep 2026 06:21:52 +0000 Subject: [PATCH 15/26] docs: trim comments and rewrite the README Comments say what is non obvious and why, and no more. The README now describes the pair, how to install and populate, and how to run every test script. --- README.md | 89 +++++++++++++++++++------------- examples/qrmi/README.md | 34 +++++++----- flux_quantum/graph.py | 17 ++---- flux_quantum/launch.py | 11 ++-- flux_quantum/populate.py | 12 ++--- flux_quantum/qresource.py | 13 ++--- flux_quantum/scout.py | 21 +++----- tests/integration/test-reload.sh | 3 +- tests/unit/test_cli_plugin.py | 3 +- 9 files changed, 93 insertions(+), 110 deletions(-) diff --git a/README.md b/README.md index d4d032c..6166c03 100644 --- a/README.md +++ b/README.md @@ -1,50 +1,65 @@ # flux-quantum -Quantum + classical coscheduling for Flux, including jobtap, validation, and client plugins. -Developed against flux-core 0.87.0 and the `add-hold` branch of -https://github.com/vsoch/flux-sched (also loads against flux-core 0.89.0). 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/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 per vendor, IBM through QRMI, AWS Braket, and a mock. 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_ -> 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 . # 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 -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 +flux python -m flux_quantum.populate ibm braket ``` -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: + +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 .../quantum.so +make -C flux_quantum/jobtap +flux jobtap load $PWD/flux_quantum/jobtap/quantum.so vendors=ibm,braket \ + total_cores=128 reserve_cores=32 protect_types=qpu preempt_after=30 ``` -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. -- `tests/integration/setup-mock-registry.sh` derives the real graph and injects `qdevice_*` - markers so the selector has vendors to discover (markers are plain graph vertices - added with `flux inject` -- no token). -- `tests/integration/test-mock-e2e.sh` runs it end to end: plugin discovery -> registry -> - `flux submit --quantum-vendor mock` (held) -> scout (mock session) -> handoff -> run. -- `tests/integration/test-admission.sh` covers the jobtap admission gate: `FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-admission.sh`. -- `tests/integration/test-preemption.sh` covers preemption of a held job: `flux start bash tests/integration/test-preemption.sh`. -- `tests/integration/test-reload.sh` unloads the jobtap plugin with a grace timer pending and reloads it with jobs running: `FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-reload.sh`. -- `tests/integration/test-select-discovery.sh` covers auto vendor discovery via `--quantum-select any`: `flux start bash tests/integration/test-select-discovery.sh`. -- `tests/integration/test-system-instance.sh` runs the same pipeline as test-mock-e2e.sh against a running system instance with fluxion already loaded: `bash tests/integration/test-system-instance.sh`. + +## 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 +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. ## License + flux-quantum is distributed under the terms of the MIT license. All new contributions must be made under this license. See [LICENSE](LICENSE), [COPYRIGHT](COPYRIGHT), and [NOTICE](NOTICE) for details. diff --git a/examples/qrmi/README.md b/examples/qrmi/README.md index 0a3d988..f1412d7 100644 --- a/examples/qrmi/README.md +++ b/examples/qrmi/README.md @@ -15,7 +15,10 @@ the scout. qiskit comes along with it and we need that for the warmup task. QRMI reads its config out of the environment. Every variable is prefixed with the resource id, so for ibm_kingston the api key lives in -`ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY`. Four of them are required. Note that this used to work for me, and it stopped. I now login with the web interface copy paste, which looks like this: +`ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY`. Four of them are required. + +The passcode login below can stop accepting credentials. If it does, use the +API key login that follows instead. ```bash ibmcloud login -a https://cloud.ibm.com -u passcode -p @@ -60,15 +63,17 @@ resource has credentials in the environment you can drop ## You need an account that can open sessions -The qiskit-runtime-service type opens a session, and IBM only allows that on plans that support sessions (Premium). On Open or Pay As You Go you can submit tasks but you cannot hold a session. You will get an error message (403). Our library will cancel the held classical job. If you have direct access instead, use `--quantum-ibm-type ibm-quantum-system`. +The qiskit-runtime-service type opens a session, which IBM only allows on plans that support sessions (Premium). Open and Pay As You Go plans can submit tasks but not hold a session, IBM returns a 403, and flux-quantum cancels the held classical job. Use `--quantum-ibm-type ibm-quantum-system` for direct access instead. ## A session is not the same as having the QPU -An IBM session goes active when its first task reaches the head of the queue. -After that, tasks in the session keep that priority. So opening a session tells -you very little on its own. After acquire we have to submit a warmup task and monitor it, -and it is small (one qubit and one measurement) and wait for it to run. Then we can at least -say that we reached the top and have priority and we release classical. If the warmup never runs we release the session and fail the submit. The classical job never starts against a QPU we do not have. +An IBM session goes active when its first task reaches the head of the queue, +and later tasks in the session inherit that priority. So opening a session +tells you little by itself. After acquiring, the scout submits a small warmup +task (one qubit, one measurement) and waits for it to run before releasing the +classical job. If the warmup never runs, the scout releases the session and +fails the submit instead of starting the classical job against a QPU it does +not have. --quantum-ibm-warmup-timeout SECONDS give up after this long. Default 0, which means wait as long as the @@ -82,17 +87,18 @@ The warmup costs one shot. ## The scout holds the session for as long as the classical runs -The fluxion allocation and the vendor session start and end -together. It does mean you are billed for the length of the classical job and not -the length of the quantum work, so set `-t`. +The fluxion allocation and the vendor session start and end together, so +you're billed for the classical job's wall time, not the quantum work. Set +`-t` accordingly. -A killed scout still releases, the release is in a finally and SIGTERM is -handled. If the node dies you are down to the vendor timeout. +A killed scout still releases, since the release runs in a finally block and +SIGTERM is handled. If the node dies, you're down to the vendor's own timeout. ## Credentials -The plugin runs in your process at submit time and reads the variables out of -your environment. The scout uses them, and it runs as you. Only the session id travels to the classical job, and it goes over the job eventlog. +The plugin runs in your process at submit time and reads the variables from +your environment. The scout runs as you and uses them too. Only the session +id reaches the classical job, over the job eventlog. ## What the classical job sees diff --git a/flux_quantum/graph.py b/flux_quantum/graph.py index b0dd5ee..9079f0e 100644 --- a/flux_quantum/graph.py +++ b/flux_quantum/graph.py @@ -64,15 +64,9 @@ def populate(handle, vendors, qpus=1, graph=None, lister=None, force=False): Idempotent, only missing vendors are added. Returns what was added. - Growing the graph while jobs hold resources corrupts fluxion. The spans of - the running jobs no longer match the graph, so every later free fails with - planner_multi_rem_span returned -1, those resources are never released, and - the instance stops scheduling anything at all. It looks like free cores that - nothing will use. - - So refuse to grow a busy graph. Populate at startup instead, before any - jobs run, which is what the vendors are there for anyway. force=True is for - a caller that knows the instance is idle. + Growing the graph while jobs hold resources corrupts fluxion. Every later + free fails with planner_multi_rem_span returned -1 and the instance stops + scheduling. So a busy graph is refused unless force is set. """ if isinstance(vendors, str): vendors = [vendors] @@ -83,10 +77,7 @@ def populate(handle, vendors, qpus=1, graph=None, lister=None, force=False): return set() if not force: - # Being unable to read the allocation count is not evidence that the - # instance is idle. Growing the graph while jobs hold resources - # corrupts fluxion irrecoverably, so an unknown answer has to be - # treated the same as a busy one. + # an unreadable allocation count is treated as busy try: busy = allocated_cores(handle, lister) except Exception as exc: diff --git a/flux_quantum/launch.py b/flux_quantum/launch.py index c8b4f22..8547f10 100644 --- a/flux_quantum/launch.py +++ b/flux_quantum/launch.py @@ -1,12 +1,7 @@ #!/usr/bin/env python3 -# Build and submit the scout job. -# -# The scout asks for a small classical foothold and the vendor device, so -# fluxion co-allocates a core and a qpu. That match is the coschedule. When the -# scout runs it opens the vendor session and releases the held classical job. -# -# The mock only replaces the vendor API. The graph match is real, so if -# qdevice_ -> qpu is missing the scout is unsatisfiable and never runs. +# Build the scout jobspec. The scout asks for a small classical foothold and +# the vendor device, so fluxion co-allocates a core and a qpu. Even with the +# mock the graph match is real, so a missing qdevice makes it unsatisfiable. import argparse import json import os diff --git a/flux_quantum/populate.py b/flux_quantum/populate.py index 00602cf..20abfab 100644 --- a/flux_quantum/populate.py +++ b/flux_quantum/populate.py @@ -5,18 +5,12 @@ # SPDX-License-Identifier: LGPL-3.0 ############################################################## -"""Add vendor devices to the fluxion graph once, at startup. - -This has to happen before any job holds resources. Growing the graph while jobs -are allocated corrupts fluxion, every later free fails with -planner_multi_rem_span returned -1, and the instance stops scheduling. Doing it -at startup means the first quantum submit never has to grow anything. +"""Add vendor devices to the fluxion graph once, at startup, before any job +holds resources. Growing the graph under running jobs corrupts fluxion. flux python -m flux_quantum.populate ibm braket mock -Run it with flux python, not the system python. The flux bindings are installed -for whichever interpreter flux was built against, so a console script on the -default python cannot import them. +Use flux python so the flux bindings are importable. """ import argparse diff --git a/flux_quantum/qresource.py b/flux_quantum/qresource.py index 152c959..879d775 100644 --- a/flux_quantum/qresource.py +++ b/flux_quantum/qresource.py @@ -29,10 +29,8 @@ def qdevice_type(vendor): def jobspec_resource(vendor, nqpus=1): """A scout jobspec resources entry for the vendor device. - The qpu must be exclusive. Fluxion only adds a planner span and emits a - leaf device into R when it is asked for exclusively, since upd_plan calls - planner_add_span only inside the excl branch. Ask for it any other way and - it is matched but then dropped from the allocation. + The qpu must be exclusive. Fluxion only writes a leaf device into R when + it is asked for exclusively, otherwise it is matched and then dropped. """ return { "type": qdevice_type(vendor), @@ -149,12 +147,7 @@ def classical_resource(live_graph, ncores=1, label="scout"): def count_cores(jobspec): - """Total cores a v1 jobspec asks for. - - Multiplies counts down the resource tree and sums the core leaves. The - jobtap plugin needs this to keep a core budget, and doing it here means it - is testable without a broker rather than parsed in C. - """ + """Total cores a v1 jobspec asks for, multiplying counts down the tree.""" def walk(entries, factor): total = 0 diff --git a/flux_quantum/scout.py b/flux_quantum/scout.py index 921803b..05b740f 100644 --- a/flux_quantum/scout.py +++ b/flux_quantum/scout.py @@ -20,16 +20,11 @@ def post_session(handle, jobid, session, rpc=None): - """Post the session id to the eventlog of the classical job as a memo. + """Memo the session id onto the classical job's eventlog. - The memo RPC is FLUX_ROLE_USER and authorized against the job owner, so - this needs no shared filesystem and no instance owner privilege. - - The released key is the durable record that this job was let go. A release - otherwise lives only in the scheduler memory, so a qmanager restart would - read the hold attribute again and park the job a second time, with nobody - left to release it and a vendor session already open and being paid for. - The eventlog survives a restart, so fluxion reads the marker from there. + The memo RPC is authorized against the job owner, so no shared filesystem + or owner privilege is needed. The released key is what qmanager reads + after a restart so it does not hold the job a second time. """ if rpc is None: rpc = handle.rpc @@ -46,12 +41,8 @@ def wait_for_job(handle, jobid, waiter=None): def abort_held(handle, jobid, why, cancel=None): - """Cancel the held classical and exit. - - Anything that fails before the release leaves the job held with its - reservation and nothing on the way to free it, so cancel it rather than - leave nodes parked on work that will never start. - """ + """Cancel the held classical and exit, so a failure before the release + does not leave it parked forever.""" if cancel is None: from flux.job import cancel try: diff --git a/tests/integration/test-reload.sh b/tests/integration/test-reload.sh index 724c9a2..e05d838 100644 --- a/tests/integration/test-reload.sh +++ b/tests/integration/test-reload.sh @@ -31,9 +31,8 @@ make -s -C "$HERE/flux_quantum/jobtap" || { echo "FAIL build"; exit 1; } flux python -m flux_quantum.populate mock >/dev/null 2>&1 \ || echo "WARNING could not populate the graph up front" -# query is by the registered name; fall back to the file name for older cores query () { - flux jobtap query quantum 2>/dev/null || flux jobtap query quantum.so + flux jobtap query quantum.so } load () { diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index c59cae2..1f5ca94 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -247,8 +247,7 @@ def rejecting_submit(handle, jobspec_json): def test_populate_failure_leaves_no_held_classical_behind(stub_flux_cli): """The graph is set up before either half is created, so a graph failure - has nothing to clean up. It used to happen after the classical was - submitted, which meant cancelling a job that had just been held.""" + has nothing to clean up.""" seen = {} def boom_populate(handle, vendors): From 813b1edf096a3caa8dfd7903e6a8571e3d221c6c Mon Sep 17 00:00:00 2001 From: vsoch Date: Wed, 9 Sep 2026 15:46:28 +0000 Subject: [PATCH 16/26] refactor: import dependencies at module level and check them on construction Vendor SDKs and the flux bindings were imported inside the functions that used them, so a missing package or credential only surfaced once a session was being opened. Every optional import now happens once at module level, bound to None when absent, and constructing a backend is the check. A missing SDK or missing credentials raise BackendError from __init__, the selector and the CLI report the message, and the scout cancels the held job instead of exiting and leaving it parked. credentials_present is gone and credential_note says where the credentials came from. The shared key names moved to keys.py so the wrapper, which runs in every classical job, no longer imports the backends package or any SDK. --- examples/qrmi/README.md | 5 +- examples/qrmi/ibm-run.sh | 4 +- flux_quantum/backends/__init__.py | 1 + flux_quantum/backends/base.py | 63 +- flux_quantum/backends/braket.py | 496 ++++++++-- flux_quantum/backends/hold/__init__.py | 0 flux_quantum/backends/hold/entry.py | 91 ++ flux_quantum/backends/mock.py | 6 - flux_quantum/backends/qrmi.py | 88 +- flux_quantum/cli.py | 24 +- flux_quantum/graph.py | 8 +- flux_quantum/keys.py | 9 + flux_quantum/launch.py | 14 +- flux_quantum/scout.py | 43 +- flux_quantum/selector.py | 12 +- flux_quantum/wrap.py | 23 +- tests/kill-all.sh | 9 + tests/probe_hold.py | 610 ++++++++++++ tests/unit/conftest.py | 8 +- tests/unit/test_backends.py | 93 +- tests/unit/test_braket.py | 41 +- tests/unit/test_braket_hold.py | 1222 ++++++++++++++++++++++++ tests/unit/test_cli_plugin.py | 10 + tests/unit/test_handoff.py | 3 - tests/unit/test_qrmi.py | 21 +- tests/unit/test_selector.py | 14 + 26 files changed, 2689 insertions(+), 229 deletions(-) create mode 100644 flux_quantum/backends/hold/__init__.py create mode 100644 flux_quantum/backends/hold/entry.py create mode 100644 flux_quantum/keys.py create mode 100755 tests/kill-all.sh create mode 100644 tests/probe_hold.py create mode 100644 tests/unit/test_braket_hold.py diff --git a/examples/qrmi/README.md b/examples/qrmi/README.md index f1412d7..b53f6bb 100644 --- a/examples/qrmi/README.md +++ b/examples/qrmi/README.md @@ -33,9 +33,10 @@ Check before you submit anything. flux python -c " from flux_quantum.backends import get_backend - print(get_backend('ibm').credentials_present())" + print(get_backend('ibm').credential_note)" -That names any variable you are missing. It never prints a value. +Constructing the backend is the check. It raises naming any variable you are +missing, and never prints a value. unset FLUX_QUANTUM_MOCK diff --git a/examples/qrmi/ibm-run.sh b/examples/qrmi/ibm-run.sh index b3d92c8..b3f1b00 100755 --- a/examples/qrmi/ibm-run.sh +++ b/examples/qrmi/ibm-run.sh @@ -88,9 +88,7 @@ echo "crn $IBM_CLOUD_CRN" >> "$OUT/run.txt" say "credentials" flux python -c " from flux_quantum.backends import get_backend -ok, msg = get_backend('ibm').credentials_present() -print(msg) -raise SystemExit(0 if ok else 1)" | tee -a "$OUT/run.txt" || die "credentials incomplete" +print(get_backend('ibm').credential_note)" | tee -a "$OUT/run.txt" || die "credentials incomplete" # 3. submit. The classical job prints and exits, so the billed window is warmup plus a few seconds. warmup_args="--quantum-ibm-warmup-timeout $WARMUP_TIMEOUT" diff --git a/flux_quantum/backends/__init__.py b/flux_quantum/backends/__init__.py index 2b9893e..6c12bca 100644 --- a/flux_quantum/backends/__init__.py +++ b/flux_quantum/backends/__init__.py @@ -9,6 +9,7 @@ from .base import ( Backend, + BackendError, # noqa: F401 Signals, register, get_backend, # noqa: F401 diff --git a/flux_quantum/backends/base.py b/flux_quantum/backends/base.py index 59fd6c6..c930ed6 100644 --- a/flux_quantum/backends/base.py +++ b/flux_quantum/backends/base.py @@ -4,12 +4,21 @@ depth, cost and availability, and it owns the vendor session. Backends run in userspace with the user credentials, both in the CLI plugin and in the scout. +Constructing a backend is the credential check. A backend whose SDK is not +installed or whose credentials cannot be found raises BackendError from +__init__, so nothing downstream ever holds an unusable backend. + Adding a vendor means adding a Backend subclass and registering it. """ from abc import ABC, abstractmethod +class BackendError(RuntimeError): + """The backend cannot be used. The message names what is missing, never + a credential value.""" + + class Signals: """Live signals for one vendor, returned by Backend.probe().""" @@ -31,6 +40,9 @@ class Backend(ABC): # vendor key like ibm, must match the qdevice_ graph type name = None + # set by __init__ to say where the credentials came from, for the log + credential_note = "" + @classmethod def add_options(cls, add_option): """Declare the flux submit options for this vendor. @@ -41,12 +53,19 @@ def add_options(cls, add_option): return def scout_options(self, args): - """Return the options for this vendor from the parsed args. + """Return the options for this vendor from the parsed args.""" + return {} - Only the selected vendor gets collected, so reading another vendor - args here is harmless. + def session_id(self, opened): + """What the classical job should be handed. + + Usually the id open_session returned. Braket publishes its token only + once the hybrid job is running, and the token rather than the job ARN + is what lets work elsewhere submit against the hold. + + Called after wait_for_priority. """ - return {} + return opened def open_session(self, options): """Open a session with the user credentials and return the id. @@ -58,8 +77,7 @@ def open_session(self, options): ) def job_environment(self, options): - """Env vars to add to the classical job, for example to tell it which - QPU it has. Defaults to none.""" + """Env vars to add to the classical job, for example which QPU it has.""" return {} def wait_for_priority(self, options=None): @@ -67,36 +85,20 @@ def wait_for_priority(self, options=None): Opening a session is not the same as having the device. IBM activates a session when its first task reaches the head of the queue, and Braket - reports a queue position and never holds anything. So this is a separate - step from open_session, and the scout does not release the classical job - until it returns ok. - - Defaults to ready, for vendors with no queue to wait on. + reports a queue position and never holds anything. The scout does not + release the classical job until this returns ok. """ return True, "not applicable" def close_session(self, session=None): - """Release the session opened by open_session. - - Called in the scout after the classical job finishes. A backend holding - vendor state, like a QRMI lock, stashes it on self and releases it - here. Defaults to doing nothing for vendors with nothing to release. - """ + """Release the session opened by open_session. Called in the scout + after the classical job finishes.""" return - @abstractmethod - def credentials_present(self): - """Return ok and a message, where ok is False if credentials are - missing. Never returns or logs the secret itself. - """ - @abstractmethod def probe(self): - """Return Signals for this vendor. - - Only called once credentials_present is ok. May raise. The selector - drops a backend that raises or is unavailable. - """ + """Return Signals for this vendor. May raise. The selector drops a + backend that raises or is unavailable.""" _REGISTRY = {} @@ -111,7 +113,10 @@ def register(cls): def get_backend(name): - """Return an instantiated backend for a vendor, or None if unknown.""" + """Return a backend for a vendor, or None if the name is unknown. + + Raises BackendError when the vendor is known but cannot be used. + """ cls = _REGISTRY.get(name) return cls() if cls else None diff --git a/flux_quantum/backends/braket.py b/flux_quantum/backends/braket.py index da536e2..661cad6 100644 --- a/flux_quantum/backends/braket.py +++ b/flux_quantum/backends/braket.py @@ -1,23 +1,36 @@ """AWS Braket backend. -Braket has no session to acquire. You submit a task, it waits in the device -queue, and it runs. So the scout submits a trivial no-op task and waits for it -to reach the front of the queue, which is as close to priority as Braket -offers. - -Note the limitation, because it is the vendor's and not ours. Reaching the -front of the queue does not reserve anything, so another user can take the -device between our task running and the classical job starting. IBM sessions do -hold the device, at the price of billing wall clock for it. Braket exposes the -queue position but no way to hold it. - -The SDK is an optional dependency, so it is imported inside the methods. +Two ways to take a device, set with --braket-hold. + +probe is the default. Submit a trivial task and wait for it to reach the front +of the queue. That reserves nothing, so another user can take the device +between our task running and the classical job starting. + +job starts a hybrid job. Braket runs one hybrid job at a time per QPU and gives +it priority, and the job publishes AMZN_BRAKET_JOB_TOKEN. A task created with +that token gets the job's priority and bills to the job. Without it, neither. + +So the hybrid job is the scout and the token is the session id. The classical +work stays on our cluster instead of the small instance Braket gives the job, +and the hold costs only that instance, about $0.12 an hour. + +The SDK is optional. The module imports without it and the backend refuses to +construct. """ +import datetime +import json import os import time -from .base import Backend, Signals, register +from .base import Backend, BackendError, Signals, register + +try: + import boto3 + from braket.aws import AwsDevice + from braket.circuits import Circuit +except ImportError: + boto3 = AwsDevice = Circuit = None # state simulator, no queue to speak of and cents per task SV1 = "arn:aws:braket:::device/quantum-simulator/amazon/sv1" @@ -55,12 +68,214 @@ def region_for(device_arn, override=None): ) +def _count(value): + """Leading number out of a queue depth field. + + Braket returns these as strings and sometimes appends a note, as in + "0 (1 prioritized hybrid job running)". Anything without a leading digit + counts as zero. + """ + import re + + m = re.match(r"\s*(\d+)", str(value or "")) + return int(m.group(1)) if m else 0 + + +def queue_depth(device_arn, cls=None): + """How many tasks are waiting on a device, split by queue. + + Braket counts Normal and Priority separately. A task submitted with a + hybrid job token lands in Priority, so the two counts are how you tell + from outside whether the token did anything. + """ + from braket.aws import AwsDevice + from braket.aws.queue_information import QueueType + + d = (cls or AwsDevice)(device_arn).queue_depth() + return { + "normal": _count(d.quantum_tasks.get(QueueType.NORMAL)), + "priority": _count(d.quantum_tasks.get(QueueType.PRIORITY)), + "jobs": _count(d.jobs), + # the note Braket sometimes appends, which says whether a prioritised + # hybrid job is running + "jobs_note": str(d.jobs or ""), + } + + +def is_open(device, now=None): + """Is the device inside an execution window right now. + + ONLINE does not mean running. Most QPUs only execute during set hours and + park everything else, which from outside looks the same as a long queue. + A window names a weekday or one of the groups Everyday, Weekdays and + Weekend, and a device can have several a day. Garnet stops between 15:30 + and 17:15 that way. + """ + now = now or datetime.datetime.now(datetime.timezone.utc) + groups = { + "Everyday": True, + "Weekdays": now.weekday() < 5, + "Weekend": now.weekday() >= 5, + now.strftime("%A"): True, + } + for w in device.properties.service.executionWindows: + if groups.get(w.executionDay.value) and ( + w.windowStartHour <= now.time() < w.windowEndHour + ): + return True + return False + + +def windows(device, local=True, now=None): + """The device execution windows, as readable tuples. + + Braket gives these in UTC. Reading them against your own clock is where + the mistakes happen, so they are converted to local time by default. A + window that crosses midnight in the conversion is marked, since the day + label no longer matches the hours. + """ + now = now or datetime.datetime.now(datetime.timezone.utc) + out = [] + for w in device.properties.service.executionWindows: + if local: + start, end = _to_local(w.windowStartHour, now), _to_local( + w.windowEndHour, now + ) + else: + start = w.windowStartHour.strftime("%H:%M") + end = w.windowEndHour.strftime("%H:%M") + out.append((w.executionDay.value, start, end)) + return out + + +def _to_local(t, now): + """A UTC clock time as local, tagged when it lands on another day. + + The day label on a window is the UTC day, so a converted time can sit + before or after it. Saying so is the difference between a readable line + and a misleading one. + """ + stamp = datetime.datetime.combine(now.date(), t, tzinfo=datetime.timezone.utc) + here = stamp.astimezone() + shift = (here.date() - stamp.date()).days + tag = {0: "", -1: " (prev day)", 1: " (next day)"}.get(shift, "") + return here.strftime("%H:%M") + tag + + +# Braket charges a flat fee per task on top of whatever the device charges +# per shot. It is not in the device properties, so it lives here. +TASK_FEE = 0.30 + + +def cost(device, shots=1): + """What one task costs on this device. + + deviceCost comes from the service, so it does not go stale here. The unit + is usually shot, sometimes minute for simulators, and the flat task fee is + on top either way. + """ + try: + c = device.properties.service.deviceCost + per, unit = float(c.price), str(c.unit).lower() + except Exception: + return None + if unit == "shot": + return TASK_FEE + per * shots + # per minute devices bill for time, so a shot count says nothing + return None + + +def shots_range(device): + """The shot count this device accepts, or None if it does not say. + + Cepheus takes 10 to 50000 and refuses 1. Finding that out from a + ValidationException means the hold is already up and billing. + """ + try: + r = device.properties.service.shotsRange + return int(r[0]), int(r[1]) + except Exception: + return None + + +def survey(minimum=0, cls=None, now=None, shots=1): + """Online QPUs with their queue and whether they are open right now. + + Open ones first, then busiest. Simulators are left out since they scale + on demand and never queue. + """ + from braket.aws import AwsDevice + + out = [] + for d in (cls or AwsDevice).get_devices(statuses=["ONLINE"]): + if "/qpu/" not in d.arn: + continue + try: + q = queue_depth(d.arn, cls=cls) + row = { + "arn": d.arn, + "open": is_open(d, now), + "windows": windows(d), + "cost": cost(d, shots), + **q, + } + except Exception: + continue + if row["normal"] >= minimum: + out.append(row) + return sorted(out, key=lambda r: (not r["open"], -r["normal"])) + + +def qpus_with_a_queue(minimum=1, cls=None, now=None): + """Devices where a priority comparison would show something. + + Both have to be true. An idle device runs a token task and a plain one at + once, and a shut device runs neither. + """ + return [d for d in survey(minimum, cls, now) if d["open"]] + + @register class BraketBackend(Backend): name = "braket" def __init__(self): + if boto3 is None or AwsDevice is None: + raise BackendError( + "braket: amazon-braket-sdk is not installed, " + "pip install amazon-braket-sdk" + ) + # resolve credentials the way boto3 normally would, and fail here + # rather than at the first API call + self._credentials = boto3.Session().get_credentials() + if self._credentials is None: + raise BackendError( + "braket: no AWS credentials. Set AWS_ACCESS_KEY_ID and " + "AWS_SECRET_ACCESS_KEY, or configure ~/.aws/credentials" + ) + method = getattr(self._credentials, "method", "unknown") + if method in ("iam-role", "instance-metadata"): + # the instance role belongs to the node and not to the user + self.credential_note = ( + "braket: using the EC2 instance role, which belongs to the " + "node and not to you. Export your own keys to keep the " + "credential in user space" + ) + else: + self.credential_note = "braket: credentials present via {}".format(method) self._task = None + self._hold_job = None + self._hold_prefix = None + self._session = None + # what this run made, recorded as it is made so cleanup never has to + # work it out afterwards + self._made = {"job": None, "bucket": None, "prefix": None} + # the device the token is good for, which is the device the hybrid job + # named and nothing else + self._token_device = None + # wall time from asking for the hold to having it + self._hold_started = None + self._hold_ready_after = None @classmethod def add_options(cls, add_option): @@ -90,6 +305,29 @@ def add_options(cls, add_option): "this queue position or closer. Default 1, next in line. Raise it " "when the classical job is slow to start", ) + add_option( + "--braket-hold", + metavar="MODE", + default=None, + help="Braket: how to take the device. probe waits for the front " + "of the queue, the default. job starts a hybrid job, which holds " + "the priority queue and hands out a token that work anywhere can " + "submit with", + ) + add_option( + "--braket-hold-instance", + metavar="TYPE", + default=None, + help="Braket: instance for the holding hybrid job. Default " + "ml.m5.large, the cheapest, since it does no work", + ) + add_option( + "--braket-hold-max-seconds", + metavar="SECONDS", + default=None, + help="Braket: give up the queue slot after this long, so a scout " + "that is never released stops billing. Default 900", + ) add_option( "--braket-queue-timeout", metavar="SECONDS", @@ -105,6 +343,12 @@ def scout_options(self, args): return { "device": device, "region": region_for(device, getattr(args, "braket_region", None)), + "hold": getattr(args, "braket_hold", None) or "probe", + "hold_instance": getattr(args, "braket_hold_instance", None) + or "ml.m5.large", + "hold_max_seconds": float( + getattr(args, "braket_hold_max_seconds", None) or 900 + ), "shots": 1 if shots is None else int(shots), "queue_timeout": 0 if timeout is None else float(timeout), "ungate_position": int( @@ -122,14 +366,20 @@ def job_environment(self, options): } def open_session(self, options): - """Submit the probe task and return its ARN, which is what the - classical job is handed. The scout calls wait_for_priority next.""" - from braket.aws import AwsDevice - from braket.circuits import Circuit + """Take the device. probe returns a task ARN, job returns a job ARN. + The scout calls wait_for_priority next. In job mode the token replaces + the ARN once the job is running. + """ + if (options.get("hold") or "probe") == "job": + return self._open_hold_job(options) device_arn = options.get("device") or SV1 region = options.get("region") or region_for(device_arn) - os.environ.setdefault("AWS_DEFAULT_REGION", region) + # the device ARN carries the region, so it wins. setdefault let a + # stale AWS_DEFAULT_REGION point the session at the wrong region, and + # then the bucket we polled for the token, and the cancel on cleanup, + # both went somewhere the job did not exist + os.environ["AWS_DEFAULT_REGION"] = region device = AwsDevice(device_arn) # identity on one qubit. Braket measures every qubit, so this is the @@ -137,18 +387,151 @@ def open_session(self, options): self._task = device.run(Circuit().i(0), shots=int(options.get("shots") or 1)) return self._task.id + def _open_hold_job(self, options): + """Start a hybrid job that does nothing but hold the queue. + + Its entry point publishes the token to S3 and waits. We cannot read + the token until the job runs, so this returns the job ARN and + wait_for_priority swaps in the token later. + """ + # only needed in job mode, so kept out of the module import + from braket.aws import AwsQuantumJob + from braket.jobs.config import InstanceConfig, StoppingCondition + + from . import hold + + device_arn = options.get("device") or SV1 + region = options.get("region") or region_for(device_arn) + # the device ARN carries the region, so it wins. setdefault let a + # stale AWS_DEFAULT_REGION point the session at the wrong region, and + # then the bucket we polled for the token, and the cancel on cleanup, + # both went somewhere the job did not exist + os.environ["AWS_DEFAULT_REGION"] = region + + self._hold_started = time.time() + # resolve the bucket now. Working it out during cleanup means a lookup + # that can fail at the one moment it has to work, and the token stays + # in the bucket when it does. + from braket.aws import AwsSession + + self._made["bucket"] = AwsSession().default_bucket() + name = "flux-quantum-hold-{}".format(int(self._hold_started)) + self._hold_prefix = "flux-quantum/{}".format(name) + self._made["prefix"] = self._hold_prefix + max_seconds = int(options.get("hold_max_seconds") or 900) + self._hold_job = AwsQuantumJob.create( + device=device_arn, + source_module=os.path.dirname(os.path.abspath(hold.__file__)), + entry_point="hold.entry:main", + job_name=name, + instance_config=InstanceConfig( + instanceType=options.get("hold_instance") or "ml.m5.large", + instanceCount=1, + volumeSizeInGb=30, + ), + # create takes no environment, so config goes in here + hyperparameters={ + "flux_quantum_prefix": self._hold_prefix, + "flux_quantum_max_seconds": str(max_seconds), + }, + # Braket stops the job itself, so a wedged entry point still + # stops billing + stopping_condition=StoppingCondition(maxRuntimeInSeconds=max_seconds + 300), + wait_until_complete=False, + ) + self._made["job"] = self._hold_job + return self._hold_job.arn + + def _wait_for_hold_job(self, options, interval, sleep): + """Ready once the job is RUNNING and its token has landed. + + Braket says RUNNING when the instance is up. The container still has + to start and publish, which takes a minute or so. RUNNING without a + token is reported as starting, not running. + """ + timeout = float(options.get("queue_timeout") or 0) + started = getattr(self, "_hold_started", None) or time.time() + deadline = started + timeout if timeout > 0 else None + s3 = boto3.client("s3") + bucket = None + while True: + state = self._hold_job.state() + waited = time.time() - started + if state in ("FAILED", "CANCELLED", "COMPLETED"): + return False, "hybrid job {} after {:.0f}s".format(state, waited) + if state == "RUNNING": + if bucket is None: + bucket = self._output_bucket() + token = self._read_token(s3, bucket) + if token: + asked = options.get("device") or SV1 + held = getattr(self, "_token_device", None) + if held and held != asked: + return False, ( + "the hold is on {} but the work is for {}. A job " + "token is only valid for the device its hybrid " + "job named, and submitting elsewhere with it " + "silently gets no priority".format(held, asked) + ) + self._session = token + self._hold_ready_after = waited + return True, "holding the device after {:.0f}s".format(waited) + said = "starting the container" + else: + said = state.lower() + print( + "braket: {} ({:.0f}s)".format(said, waited), + flush=True, + ) + if deadline is not None and time.time() + interval >= deadline: + return False, "gave up after {:.0f}s, still {}".format(waited, said) + sleep(interval) + + def _output_bucket(self): + """The bucket the job writes to. + + We pass no code_location or output_data_config, so the job uses the + session default. That is what the container sees as + AMZN_BRAKET_OUT_S3_BUCKET. + """ + from braket.aws import AwsSession + + return AwsSession().default_bucket() + + def _read_token(self, s3, bucket): + """The token the job published, or None if it is not there yet. + + The device is kept alongside it. A token is only valid for the device + its hybrid job named, so it has to travel with one. + """ + try: + body = s3.get_object( + Bucket=bucket, Key="{}/token.json".format(self._hold_prefix) + )["Body"].read() + payload = json.loads(body) + except Exception: + return None + self._token_device = payload.get("device") # noqa: E501 + return payload.get("token") + + def session_id(self, opened): + """The token replaces the job ARN. The token is what a task elsewhere + submits with.""" + return self._session or opened + def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep): """Poll until the probe task is next in line. Ready at the requested position or closer, and also once the task is RUNNING or COMPLETED, because a task that left the queue reports no - position and the value we wait for can never arrive. FAILED and - CANCELLED are a failure and not readiness, so the caller can cancel the - classical job rather than start it against nothing. + position. FAILED and CANCELLED are a failure and not readiness, so the + caller cancels the classical job rather than starting it. Returns (ok, reason). """ options = options or {} + if (options.get("hold") or "probe") == "job": + return self._wait_for_hold_job(options, interval, sleep) timeout = float(options.get("queue_timeout") or 0) position = int(options.get("ungate_position") or 1) deadline = time.time() + timeout if timeout > 0 else None @@ -159,8 +542,6 @@ def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep): if state in RAN: return True, state pos = self.queue_position() - # positions over 2000 come back as the string >2000, so compare as - # a string rather than an int if position_at_most(pos, position): return True, "queue position {}".format(pos) print("braket: queued at position {}".format(pos), flush=True) @@ -178,43 +559,50 @@ def queue_position(self): return None def close_session(self, session=None): - """Nothing to release. The probe task finishes or is already terminal.""" - return + """Undo whatever this run made, in reverse. - def credentials_present(self): - """Ask boto3 to resolve credentials the normal way, and say where they - came from. An EC2 instance role works but belongs to the node rather - than to you, so it is called out.""" - try: - import boto3 - except ImportError: - return ( - False, - "braket: boto3 is not installed, pip install amazon-braket-sdk", - ) - creds = boto3.Session().get_credentials() - if creds is None: - return False, ( - "braket: no AWS credentials. Set AWS_ACCESS_KEY_ID and " - "AWS_SECRET_ACCESS_KEY, or configure ~/.aws/credentials" - ) - method = getattr(creds, "method", "unknown") - if method in ("iam-role", "instance-metadata"): - return True, ( - "braket: using the EC2 instance role, which belongs to the node " - "and not to you. Export your own keys to keep the credential in " - "user space" + Driven by what was recorded at the time, so a run that failed halfway + still cleans up the half it made. Safe to call twice, and safe to call + when nothing was made at all. + + The release marker goes before the cancel, or a container polling in + between never sees it. Each step is best effort, because a later step + still needs to run when an earlier one fails. + """ + made, self._made = self._made, {"job": None, "bucket": None, "prefix": None} + bucket, prefix, job = made["bucket"], made["prefix"], made["job"] + + if bucket and prefix: + self._s3_call( + "put_object", + Bucket=bucket, + Key="{}/release".format(prefix), + Body=b"released", ) - return True, "braket: credentials present via {}".format(method) + if job is not None: + try: + if job.state() not in ("COMPLETED", "FAILED", "CANCELLED"): + job.cancel() + except Exception as e: + print("braket: could not cancel the hold job: {}".format(e)) + if bucket and prefix: + for key in ("token.json", "release"): + self._s3_call( + "delete_object", Bucket=bucket, Key="{}/{}".format(prefix, key) + ) + + @staticmethod + def _s3_call(op, **kw): + """One S3 call that must not raise. Cleanup that gives up on the first + failure leaves more behind than it removes.""" + try: + getattr(boto3.client("s3"), op)(**kw) + except Exception as e: + print("braket: {} on {} failed: {}".format(op, kw.get("Key"), e)) def probe(self): """Report the device status. Cheap, it is a metadata call.""" - ok, _ = self.credentials_present() - if not ok: - return Signals(available=False) try: - from braket.aws import AwsDevice - device = AwsDevice(SV1) return Signals(available=device.status == "ONLINE", detail={"device": SV1}) except Exception as e: diff --git a/flux_quantum/backends/hold/__init__.py b/flux_quantum/backends/hold/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/flux_quantum/backends/hold/entry.py b/flux_quantum/backends/hold/entry.py new file mode 100644 index 0000000..e50d651 --- /dev/null +++ b/flux_quantum/backends/hold/entry.py @@ -0,0 +1,91 @@ +"""The hybrid job that holds device priority. + +Runs inside the Braket container. It does no quantum work. It exists so the +job is running, because a running hybrid job holds the priority queue for its +device, and it publishes the token that lets work elsewhere use that priority. + +AMZN_BRAKET_JOB_TOKEN goes to CreateQuantumTask as jobToken. Tasks without it +get no priority and bill as standalone, so the token is what we carry out. + +Then it waits for a release marker in S3, or its own deadline. That is the +hold. The classical work runs on our cluster while this keeps the place. +""" + +import json +import os +import time + + +def _s3(): + import boto3 + + return boto3.client("s3") + + +def publish(bucket, prefix, payload, client=None): + """Write the token where the scout can read it.""" + client = client or _s3() + client.put_object( + Bucket=bucket, + Key="{}/token.json".format(prefix.rstrip("/")), + Body=json.dumps(payload).encode(), + ) + + +def released(bucket, prefix, client=None): + """True once the release marker exists.""" + client = client or _s3() + try: + client.head_object(Bucket=bucket, Key="{}/release".format(prefix.rstrip("/"))) + return True + except Exception: + return False + + +def _hyperparameters(): + """What the scout passed in. create takes no environment, so config comes + through as hyperparameters.""" + try: + from braket.jobs.environment_variables import get_hyperparameters + + return get_hyperparameters() + except Exception: + return {} + + +def main(env=None, sleep=time.sleep, client=None, params=None): + env = env or os.environ + params = _hyperparameters() if params is None else params + bucket = env["AMZN_BRAKET_OUT_S3_BUCKET"] + prefix = ( + params.get("flux_quantum_prefix") + or env.get("FLUX_QUANTUM_PREFIX") + or "flux-quantum/{}".format(env.get("AMZN_BRAKET_JOB_NAME", "job")) + ) + payload = { + "token": env.get("AMZN_BRAKET_JOB_TOKEN"), + "device": env.get("AMZN_BRAKET_DEVICE_ARN"), + "job_arn": env.get("AMZN_BRAKET_JOB_ARN"), + "published": time.time(), + } + publish(bucket, prefix, payload, client=client) + print("hold: published token for {}".format(payload["device"]), flush=True) + + # Hold until released or out of time. The instance bills by the minute, + # so a scout that never releases needs a stop. + deadline = time.time() + float( + params.get("flux_quantum_max_seconds") + or env.get("FLUX_QUANTUM_MAX_SECONDS", 900) + ) + interval = float(env.get("FLUX_QUANTUM_POLL", 5)) + while time.time() < deadline: + if released(bucket, prefix, client=client): + print("hold: released", flush=True) + return 0 + sleep(interval) + print("hold: deadline reached, giving up the queue slot", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/flux_quantum/backends/mock.py b/flux_quantum/backends/mock.py index e89cf10..28c9fef 100644 --- a/flux_quantum/backends/mock.py +++ b/flux_quantum/backends/mock.py @@ -137,9 +137,6 @@ def close_session(self, session=None): self._closed = session return - def credentials_present(self): - return True, "mock: no credentials required" - def probe(self): q = os.environ.get("FLUX_QUANTUM_MOCK_QUEUE") c = os.environ.get("FLUX_QUANTUM_MOCK_COST") @@ -159,8 +156,5 @@ def close_session(self, session=None): self._closed = session return - def credentials_present(self): - return True, "mock_busy: no credentials required" - def probe(self): return Signals(available=True, queue_depth=9, cost=5.0, detail={"mock": True}) diff --git a/flux_quantum/backends/qrmi.py b/flux_quantum/backends/qrmi.py index e392ae3..31e6477 100644 --- a/flux_quantum/backends/qrmi.py +++ b/flux_quantum/backends/qrmi.py @@ -12,15 +12,26 @@ Credentials stay in the user environment. We only ever check that a variable is set and report its name, never its value. -qrmi is an optional dependency and needs python 3.11 or newer, so it is -imported inside the methods that use it. +qrmi is optional and needs python 3.11 or newer. The module imports without +it and the backend refuses to construct. """ import json import os import time -from .base import Backend, Signals +from .base import Backend, BackendError, Signals + +try: + from qrmi import Payload, QuantumResource, ResourceType, TaskStatus +except ImportError: + Payload = QuantumResource = ResourceType = TaskStatus = None + +# only the warmup task needs qiskit +try: + from qiskit import QuantumCircuit, qasm3 +except ImportError: + QuantumCircuit = qasm3 = None # type strings QRMI itself uses in qrmi_config.json and in QRMI_JOB_QPU_TYPES, # mapped to the ResourceType attribute @@ -61,8 +72,8 @@ def resource_type(name): """Return the QRMI ResourceType for a type string like ibm-quantum-system.""" - from qrmi import ResourceType - + if ResourceType is None: + raise BackendError("qrmi is not installed, pip install 'qrmi[ibm]'") attr = RESOURCE_TYPES.get(name, name) if not hasattr(ResourceType, attr): known = ", ".join(sorted(RESOURCE_TYPES)) @@ -90,10 +101,7 @@ def missing_env(resource, rtype): return [s for s in REQUIRED_ENV.get(rtype, ()) if not os.environ.get(resource + s)] -# QRMI passes the acquisition token to the job in this variable, prefixed with -# the resource id. The Slurm and LSF plugins set it, so we set it too and a -# workload written for either runs here unchanged. -ACQUISITION_TOKEN = "_QRMI_JOB_ACQUISITION_TOKEN" +from ..keys import ACQUISITION_TOKEN # noqa: E402,F401 def warmup_payload(shots=1): @@ -103,9 +111,7 @@ def warmup_payload(shots=1): qiskit generates the qasm3 so we are not hand writing a circuit. A measure only circuit needs no basis gate translation, so it is ISA valid anywhere. """ - try: - from qiskit import QuantumCircuit, qasm3 - except ImportError: + if QuantumCircuit is None: raise RuntimeError( "qiskit is needed to build the warmup task. Install it with " "pip install 'qrmi[ibm]', or pass --quantum-ibm-skip-warmup to " @@ -168,6 +174,36 @@ class QRMIBackend(Backend): default_type = None def __init__(self): + # The resource id is only known once the submit options are parsed, + # so this checks that at least one resource of the default type is + # fully configured. open_session checks the one actually chosen. + found = resources_in_environment(self.default_type) + if not found: + wanted = ", ".join( + "" + s for s in REQUIRED_ENV.get(self.default_type, ()) + ) + raise BackendError( + "{}: no QRMI credentials in the environment, set {}".format( + self.name, wanted + ) + ) + incomplete = {r: missing_env(r, self.default_type) for r in found} + usable = [r for r, m in incomplete.items() if not m] + if not usable: + r = found[0] + raise BackendError( + "{}: {} is missing {}".format( + self.name, r, ", ".join(r + s for s in incomplete[r]) + ) + ) + self.credential_note = "{}: credentials present for {}".format( + self.name, " ".join(usable) + ) + if QuantumResource is None: + raise BackendError( + "{}: qrmi is not installed, pip install 'qrmi[ibm]' " + "(needs python 3.11 or newer)".format(self.name) + ) self._resource = None self._lock = None @@ -256,10 +292,6 @@ def open_session(self, options): self.name, resource, ", ".join(resource + s for s in missing) ) ) - # deferred so a missing resource or variable is reported as a - # ValueError above even when qrmi itself is not installed - from qrmi import QuantumResource - self._resource = QuantumResource(resource, resource_type(rtype)) try: self._lock = self._resource.acquire() @@ -317,8 +349,6 @@ def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep, shots= Returns (ok, reason). A timeout of 0 waits as long as the scout job is allowed to live. """ - from qrmi import Payload, TaskStatus - options = options or {} if options.get("skip_warmup"): return True, "warmup skipped, priority unconfirmed" @@ -348,26 +378,6 @@ def close_session(self, session=None): self._resource = None self._lock = None - def credentials_present(self): - found = resources_in_environment(self.default_type) - if not found: - wanted = ", ".join( - "" + s for s in REQUIRED_ENV.get(self.default_type, ()) - ) - return False, "{}: no QRMI credentials in the environment, set {}".format( - self.name, wanted - ) - incomplete = {r: missing_env(r, self.default_type) for r in found} - usable = [r for r, m in incomplete.items() if not m] - if not usable: - r = found[0] - return False, "{}: {} is missing {}".format( - self.name, r, ", ".join(r + s for s in incomplete[r]) - ) - return True, "{}: credentials present for {}".format( - self.name, " ".join(usable) - ) - def probe(self): """Ask QRMI whether the resource is reachable. @@ -383,8 +393,6 @@ def probe(self): if len(found) != 1: return Signals(available=bool(found), detail={"resources": found}) try: - from qrmi import QuantumResource - qr = QuantumResource(found[0], resource_type(self.default_type)) return Signals( available=bool(qr.is_accessible()), detail={"resource": found[0]} diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index 51a6ac6..59e8e88 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -21,10 +21,11 @@ import flux from flux.cli.plugin import CLIPlugin +from flux.job import cancel, submit from . import graph, qresource from .selector import select_vendor, SelectionError, discover_registry_vendors -from .backends import get_backend, backend_classes +from .backends import BackendError, get_backend, backend_classes from .launch import build_scout_jobspec # subcommands where quantum submission makes sense @@ -115,12 +116,8 @@ def prepare_pair( flux submits whatever is left in the jobspec, so it submits the scout. The flux calls are injectable so this is testable without a broker. """ - # flux.job is imported here and not at the top so the unit tests can drive - # this without flux installed - if submit_fn is None: - from flux.job import submit as submit_fn - if cancel_fn is None: - from flux.job import cancel as cancel_fn + submit_fn = submit_fn or submit + cancel_fn = cancel_fn or cancel populate_fn = populate_fn or graph.populate get_graph_fn = get_graph_fn or graph.get_live_graph if wrap_path is None: @@ -302,7 +299,10 @@ def modify_jobspec(self, args, jobspec): if not vendor: return # not a quantum submit - backend = get_backend(vendor) + try: + backend = get_backend(vendor) + except BackendError as e: + raise SystemExit("flux quantum: {}".format(e)) options = backend.scout_options(args) if backend else {} job_env = backend.job_environment(options) if backend else {} handle = flux.Flux() @@ -322,9 +322,9 @@ def validate(self, jobspec): vendor = jobspec.getattr("system.quantum.vendor") except KeyError: return # not a quantum job - backend = get_backend(vendor) + try: + backend = get_backend(vendor) + except BackendError as e: + raise ValueError(str(e)) if backend is None: raise ValueError("quantum: no backend for vendor '{}'".format(vendor)) - ok, msg = backend.credentials_present() - if not ok: - raise ValueError(msg) diff --git a/flux_quantum/graph.py b/flux_quantum/graph.py index 9079f0e..2c68cf2 100644 --- a/flux_quantum/graph.py +++ b/flux_quantum/graph.py @@ -17,6 +17,11 @@ from . import qresource +try: + from flux.resource import resource_list +except ImportError: + resource_list = None + def get_live_graph(handle, criteria="status=up"): """Return the live fluxion graph of nodes and edges from the find RPC.""" @@ -54,8 +59,9 @@ def vendors_present(graph): def allocated_cores(handle, lister=None): """Cores the scheduler currently has allocated.""" + lister = lister or resource_list if lister is None: - from flux.resource import resource_list as lister + raise RuntimeError("the flux bindings are not importable") return int(lister(handle).get().allocated.ncores) diff --git a/flux_quantum/keys.py b/flux_quantum/keys.py new file mode 100644 index 0000000..c940343 --- /dev/null +++ b/flux_quantum/keys.py @@ -0,0 +1,9 @@ +"""Names shared by the scout, the wrapper and the backends. Kept apart so the +wrapper, which runs in every classical job, imports no vendor SDK.""" + +# memo key carrying the session id to the classical job +SESSION_KEY = "quantum_session" + +# QRMI reads the acquisition token from + this. The Slurm and LSF +# plugins set it, so a workload written for either runs here unchanged. +ACQUISITION_TOKEN = "_QRMI_JOB_ACQUISITION_TOKEN" diff --git a/flux_quantum/launch.py b/flux_quantum/launch.py index 8547f10..2316b84 100644 --- a/flux_quantum/launch.py +++ b/flux_quantum/launch.py @@ -8,6 +8,12 @@ from flux_quantum import qresource +try: + import flux + from flux.job import submit +except ImportError: + flux = submit = None + def build_scout_jobspec( vendor, @@ -111,9 +117,11 @@ def main(): print(json.dumps(jobspec, indent=2)) return - import flux - from flux.job import submit - + if flux is None: + raise SystemExit( + "quantum-scout-launch: the flux bindings are not importable, run " + "under flux python" + ) h = flux.Flux() print(submit(h, json.dumps(jobspec))) diff --git a/flux_quantum/scout.py b/flux_quantum/scout.py index 05b740f..a83e3ea 100644 --- a/flux_quantum/scout.py +++ b/flux_quantum/scout.py @@ -13,10 +13,17 @@ import signal import sys -from flux_quantum.backends import get_backend +from flux_quantum.backends import BackendError, get_backend +from flux_quantum.keys import SESSION_KEY # noqa: F401 (re-exported) -# memo key carrying the session id to the classical job -SESSION_KEY = "quantum_session" +# the flux bindings are only importable under flux python. Every flux call +# takes an injectable so the unit tests can run without them. +try: + import flux + from flux.job import JobID, event_wait + from flux.job import cancel as flux_cancel +except ImportError: + flux = JobID = event_wait = flux_cancel = None def post_session(handle, jobid, session, rpc=None): @@ -34,8 +41,7 @@ def post_session(handle, jobid, session, rpc=None): def wait_for_job(handle, jobid, waiter=None): """Block until the job reaches clean.""" - if waiter is None: - from flux.job import event_wait as waiter + waiter = waiter or event_wait # a failed job still reaches clean, and we close the session either way waiter(handle, jobid, "clean", raiseJobException=False) @@ -43,8 +49,7 @@ def wait_for_job(handle, jobid, waiter=None): def abort_held(handle, jobid, why, cancel=None): """Cancel the held classical and exit, so a failure before the release does not leave it parked forever.""" - if cancel is None: - from flux.job import cancel + cancel = cancel or flux_cancel try: cancel(handle, jobid, why) print("quantum-scout: cancelled held job {}".format(jobid), file=sys.stderr) @@ -91,10 +96,10 @@ def main(): ) args = ap.parse_args() - # flux is imported here and not at the top so the unit tests can import - # this module without flux installed - import flux - from flux.job import JobID + if flux is None: + sys.exit( + "quantum-scout: the flux bindings are not importable, run under flux python" + ) jobid = int(JobID(args.job)) h = flux.Flux() @@ -104,11 +109,16 @@ def main(): if args.session: session = args.session else: - backend = get_backend(args.vendor) + try: + backend = get_backend(args.vendor) + except BackendError as e: + abort_held(h, jobid, str(e)) if backend is None: - sys.exit( - "quantum-scout: no backend registered for vendor {}, " - "set FLUX_QUANTUM_MOCK for the mock vendor".format(args.vendor) + abort_held( + h, + jobid, + "no backend registered for vendor {}, set FLUX_QUANTUM_MOCK for " + "the mock vendor".format(args.vendor), ) opts = json.loads(args.options) try: @@ -129,6 +139,9 @@ def main(): backend.close_session(session) abort_held(h, jobid, "{} never became ours: {}".format(args.vendor, why)) print("quantum-scout: {}".format(why)) + # what the classical job needs may only exist once the device is + # held. Braket publishes its token after the hybrid job starts. + session = backend.session_id(session) # must land before the release, or the job could start with no session try: diff --git a/flux_quantum/selector.py b/flux_quantum/selector.py index c8ec8e3..89d5d8b 100644 --- a/flux_quantum/selector.py +++ b/flux_quantum/selector.py @@ -11,7 +11,7 @@ import json from . import qresource -from .backends import get_backend, known_vendors +from .backends import BackendError, get_backend, known_vendors class SelectionError(Exception): @@ -43,14 +43,14 @@ def select_vendor(candidates=None, policy=None): usable = [] log = [] for name in names: - backend = get_backend(name) + try: + backend = get_backend(name) + except BackendError as e: # no SDK or no credentials, said in e + log.append(str(e)) + continue if backend is None: log.append("{}: no backend registered".format(name)) continue - ok, msg = backend.credentials_present() - if not ok: - log.append(msg) - continue try: sig = backend.probe() except Exception as e: # a raising backend is simply not a candidate diff --git a/flux_quantum/wrap.py b/flux_quantum/wrap.py index 78c6678..e995ae9 100644 --- a/flux_quantum/wrap.py +++ b/flux_quantum/wrap.py @@ -11,8 +11,13 @@ import signal import sys -from flux_quantum.backends.qrmi import ACQUISITION_TOKEN -from flux_quantum.scout import SESSION_KEY +from flux_quantum.keys import ACQUISITION_TOKEN, SESSION_KEY + +try: + import flux + from flux.job import JobID, event_watch +except ImportError: + flux = JobID = event_watch = None class _Timeout(Exception): @@ -21,8 +26,7 @@ class _Timeout(Exception): def read_session(handle, jobid, timeout=60.0, watcher=None): """Return the session id the scout put on the eventlog for this job.""" - if watcher is None: - from flux.job import event_watch as watcher + watcher = watcher or event_watch def _alarm(signum, frame): raise _Timeout() @@ -47,9 +51,8 @@ def session_environment(session, resources=None): """Return the env the wrapped program should see. QRMI reads the acquisition token from _QRMI_JOB_ACQUISITION_TOKEN, - the same variable the Slurm and LSF plugins set, so a workload written for - either runs here unchanged. The resource names arrive at submit time, the - token only exists once the scout has acquired. + the same variable the Slurm and LSF plugins set. The resource names arrive + at submit time, the token only exists once the scout has acquired. """ if resources is None: resources = os.environ.get("QRMI_JOB_QPU_RESOURCES", "") @@ -67,8 +70,10 @@ def main(): ap.add_argument("command", nargs=argparse.REMAINDER) args = ap.parse_args() - import flux - from flux.job import JobID + if flux is None: + sys.exit( + "quantum-wrap: the flux bindings are not importable, run under flux python" + ) fjid = os.environ.get("FLUX_JOB_ID") if not fjid: diff --git a/tests/kill-all.sh b/tests/kill-all.sh new file mode 100755 index 0000000..73296db --- /dev/null +++ b/tests/kill-all.sh @@ -0,0 +1,9 @@ +#!/bin/bash +for r in us-west-1 us-east-1 eu-north-1; do + for a in $(aws braket search-jobs --region $r --filters \ + '[{"name":"jobName","operator":"CONTAINS","values":["flux-quantum-hold"]}]' \ + --query 'jobs[?status==`QUEUED`||status==`RUNNING`].jobArn' --output text); do + echo "cancelling $a" + aws braket cancel-job --region $r --job-arn "$a" + done +done diff --git a/tests/probe_hold.py b/tests/probe_hold.py new file mode 100644 index 0000000..e3fcd07 --- /dev/null +++ b/tests/probe_hold.py @@ -0,0 +1,610 @@ +#!/usr/bin/env python3 +"""Does a hybrid job token give priority to work submitted elsewhere. + +That is what this branch turns on. If the token travels, the hybrid job is the +scout and the classical work stays on our cluster. If not, the job has to +submit for us, which is a different design. + +Run against SV1 first. A simulator runs tasks straight away so this says +nothing about ordering, but it does say whether the API takes the token from +outside the container and links the task to the job. That costs a few cents. +Then run it against a QPU, where queuePriority is the direct answer. + + python3 probe_hold.py # SV1 + python3 probe_hold.py --device arn:... # a real QPU, costs money + +The hold is closed in a finally, so nothing is left running. +""" + +import argparse +import datetime +import json +import time +import uuid + +import boto3 +from braket.aws import AwsSession + +from flux_quantum.backends.braket import ( + SV1, + BraketBackend, + is_open, + queue_depth, + shots_range, + survey, + windows, +) + +# one qubit, measured. The cheapest thing that is still a real task. +BELL = """OPENQASM 3; +qubit[1] q; +bit[1] c; +h q[0]; +c[0] = measure q[0]; +""" + + +def submit(braket, device, bucket, token=None, shots=1): + """Create a task, with or without the token. + + Without it the docs say the task gets no priority and bills standalone. + The pair of calls is the comparison. + """ + args = { + "clientToken": str(uuid.uuid4()), + "deviceArn": device, + "shots": shots, + "outputS3Bucket": bucket, + "outputS3KeyPrefix": "flux-quantum-probe", + "action": json.dumps( + { + "braketSchemaHeader": { + "name": "braket.ir.openqasm.program", + "version": "1", + }, + "source": BELL, + } + ), + } + if token: + args["jobToken"] = token + return braket.create_quantum_task(**args)["quantumTaskArn"] + + +TERMINAL = ("COMPLETED", "FAILED", "CANCELLED") + + +def describe(braket, arn): + """What the service says about the task.""" + t = braket.get_quantum_task(quantumTaskArn=arn) + q = t.get("queueInfo") or {} + return { + "status": t.get("status"), + "job": t.get("jobArn"), + "queue": q.get("queue"), + "position": q.get("position"), + "priority": q.get("queuePriority"), + } + + +def settle(braket, arns, timeout=180, interval=1.0, sleep=time.sleep): + """Follow both tasks to the end, keeping the best view of each. + + jobArn is set when the task is created and is still there afterwards, so + the verdict never has to catch anything in flight. queueInfo is the + opposite. It describes a task that is still waiting and is gone once it + runs, so we record it if we see it and never wait for it. A one qubit task + on a simulator can finish before the first poll returns. + """ + best = {a: {} for a in arns} + deadline = time.time() + timeout + while time.time() < deadline: + pending = False + for a in arns: + seen = describe(braket, a) + for k, v in seen.items(): + # keep the first real value we see + if v and not best[a].get(k): + best[a][k] = v + best[a]["status"] = seen["status"] + if seen["status"] not in TERMINAL: + pending = True + if not pending: + break + sleep(interval) + return best + + +def make_queue(braket, device, bucket, count, shots=1): + """Submit filler tasks so there is a queue to jump. + + A device is only busy while it is shut, and drains as soon as a window + opens, so waiting for someone else's backlog means waiting for a window to + open with work still in front of it. Making our own removes that. + + These go in without a token so they sit in the Normal queue, and they go + in after the hold is ready, because a fast device empties during the two + minutes the hold takes to provision. + + They cost the task fee each, so keep the count small. + """ + arns = [] + for _ in range(count): + arns.append(submit(braket, device, bucket, None, shots)) + return arns + + +def drain_queue(braket, arns): + """Cancel the filler, so it does not run and bill for shots. + + A task that has already started cannot be cancelled, which is fine. The + task fee is spent either way, the shots are what this saves. + """ + for a in arns: + try: + # clientToken is required on this call, not optional as on most + braket.cancel_quantum_task(quantumTaskArn=a, clientToken=str(uuid.uuid4())) + except Exception as e: + # a task that already ran cannot be cancelled, which is expected + # on a fast device and not worth a wall of text + if "COMPLETED" not in str(e): + print("could not cancel %s: %s" % (a.split("/")[-1], e)) + + +def wait_for_queue(device, timeout=20, interval=0.5, sleep=time.sleep): + """Wait until the filler actually shows up in the Normal queue. + + Submitting returns before the service counts the task, and on a fast + device the window between counted and finished is short. Reading the + depth straight after the submit can miss it in either direction. + """ + deadline = time.time() + timeout + last = queue_depth(device) + while time.time() < deadline: + if last["normal"] > 0: + return last + sleep(interval) + last = queue_depth(device) + return last + + +def calibration(device): + """A summary of the device's calibration state, not the whole dump. + + Which edges were usable and when they were last characterised is what a + result has to be read against. A run against a device with a sixth of its + couplers dead is not the same measurement as one against a healthy one. + """ + try: + specs = device.properties.provider.specs + except Exception: + return None + edges, stamps = [], [] + for group in _walk_specs(specs): + ids, value, when = group + edges.append((ids, value)) + if when: + stamps.append(when) + live = [v for _, v in edges if v is not None and v > 0.6] + if not live: + return None + return { + "live": len(live), + "dead": len(edges) - len(live), + "error_min": round(1 - max(live), 6), + "error_max": round(1 - min(live), 6), + "calibrated_from": min(stamps) if stamps else None, + "calibrated_to": max(stamps) if stamps else None, + } + + +def _walk_specs(node): + """fCZ entries wherever the provider put them.""" + if isinstance(node, dict): + if "node_ids" in node and "characteristics" in node: + for c in node["characteristics"]: + if c.get("name") == "fCZ": + yield tuple(node["node_ids"]), c.get("value"), c.get("timestamp") + return + for v in node.values(): + yield from _walk_specs(v) + elif isinstance(node, list): + for v in node: + yield from _walk_specs(v) + + +def used_qubits(braket_task_arn): + """Which physical qubits the task actually ran on. + + The compiler places the circuit, so what you asked for and what ran are + not the same thing. Without this a result cannot be read against the + calibration data for the edges it used. + """ + try: + from braket.aws import AwsQuantumTask + + task = AwsQuantumTask(braket_task_arn) + # result() polls until the task finishes, five days by default. On a + # task that is still queued that is a hang, not a wait. + if task.state() != "COMPLETED": + return {"state": task.state(), "note": "not finished, nothing to read"} + r = task.result() + out = {"measured_qubits": list(r.measured_qubits or [])} + meta = getattr(r, "additional_metadata", None) + prog = getattr(getattr(meta, "rigettiMetadata", None), "compiledProgram", None) + if prog: + out["compiled"] = prog + return out + except Exception as e: + return {"error": str(e)} + + +def record(path, row): + """One JSON line per run. + + Printing to a terminal that is then closed is how a campaign ends up with + numbers nobody can reproduce. The token is left out on purpose: it grants + priority and bills to the job, so it does not belong in a log. + """ + import json as _json + + with open(path, "a") as fh: + fh.write(_json.dumps(row, default=str) + "\n") + + +def settled_depth(device, timeout=30, interval=2.0, sleep=time.sleep): + """Read the queue once it has stopped moving. + + The counters lag submission by a few seconds. Reading straight after a + submit gave normal=1 right after six filler tasks went in, then 7 a + moment later, so a token task's arrival in Priority was attributed to the + filler catching up. Wait for two readings in a row to agree. + """ + last = queue_depth(device) + deadline = time.time() + timeout + while time.time() < deadline: + sleep(interval) + now = queue_depth(device) + if all(now.get(k) == last.get(k) for k in ("normal", "priority")): + return now + last = now + return last + + +def moved(before, after): + """Which queue grew between two readings.""" + return { + k: after.get(k, 0) - v + for k, v in before.items() + if isinstance(v, int) and isinstance(after.get(k), int) + } + + +def _tzname(): + """What the times below are in, so nobody has to guess.""" + import time as _t + + return _t.tzname[_t.daylight and _t.localtime().tm_isdst > 0] or "local time" + + +def show_survey(shots=1): + """What is worth running against right now.""" + rows = survey() + if not rows: + print("no online QPUs found") + return + print("%-70s %-6s %-7s %s" % ("device", "open", "normal", "priority")) + for r in rows: + print( + "%-70s %-6s %-7d %d" + % (r["arn"], "yes" if r["open"] else "no", r["normal"], r["priority"]) + ) + usable = [r for r in rows if r["open"] and r["normal"] > 0] + + # every device, with when it runs. A device being open and empty now says + # nothing about whether it is worth coming back to. + print("\n windows in %s" % _tzname()) + for r in rows: + print( + " %s (%d queued, %s)" + % (r["arn"], r["normal"], "open" if r["open"] else "shut") + ) + for w in r["windows"]: + print(" %-10s %s to %s" % w) + + if usable: + # busiest first is the best demonstration, but cheapest is the one + # worth naming when they are all queued + cheapest = min( + (r for r in usable if r["cost"] is not None), + key=lambda r: r["cost"], + default=usable[0], + ) + print("\nbusiest queue: %s" % usable[0]["arn"]) + if cheapest is not usable[0] and cheapest["cost"] is not None: + print( + "cheapest probe: %s at $%.2f" % (cheapest["arn"], 2 * cheapest["cost"]) + ) + + +def verdict( + with_token, without, device, by_token=None, by_control=None, after_token=None +): + """What the two tasks say, strongest evidence first. + + jobArn is the one that matters. AWS says a task without the token gets no + priority and bills standalone, so association is what the token carries. + + queuePriority only exists while a task waits, so it is reported when we + saw it and not otherwise. A simulator usually shows nothing here. + """ + out = [] + if with_token.get("job"): + out.append("ASSOCIATED. The token travels: the task belongs to the job.") + else: + out.append( + "NOT ASSOCIATED. The call was accepted but the task is standalone,\n" + "which bills separately and gets no priority. Check the bill rather\n" + "than the return code before believing otherwise." + ) + if without.get("job"): + out.append( + "WARNING: the control task is also associated, so this comparison\n" + "says nothing. Something other than the token is linking them." + ) + + # Queue depth is counted per device and per queue, so a task landing in + # Priority shows up without having to catch it in flight. + settled = False + + # The count itself beats the difference between two counts. Priority has + # read zero on every survey of these devices, so one task sitting in it + # after we submitted exactly one token task is the plainest evidence + # there is. Movement missed this because the counters lag a submit. + if after_token and after_token.get("priority", 0) > 0: + out.append( + "IN THE PRIORITY QUEUE. %d task%s waiting there against %d in " + "Normal, and the only token task submitted was ours." + % ( + after_token["priority"], + "" if after_token["priority"] == 1 else "s", + after_token.get("normal", 0), + ) + ) + settled = True + + # A queue that emptied while the hold was coming up is not the same as a + # device that never had one. Taking the hold costs a couple of minutes, + # and a fast device clears its backlog in that time. + drained = by_token and by_token.get("normal", 0) < 0 + if drained: + out.append( + "The Normal queue lost %d tasks while the hold was provisioning, " + "so there was nothing left to jump by the time the token task " + "went in. Use --make-queue to build one after the hold is up." + % -by_token["normal"] + ) + settled = True + if not settled and by_token and by_control: + if by_token.get("priority", 0) > 0 and by_control.get("normal", 0) > 0: + out.append( + "PRIORITISED. The token submit grew the Priority queue and " + "the control grew Normal." + ) + settled = True + elif by_token.get("normal", 0) > 0: + out.append( + "NOT PRIORITISED. The token submit grew the Normal queue, so " + "the task queues like any other." + ) + settled = True + + if not settled: + a, c = with_token.get("priority"), without.get("priority") + if a and c and a != c: + out.append("PRIORITISED. {} against {} without the token.".format(a, c)) + elif a and c: + out.append( + "Same queuePriority either way ({}), so the token associates " + "but does not lift.".format(a) + ) + else: + out.append( + "Nothing was waiting on the device, so priority cannot be " + "seen. Rerun against a QPU with a queue." + ) + + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--device", default=SV1) + ap.add_argument("--shots", type=int, default=1) + ap.add_argument( + "--max-seconds", + type=int, + default=900, + help="ceiling on the hold, so a mistake stops billing", + ) + ap.add_argument( + "--record", + default="probe-runs.jsonl", + metavar="PATH", + help="append a JSON line per run, so a result can be read back " + "against the device state it was taken on", + ) + ap.add_argument( + "--survey", + action="store_true", + help="list online QPUs with their queues and windows, then stop", + ) + ap.add_argument( + "--filler-shots", + type=int, + default=500, + metavar="N", + help="shots per filler task. Shots are what keep a task occupying the " + "device: ten of them finish before the queue can be read, five " + "hundred hold it long enough to queue behind", + ) + ap.add_argument( + "--make-queue", + type=int, + default=0, + metavar="N", + help="submit N filler tasks first so there is a Normal queue to jump. " + "Costs the task fee each, and they are cancelled afterwards", + ) + ap.add_argument( + "--force", + action="store_true", + help="run even if the device is shut or has nothing queued", + ) + args = ap.parse_args() + + if args.survey: + show_survey(args.shots) + return + + opts = { + "hold": "job", + "device": args.device, + "hold_instance": "ml.m5.large", + "hold_max_seconds": args.max_seconds, + "queue_timeout": args.max_seconds, + } + + run = { + "started": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "device": args.device, + "shots": args.shots, + "filler_tasks": args.make_queue, + "filler_shots": args.filler_shots, + } + + b = BraketBackend() + print(b.credential_note) + print("device:", args.device) + + # A shut device or an empty queue means the run cannot show anything, and + # finding that out after provisioning wastes time and money. + q0 = queue_depth(args.device) + print("device queues:", q0) + if args.device != SV1 and not args.force: + from braket.aws import AwsDevice + + d = AwsDevice(args.device) + run["calibration"] = calibration(d) + run["windows"] = windows(d) + limits = shots_range(d) + for what, n in (("--shots", args.shots), ("--filler-shots", args.filler_shots)): + if limits and not (limits[0] <= n <= limits[1]): + raise SystemExit( + "this device takes {} to {} shots and {} is {}. " + "Submitting would fail after the hold is up and " + "billing.".format(limits[0], limits[1], what, n) + ) + if not is_open(d): + print("execution windows, {}:".format(_tzname())) + for w in windows(d): + print(" %-10s %s to %s" % w) + raise SystemExit( + "the device is outside its execution window, so the job would " + "just sit there. Wait for a window, or pass --force." + ) + if q0["normal"] == 0 and not args.make_queue: + raise SystemExit( + "nothing is waiting on this device, so a priority comparison " + "would say nothing. Pick a busier one, pass --make-queue N to " + "build a backlog, or pass --force." + ) + + arn = None + filler = [] + braket = boto3.client("braket") + bucket = AwsSession().default_bucket() + try: + arn = b.open_session(opts) + print("hybrid job:", arn) + ok, why = b.wait_for_priority(opts) + print(why) + if not ok: + raise SystemExit("never got the hold, nothing submitted, nothing spent") + + run["hold_job"] = arn + run["hold_ready_after_s"] = getattr(b, "_hold_ready_after", None) + token = b.session_id(arn) + print("token:", token[:12], "...") + print("job state at submit time:", b._hold_job.state()) + + # Build the backlog now, not before the hold. Taking the hold costs a + # couple of minutes, and a fast device empties in that time: four ten + # shot tasks on Cepheus were all COMPLETED before the hold was ready, + # so there was nothing left to jump. + if args.make_queue: + filler = make_queue( + braket, args.device, bucket, args.make_queue, args.filler_shots + ) + q0 = settled_depth(args.device) + print("queued %d filler tasks, device now %s" % (len(filler), q0)) + if q0["normal"] == 0: + print( + "the filler ran before it could queue. %d tasks of %d " + "shots is not enough to occupy this device, so raise " + "--filler-shots." % (args.make_queue, args.filler_shots) + ) + + with_token = submit(braket, args.device, bucket, token, args.shots) + q1 = settled_depth(args.device) + print("\nwith the token:", with_token, "queues", q1) + + without = submit(braket, args.device, bucket, None, args.shots) + q2 = settled_depth(args.device) + print("without it: ", without, "queues", q2) + + by_token, by_control = moved(q0, q1), moved(q1, q2) + print("\nmovement, token submit: ", by_token) + print("movement, control submit:", by_control) + + seen = settle(braket, [with_token, without]) + a, c = seen[with_token], seen[without] + print("\n %-10s %-30s %s" % ("", "with token", "without")) + for k in ("status", "job", "queue", "position", "priority"): + print(" %-10s %-30s %s" % (k, a.get(k), c.get(k))) + + run.update( + { + "queues": {"before": q0, "after_token": q1, "after_control": q2}, + "moved": {"token": by_token, "control": by_control}, + "with_token": {"arn": with_token, **a}, + "without": {"arn": without, **c}, + } + ) + lines = verdict(a, c, args.device, by_token, by_control, q1) + run["verdict"] = lines + for line in lines: + print(line) + + # which physical qubits ran, so the result can be read against the + # calibration for those edges + run["with_token"]["qubits"] = used_qubits(with_token) + run["without"]["qubits"] = used_qubits(without) + finally: + run["ended"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + try: + record(args.record, run) + print("recorded to %s" % args.record) + except Exception as e: + print("could not record this run: %s" % e) + if filler: + print("\ncancelling %d filler tasks" % len(filler)) + drain_queue(braket, filler) + if arn: + print("releasing the hold") + b.close_session() + print("job state:", b._hold_job.state() if b._hold_job else "none") + + +if __name__ == "__main__": + main() diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 8744531..87a2dd2 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,7 +1,13 @@ -import sys import importlib +import os +import sys + import pytest +# probe_hold lives in tests/, one level up, and pytest only puts the test +# file's own directory on the path +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + def _fresh(mock, monkeypatch): """Import flux_quantum.backends fresh with FLUX_QUANTUM_MOCK set/unset.""" diff --git a/tests/unit/test_backends.py b/tests/unit/test_backends.py index cf2923a..9d28a1a 100644 --- a/tests/unit/test_backends.py +++ b/tests/unit/test_backends.py @@ -1,3 +1,8 @@ +import sys +import types + +import pytest + IBM_CREDS = { "ibm_kingston_QRMI_IBM_QRS_ENDPOINT": "supersecret", "ibm_kingston_QRMI_IBM_QRS_IAM_ENDPOINT": "supersecret", @@ -23,40 +28,59 @@ def test_mock_present_with_env(backends_mock): assert {"mock", "mock_busy"} <= backends_mock.known_vendors() -def test_missing_creds_reported(backends_real): - ok, msg = backends_real.get_backend("ibm").credentials_present() - # names the variables it wants, never a value - assert not ok and "_QRMI_IBM_QRS_IAM_APIKEY" in msg +def _fake_qrmi(monkeypatch): + """The backend binds qrmi when its module is imported, so the stub has to + be in sys.modules before the package is imported fresh.""" + mod = types.ModuleType("qrmi") + mod.QuantumResource = object + mod.ResourceType = type("ResourceType", (), {"IBMQiskitRuntimeService": "x"}) + mod.Payload = mod.TaskStatus = object + monkeypatch.setitem(sys.modules, "qrmi", mod) + + +def test_missing_creds_fail_at_construction(backends_real): + """Constructing the backend is the check. The message names the variables + it wants, never a value.""" + with pytest.raises(backends_real.BackendError) as e: + backends_real.get_backend("ibm") + assert "_QRMI_IBM_QRS_IAM_APIKEY" in str(e.value) + + +def test_creds_present_when_env_set(fresh, monkeypatch): + _fake_qrmi(monkeypatch) + b = fresh(mock=False, env=IBM_CREDS) + assert "ibm_kingston" in b.get_backend("ibm").credential_note -def test_creds_present_when_env_set(fresh): +def test_missing_sdk_is_reported_after_credentials(fresh): + """With credentials in place but no qrmi, say what to install.""" b = fresh(mock=False, env=IBM_CREDS) - ok, msg = b.get_backend("ibm").credentials_present() - assert ok and "ibm_kingston" in msg + with pytest.raises(b.BackendError, match="qrmi is not installed"): + b.get_backend("ibm") def test_partial_creds_name_only_the_missing_variable(fresh): partial = dict(IBM_CREDS) del partial["ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN"] b = fresh(mock=False, env=partial) - ok, msg = b.get_backend("ibm").credentials_present() - assert not ok + with pytest.raises(b.BackendError) as e: + b.get_backend("ibm") + msg = str(e.value) assert "ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN" in msg # the api key value must never be echoed back assert "supersecret" not in msg def test_mock_needs_no_creds(backends_mock): - ok, _ = backends_mock.get_backend("mock").credentials_present() - assert ok + assert backends_mock.get_backend("mock") is not None -def test_backend_declares_options_and_scout_options_roundtrip(): +def test_backend_declares_options_and_scout_options_roundtrip(fresh, monkeypatch): """A vendor backend declares CLI options and extracts them from args.""" - from flux_quantum.backends import get_backend - import flux_quantum.backends.ibm # ensure ibm registered + _fake_qrmi(monkeypatch) + b = fresh(mock=False, env=IBM_CREDS) - ibm = get_backend("ibm") + ibm = b.get_backend("ibm") declared = [] ibm.add_options(lambda name, **kw: declared.append(name)) assert "--ibm-resource" in declared and "--ibm-type" in declared @@ -70,8 +94,9 @@ class _Args: assert opts["type"] == "qiskit-runtime-service" -def test_scout_options_infer_the_resource_from_the_environment(fresh): +def test_scout_options_infer_the_resource_from_the_environment(fresh, monkeypatch): """One resource configured is unambiguous, so the id need not be repeated.""" + _fake_qrmi(monkeypatch) b = fresh(mock=False, env=IBM_CREDS) class _Args: @@ -81,9 +106,10 @@ class _Args: assert b.get_backend("ibm").scout_options(_Args())["resource"] == "ibm_kingston" -def test_job_environment_follows_the_qrmi_convention(fresh): +def test_job_environment_follows_the_qrmi_convention(fresh, monkeypatch): """The classical job gets the same variables Slurm and LSF set, so user code can call get_job_qpu_resources_and_types unchanged.""" + _fake_qrmi(monkeypatch) b = fresh(mock=False, env=IBM_CREDS) env = b.get_backend("ibm").job_environment( {"resource": "ibm_kingston", "type": "qiskit-runtime-service"} @@ -105,20 +131,23 @@ def test_mock_open_session_uses_options(monkeypatch): assert b.open_session({}).startswith("mock-session-") -def test_open_session_without_a_resource_is_explicit(): - from flux_quantum.backends import get_backend - import flux_quantum.backends.ibm # noqa - - with __import__("pytest").raises(ValueError, match="no QRMI resource id"): - get_backend("ibm").open_session({}) - - -def test_open_session_names_the_missing_variable(fresh): - partial = dict(IBM_CREDS) - del partial["ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY"] - b = fresh(mock=False, env=partial) - with __import__("pytest").raises(ValueError) as e: +def test_open_session_without_a_resource_is_explicit(fresh, monkeypatch): + _fake_qrmi(monkeypatch) + b = fresh(mock=False, env=IBM_CREDS) + with pytest.raises(ValueError, match="no QRMI resource id"): + b.get_backend("ibm").open_session({}) + + +def test_open_session_names_the_missing_variable(fresh, monkeypatch): + """Two resources, one complete. Construction passes on the complete one, + and asking for the other names what it lacks.""" + _fake_qrmi(monkeypatch) + env = dict(IBM_CREDS) + env["ibm_torino_QRMI_IBM_QRS_ENDPOINT"] = "supersecret" + b = fresh(mock=False, env=env) + with pytest.raises(ValueError) as e: b.get_backend("ibm").open_session( - {"resource": "ibm_kingston", "type": "qiskit-runtime-service"} + {"resource": "ibm_torino", "type": "qiskit-runtime-service"} ) - assert "ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY" in str(e.value) + assert "ibm_torino_QRMI_IBM_QRS_IAM_APIKEY" in str(e.value) + assert "supersecret" not in str(e.value) diff --git a/tests/unit/test_braket.py b/tests/unit/test_braket.py index 68d9b92..5bf103e 100644 --- a/tests/unit/test_braket.py +++ b/tests/unit/test_braket.py @@ -55,20 +55,57 @@ def i(self, q): circuits.Circuit = FakeCircuit pkg = types.ModuleType("braket") + + # the backend resolves credentials through boto3 when it is constructed + class FakeSession: + def get_credentials(self): + return types.SimpleNamespace(method="env") + + boto3 = types.ModuleType("boto3") + boto3.Session = FakeSession monkeypatch.setitem(sys.modules, "braket", pkg) monkeypatch.setitem(sys.modules, "braket.aws", aws) monkeypatch.setitem(sys.modules, "braket.circuits", circuits) + monkeypatch.setitem(sys.modules, "boto3", boto3) + # the SDK is bound when the backend module is imported, so import it again + for name in [m for m in sys.modules if m.startswith("flux_quantum")]: + del sys.modules[name] return submitted def _backend(): from flux_quantum.backends import get_backend - import flux_quantum.backends.braket # noqa return get_backend("braket") -def test_defaults_to_the_sv1_simulator(): +def test_missing_sdk_fails_at_construction(monkeypatch): + """No SDK, no backend. The message says what to install.""" + monkeypatch.setitem(sys.modules, "boto3", None) + monkeypatch.setitem(sys.modules, "braket", None) + for name in [m for m in sys.modules if m.startswith("flux_quantum")]: + del sys.modules[name] + from flux_quantum.backends import BackendError, get_backend + + with pytest.raises(BackendError, match="amazon-braket-sdk"): + get_backend("braket") + + +def test_missing_credentials_fail_at_construction(fake_braket, monkeypatch): + import boto3 + + monkeypatch.setattr(boto3.Session, "get_credentials", lambda self: None) + from flux_quantum.backends import BackendError + + with pytest.raises(BackendError, match="no AWS credentials"): + _backend() + + +def test_credential_note_says_where_they_came_from(fake_braket): + assert "via env" in _backend().credential_note + + +def test_defaults_to_the_sv1_simulator(fake_braket): """SV1 is cents per task, so it is the sane default for a probe.""" from flux_quantum.backends.braket import SV1 diff --git a/tests/unit/test_braket_hold.py b/tests/unit/test_braket_hold.py new file mode 100644 index 0000000..8bc1b32 --- /dev/null +++ b/tests/unit/test_braket_hold.py @@ -0,0 +1,1222 @@ +"""The hybrid job hold, tested without AWS. + +A running hybrid job holds the priority queue for its device and publishes +AMZN_BRAKET_JOB_TOKEN. Work submitted anywhere with that token gets the job's +priority. Without it the docs say it gets none and bills standalone. So the +token, not the job ARN, is what the classical job has to be handed. +""" + +import json + + +def test_entry_publishes_the_token_then_waits_for_release(): + from flux_quantum.backends.hold import entry + + written = {} + released = {"yet": False} + + class FakeS3: + def put_object(self, Bucket, Key, Body): + written[Key] = Body + + def head_object(self, Bucket, Key): + if not released["yet"]: + raise RuntimeError("not there") + return {} + + def sleep(_): + released["yet"] = True + + env = { + "AMZN_BRAKET_OUT_S3_BUCKET": "bucket", + "AMZN_BRAKET_JOB_TOKEN": "tok-123", + "AMZN_BRAKET_DEVICE_ARN": "arn:aws:braket:::device/quantum-simulator/amazon/sv1", + "AMZN_BRAKET_JOB_ARN": "arn:aws:braket:us-east-1:1:job/x", + "FLUX_QUANTUM_PREFIX": "flux-quantum/x", + "FLUX_QUANTUM_MAX_SECONDS": "60", + "FLUX_QUANTUM_POLL": "0", + } + rc = entry.main(env=env, sleep=sleep, client=FakeS3()) + assert rc == 0 + body = json.loads(written["flux-quantum/x/token.json"].decode()) + assert body["token"] == "tok-123" + assert body["device"].endswith("sv1") + + +def test_entry_gives_up_rather_than_billing_forever(): + """A scout that is never released has to stop. The instance bills by the + minute for as long as it sits there.""" + from flux_quantum.backends.hold import entry + + class FakeS3: + def put_object(self, Bucket, Key, Body): + pass + + def head_object(self, Bucket, Key): + raise RuntimeError("never released") + + calls = {"n": 0} + + def sleep(_): + calls["n"] += 1 + if calls["n"] > 50: + raise AssertionError("did not respect the deadline") + + env = { + "AMZN_BRAKET_OUT_S3_BUCKET": "bucket", + "AMZN_BRAKET_JOB_TOKEN": "tok", + "FLUX_QUANTUM_PREFIX": "p", + "FLUX_QUANTUM_MAX_SECONDS": "0", + "FLUX_QUANTUM_POLL": "0", + } + assert entry.main(env=env, sleep=sleep, client=FakeS3()) == 0 + + +def _backend(): + """Skip __init__, which wants the SDK and credentials. Nothing here needs + either.""" + from flux_quantum.backends.braket import BraketBackend + + b = BraketBackend.__new__(BraketBackend) + b._session = None + return b + + +def test_scout_options_carry_the_hold_mode(): + + class Args: + braket_device = None + braket_region = None + braket_shots = None + braket_queue_timeout = None + braket_ungate_position = None + braket_hold = "job" + braket_hold_instance = None + braket_hold_max_seconds = None + + opts = _backend().scout_options(Args()) + assert opts["hold"] == "job" + # the cheapest instance, since it does no work + assert opts["hold_instance"] == "ml.m5.large" + assert opts["hold_max_seconds"] == 900 + + +def test_probe_stays_the_default(): + """The older mode keeps working for anyone not asking for a hybrid job.""" + + class Args: + braket_device = None + braket_region = None + braket_shots = None + braket_queue_timeout = None + braket_ungate_position = None + + assert _backend().scout_options(Args())["hold"] == "probe" + + +def test_the_token_replaces_the_job_arn_as_the_session(): + """open_session returns the job ARN. The classical job needs the token, + which only exists once the job runs.""" + b = _backend() + assert b.session_id("arn:job") == "arn:job" + b._session = "tok-abc" + assert b.session_id("arn:job") == "tok-abc" + + +def test_base_session_id_is_a_passthrough(): + """The default hands back what open_session returned.""" + from flux_quantum.backends.base import Backend + + class Plain(Backend): + name = "plain" + + def probe(self): + return None + + assert Plain().session_id("whatever") == "whatever" + + +def _held(states, tokens): + """A backend whose job walks through states and whose token shows up when + tokens says so.""" + import types + + import flux_quantum.backends.braket as bk + + b = bk.BraketBackend.__new__(bk.BraketBackend) + b._hold_prefix = "p" + b._session = None + b._hold_ready_after = None + seq = list(states) + b._hold_job = types.SimpleNamespace(state=lambda: seq.pop(0)) + b._output_bucket = lambda: "bucket" + toks = list(tokens) + b._read_token = lambda s3, bucket: toks.pop(0) + return b, bk + + +def test_running_without_a_token_is_reported_as_starting(capsys): + """Braket says RUNNING when the instance is up. The device is not ours + until the container publishes its token, so that stretch reads as + starting.""" + import types + + b, bk = _held(["RUNNING", "RUNNING"], [None, "tok"]) + b._hold_started = __import__("time").time() + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: None) + try: + ok, why = b._wait_for_hold_job({}, 0, lambda _: None) + finally: + bk.boto3 = original + + assert ok + assert "holding the device" in why + assert "starting the container" in capsys.readouterr().out + + +def test_the_wait_is_timed(): + """How long from asking for the hold to having it. Minutes rather than + milliseconds decides whether a scout can be started per pair.""" + import time + import types + + b, bk = _held(["RUNNING"], ["tok"]) + b._hold_started = time.time() - 42 + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: None) + try: + ok, why = b._wait_for_hold_job({}, 0, lambda _: None) + finally: + bk.boto3 = original + + assert ok and "42s" in why + assert 41 < b._hold_ready_after < 44 + + +def test_a_job_that_died_says_so_and_for_how_long(): + import time + import types + + b, bk = _held(["FAILED"], [None]) + b._hold_started = time.time() - 10 + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: None) + try: + ok, why = b._wait_for_hold_job({}, 0, lambda _: None) + finally: + bk.boto3 = original + + assert not ok and "FAILED" in why and "10s" in why + + +def test_the_verdict_survives_a_task_that_finishes_instantly(): + """jobArn is set at creation and is still there afterwards, so the check + never has to win a race. A one qubit task on a simulator can finish before + the first poll returns.""" + import probe_hold + + class Instant: + def get_quantum_task(self, quantumTaskArn): + return { + "status": "COMPLETED", + "jobArn": "arn:job" if quantumTaskArn == "t1" else None, + } + + seen = probe_hold.settle(Instant(), ["t1", "t2"], sleep=lambda _: None) + assert seen["t1"]["job"] == "arn:job" + assert not seen["t2"].get("job") + + +def test_a_queue_glimpse_is_kept_but_not_required(): + """queueInfo describes a waiting task and is gone once it runs. Record it + if it shows up, never depend on it.""" + import probe_hold + + class Fleeting: + def __init__(self): + self.n = 0 + + def get_quantum_task(self, quantumTaskArn): + self.n += 1 + if self.n <= 2: + return { + "status": "QUEUED", + "jobArn": "arn:job", + "queueInfo": {"queuePriority": "Priority", "position": "1"}, + } + return {"status": "COMPLETED", "jobArn": "arn:job"} + + seen = probe_hold.settle(Fleeting(), ["t1"], sleep=lambda _: None) + assert seen["t1"]["priority"] == "Priority" + assert seen["t1"]["status"] == "COMPLETED" + + +def test_settle_gives_up_rather_than_hanging(): + import probe_hold + + class Stuck: + def get_quantum_task(self, quantumTaskArn): + return {"status": "QUEUED", "jobArn": "arn:job"} + + calls = {"n": 0} + + def sleep(_): + calls["n"] += 1 + if calls["n"] > 500: + raise AssertionError("did not respect the timeout") + + seen = probe_hold.settle(Stuck(), ["t1"], timeout=0, sleep=sleep) + assert seen["t1"] == {} + + +def test_verdict_handles_a_queue_that_was_never_visible(): + """The ordinary simulator case. settle only records fields it saw, so + reading priority unconditionally raised KeyError on the very run this + probe is for.""" + from probe_hold import SV1, verdict + + out = "\n".join(verdict({"job": "arn:job"}, {}, SV1)) + assert "ASSOCIATED" in out + assert "cannot be seen" in out + + +def test_verdict_reports_priority_when_both_were_seen(): + from probe_hold import verdict + + out = "\n".join( + verdict( + {"job": "arn:job", "priority": "Priority"}, {"priority": "Normal"}, "qpu" + ) + ) + assert "PRIORITISED" in out + + +def test_verdict_calls_out_a_control_that_is_also_associated(): + """If the control is linked too, the comparison says nothing.""" + from probe_hold import verdict + + out = "\n".join(verdict({"job": "a"}, {"job": "a"}, "qpu")) + assert "WARNING" in out + + +def test_verdict_says_when_the_token_did_not_travel(): + from probe_hold import SV1, verdict + + assert "NOT ASSOCIATED" in "\n".join(verdict({}, {}, SV1)) + + +class _FakeWindow: + def __init__(self, day, start, end): + import datetime + + self.executionDay = type("D", (), {"value": day})() + self.windowStartHour = datetime.time(*start) + self.windowEndHour = datetime.time(*end) + + +ALWAYS = [_FakeWindow("Everyday", (0, 0), (23, 59))] + + +class _FakeDevice: + """Enough of AwsDevice to exercise queue and window logic without AWS. + + fleet maps an arn to (normal, priority) or (normal, priority, windows). + """ + + fleet = {} + + def __init__(self, arn): + self.arn = arn + wins = self.fleet[arn][2] if len(self.fleet[arn]) > 2 else ALWAYS + self.properties = type( + "P", (), {"service": type("S", (), {"executionWindows": wins})()} + )() + + def queue_depth(self): + from braket.aws.queue_information import QueueDepthInfo, QueueType + + n, p = self.fleet[self.arn][:2] + return QueueDepthInfo( + quantum_tasks={QueueType.NORMAL: str(n), QueueType.PRIORITY: str(p)}, + jobs="0", + ) + + @classmethod + def get_devices(cls, statuses=None): + return [cls(a) for a in cls.fleet] + + +def test_queue_depth_splits_normal_from_priority(): + from flux_quantum.backends.braket import queue_depth + + _FakeDevice.fleet = {"arn:aws:braket:::device/qpu/x/y": (7, 3)} + q = queue_depth("arn:aws:braket:::device/qpu/x/y", cls=_FakeDevice) + assert (q["normal"], q["priority"], q["jobs"]) == (7, 3, 0) + + +def test_only_qpus_with_work_waiting_are_offered(): + """Priority means nothing on an idle device, and simulators never queue.""" + from flux_quantum.backends.braket import qpus_with_a_queue + + _FakeDevice.fleet = { + "arn:aws:braket:::device/qpu/quera/Aquila": (59, 0), + "arn:aws:braket:::device/qpu/iqm/Emerald": (1, 0), + "arn:aws:braket:::device/qpu/iqm/Garnet": (0, 0), + "arn:aws:braket:::device/quantum-simulator/amazon/sv1": (99, 0), + } + got = qpus_with_a_queue(cls=_FakeDevice) + assert [d["arn"].split("/")[-1] for d in got] == ["Aquila", "Emerald"] + + +def test_a_device_that_cannot_be_read_is_skipped(): + from flux_quantum.backends.braket import qpus_with_a_queue + + class Broken(_FakeDevice): + def queue_depth(self): + raise RuntimeError("AccessDenied in this region") + + Broken.fleet = {"arn:aws:braket:::device/qpu/x/y": (5, 0)} + assert qpus_with_a_queue(cls=Broken) == [] + + +def test_priority_is_read_from_which_queue_grew(): + from probe_hold import verdict + + out = "\n".join( + verdict( + {"job": "arn:job"}, + {}, + "qpu", + {"normal": 0, "priority": 1}, + {"normal": 1, "priority": 0}, + ) + ) + assert "PRIORITISED" in out and "NOT PRIORITISED" not in out + + +def test_a_token_that_queues_normally_is_called_out(): + from probe_hold import verdict + + out = "\n".join( + verdict( + {"job": "arn:job"}, + {}, + "qpu", + {"normal": 1, "priority": 0}, + {"normal": 1, "priority": 0}, + ) + ) + assert "NOT PRIORITISED" in out + + +def test_an_idle_device_gives_one_verdict_not_two(): + from probe_hold import SV1, verdict + + out = verdict( + {"job": "arn:job"}, + {}, + SV1, + {"normal": 0, "priority": 0}, + {"normal": 0, "priority": 0}, + ) + assert len([x for x in out if "PRIORITIS" in x or "cannot be seen" in x]) == 1 + + +def test_a_maintenance_gap_counts_as_shut(): + """Garnet runs weekdays but stops between 15:30 and 17:15. From outside + that looks the same as a long queue.""" + import datetime + + from flux_quantum.backends.braket import is_open + + wins = [ + _FakeWindow("Weekdays", (3, 15), (15, 29)), + _FakeWindow("Weekdays", (17, 15), (23, 59)), + ] + _FakeDevice.fleet = {"arn:garnet": (7, 0, wins)} + d = _FakeDevice("arn:garnet") + monday = datetime.datetime(2026, 9, 28, tzinfo=datetime.timezone.utc) + + assert not is_open(d, monday.replace(hour=16, minute=23)) + assert is_open(d, monday.replace(hour=18)) + # Weekdays must not fire at the weekend + assert not is_open(d, monday.replace(day=26, hour=18)) + + +def test_everyday_windows_are_understood(): + import datetime + + from flux_quantum.backends.braket import is_open + + _FakeDevice.fleet = {"arn:x": (0, 0, ALWAYS)} + assert is_open( + _FakeDevice("arn:x"), + datetime.datetime(2026, 9, 27, 3, tzinfo=datetime.timezone.utc), + ) + + +def test_survey_puts_open_devices_first_and_skips_shut_ones(): + """The real case from a Monday afternoon. Aquila had 60 waiting and Garnet + 7, but both were shut, so only Forte was worth paying for.""" + import datetime + + from flux_quantum.backends.braket import qpus_with_a_queue, survey + + shut = [_FakeWindow("Friday", (4, 0), (11, 59))] + _FakeDevice.fleet = { + "arn:aws:braket:::device/qpu/quera/Aquila": (60, 0, shut), + "arn:aws:braket:::device/qpu/iqm/Garnet": (7, 0, shut), + "arn:aws:braket:::device/qpu/ionq/Forte": (1, 0, ALWAYS), + "arn:aws:braket:::device/quantum-simulator/amazon/sv1": (99, 0, ALWAYS), + } + monday = datetime.datetime(2026, 9, 28, 16, tzinfo=datetime.timezone.utc) + + names = [r["arn"].split("/")[-1] for r in survey(cls=_FakeDevice, now=monday)] + assert names[0] == "Forte", names + assert "sv1" not in names + assert [ + r["arn"].split("/")[-1] for r in qpus_with_a_queue(cls=_FakeDevice, now=monday) + ] == ["Forte"] + + +def test_windows_are_reported_readably(): + from flux_quantum.backends.braket import windows + + _FakeDevice.fleet = {"arn:x": (0, 0, [_FakeWindow("Weekdays", (17, 15), (23, 59))])} + assert windows(_FakeDevice("arn:x"), local=False) == [ + ("Weekdays", "17:15", "23:59") + ] + + +def test_queue_depth_survives_the_note_braket_appends(): + """The jobs field is a string and sometimes carries a message, as in + "0 (1 prioritized hybrid job running)". int() on that raised mid run, + after a task had already been submitted and paid for.""" + from braket.aws.queue_information import QueueDepthInfo, QueueType + + from flux_quantum.backends.braket import queue_depth + + class Chatty: + def __init__(self, arn): + pass + + def queue_depth(self): + return QueueDepthInfo( + quantum_tasks={QueueType.NORMAL: "7", QueueType.PRIORITY: "0"}, + jobs="0 (1 prioritized hybrid job running)", + ) + + q = queue_depth("arn:x", cls=Chatty) + assert q["normal"] == 7 + assert q["jobs"] == 0 + assert "prioritized hybrid job running" in q["jobs_note"] + + +def test_moved_ignores_the_note(): + from probe_hold import moved + + before = {"normal": 1, "priority": 0, "jobs": 0, "jobs_note": "0 (1 running)"} + after = {"normal": 1, "priority": 1, "jobs": 0, "jobs_note": "0 (1 running)"} + assert moved(before, after) == {"normal": 0, "priority": 1, "jobs": 0} + + +def _device(wins): + return type( + "D", + (), + { + "properties": type( + "P", (), {"service": type("S", (), {"executionWindows": wins})()} + )() + }, + )() + + +def test_windows_come_back_in_local_time(): + """Braket gives these in UTC. Reading them against your own clock is where + the mistakes happen.""" + import datetime + import os + import time + + from flux_quantum.backends.braket import windows + + os.environ["TZ"] = "Europe/Berlin" + time.tzset() + try: + d = _device([_FakeWindow("Weekdays", (3, 15), (15, 29))]) + now = datetime.datetime(2026, 9, 28, tzinfo=datetime.timezone.utc) + assert windows(d, now=now) == [("Weekdays", "05:15", "17:29")] + assert windows(d, local=False, now=now) == [("Weekdays", "03:15", "15:29")] + finally: + del os.environ["TZ"] + time.tzset() + + +def test_a_window_that_lands_on_another_day_says_so(): + """The day label is the UTC day, so a converted time can sit outside it.""" + import datetime + import os + import time + + from flux_quantum.backends.braket import windows + + os.environ["TZ"] = "America/Los_Angeles" + time.tzset() + try: + d = _device([_FakeWindow("Weekdays", (3, 15), (15, 29))]) + now = datetime.datetime(2026, 9, 28, tzinfo=datetime.timezone.utc) + assert windows(d, now=now) == [("Weekdays", "20:15 (prev day)", "08:29")] + finally: + del os.environ["TZ"] + time.tzset() + + +def test_cleanup_undoes_a_run_that_failed_halfway(): + """The job was made but nothing else. Cleanup should still cancel it and + clear the keys it would have used.""" + import types + + import flux_quantum.backends.braket as bk + + calls = [] + + class S3: + def put_object(self, Bucket, Key, Body): + calls.append(("put", Key)) + + def delete_object(self, Bucket, Key): + calls.append(("delete", Key)) + + b = bk.BraketBackend.__new__(bk.BraketBackend) + b._made = { + "job": types.SimpleNamespace( + state=lambda: "RUNNING", cancel=lambda: calls.append(("cancel", None)) + ), + "bucket": "amazon-braket-1", + "prefix": "flux-quantum/j", + } + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: S3()) + try: + b.close_session() + # calling twice must not repeat the work + b.close_session() + finally: + bk.boto3 = original + + assert calls == [ + ("put", "flux-quantum/j/release"), + ("cancel", None), + ("delete", "flux-quantum/j/token.json"), + ("delete", "flux-quantum/j/release"), + ], calls + + +def test_cleanup_continues_when_a_step_fails(): + """A cancel that fails must not stop the keys being deleted.""" + import types + + import flux_quantum.backends.braket as bk + + deleted = [] + + class S3: + def put_object(self, **kw): + pass + + def delete_object(self, Bucket, Key): + deleted.append(Key) + + def boom(): + raise RuntimeError("AccessDenied") + + b = bk.BraketBackend.__new__(bk.BraketBackend) + b._made = { + "job": types.SimpleNamespace(state=lambda: "RUNNING", cancel=boom), + "bucket": "bucket", + "prefix": "p", + } + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: S3()) + try: + b.close_session() + finally: + bk.boto3 = original + assert deleted == ["p/token.json", "p/release"] + + +def test_cleanup_with_nothing_made_is_a_noop(): + import flux_quantum.backends.braket as bk + + b = bk.BraketBackend.__new__(bk.BraketBackend) + b._made = {"job": None, "bucket": None, "prefix": None} + b.close_session() + + +def _holding(token_payload, device): + """A backend whose hybrid job is running and has published a token.""" + import json + import types + + import flux_quantum.backends.braket as bk + + b = bk.BraketBackend.__new__(bk.BraketBackend) + b._hold_prefix = "p" + b._session = None + b._token_device = None + b._hold_ready_after = None + b._hold_started = __import__("time").time() + b._hold_job = types.SimpleNamespace(state=lambda: "RUNNING") + b._output_bucket = lambda: "bucket" + + class Body: + @staticmethod + def read(): + return json.dumps(token_payload).encode() + + class S3: + def get_object(self, Bucket, Key): + return {"Body": Body} + + return b, bk, S3(), {"device": device} + + +def test_a_token_for_another_device_is_refused(): + """Seb: the device ARN in AwsQuantumJob.create decides which device the + token is valid for. Using it elsewhere silently gets no priority, so it + has to fail loudly instead.""" + import types + + b, bk, s3, opts = _holding( + { + "token": "tok", + "device": "arn:aws:braket:::device/quantum-simulator/amazon/sv1", + }, + "arn:aws:braket:eu-north-1::device/qpu/iqm/Garnet", + ) + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: s3) + try: + ok, why = b._wait_for_hold_job(opts, 0, lambda _: None) + finally: + bk.boto3 = original + + assert not ok + assert "only valid for the device" in why + assert b._session is None + + +def test_a_token_for_the_right_device_is_accepted(): + import types + + garnet = "arn:aws:braket:eu-north-1::device/qpu/iqm/Garnet" + b, bk, s3, opts = _holding({"token": "tok", "device": garnet}, garnet) + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: s3) + try: + ok, why = b._wait_for_hold_job(opts, 0, lambda _: None) + finally: + bk.boto3 = original + + assert ok and b._session == "tok" + + +def test_a_token_without_a_device_is_not_second_guessed(): + """An older hold published no device. Refusing on that would break a + working setup for the sake of a check we cannot make.""" + import types + + b, bk, s3, opts = _holding({"token": "tok"}, "arn:whatever") + original = bk.boto3 + bk.boto3 = types.SimpleNamespace(client=lambda _: s3) + try: + ok, _ = b._wait_for_hold_job(opts, 0, lambda _: None) + finally: + bk.boto3 = original + assert ok and b._session == "tok" + + +def _priced(price, unit="shot"): + return type( + "D", + (), + { + "properties": type( + "P", + (), + { + "service": type( + "S", + (), + { + "deviceCost": type( + "C", (), {"price": price, "unit": unit} + )(), + "executionWindows": [], + }, + )() + }, + )() + }, + )() + + +def test_cost_is_the_task_fee_plus_shots(): + """Rigetti is $0.00090 a shot on top of the flat $0.30 task fee.""" + from flux_quantum.backends.braket import cost + + assert round(cost(_priced(0.00090), shots=1), 5) == 0.3009 + assert round(cost(_priced(0.00090), shots=100), 4) == 0.39 + + +def test_per_minute_devices_have_no_per_shot_price(): + """Simulators bill for time, so a shot count says nothing about cost.""" + from flux_quantum.backends.braket import cost + + assert cost(_priced(0.075, unit="minute")) is None + + +def test_a_device_with_no_price_does_not_break_the_survey(): + from flux_quantum.backends.braket import cost + + assert cost(type("D", (), {})()) is None + + +def _api_bound_client(seen): + """A Braket client that rejects anything the real API would.""" + import botocore.session + + m = botocore.session.get_session().get_service_model("braket") + create = m.operation_model("CreateQuantumTask").input_shape + cancel = m.operation_model("CancelQuantumTask").input_shape + + class Client: + def create_quantum_task(self, **kw): + assert not set(create.required_members) - set(kw) + assert not set(kw) - set(create.members) + seen.setdefault("tokens", []).append(kw.get("jobToken")) + seen.setdefault("shots", []).append(kw.get("shots")) + return {"quantumTaskArn": "arn:task:%d" % len(seen["tokens"])} + + def cancel_quantum_task(self, **kw): + assert not set(cancel.required_members) - set(kw) + assert not set(kw) - set(cancel.members) + seen.setdefault("cancelled", []).append(kw["quantumTaskArn"]) + + return Client() + + +def test_filler_goes_in_without_a_token(): + """It is there to sit in the Normal queue. A token would put it in + Priority, which is the queue we are trying to jump.""" + from probe_hold import make_queue + + seen = {} + arns = make_queue(_api_bound_client(seen), "arn:dev", "bucket", 3) + assert len(arns) == 3 + assert seen["tokens"] == [None, None, None] + + +def test_filler_is_cancelled_with_a_client_token(): + """CancelQuantumTask requires clientToken, unlike most calls. Without it + the cleanup fails at the moment it matters.""" + from probe_hold import drain_queue, make_queue + + seen = {} + client = _api_bound_client(seen) + drain_queue(client, make_queue(client, "arn:dev", "bucket", 2)) + assert len(seen["cancelled"]) == 2 + + +def test_a_cancel_that_fails_does_not_stop_the_rest(): + """A task that already started cannot be cancelled, and the others still + need clearing.""" + from probe_hold import drain_queue + + done = [] + + class Client: + def cancel_quantum_task(self, quantumTaskArn, clientToken): + if quantumTaskArn == "arn:b": + raise RuntimeError("already RUNNING") + done.append(quantumTaskArn) + + drain_queue(Client(), ["arn:a", "arn:b", "arn:c"]) + assert done == ["arn:a", "arn:c"] + + +def test_shots_range_comes_from_the_device(): + """Cepheus takes 10 to 50000 and refuses 1. Learning that from a + ValidationException means the hold is already up and billing.""" + from flux_quantum.backends.braket import shots_range + + d = type( + "D", + (), + { + "properties": type( + "P", (), {"service": type("S", (), {"shotsRange": (10, 50000)})()} + )() + }, + )() + assert shots_range(d) == (10, 50000) + + +def test_a_device_that_does_not_say_is_not_guessed_at(): + from flux_quantum.backends.braket import shots_range + + assert shots_range(type("D", (), {})()) is None + + +def test_the_device_region_wins_over_the_environment(): + """A stale AWS_DEFAULT_REGION pointed the session at the wrong region, so + the bucket polled for the token was the wrong one and the hold looked like + it never started.""" + import os + import types + + import flux_quantum.backends.braket as bk + + os.environ["AWS_DEFAULT_REGION"] = "eu-north-1" + seen = {} + + def fake_create(**kw): + seen["region"] = os.environ["AWS_DEFAULT_REGION"] + job = types.SimpleNamespace(arn="arn:job") + return job + + import braket.aws as aws_mod + + original, original_session = aws_mod.AwsQuantumJob.create, aws_mod.AwsSession + aws_mod.AwsQuantumJob.create = staticmethod(fake_create) + aws_mod.AwsSession = lambda: types.SimpleNamespace( + default_bucket=lambda: "amazon-braket-us-west-1-1" + ) + try: + b = bk.BraketBackend.__new__(bk.BraketBackend) + b._hold_job = b._hold_prefix = b._session = None + b._made = {"job": None, "bucket": None, "prefix": None} + b._open_hold_job( + {"device": "arn:aws:braket:us-west-1::device/qpu/rigetti/Cepheus-1-108Q"} + ) + finally: + aws_mod.AwsQuantumJob.create = original + aws_mod.AwsSession = original_session + del os.environ["AWS_DEFAULT_REGION"] + + assert seen["region"] == "us-west-1" + + +def test_wait_for_queue_returns_as_soon_as_the_filler_registers(): + """Submitting returns before the service counts the task, so reading the + depth straight after can miss it.""" + import probe_hold + + depths = [ + {"normal": 0, "priority": 0}, + {"normal": 0, "priority": 0}, + {"normal": 3, "priority": 0}, + ] + probe_hold.queue_depth = lambda d: depths.pop(0) if depths else {"normal": 3} + try: + assert probe_hold.wait_for_queue("arn:dev", sleep=lambda _: None)["normal"] == 3 + finally: + from flux_quantum.backends.braket import queue_depth + + probe_hold.queue_depth = queue_depth + + +def test_wait_for_queue_gives_up_on_a_device_too_fast_to_queue(): + """Cepheus ran four ten shot tasks before the hold was ready. Waiting + forever for a queue that will never form helps nobody.""" + import probe_hold + + probe_hold.queue_depth = lambda d: {"normal": 0, "priority": 0} + calls = {"n": 0} + + def sleep(_): + calls["n"] += 1 + if calls["n"] > 200: + raise AssertionError("did not respect the timeout") + + try: + assert ( + probe_hold.wait_for_queue("arn:dev", timeout=0, sleep=sleep)["normal"] == 0 + ) + finally: + from flux_quantum.backends.braket import queue_depth + + probe_hold.queue_depth = queue_depth + + +def test_a_completed_task_is_not_reported_as_a_cancel_failure(): + """On a fast device every filler task has already run, and four lines of + ConflictException are noise.""" + import probe_hold + + said = [] + + class Client: + def cancel_quantum_task(self, quantumTaskArn, clientToken): + raise RuntimeError( + "ConflictException: cannot cancel a quantum task in the " + "COMPLETED status" + ) + + import builtins + + real_print = builtins.print + builtins.print = lambda *a, **k: said.append(" ".join(str(x) for x in a)) + try: + probe_hold.drain_queue(Client(), ["arn:a"]) + finally: + builtins.print = real_print + assert said == [] + + +def test_filler_shots_are_separate_from_probe_shots(): + """They want opposite things. Filler needs shots to occupy the device, + probe tasks want few so the run stays cheap. Ten shot filler finished + before the queue could be read, twice.""" + from probe_hold import make_queue + + seen = {} + client = _api_bound_client(seen) + make_queue(client, "arn:dev", "bucket", 2, 500) + assert seen["shots"] == [500, 500] + + +def test_calibration_summarises_rather_than_dumps(): + """What a result has to be read against is which edges were usable and + when they were characterised, not a thousand lines of specs.""" + from probe_hold import calibration + + specs = { + "instructions": [ + { + "node_ids": [88, 89], + "characteristics": [ + { + "name": "fCZ", + "value": 0.9967, + "timestamp": "2026-09-28T15:21:58+00:00", + } + ], + }, + { + "node_ids": [97, 106], + "characteristics": [ + { + "name": "fCZ", + "value": 0.8857, + "timestamp": "2026-09-28T15:41:47+00:00", + } + ], + }, + { + "node_ids": [64, 73], + "characteristics": [ + { + "name": "fCZ", + "value": 0.5, + "timestamp": "2026-09-28T15:41:46+00:00", + } + ], + }, + ] + } + d = type( + "D", + (), + { + "properties": type( + "P", (), {"provider": type("V", (), {"specs": specs})()} + )() + }, + )() + c = calibration(d) + assert c["live"] == 2 and c["dead"] == 1 + assert round(c["error_min"], 4) == 0.0033 + assert round(c["error_max"], 4) == 0.1143 + assert c["calibrated_from"].endswith("15:21:58+00:00") + + +def test_the_record_never_holds_the_token(): + """It grants priority and bills to the job. A log is not the place.""" + import json + import tempfile + + from probe_hold import record + + path = tempfile.mktemp(suffix=".jsonl") + row = {"device": "arn:dev", "hold_job": "arn:job", "verdict": ["ASSOCIATED"]} + record(path, row) + record(path, row) + lines = open(path).read().strip().split("\n") + assert len(lines) == 2 + assert "token" not in json.loads(lines[0]) + + +def test_a_calibration_we_cannot_read_is_not_invented(): + from probe_hold import calibration + + assert calibration(type("D", (), {})()) is None + + +def test_a_queue_that_drained_during_the_hold_says_so(): + """Seen four times: the device has a backlog when the run starts and none + by the time the hold is up. That is not the same as never having one.""" + from probe_hold import verdict + + out = "\n".join( + verdict({"job": "j"}, {}, "qpu", {"normal": -7, "priority": 0}, {"normal": 0}) + ) + assert "lost 7 tasks while the hold was provisioning" in out + assert "Nothing was waiting" not in out + + +def test_an_idle_device_still_reads_as_idle(): + from probe_hold import verdict + + out = "\n".join( + verdict({"job": "j"}, {}, "qpu", {"normal": 0, "priority": 0}, {"normal": 0}) + ) + assert "Nothing was waiting" in out + assert "provisioning" not in out + + +def test_used_qubits_does_not_block_on_an_unfinished_task(): + """result() polls until the task completes, five days by default. On a + queued task that is a hang, and it hung a live run.""" + import probe_hold + + class Task: + def __init__(self, arn): + pass + + def state(self): + return "QUEUED" + + def result(self): + raise AssertionError("must not be called on an unfinished task") + + import braket.aws as aws + + original = aws.AwsQuantumTask + aws.AwsQuantumTask = Task + try: + got = probe_hold.used_qubits("arn:task") + finally: + aws.AwsQuantumTask = original + assert got["state"] == "QUEUED" + + +def test_every_device_gets_its_windows(): + """A device open and empty now says nothing about whether it is worth + coming back to, so all of them are listed.""" + import builtins + + import probe_hold + + rows = [ + { + "arn": "arn:open", + "open": True, + "normal": 0, + "priority": 0, + "cost": 0.3, + "windows": [("Everyday", "02:00", "12:00")], + }, + { + "arn": "arn:shut", + "open": False, + "normal": 60, + "priority": 0, + "cost": 0.3, + "windows": [("Monday", "18:00", "20:59")], + }, + ] + real, probe_hold.survey = probe_hold.survey, lambda shots=1: rows + said = [] + rp, builtins.print = builtins.print, lambda *a, **k: said.append( + " ".join(str(x) for x in a) + ) + try: + probe_hold.show_survey() + finally: + builtins.print = rp + probe_hold.survey = real + + out = "\n".join(said) + assert "arn:open (0 queued, open)" in out + assert "arn:shut (60 queued, shut)" in out + assert "02:00 to 12:00" in out and "18:00 to 20:59" in out + + +def test_a_task_sitting_in_priority_is_the_plainest_evidence(): + """Priority read zero on every survey of these devices. One task in it + after we submitted exactly one token task is the answer, and movement + between two laggy counts is not.""" + from probe_hold import verdict + + out = "\n".join( + verdict( + {"job": "j"}, + {}, + "qpu", + {"normal": 6, "priority": 0}, + {"normal": 0}, + after_token={"normal": 8, "priority": 1}, + ) + ) + assert "IN THE PRIORITY QUEUE" in out + assert "NOT PRIORITISED" not in out + + +def test_an_empty_priority_queue_falls_back_to_movement(): + from probe_hold import verdict + + out = "\n".join( + verdict( + {"job": "j"}, + {}, + "qpu", + {"normal": 1, "priority": 0}, + {"normal": 1}, + after_token={"normal": 8, "priority": 0}, + ) + ) + assert "NOT PRIORITISED" in out + + +def test_a_reading_is_taken_only_once_it_stops_moving(): + """The counters lag a submit. Six filler tasks read as one, then seven a + moment later, and a token task's arrival in Priority was credited to the + filler catching up.""" + import probe_hold + + readings = [ + {"normal": 1, "priority": 0}, + {"normal": 7, "priority": 0}, + {"normal": 7, "priority": 1}, + {"normal": 7, "priority": 1}, + ] + real, probe_hold.queue_depth = probe_hold.queue_depth, lambda d: readings.pop(0) + try: + got = probe_hold.settled_depth("arn:dev", sleep=lambda _: None) + finally: + probe_hold.queue_depth = real + assert got == {"normal": 7, "priority": 1} + + +def test_settling_gives_up_rather_than_waiting_forever(): + import probe_hold + + flip = [{"normal": i % 2, "priority": 0} for i in range(200)] + real, probe_hold.queue_depth = probe_hold.queue_depth, lambda d: flip.pop(0) + try: + probe_hold.settled_depth("arn:dev", timeout=0, sleep=lambda _: None) + finally: + probe_hold.queue_depth = real diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index 1f5ca94..0d85e61 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -26,9 +26,19 @@ def add_option(self, name, **kw): m_p.CLIPlugin = CLIPlugin m_flux.cli = m_cli m_cli.plugin = m_p + + # cli.py imports submit and cancel at module level. The tests inject + # their own, so these only have to exist. + def _stub(*a, **k): + raise AssertionError("flux.job is a stub in the unit tests") + + m_job = types.ModuleType("flux.job") + m_job.submit = m_job.cancel = _stub + m_flux.job = m_job monkeypatch.setitem(sys.modules, "flux", m_flux) monkeypatch.setitem(sys.modules, "flux.cli", m_cli) monkeypatch.setitem(sys.modules, "flux.cli.plugin", m_p) + monkeypatch.setitem(sys.modules, "flux.job", m_job) monkeypatch.setenv("FLUX_QUANTUM_MOCK", "1") for name in [m for m in sys.modules if m.startswith("flux_quantum")]: del sys.modules[name] diff --git a/tests/unit/test_handoff.py b/tests/unit/test_handoff.py index 6af1d87..14f03f2 100644 --- a/tests/unit/test_handoff.py +++ b/tests/unit/test_handoff.py @@ -91,9 +91,6 @@ def test_backend_close_session_defaults_to_noop(): class _Bare(Backend): name = "bare" - def credentials_present(self): - return True, "ok" - def probe(self): return None diff --git a/tests/unit/test_qrmi.py b/tests/unit/test_qrmi.py index b59e527..baffe64 100644 --- a/tests/unit/test_qrmi.py +++ b/tests/unit/test_qrmi.py @@ -73,6 +73,9 @@ def fake_qrmi(monkeypatch): monkeypatch.setitem(sys.modules, "qrmi", mod) for k, v in CREDS.items(): monkeypatch.setenv(k, v) + # qrmi is bound when the backend module is imported, so import it again + for name in [m for m in sys.modules if m.startswith("flux_quantum")]: + del sys.modules[name] return FakeResource @@ -141,9 +144,12 @@ def test_unknown_resource_type_is_rejected(fake_qrmi): def test_secrets_never_appear_in_messages(fake_qrmi, monkeypatch): """Credential checks report variable names and never values.""" + from flux_quantum.backends import BackendError + monkeypatch.delenv(RESOURCE + "_QRMI_IBM_QRS_IAM_APIKEY") - ok, msg = _ibm().credentials_present() - assert not ok + with pytest.raises(BackendError) as e: + _ibm() + msg = str(e.value) assert RESOURCE + "_QRMI_IBM_QRS_IAM_APIKEY" in msg assert "supersecret" not in msg @@ -301,17 +307,10 @@ def test_warmup_payload_is_a_single_shot_measure(fake_qrmi): def test_missing_qiskit_says_what_to_install(fake_qrmi, monkeypatch): """The warmup needs qiskit. Say so, and offer the way out.""" - import builtins - - real = builtins.__import__ - - def no_qiskit(name, *a, **k): - if name == "qiskit": - raise ImportError("no qiskit") - return real(name, *a, **k) + from flux_quantum.backends import qrmi as qrmimod b = _acquired(fake_qrmi, ["Queued"]) - monkeypatch.setattr(builtins, "__import__", no_qiskit) + monkeypatch.setattr(qrmimod, "QuantumCircuit", None) with pytest.raises(RuntimeError) as e: b.wait_for_priority({}) assert "qrmi[ibm]" in str(e.value) diff --git a/tests/unit/test_selector.py b/tests/unit/test_selector.py index 0c99e4f..f3e5c79 100644 --- a/tests/unit/test_selector.py +++ b/tests/unit/test_selector.py @@ -7,6 +7,9 @@ import importlib +import sys +import types + import pytest @@ -14,6 +17,15 @@ def _selector(): return importlib.import_module("flux_quantum.selector") +def _fake_qrmi(monkeypatch): + """Enough of qrmi for the ibm backend to construct.""" + mod = types.ModuleType("qrmi") + mod.QuantumResource = object + mod.ResourceType = type("ResourceType", (), {"IBMQiskitRuntimeService": "x"}) + mod.Payload = mod.TaskStatus = object + monkeypatch.setitem(sys.modules, "qrmi", mod) + + def test_no_usable_vendor_raises(backends_real): sel = _selector() with pytest.raises(sel.SelectionError) as e: @@ -25,6 +37,7 @@ def test_no_usable_vendor_raises(backends_real): def test_select_ibm_when_reachable(fresh, monkeypatch): """Credentials alone are not enough. The probe asks QRMI whether the resource is reachable, so stub that rather than calling out to IBM.""" + _fake_qrmi(monkeypatch) b = fresh(mock=False, env=IBM_CREDS) from flux_quantum.backends.base import Signals @@ -37,6 +50,7 @@ def test_select_ibm_when_reachable(fresh, monkeypatch): def test_unreachable_ibm_is_not_a_candidate(fresh, monkeypatch): + _fake_qrmi(monkeypatch) b = fresh(mock=False, env=IBM_CREDS) from flux_quantum.backends.base import Signals From 2c76106237715688bff32c2c2f73f9b063eee38c Mon Sep 17 00:00:00 2001 From: vsoch Date: Thu, 1 Oct 2026 17:00:17 +0000 Subject: [PATCH 17/26] test: read priority from the task and race the token from behind The probe never asked GetQuantumTask for QueueInfo, so it never saw a task's queue or position and fell back to the device counters, which lag a submit by minutes. It also gave up after a fixed three minutes whatever the hold was, and submitted the token task first, so even a correct finishing order proved nothing. Now the control goes in first and the token task after it, both are polled with QueueInfo until the token task finishes with the hold up for the whole --max-seconds, and the verdict reads only per-task evidence: the job association, which queue each waited in and where, and the finishing order from the service timestamps. A device that served both on arrival says NO CONTENTION rather than claiming an order. Also: --log tees everything to a file, --inspect reads old tasks back, --pace measures whether a device serves one task at a time, Aquila gets the analog program it needs with --check-program to validate it locally, and a heartbeat line says the wait is a wait and not a hang. On IonQ Forte this showed a token task submitted after a control queue at Priority position 1 against Normal position 3, and finish first. Signed-off-by: vsoch --- tests/probe_hold.py | 795 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 618 insertions(+), 177 deletions(-) diff --git a/tests/probe_hold.py b/tests/probe_hold.py index e3fcd07..7b71ef3 100644 --- a/tests/probe_hold.py +++ b/tests/probe_hold.py @@ -1,24 +1,44 @@ #!/usr/bin/env python3 -"""Does a hybrid job token give priority to work submitted elsewhere. +"""Does a hybrid job token give priority to work submitted from outside. That is what this branch turns on. If the token travels, the hybrid job is the scout and the classical work stays on our cluster. If not, the job has to submit for us, which is a different design. +The measurement is a race the token task has to win from behind. + + open the hold, a hybrid job that does nothing but publish its token + submit a control task without the token + submit the token task after it, so it is behind the control + watch both until the token task finishes, with the hold still up + +Two things say whether the token bought anything. Braket reports which queue +a task is waiting in, Normal or Priority, and its position there, but only +when GetQuantumTask is asked for QueueInfo. Earlier runs never asked, so +they saw nothing and fell back to the device counters, which lag a submit by +minutes. Then the service timestamps say whether the token task finished +ahead of a control that was submitted before it. Both are read straight from +the tasks, so neither depends on the counters. + Run against SV1 first. A simulator runs tasks straight away so this says nothing about ordering, but it does say whether the API takes the token from outside the container and links the task to the job. That costs a few cents. -Then run it against a QPU, where queuePriority is the direct answer. +Then run it against a QPU that has a queue, where the queue is the answer. - python3 probe_hold.py # SV1 - python3 probe_hold.py --device arn:... # a real QPU, costs money + python3 tests/probe_hold.py --survey # who has a queue + python3 tests/probe_hold.py # SV1 + python3 tests/probe_hold.py --device arn:... --shots 100 --max-seconds 1800 -The hold is closed in a finally, so nothing is left running. +--max-seconds is the hold, and the token task has to finish inside it. On a +device where each task takes minutes, give it half an hour. The hold is +closed in a finally, so nothing is left running. """ import argparse import datetime import json +import os +import sys import time import uuid @@ -28,8 +48,10 @@ from flux_quantum.backends.braket import ( SV1, BraketBackend, + cost, is_open, queue_depth, + region_for, shots_range, survey, windows, @@ -43,8 +65,77 @@ c[0] = measure q[0]; """ +TERMINAL = ("COMPLETED", "FAILED", "CANCELLED") + -def submit(braket, device, bucket, token=None, shots=1): +def qasm_action(): + """The one qubit circuit as a task action.""" + return json.dumps( + { + "braketSchemaHeader": { + "name": "braket.ir.openqasm.program", + "version": "1", + }, + "source": BELL, + } + ) + + +def ahs_program(): + """The smallest analog program Aquila will take. + + Aquila runs Hamiltonians, not circuits, so a circuit is refused. One + atom, one microsecond, every field held at zero. Built with the SDK so + the schema is the SDK's and not ours. + """ + from braket.ahs.analog_hamiltonian_simulation import AnalogHamiltonianSimulation + from braket.ahs.atom_arrangement import AtomArrangement + from braket.ahs.driving_field import DrivingField + from braket.timings.time_series import TimeSeries + + zero = TimeSeries().put(0.0, 0.0).put(1e-6, 0.0) + return AnalogHamiltonianSimulation( + register=AtomArrangement().add((0.0, 0.0)), + hamiltonian=DrivingField(amplitude=zero, phase=zero, detuning=zero), + ) + + +def ahs_action(): + ir = ahs_program().to_ir() + dump = getattr(ir, "model_dump_json", None) or ir.json + return dump() + + +def program_action(device): + """The task action for this device: analog for QuEra, a circuit otherwise.""" + if "/quera/" in device: + return ahs_action() + return qasm_action() + + +def check_program(device): + """Run the program on the local simulator, which validates it for free. + + The analog simulator checks the program against Aquila's own limits + before it simulates, so a program it accepts is one Aquila accepts, and + a program it refuses would have been refused after the hold was up. + """ + from braket.devices import LocalSimulator + + if "/quera/" in device: + LocalSimulator("braket_ahs").run(ahs_program(), shots=1).result() + else: + from braket.ir.openqasm import Program + + LocalSimulator().run(Program(source=BELL), shots=1).result() + + +# the ml.m5.large the hold runs on, per hour. Braket bills the instance by +# the minute and nothing else for a job that submits no tasks of its own +HOLD_PER_HOUR = 0.115 + + +def submit(braket, device, bucket, token=None, shots=1, action=None): """Create a task, with or without the token. Without it the docs say the task gets no priority and bills standalone. @@ -56,75 +147,142 @@ def submit(braket, device, bucket, token=None, shots=1): "shots": shots, "outputS3Bucket": bucket, "outputS3KeyPrefix": "flux-quantum-probe", - "action": json.dumps( - { - "braketSchemaHeader": { - "name": "braket.ir.openqasm.program", - "version": "1", - }, - "source": BELL, - } - ), + "action": action or program_action(device), } if token: args["jobToken"] = token return braket.create_quantum_task(**args)["quantumTaskArn"] -TERMINAL = ("COMPLETED", "FAILED", "CANCELLED") - - def describe(braket, arn): - """What the service says about the task.""" - t = braket.get_quantum_task(quantumTaskArn=arn) + """What the service says about the task, queue included. + + queueInfo only comes back when the call asks for it. Without + additionalAttributeNames the response has no queue, no position and no + queuePriority, which is what every run before this one saw. + """ + t = braket.get_quantum_task( + quantumTaskArn=arn, additionalAttributeNames=["QueueInfo"] + ) q = t.get("queueInfo") or {} + position = q.get("position") + if position in ("None", ""): + position = None return { "status": t.get("status"), "job": t.get("jobArn"), "queue": q.get("queue"), - "position": q.get("position"), + "position": position, "priority": q.get("queuePriority"), + "message": q.get("message"), + "created": t.get("createdAt"), + "ended": t.get("endedAt"), } -def settle(braket, arns, timeout=180, interval=1.0, sleep=time.sleep): - """Follow both tasks to the end, keeping the best view of each. +def _stamp(now): + return datetime.datetime.fromtimestamp(now, datetime.timezone.utc).isoformat() + + +def follow( + braket, + arns, + timeout=1800, + interval=5.0, + sleep=time.sleep, + clock=time.time, + done=None, + views=None, + say=None, + heartbeat=60, +): + """Poll the tasks until done says so, keeping what each one showed. + + The queue fields describe a task that is waiting and are gone once it + runs, so the first sighting of each is kept as the task's queue, and the + timeline records every change so the run can be read back later. A one + qubit task on a simulator can finish before the first poll returns, and + then there is simply no queue entry, which the verdict allows for. - jobArn is set when the task is created and is still there afterwards, so - the verdict never has to catch anything in flight. queueInfo is the - opposite. It describes a task that is still waiting and is gone once it - runs, so we record it if we see it and never wait for it. A one qubit task - on a simulator can finish before the first poll returns. + done takes the views and says whether to stop. The default is every task + terminal. Pass views back in to keep following after a stop. """ - best = {a: {} for a in arns} - deadline = time.time() + timeout - while time.time() < deadline: - pending = False + if done is None: + + def done(v): + return all(v[a].get("status") in TERMINAL for a in arns) + + views = views if views is not None else {} + for a in arns: + views.setdefault(a, {"timeline": []}) + deadline = clock() + timeout + started = last_word = clock() + while True: + now = clock() for a in arns: + v = views[a] + if v.get("status") in TERMINAL: + continue seen = describe(braket, a) - for k, v in seen.items(): - # keep the first real value we see - if v and not best[a].get(k): - best[a][k] = v - best[a]["status"] = seen["status"] - if seen["status"] not in TERMINAL: - pending = True - if not pending: - break + for k in ("job", "created", "ended", "queue", "message"): + if seen.get(k) and not v.get(k): + v[k] = seen[k] + # which queue the task waited in, and where it started + if seen["priority"] and not v.get("priority"): + v["priority"] = seen["priority"] + v["position"] = seen["position"] + if seen["position"]: + v["position_last"] = seen["position"] + v["status"] = seen["status"] + change = (seen["status"], seen["priority"], seen["position"]) + if not v["timeline"] or v["timeline"][-1][1:] != list(change): + v["timeline"].append([_stamp(now), *change]) + if say: + say(a, seen) + last_word = now + if done(views) or clock() + interval > deadline: + return views + # a device that takes minutes per task prints nothing for minutes, + # which looks like a hang. Say we are still here, and what we see + if say and heartbeat and now - last_word >= heartbeat: + print( + " still waiting after %.0fs, %s" + % ( + now - started, + ", ".join( + "%s %s%s" + % ( + a.split("/")[-1][:8], + views[a].get("status"), + ( + " at %s" % views[a]["position_last"] + if views[a].get("position_last") + else "" + ), + ) + for a in arns + ), + ), + flush=True, + ) + last_word = now sleep(interval) - return best + + +def seconds_between(earlier, later): + """later minus earlier in seconds, or None if either is missing.""" + if not earlier or not later: + return None + return (later - earlier).total_seconds() def make_queue(braket, device, bucket, count, shots=1): """Submit filler tasks so there is a queue to jump. - A device is only busy while it is shut, and drains as soon as a window - opens, so waiting for someone else's backlog means waiting for a window to - open with work still in front of it. Making our own removes that. - - These go in without a token so they sit in the Normal queue, and they go - in after the hold is ready, because a fast device empties during the two - minutes the hold takes to provision. + A device that is idle runs the control and the token task at once and + shows nothing. These go in without a token so they sit in Normal, and + they go in after the hold is ready, because a fast device empties during + the two minutes the hold takes to provision. They cost the task fee each, so keep the count small. """ @@ -135,7 +293,7 @@ def make_queue(braket, device, bucket, count, shots=1): def drain_queue(braket, arns): - """Cancel the filler, so it does not run and bill for shots. + """Cancel what is still waiting, so it does not run and bill for shots. A task that has already started cannot be cancelled, which is fine. The task fee is spent either way, the shots are what this saves. @@ -151,23 +309,6 @@ def drain_queue(braket, arns): print("could not cancel %s: %s" % (a.split("/")[-1], e)) -def wait_for_queue(device, timeout=20, interval=0.5, sleep=time.sleep): - """Wait until the filler actually shows up in the Normal queue. - - Submitting returns before the service counts the task, and on a fast - device the window between counted and finished is short. Reading the - depth straight after the submit can miss it in either direction. - """ - deadline = time.time() + timeout - last = queue_depth(device) - while time.time() < deadline: - if last["normal"] > 0: - return last - sleep(interval) - last = queue_depth(device) - return last - - def calibration(device): """A summary of the device's calibration state, not the whole dump. @@ -239,6 +380,41 @@ def used_qubits(braket_task_arn): return {"error": str(e)} +class Tee: + """Everything printed goes to the terminal and to a file. + + A terminal that is closed takes the run with it. The backend prints its + progress too, and a traceback goes to stderr, so both streams are teed. + """ + + def __init__(self, path, stream): + self.fh = open(path, "a") + self.stream = stream + + def write(self, s): + self.stream.write(s) + self.fh.write(s) + self.fh.flush() + return len(s) + + def flush(self): + self.stream.flush() + self.fh.flush() + + +def log_to(path, argv=None): + """Tee stdout and stderr to path, and write a header naming the run.""" + sys.stdout = Tee(path, sys.stdout) + sys.stderr = Tee(path, sys.stderr) + print( + "\n=== %s %s" + % ( + datetime.datetime.now(datetime.timezone.utc).isoformat(), + " ".join(argv if argv is not None else sys.argv), + ) + ) + + def record(path, row): """One JSON line per run. @@ -252,32 +428,94 @@ def record(path, row): fh.write(_json.dumps(row, default=str) + "\n") -def settled_depth(device, timeout=30, interval=2.0, sleep=time.sleep): - """Read the queue once it has stopped moving. +def inspect(braket, arns): + """What the service says about tasks from earlier runs. Free. - The counters lag submission by a few seconds. Reading straight after a - submit gave normal=1 right after six filler tasks went in, then 7 a - moment later, so a token task's arrival in Priority was attributed to the - filler catching up. Wait for two readings in a row to agree. + createdAt and endedAt give how long a device took to serve a task, which + is what decides whether a queue on it is real. The queue label survives + completion, so the Priority against Normal reading is still there too. """ - last = queue_depth(device) - deadline = time.time() + timeout - while time.time() < deadline: - sleep(interval) - now = queue_depth(device) - if all(now.get(k) == last.get(k) for k in ("normal", "priority")): - return now - last = now - return last + rows = [describe(braket, a) for a in arns] + print( + " %-14s %-10s %-9s %-26s %-26s %s" + % ("task", "status", "queue", "created", "ended", "took") + ) + for a, r in zip(arns, rows): + took = seconds_between(r.get("created"), r.get("ended")) + print( + " %-14s %-10s %-9s %-26s %-26s %s" + % ( + a.split("/")[-1][:12], + r.get("status"), + r.get("priority") or "-", + r.get("created") or "-", + r.get("ended") or "-", + "-" if took is None else "%.0fs" % took, + ) + ) + return rows -def moved(before, after): - """Which queue grew between two readings.""" - return { - k: after.get(k, 0) - v - for k, v in before.items() - if isinstance(v, int) and isinstance(after.get(k), int) - } +def pace(braket, device, bucket, count, shots, timeout=600, **kw): + """How fast a device serves plain tasks, with no hold and no token. + + Whether filler can build a queue depends on this alone. Cepheus served + eight tasks in under a second, Forte held two for six minutes. Tasks go + in back to back and each is followed to its end, and the report says + how long each took and whether they ran one after another or all at + once. Costs the task fee and shots each, nothing else. + """ + arns = make_queue(braket, device, bucket, count, shots) + views = follow(braket, arns, timeout=timeout, **kw) + rows = [] + for a in arns: + v = views[a] + rows.append( + { + "arn": a, + "status": v.get("status"), + "created": v.get("created"), + "ended": v.get("ended"), + "took": seconds_between(v.get("created"), v.get("ended")), + "position": v.get("position"), + } + ) + return rows + + +def pace_report(rows): + """Serial or not, and how long each task took.""" + out = [] + for r in rows: + out.append( + " %-14s %-10s queued at %-5s took %s" + % ( + r["arn"].split("/")[-1][:12], + r["status"], + r["position"] or "-", + "-" if r["took"] is None else "%.1fs" % r["took"], + ) + ) + done = [r for r in rows if r["created"] and r["ended"]] + if len(done) >= 2: + # created to ended includes the wait, so it says nothing on its own. + # The gap between one finish and the next is the service time, and + # a queue is only visible when that is longer than a poll + ends = sorted(r["ended"] for r in done) + gaps = [seconds_between(a, b) for a, b in zip(ends, ends[1:])] + if min(gaps) >= 5: + out.append( + "SERIAL. %d tasks finished %.0fs apart at the least, about " + "%.0fs each, so filler will hold this device." + % (len(done), min(gaps), sum(gaps) / len(gaps)) + ) + else: + out.append( + "OVERLAPPING. %d tasks finished within %.1fs of each other, so " + "this device runs them together or in under a poll, and " + "filler will not build a queue here." % (len(done), min(gaps)) + ) + return out def _tzname(): @@ -327,16 +565,28 @@ def show_survey(shots=1): ) -def verdict( - with_token, without, device, by_token=None, by_control=None, after_token=None -): +def _where(view): + """A task's queue and position as words, for the verdict.""" + if not view.get("priority"): + return None + out = view["priority"] + if view.get("position"): + out += " at position %s" % view["position"] + return out + + +def verdict(with_token, without, device, queues=None): """What the two tasks say, strongest evidence first. jobArn is the one that matters. AWS says a task without the token gets no priority and bills standalone, so association is what the token carries. - queuePriority only exists while a task waits, so it is reported when we - saw it and not otherwise. A simulator usually shows nothing here. + Then the queue each task waited in, read from the task itself, and then + the order they finished in. The token task went in after the control, so + finishing before it means it was served out of order. + + The device counters are not consulted. They lag a submit by minutes and + are the reason earlier runs read nothing. """ out = [] if with_token.get("job"): @@ -353,70 +603,110 @@ def verdict( "says nothing. Something other than the token is linking them." ) - # Queue depth is counted per device and per queue, so a task landing in - # Priority shows up without having to catch it in flight. - settled = False - - # The count itself beats the difference between two counts. Priority has - # read zero on every survey of these devices, so one task sitting in it - # after we submitted exactly one token task is the plainest evidence - # there is. Movement missed this because the counters lag a submit. - if after_token and after_token.get("priority", 0) > 0: + tq, cq = _where(with_token), _where(without) + if with_token.get("priority") == "Priority": + line = "IN THE PRIORITY QUEUE. The token task waited as %s" % tq + if cq: + line += ", the control as %s" % cq + out.append(line + ".") + elif with_token.get("priority") == "Normal": out.append( - "IN THE PRIORITY QUEUE. %d task%s waiting there against %d in " - "Normal, and the only token task submitted was ours." - % ( - after_token["priority"], - "" if after_token["priority"] == 1 else "s", - after_token.get("normal", 0), - ) + "NOT PRIORITISED. The token task waited in the Normal queue like " + "any other%s." % (", the control as %s" % cq if cq else "") + ) + elif not cq: + # neither task was ever seen waiting, which is what an idle device or + # a simulator looks like. Ordering says nothing then either. + out.append( + "Nothing was waiting on the device, so priority cannot be seen. " + "Rerun against a QPU with a queue." + ) + return out + else: + out.append( + "The token task was never seen waiting, while the control waited " + "as %s." % cq ) - settled = True - # A queue that emptied while the hold was coming up is not the same as a - # device that never had one. Taking the hold costs a couple of minutes, - # and a fast device clears its backlog in that time. - drained = by_token and by_token.get("normal", 0) < 0 - if drained: + lag = seconds_between(without.get("created"), with_token.get("created")) + if lag is not None and lag < 0: out.append( - "The Normal queue lost %d tasks while the hold was provisioning, " - "so there was nothing left to jump by the time the token task " - "went in. Use --make-queue to build one after the hold is up." - % -by_token["normal"] + "WARNING: the token task was created before the control, so the " + "finishing order says nothing." ) - settled = True - if not settled and by_token and by_control: - if by_token.get("priority", 0) > 0 and by_control.get("normal", 0) > 0: - out.append( - "PRIORITISED. The token submit grew the Priority queue and " - "the control grew Normal." - ) - settled = True - elif by_token.get("normal", 0) > 0: - out.append( - "NOT PRIORITISED. The token submit grew the Normal queue, so " - "the task queues like any other." - ) - settled = True + return out + after = "" if lag is None else " %.0fs after the control" % lag - if not settled: - a, c = with_token.get("priority"), without.get("priority") - if a and c and a != c: - out.append("PRIORITISED. {} against {} without the token.".format(a, c)) - elif a and c: + ts, cs = with_token.get("status"), without.get("status") + waited = any( + v.get("position") or v.get("position_last") for v in (with_token, without) + ) + if ts == "COMPLETED" and not waited: + # the queue labels survive completion, but a position is only ever + # reported while a task waits. Neither had one, so neither waited, + # and finishing order on an idle device is submission order + took = seconds_between(with_token.get("created"), with_token.get("ended")) + out.append( + "NO CONTENTION. Neither task was ever seen waiting%s, so there " + "was no queue to jump and the finishing order says nothing. The " + "queue labels above are still the service's own. Rerun against a " + "device with a backlog." + % ("" if took is None else " and the token task ran in %.1fs" % took) + ) + return out + if ts in ("FAILED", "CANCELLED"): + out.append( + "The token task %s, so it never ran and the order says nothing. " + "Check the task's failure reason in the console." % ts + ) + return out + te = with_token.get("ended") + # a cancelled or failed control has an end time but never ran, so it + # counts as still waiting, not as finished + ce = without.get("ended") if cs not in ("FAILED", "CANCELLED") else None + if te and ce: + gap = seconds_between(te, ce) + if gap > 0: out.append( - "Same queuePriority either way ({}), so the token associates " - "but does not lift.".format(a) + "SERVED FIRST. Submitted%s and finished %.0fs before it." % (after, gap) ) else: out.append( - "Nothing was waiting on the device, so priority cannot be " - "seen. Rerun against a QPU with a queue." + "SERVED IN ORDER. Submitted%s and finished %.0fs after it, so " + "the token bought no place." % (after, -gap) ) - + elif te: + pos = without.get("position_last") + where = " at position %s" % pos if pos else "" + if cs == "CANCELLED": + tail = "still waiting%s, and was cancelled afterwards" % where + else: + tail = "still %s%s" % (cs or "waiting", where) + out.append( + "SERVED FIRST. Submitted%s and finished while the control was %s." + % (after, tail) + ) + else: + out.append( + "The token task did not finish while the hold was up (last seen " + "%s), so the order is unknown. Raise --max-seconds." + % (with_token.get("status") or "unknown") + ) return out +def estimate(device, shots, filler, filler_shots, max_seconds): + """What the run will cost, before anything is created.""" + per = cost(device, shots) + if per is None: + return None + total = 2 * per + if filler: + total += filler * (cost(device, filler_shots) or 0) + total += HOLD_PER_HOUR * max_seconds / 3600.0 + return total + + def main(): ap = argparse.ArgumentParser() ap.add_argument("--device", default=SV1) @@ -425,7 +715,22 @@ def main(): "--max-seconds", type=int, default=900, - help="ceiling on the hold, so a mistake stops billing", + help="how long the hold may last. The token task has to finish inside " + "it, so give a slow device half an hour. A mistake stops billing here", + ) + ap.add_argument( + "--follow-seconds", + type=int, + default=None, + help="after the token task finishes and the hold is released, keep " + "watching the control for this long so its finish time is on record. " + "Default is --max-seconds", + ) + ap.add_argument( + "--cancel-control", + action="store_true", + help="cancel the control once the token task has finished instead of " + "waiting for it. Saves its shots, loses its finish time", ) ap.add_argument( "--record", @@ -454,18 +759,63 @@ def main(): default=0, metavar="N", help="submit N filler tasks first so there is a Normal queue to jump. " - "Costs the task fee each, and they are cancelled afterwards", + "Costs the task fee each, and they are cancelled afterwards. Only " + "works on a device slow enough to queue: Cepheus ran six 500 shot " + "tasks in under a second, so there it buys nothing", ) ap.add_argument( "--force", action="store_true", help="run even if the device is shut or has nothing queued", ) + ap.add_argument( + "--inspect", + nargs="+", + metavar="TASK_ARN", + help="print what the service says about these tasks and stop. Free. " + "How long a device took to serve a task says whether a queue on it " + "is real, and the queue label is still on a finished task", + ) + ap.add_argument( + "--pace", + type=int, + default=0, + metavar="N", + help="submit N plain tasks with no hold, at --filler-shots, follow " + "them to the end and report how long each took and whether they ran " + "one at a time. Says whether filler can build a queue here, for the " + "task fee and shots each", + ) + ap.add_argument( + "--check-program", + action="store_true", + help="run the task program on the local simulator and stop. Free, " + "and the analog simulator checks Aquila's limits, so a refusal " + "here is one that would otherwise come after the hold is billing", + ) + ap.add_argument( + "--log", + default="probe-hold.log", + metavar="PATH", + help="append everything printed, including the backend's progress " + "and any traceback, to this file. Empty to disable", + ) args = ap.parse_args() + if args.log: + log_to(args.log) + if args.survey: show_survey(args.shots) return + if args.inspect: + os.environ["AWS_DEFAULT_REGION"] = region_for(args.inspect[0]) + inspect(boto3.client("braket"), args.inspect) + return + if args.check_program: + check_program(args.device) + print("the program for %s runs on the local simulator" % args.device) + return opts = { "hold": "job", @@ -481,6 +831,7 @@ def main(): "shots": args.shots, "filler_tasks": args.make_queue, "filler_shots": args.filler_shots, + "max_seconds": args.max_seconds, } b = BraketBackend() @@ -491,6 +842,7 @@ def main(): # finding that out after provisioning wastes time and money. q0 = queue_depth(args.device) print("device queues:", q0) + d = None if args.device != SV1 and not args.force: from braket.aws import AwsDevice @@ -513,17 +865,60 @@ def main(): "the device is outside its execution window, so the job would " "just sit there. Wait for a window, or pass --force." ) - if q0["normal"] == 0 and not args.make_queue: + if q0["normal"] == 0 and not (args.make_queue or args.pace): raise SystemExit( "nothing is waiting on this device, so a priority comparison " "would say nothing. Pick a busier one, pass --make-queue N to " "build a backlog, or pass --force." ) + total = estimate( + d, args.shots, args.make_queue, args.filler_shots, args.max_seconds + ) + if total is not None and not args.pace: + print("this run costs about $%.2f if the hold runs its full time" % total) + + if args.pace: + os.environ["AWS_DEFAULT_REGION"] = region_for(args.device) + per = cost(d, args.filler_shots) if d is not None else None + if per is not None: + print("this costs about $%.2f" % (per * args.pace)) + rows = pace( + boto3.client("braket"), + args.device, + AwsSession().default_bucket(), + args.pace, + args.filler_shots, + sleep=time.sleep, + clock=time.time, + ) + lines = pace_report(rows) + for line in lines: + print(line) + run.update({"pace": rows, "verdict": lines}) + run["ended"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + record(args.record, run) + print("recorded to %s" % args.record) + return arn = None filler = [] + # the device ARN names the region. The backend sets this too, but only + # once the hold is opened, and the client and bucket are made before that + os.environ["AWS_DEFAULT_REGION"] = region_for(args.device) braket = boto3.client("braket") bucket = AwsSession().default_bucket() + views = {} + control = with_token = None + + def say(a, seen): + who = "token " if a == with_token else "control" + where = "" + if seen["priority"]: + where = " %s queue" % seen["priority"] + if seen["position"]: + where += " position %s" % seen["position"] + print(" %s %-10s%s" % (who, seen["status"], where), flush=True) + try: arn = b.open_session(opts) print("hybrid job:", arn) @@ -532,6 +927,9 @@ def main(): if not ok: raise SystemExit("never got the hold, nothing submitted, nothing spent") + # the container gives up the hold this long after it published the + # token, and everything the token task needs has to happen before then + hold_until = time.time() + args.max_seconds - 30 run["hold_job"] = arn run["hold_ready_after_s"] = getattr(b, "_hold_ready_after", None) token = b.session_id(arn) @@ -539,57 +937,100 @@ def main(): print("job state at submit time:", b._hold_job.state()) # Build the backlog now, not before the hold. Taking the hold costs a - # couple of minutes, and a fast device empties in that time: four ten - # shot tasks on Cepheus were all COMPLETED before the hold was ready, - # so there was nothing left to jump. + # couple of minutes, and a fast device empties in that time. if args.make_queue: filler = make_queue( braket, args.device, bucket, args.make_queue, args.filler_shots ) - q0 = settled_depth(args.device) - print("queued %d filler tasks, device now %s" % (len(filler), q0)) - if q0["normal"] == 0: - print( - "the filler ran before it could queue. %d tasks of %d " - "shots is not enough to occupy this device, so raise " - "--filler-shots." % (args.make_queue, args.filler_shots) - ) - + run["filler_arns"] = filler + print("queued %d filler tasks" % len(filler)) + for a in filler: + print(" filler: ", a) + + # control first, so the token task has something to overtake. + # Finishing first after being submitted first would prove nothing. + control = submit(braket, args.device, bucket, None, args.shots) with_token = submit(braket, args.device, bucket, token, args.shots) - q1 = settled_depth(args.device) - print("\nwith the token:", with_token, "queues", q1) + q1 = queue_depth(args.device) + print("\ncontrol: ", control) + print("with token: ", with_token) + print("device queues after both (these lag):", q1) - without = submit(braket, args.device, bucket, None, args.shots) - q2 = settled_depth(args.device) - print("without it: ", without, "queues", q2) - - by_token, by_control = moved(q0, q1), moved(q1, q2) - print("\nmovement, token submit: ", by_token) - print("movement, control submit:", by_control) + print( + "\nwatching until the token task finishes, hold up for %ds" + % (hold_until - time.time()) + ) + follow( + braket, + [control, with_token], + timeout=hold_until - time.time(), + done=lambda v: v[with_token].get("status") in TERMINAL, + views=views, + say=say, + sleep=time.sleep, + clock=time.time, + ) + q2 = queue_depth(args.device) + run["queues"] = {"before": q0, "after_submit": q1, "after_token": q2} + + # the hold has done its work, stop paying for it before waiting on + # the control + print("\nreleasing the hold") + b.close_session() + arn = None + + if views[control].get("status") not in TERMINAL: + if args.cancel_control: + print("cancelling the control") + drain_queue(braket, [control]) + # one more look, so the record says cancelled and not queued + follow( + braket, + [control], + timeout=0, + views=views, + say=say, + sleep=time.sleep, + clock=time.time, + ) + else: + left = ( + args.max_seconds + if args.follow_seconds is None + else args.follow_seconds + ) + print("watching the control for up to %ds more" % left) + follow( + braket, + [control], + timeout=left, + views=views, + say=say, + sleep=time.sleep, + clock=time.time, + ) - seen = settle(braket, [with_token, without]) - a, c = seen[with_token], seen[without] - print("\n %-10s %-30s %s" % ("", "with token", "without")) - for k in ("status", "job", "queue", "position", "priority"): + a, c = views[with_token], views[control] + print("\n %-10s %-30s %s" % ("", "with token", "control")) + for k in ("status", "job", "priority", "position", "created", "ended"): print(" %-10s %-30s %s" % (k, a.get(k), c.get(k))) run.update( { - "queues": {"before": q0, "after_token": q1, "after_control": q2}, - "moved": {"token": by_token, "control": by_control}, "with_token": {"arn": with_token, **a}, - "without": {"arn": without, **c}, + "without": {"arn": control, **c}, } ) - lines = verdict(a, c, args.device, by_token, by_control, q1) + lines = verdict(a, c, args.device) run["verdict"] = lines + print() for line in lines: print(line) # which physical qubits ran, so the result can be read against the # calibration for those edges run["with_token"]["qubits"] = used_qubits(with_token) - run["without"]["qubits"] = used_qubits(without) + run["without"]["qubits"] = used_qubits(control) finally: run["ended"] = datetime.datetime.now(datetime.timezone.utc).isoformat() try: From 07afa5cb2abb6971ca17bb7a83e9f5cecf91398a Mon Sep 17 00:00:00 2001 From: vsoch Date: Thu, 1 Oct 2026 17:00:17 +0000 Subject: [PATCH 18/26] refactor: one set of submit options for every vendor, and vendor packages A user picks the vendor and the device and nothing else changes between vendors. The plugin offers --quantum-device, --quantum-hold, --quantum-hold-max, --quantum-wait and --quantum-dry-run, collects them once, applies the dry run, and each backend maps them to its own terms in scout_options. The session hold is Braket's hybrid job and IonQ's and IBM's session; probe is a front of queue job on any vendor. A hold a vendor cannot take fails at submit, before anything is held. The hold limit follows the job's duration, so a walltime sets it, and FLUX_QUANTUM_MOCK implies a dry run. The twenty or so vendor flags are gone. What was tuning rather than a description of the job, the Braket hold instance, warm-up shots, the QRMI resource type, the mock's queue knobs, is a FLUX_QUANTUM_* variable listed in each backend's docstring. The base class carries the common logic: a simulator attribute, the holds a vendor supports, dry_run and check_hold. Vendors with more than one file are packages: braket/ holds the backend and the hybrid job entry point, which uploads as a single file so no module named braket lands in the container next to the SDK. Tests follow, under tests/unit/vendors/, since a test package named braket on the test path shadows the SDK. Signed-off-by: vsoch --- examples/qrmi/README.md | 16 +- examples/qrmi/env.example | 2 +- examples/qrmi/ibm-run.sh | 7 +- flux_quantum/backends/base.py | 87 +- .../{braket.py => braket/__init__.py} | 119 +- .../{hold/entry.py => braket/hold.py} | 0 flux_quantum/backends/mock.py | 92 +- flux_quantum/backends/qrmi.py | 75 +- flux_quantum/cli.py | 97 +- tests/integration/test-mock-e2e.sh | 9 +- tests/integration/test-select-discovery.sh | 2 +- tests/integration/test-system-instance.sh | 2 +- tests/unit/test_cli_plugin.py | 103 +- tests/unit/test_mock_queue.py | 27 +- .../hold => tests/unit/vendors}/__init__.py | 0 tests/unit/vendors/braket/__init__.py | 0 .../braket/test_backend.py} | 9 +- .../braket/test_hold.py} | 1123 +++++++++++++---- tests/unit/vendors/ibm/__init__.py | 0 tests/unit/{ => vendors/ibm}/test_qrmi.py | 30 +- 20 files changed, 1277 insertions(+), 523 deletions(-) rename flux_quantum/backends/{braket.py => braket/__init__.py} (84%) rename flux_quantum/backends/{hold/entry.py => braket/hold.py} (100%) rename {flux_quantum/backends/hold => tests/unit/vendors}/__init__.py (100%) create mode 100644 tests/unit/vendors/braket/__init__.py rename tests/unit/{test_braket.py => vendors/braket/test_backend.py} (97%) rename tests/unit/{test_braket_hold.py => vendors/braket/test_hold.py} (52%) create mode 100644 tests/unit/vendors/ibm/__init__.py rename tests/unit/{ => vendors/ibm}/test_qrmi.py (92%) diff --git a/examples/qrmi/README.md b/examples/qrmi/README.md index b53f6bb..ccf2ace 100644 --- a/examples/qrmi/README.md +++ b/examples/qrmi/README.md @@ -41,8 +41,8 @@ missing, and never prints a value. unset FLUX_QUANTUM_MOCK flux submit -t 30m \ - --quantum-vendor ibm --quantum-ibm-resource \ - --quantum-ibm-warmup-timeout 600 \ + --quantum-vendor ibm --quantum-device \ + --quantum-wait 600 \ -n4 flux python /opt/flux-quantum/examples/qrmi/workload.py flux jobs -a @@ -51,12 +51,12 @@ flux job eventlog | grep memo ## Submit - flux submit -t 30m --quantum-vendor ibm --quantum-ibm-resource ibm_kingston \ + flux submit -t 30m --quantum-vendor ibm --quantum-device ibm_kingston \ -n4 flux python examples/qrmi/workload.py stdout is the scout id. The classical job id comes back on stderr. If only one resource has credentials in the environment you can drop -`--quantum-ibm-resource` and we work it out. +`--quantum-device` and we work it out. flux jobs -a flux job eventlog | grep memo @@ -64,7 +64,7 @@ resource has credentials in the environment you can drop ## You need an account that can open sessions -The qiskit-runtime-service type opens a session, which IBM only allows on plans that support sessions (Premium). Open and Pay As You Go plans can submit tasks but not hold a session, IBM returns a 403, and flux-quantum cancels the held classical job. Use `--quantum-ibm-type ibm-quantum-system` for direct access instead. +The qiskit-runtime-service type opens a session, which IBM only allows on plans that support sessions (Premium). Open and Pay As You Go plans can submit tasks but not hold a session, IBM returns a 403, and flux-quantum cancels the held classical job. Export `FLUX_QUANTUM_IBM_TYPE=ibm-quantum-system` for direct access instead. ## A session is not the same as having the QPU @@ -76,12 +76,12 @@ classical job. If the warmup never runs, the scout releases the session and fails the submit instead of starting the classical job against a QPU it does not have. - --quantum-ibm-warmup-timeout SECONDS give up after this long. Default 0, + --quantum-wait SECONDS give up after this long. Default 0, which means wait as long as the scout job is allowed to live - --quantum-ibm-skip-warmup release as soon as the session opens, + FLUX_QUANTUM_IBM_SKIP_WARMUP=1 release as soon as the session opens, no priority guarantee - --quantum-ibm-ready-timeout SECONDS default 120, how long to wait for the + FLUX_QUANTUM_IBM_READY_TIMEOUT=N default 120, how long to wait for the backend to report itself up The warmup costs one shot. diff --git a/examples/qrmi/env.example b/examples/qrmi/env.example index bb0d669..cfc6a40 100644 --- a/examples/qrmi/env.example +++ b/examples/qrmi/env.example @@ -13,7 +13,7 @@ ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN=REPLACE_ME # ibm_kingston_QRMI_IBM_QRS_SESSION_MODE=batch # For a direct access system use the QS prefix and -# --quantum-ibm-type ibm-quantum-system instead: +# FLUX_QUANTUM_IBM_TYPE=ibm-quantum-system instead: # # ibm_kingston_QRMI_IBM_QS_ENDPOINT= # ibm_kingston_QRMI_IBM_QS_IAM_ENDPOINT=https://iam.cloud.ibm.com diff --git a/examples/qrmi/ibm-run.sh b/examples/qrmi/ibm-run.sh index b3f1b00..d5b964f 100755 --- a/examples/qrmi/ibm-run.sh +++ b/examples/qrmi/ibm-run.sh @@ -91,14 +91,15 @@ from flux_quantum.backends import get_backend print(get_backend('ibm').credential_note)" | tee -a "$OUT/run.txt" || die "credentials incomplete" # 3. submit. The classical job prints and exits, so the billed window is warmup plus a few seconds. -warmup_args="--quantum-ibm-warmup-timeout $WARMUP_TIMEOUT" -[ "$SKIP_WARMUP" = 1 ] && warmup_args="--quantum-ibm-skip-warmup" +# the warm-up wait is the common --quantum-wait, skipping it is operator tuning +warmup_args="--quantum-wait $WARMUP_TIMEOUT" +[ "$SKIP_WARMUP" = 1 ] && export FLUX_QUANTUM_IBM_SKIP_WARMUP=1 say "submitting" t0=$(date +%s) err=$(mktemp) scout=$(flux submit -t "$WALLTIME" --quantum-vendor ibm \ - --quantum-ibm-resource "$RESOURCE" $warmup_args -n"$NTASKS" \ + --quantum-device "$RESOURCE" $warmup_args -n"$NTASKS" \ -- sh -c "echo classical got \$QUANTUM_SESSION_ID; sleep $SLEEP" 2>"$err") tee -a "$OUT/run.txt" < "$err" main=$(grep -oE 'held classical job [0-9]+' "$err" | awk '{print $NF}') diff --git a/flux_quantum/backends/base.py b/flux_quantum/backends/base.py index c930ed6..576d511 100644 --- a/flux_quantum/backends/base.py +++ b/flux_quantum/backends/base.py @@ -11,8 +11,45 @@ Adding a vendor means adding a Backend subclass and registering it. """ +import os from abc import ABC, abstractmethod +# The submit options every vendor is driven by. The CLI collects them once, +# and each backend maps them to its own terms in scout_options. A user picks +# the vendor and the device and nothing else changes between vendors. +# +# device the vendor's own name for the device: a Braket ARN, an IonQ +# backend, a QRMI resource id. Each vendor has a default +# hold 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 +# hold_max seconds the hold may last, so a scout that is never released +# stops costing +# wait seconds to wait for the hold to be ours before giving up and +# cancelling the held job. 0 means as long as the scout may run +# dry_run run on the vendor's simulator +COMMON_OPTIONS = ("device", "hold", "hold_max", "wait", "dry_run") +HOLDS = ("session", "probe") + + +def tuning(name, default=None): + """An operator setting from the environment, FLUX_QUANTUM_. + + Things like the instance the Braket hold runs on or the shot count of a + warm-up job are tuning for whoever operates the installation, not a + choice for someone submitting a job, so they are not submit options. + """ + return os.environ.get("FLUX_QUANTUM_" + name.upper(), default) + + +def truthy(value): + """An environment value read as a flag.""" + if value is None: + return False + if isinstance(value, str): + return value.strip().lower() not in ("", "0", "false", "no") + return bool(value) + class BackendError(RuntimeError): """The backend cannot be used. The message names what is missing, never @@ -40,21 +77,51 @@ class Backend(ABC): # vendor key like ibm, must match the qdevice_ graph type name = None + # what a dry run targets, or None when the vendor has no simulator + simulator = None + + # the holds this vendor can take, out of HOLDS + holds = ("session",) + # set by __init__ to say where the credentials came from, for the log credential_note = "" - @classmethod - def add_options(cls, add_option): - """Declare the flux submit options for this vendor. + def scout_options(self, common): + """Map the common submit options onto this vendor's terms. - Names get a quantum prefix, so --ibm-backend becomes - --quantum-ibm-backend. Namespace them by vendor to avoid collisions. + common has the COMMON_OPTIONS keys. What comes back is whatever + open_session and wait_for_priority need, and it travels to the + scout as JSON. The default passes the common options through. """ - return - - def scout_options(self, args): - """Return the options for this vendor from the parsed args.""" - return {} + self.check_hold(common.get("hold")) + return dict(common) + + def check_hold(self, hold): + """Refuse a hold this vendor cannot take, at submit time, before + anything is held.""" + hold = hold or "session" + if hold not in HOLDS: + raise BackendError( + "{}: unknown hold {}. The holds are {}".format( + self.name, hold, " and ".join(HOLDS) + ) + ) + if hold not in self.holds: + raise BackendError( + "{}: this vendor has no {} hold, only {}".format( + self.name, hold, " and ".join(self.holds) + ) + ) + return hold + + def dry_run(self, common): + """The common options for a dry run: the device becomes the + simulator. A vendor whose simulator needs more, such as a noise + model standing in for the hardware, adds it here.""" + out = dict(common, dry_run=True) + if self.simulator: + out["device"] = self.simulator + return out def session_id(self, opened): """What the classical job should be handed. diff --git a/flux_quantum/backends/braket.py b/flux_quantum/backends/braket/__init__.py similarity index 84% rename from flux_quantum/backends/braket.py rename to flux_quantum/backends/braket/__init__.py index 661cad6..a320792 100644 --- a/flux_quantum/backends/braket.py +++ b/flux_quantum/backends/braket/__init__.py @@ -1,14 +1,22 @@ """AWS Braket backend. -Two ways to take a device, set with --braket-hold. +Two ways to take a device, set with --quantum-hold. -probe is the default. Submit a trivial task and wait for it to reach the front -of the queue. That reserves nothing, so another user can take the device -between our task running and the classical job starting. +session, the default, starts a hybrid job. Braket runs one hybrid job at a +time per QPU and gives it priority, and the job publishes +AMZN_BRAKET_JOB_TOKEN. A task created with that token gets the job's priority +and bills to the job. Without it, neither. Internally this is hold=job. -job starts a hybrid job. Braket runs one hybrid job at a time per QPU and gives -it priority, and the job publishes AMZN_BRAKET_JOB_TOKEN. A task created with -that token gets the job's priority and bills to the job. Without it, neither. +probe submits a trivial task and waits for it to reach the front of the +queue. That reserves nothing, so another user can take the device between our +task running and the classical job starting. + +Tuning, from the environment: + + 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_BRAKET_UNGATE_POSITION probe: release at this queue position, 1 + FLUX_QUANTUM_BRAKET_REGION region, otherwise the one in the ARN So the hybrid job is the scout and the token is the session id. The classical work stays on our cluster instead of the small instance Braket gives the job, @@ -23,7 +31,7 @@ import os import time -from .base import Backend, BackendError, Signals, register +from ..base import Backend, BackendError, Signals, register, tuning try: import boto3 @@ -277,84 +285,23 @@ def __init__(self): self._hold_started = None self._hold_ready_after = None - @classmethod - def add_options(cls, add_option): - add_option( - "--braket-device", - metavar="ARN", - default=None, - help="Braket: device ARN, defaults to the SV1 simulator", - ) - add_option( - "--braket-region", - metavar="REGION", - default=None, - help="Braket: AWS region, defaults to the one in the device ARN", - ) - add_option( - "--braket-shots", - metavar="N", - default=None, - help="Braket: shots for the queue probe task, default 1", - ) - add_option( - "--braket-ungate-position", - metavar="N", - default=None, - help="Braket: release the classical job once the probe task is at " - "this queue position or closer. Default 1, next in line. Raise it " - "when the classical job is slow to start", - ) - add_option( - "--braket-hold", - metavar="MODE", - default=None, - help="Braket: how to take the device. probe waits for the front " - "of the queue, the default. job starts a hybrid job, which holds " - "the priority queue and hands out a token that work anywhere can " - "submit with", - ) - add_option( - "--braket-hold-instance", - metavar="TYPE", - default=None, - help="Braket: instance for the holding hybrid job. Default " - "ml.m5.large, the cheapest, since it does no work", - ) - add_option( - "--braket-hold-max-seconds", - metavar="SECONDS", - default=None, - help="Braket: give up the queue slot after this long, so a scout " - "that is never released stops billing. Default 900", - ) - add_option( - "--braket-queue-timeout", - metavar="SECONDS", - default=None, - help="Braket: give up if the probe task is still queued after this " - "long. Default 0, meaning wait as long as the scout may run", - ) + simulator = SV1 + holds = ("session", "probe") - def scout_options(self, args): - device = getattr(args, "braket_device", None) or SV1 - shots = getattr(args, "braket_shots", None) - timeout = getattr(args, "braket_queue_timeout", None) + def scout_options(self, common): + hold = self.check_hold(common.get("hold")) + device = common.get("device") or SV1 return { "device": device, - "region": region_for(device, getattr(args, "braket_region", None)), - "hold": getattr(args, "braket_hold", None) or "probe", - "hold_instance": getattr(args, "braket_hold_instance", None) - or "ml.m5.large", - "hold_max_seconds": float( - getattr(args, "braket_hold_max_seconds", None) or 900 - ), - "shots": 1 if shots is None else int(shots), - "queue_timeout": 0 if timeout is None else float(timeout), - "ungate_position": int( - getattr(args, "braket_ungate_position", None) - or os.environ.get("QUANTUM_BRAKET_UNGATE_POSITION", 1) - ), + "region": region_for(device, tuning("braket_region")), + # the hybrid job is Braket's session + "hold": "job" if hold == "session" else "probe", + "hold_instance": tuning("braket_hold_instance", "ml.m5.large"), + "hold_max_seconds": float(common.get("hold_max") or 900), + "shots": int(tuning("braket_shots", 1)), + "queue_timeout": float(common.get("wait") or 0), + "ungate_position": int(tuning("braket_ungate_position", 1)), + "dry_run": bool(common.get("dry_run")), } def job_environment(self, options): @@ -421,8 +368,10 @@ def _open_hold_job(self, options): max_seconds = int(options.get("hold_max_seconds") or 900) self._hold_job = AwsQuantumJob.create( device=device_arn, - source_module=os.path.dirname(os.path.abspath(hold.__file__)), - entry_point="hold.entry:main", + # the one file is the whole job. Uploading the package would put + # a module named braket in the container, next to the SDK's + source_module=os.path.abspath(hold.__file__), + entry_point="hold:main", job_name=name, instance_config=InstanceConfig( instanceType=options.get("hold_instance") or "ml.m5.large", diff --git a/flux_quantum/backends/hold/entry.py b/flux_quantum/backends/braket/hold.py similarity index 100% rename from flux_quantum/backends/hold/entry.py rename to flux_quantum/backends/braket/hold.py diff --git a/flux_quantum/backends/mock.py b/flux_quantum/backends/mock.py index 28c9fef..dfa8dcb 100644 --- a/flux_quantum/backends/mock.py +++ b/flux_quantum/backends/mock.py @@ -6,7 +6,8 @@ mock queue_depth 0, wins --select any mock_busy queue_depth 9, so --select queue prefers mock -FLUX_QUANTUM_MOCK_QUEUE and FLUX_QUANTUM_MOCK_COST override the mock vendor. +FLUX_QUANTUM_MOCK_QUEUE and FLUX_QUANTUM_MOCK_COST are what probe reports, and +the simulated wait is set from FLUX_QUANTUM_MOCK_* too, see scout_options. The mock also simulates a vendor queue, which is the only way to control queue depth as an experiment variable. No real vendor lets you do that. The wait is @@ -22,74 +23,41 @@ import random import time -from .base import Backend, Signals, register +from .base import Backend, Signals, register, tuning @register class MockBackend(Backend): name = "mock" - @classmethod - def add_options(cls, add_option): - add_option( - "--mock-session", - metavar="ID", - default=None, - help="mock: force this session id (testing)", - ) - add_option( - "--mock-latency", - metavar="SECONDS", - default=None, - help="mock: delay this many seconds before opening a session", - ) - add_option( - "--mock-queue-depth", - metavar="N", - default=None, - help="mock: tasks ahead of us in the simulated vendor queue", - ) - add_option( - "--mock-service-time", - metavar="SECONDS", - default=None, - help="mock: seconds per task ahead of us, default 0.1 so a deep " - "queue still finishes in a runnable time", - ) - add_option( - "--mock-base-overhead", - metavar="SECONDS", - default=None, - help="mock: fixed wait even at depth 0, default 10 which is what " - "ibm_marrakesh measured", - ) - add_option( - "--mock-jitter", - metavar="FRACTION", - default=None, - help="mock: multiplicative noise on the wait, 0 for a " - "deterministic run which is what an experiment wants", - ) - add_option( - "--mock-seed", - metavar="N", - default=None, - help="mock: seed for the jitter, so a run can be repeated", - ) - - def scout_options(self, args): - def num(name, default): - v = getattr(args, name, None) - return default if v is None else float(v) - + holds = ("session", "probe") + + def scout_options(self, common): + """The simulated queue is an experiment instrument, so its knobs are + environment variables and not submit options: + + FLUX_QUANTUM_MOCK_SESSION force this session id + FLUX_QUANTUM_MOCK_LATENCY seconds before opening a session + FLUX_QUANTUM_MOCK_QUEUE tasks ahead of us, also what + probe reports + FLUX_QUANTUM_MOCK_SERVICE_TIME seconds per task ahead, 0.1 + FLUX_QUANTUM_MOCK_BASE_OVERHEAD fixed wait at depth 0, 10, which + is what ibm_marrakesh measured + FLUX_QUANTUM_MOCK_JITTER multiplicative noise, 0 for a + deterministic run + FLUX_QUANTUM_MOCK_SEED seed for the jitter + """ + self.check_hold(common.get("hold")) return { - "session": getattr(args, "mock_session", None), - "latency": getattr(args, "mock_latency", None), - "queue_depth": int(num("mock_queue_depth", 0)), - "service_time": num("mock_service_time", 0.1), - "base_overhead": num("mock_base_overhead", 10.0), - "jitter": num("mock_jitter", 0.0), - "seed": getattr(args, "mock_seed", None), + "device": common.get("device") or "mock", + "session": tuning("mock_session"), + "latency": tuning("mock_latency"), + "queue_depth": int(tuning("mock_queue", 0)), + "service_time": float(tuning("mock_service_time", 0.1)), + "base_overhead": float(tuning("mock_base_overhead", 10.0)), + "jitter": float(tuning("mock_jitter", 0.0)), + "seed": tuning("mock_seed"), + "dry_run": bool(common.get("dry_run")), } def queue_wait(self, options): diff --git a/flux_quantum/backends/qrmi.py b/flux_quantum/backends/qrmi.py index 31e6477..818c838 100644 --- a/flux_quantum/backends/qrmi.py +++ b/flux_quantum/backends/qrmi.py @@ -20,7 +20,7 @@ import os import time -from .base import Backend, BackendError, Signals +from .base import Backend, BackendError, Signals, truthy, tuning try: from qrmi import Payload, QuantumResource, ResourceType, TaskStatus @@ -114,7 +114,7 @@ def warmup_payload(shots=1): if QuantumCircuit is None: raise RuntimeError( "qiskit is needed to build the warmup task. Install it with " - "pip install 'qrmi[ibm]', or pass --quantum-ibm-skip-warmup to " + "pip install 'qrmi[ibm]', or export FLUX_QUANTUM_IBM_SKIP_WARMUP=1 to " "release without confirming priority" ) @@ -143,7 +143,7 @@ def explain_acquire_failure(error, resource, rtype): "sessions on plans that support them, such as Premium. On an Open " "or Pay As You Go plan you can submit tasks but not hold a " "session, so there is nothing for the scout to co-allocate.\n" - " use --quantum-ibm-type ibm-quantum-system if you have direct " + " export FLUX_QUANTUM_IBM_TYPE=ibm-quantum-system if you have direct " "access, or --quantum-vendor mock to test the pipeline.\n" " QRMI said: {e}".format(r=resource, t=rtype, e=text) ) @@ -207,63 +207,32 @@ def __init__(self): self._resource = None self._lock = None - @classmethod - def add_options(cls, add_option): - add_option( - "--{}-resource".format(cls.name), - metavar="ID", - default=None, - help="{}: QRMI resource id, for example ibm_kingston".format(cls.name), - ) - add_option( - "--{}-ready-timeout".format(cls.name), - metavar="SECONDS", - default=None, - help="{}: wait this long after acquiring for the resource to report " - "itself accessible, 0 to skip. Default 120".format(cls.name), - ) - add_option( - "--{}-skip-warmup".format(cls.name), - action="store_true", - help="{}: release the classical job as soon as the session opens, " - "without waiting for a warmup task to reach the head of the " - "queue. Faster, but priority is not confirmed".format(cls.name), - ) - add_option( - "--{}-warmup-timeout".format(cls.name), - metavar="SECONDS", - default=None, - help="{}: give up if the warmup task is still queued after this " - "long. Default 0, meaning wait for as long as the scout job " - "is allowed to run".format(cls.name), - ) - add_option( - "--{}-type".format(cls.name), - metavar="TYPE", - default=None, - help="{}: QRMI resource type, default {}".format( - cls.name, cls.default_type - ), - ) + def scout_options(self, common): + """QRMI's terms. The device is the resource id, and the type, the + readiness wait and the warm-up are operator tuning: - def scout_options(self, args): - rtype = getattr(args, "{}_type".format(self.name), None) or self.default_type - resource = getattr(args, "{}_resource".format(self.name), None) + FLUX_QUANTUM__TYPE resource type + FLUX_QUANTUM__READY_TIMEOUT seconds for the resource to + report accessible, 120 + FLUX_QUANTUM__SKIP_WARMUP release as soon as the + session opens, no priority + confirmed + """ + self.check_hold(common.get("hold")) + rtype = tuning(self.name + "_type") or self.default_type + resource = common.get("device") if not resource: # only one resource configured in the environment is unambiguous found = resources_in_environment(rtype) if len(found) == 1: resource = found[0] - ready = getattr(args, "{}_ready_timeout".format(self.name), None) - warmup = getattr(args, "{}_warmup_timeout".format(self.name), None) return { "resource": resource, "type": rtype, - "ready_timeout": 120 if ready is None else float(ready), - "skip_warmup": bool( - getattr(args, "{}_skip_warmup".format(self.name), False) - ), - "warmup_timeout": 0 if warmup is None else float(warmup), + "ready_timeout": float(tuning(self.name + "_ready_timeout", 120)), + "skip_warmup": truthy(tuning(self.name + "_skip_warmup")), + "warmup_timeout": float(common.get("wait") or 0), + "dry_run": bool(common.get("dry_run")), } def job_environment(self, options): @@ -282,8 +251,8 @@ def open_session(self, options): rtype = options.get("type") or self.default_type if not resource: raise ValueError( - "{}: no QRMI resource id. Pass --quantum-{}-resource or export " - "the credentials for exactly one resource".format(self.name, self.name) + "{}: no QRMI resource id. Pass --quantum-device or export " + "the credentials for exactly one resource".format(self.name) ) missing = missing_env(resource, rtype) if missing: diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index 59e8e88..e56623b 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -25,7 +25,7 @@ from . import graph, qresource from .selector import select_vendor, SelectionError, discover_registry_vendors -from .backends import BackendError, get_backend, backend_classes +from .backends import BackendError, get_backend from .launch import build_scout_jobspec # subcommands where quantum submission makes sense @@ -215,6 +215,56 @@ def prepare_pair( return main_id +def _opt(args, name): + """A plugin option off the parsed args, by its prefixed name first. + + flux proxies args.device to args.quantum_device inside plugin callbacks, + but flux has options of its own and a bare name could be one of them. + """ + for attr in ("quantum_" + name, name): + value = getattr(args, attr, None) + if value is not None: + return value + return None + + +def common_options(args, jobspec=None, environ=None): + """The COMMON_OPTIONS from the parsed args, with their defaults. + + The hold limit follows the job's duration when it has one, so setting a + walltime sets it. FLUX_QUANTUM_MOCK makes every submit a dry run, since + nothing in a mock run should reach a billed device. + """ + environ = os.environ if environ is None else environ + duration = 0 + if jobspec: + duration = ( + jobspec.get("attributes", {}).get("system", {}).get("duration", 0) or 0 + ) + hold_max = _opt(args, "hold_max") + wait = _opt(args, "wait") + return { + "device": _opt(args, "device") or None, + "hold": _opt(args, "hold") or "session", + "hold_max": ( + float(hold_max) if hold_max else float(duration + 300 if duration else 900) + ), + "wait": float(wait) if wait else 0.0, + "dry_run": bool(_opt(args, "dry_run")) + or bool(environ.get("FLUX_QUANTUM_MOCK")), + } + + +def scout_options(backend, args, jobspec=None): + """What the scout is given for this vendor: the common options, made a + dry run if asked, then mapped by the backend. Raises BackendError for a + hold the vendor cannot take.""" + common = common_options(args, jobspec) + if common["dry_run"]: + common = backend.dry_run(common) + return backend.scout_options(common) + + class QuantumCLIPlugin(CLIPlugin): """Select a quantum vendor at submit time and prepare the held job.""" @@ -234,11 +284,44 @@ def __init__(self, prog): default=None, help="auto-select vendor: any | queue | cost", ) - # let each registered vendor backend contribute its own options - # (namespaced, e.g. --quantum-ibm-backend, --quantum-mock-session), so a - # quantum submit can carry vendor-specific parameters for the scout. - for backend_cls in backend_classes(): - backend_cls.add_options(self.add_option) + # the same options whichever vendor it is. Each backend maps them to + # its own terms, and operator tuning lives in FLUX_QUANTUM_* variables + self.add_option( + "--device", + metavar="NAME", + default=None, + help="the vendor's name for the device: a Braket ARN, an IonQ " + "backend, a QRMI resource id. Each vendor has a default", + ) + self.add_option( + "--hold", + metavar="MODE", + default=None, + help="session takes the vendor's real hold, a hybrid job on " + "Braket, a session on IonQ and IBM, the default. probe submits a " + "front of queue job and holds nothing", + ) + self.add_option( + "--hold-max", + metavar="SECONDS", + default=None, + help="give the hold up after this long, so a scout that is never " + "released stops costing. Default is the job's duration plus " + "300, or 900 when it has none", + ) + self.add_option( + "--wait", + metavar="SECONDS", + default=None, + help="give up if the hold is not ours after this long, and cancel " + "the held job. Default 0, as long as the scout may run", + ) + self.add_option( + "--dry-run", + action="store_true", + default=False, + help="run on the vendor's simulator. FLUX_QUANTUM_MOCK implies it", + ) self._chosen = None def _is_quantum(self, args): @@ -301,9 +384,9 @@ def modify_jobspec(self, args, jobspec): try: backend = get_backend(vendor) + options = scout_options(backend, args, jobspec.jobspec) if backend else {} except BackendError as e: raise SystemExit("flux quantum: {}".format(e)) - options = backend.scout_options(args) if backend else {} job_env = backend.job_environment(options) if backend else {} handle = flux.Flux() main_id = prepare_pair( diff --git a/tests/integration/test-mock-e2e.sh b/tests/integration/test-mock-e2e.sh index 6bb7918..ddd398a 100755 --- a/tests/integration/test-mock-e2e.sh +++ b/tests/integration/test-mock-e2e.sh @@ -38,10 +38,11 @@ echo "" echo "=== 2. ONE quantum submit -> plugin submits held classical AND scout ===" # the plugin wraps the user command, so do not pre-wrap it here. stdout is the # scout id and stderr carries the held classical job id -# a vendor option, so this value must flow from the CLI through the plugin -# and scout to the backend and into the classical environment +# a mock tuning variable, so this value must flow from the submit +# environment through the plugin and scout to the backend and into the +# classical environment FORCED="mocksess-$$" -scout_id=$(flux submit --quantum-vendor mock --quantum-mock-session "$FORCED" \ +scout_id=$(FLUX_QUANTUM_MOCK_SESSION="$FORCED" flux submit --quantum-vendor mock \ -n1 \ -- sh -c 'echo QUANTUM_SESSION=$QUANTUM_SESSION_ID' 2>"$ERR") cat "$ERR" >&2 @@ -58,7 +59,7 @@ if [ -n "$main_id" ]; then out=$(flux job attach "$main_id" &1) if echo "$out" | grep -q "QUANTUM_SESSION=$FORCED"; then echo "PASS: classical ran with the VENDOR-SUPPLIED session ($FORCED)" - echo " (--quantum-mock-session flowed CLI -> backend.open_session)" + echo " (FLUX_QUANTUM_MOCK_SESSION flowed submit env -> backend.open_session)" elif echo "$out" | grep -qE 'QUANTUM_SESSION=.+'; then echo "FAIL: got a session but NOT the vendor option value ($FORCED):" echo "$out" | grep 'QUANTUM_SESSION='; rc=1 diff --git a/tests/integration/test-select-discovery.sh b/tests/integration/test-select-discovery.sh index f3aedcf..7143c27 100755 --- a/tests/integration/test-select-discovery.sh +++ b/tests/integration/test-select-discovery.sh @@ -24,7 +24,7 @@ VENDORS="mock ibm braket" bash "$SELF/setup-mock-registry.sh" || exit 1 echo "" echo "=== 2. submit with --quantum-select any (NO explicit vendor) ===" # the plugin wraps the command itself, so do not pass wrap.py here -scout_id=$(flux submit --quantum-select any --quantum-mock-session AUTOSESS456 -n1 \ +scout_id=$(FLUX_QUANTUM_MOCK_SESSION=AUTOSESS456 flux submit --quantum-select any -n1 \ -- sh -c 'echo QUANTUM_SESSION=$QUANTUM_SESSION_ID' 2>"$ERR") main_id=$(grep -oE 'held classical job [0-9]+' "$ERR" | awk '{print $NF}') echo "scout=$scout_id classical=$main_id" diff --git a/tests/integration/test-system-instance.sh b/tests/integration/test-system-instance.sh index 038a24b..fe9f00f 100755 --- a/tests/integration/test-system-instance.sh +++ b/tests/integration/test-system-instance.sh @@ -38,7 +38,7 @@ echo "" echo "=== 4. one quantum submit -> held classical + scout ===" SID="systest-$(date +%s)" ERR=$(mktemp) -scout=$(flux submit --quantum-vendor mock --quantum-mock-session "$SID" -n1 \ +scout=$(FLUX_QUANTUM_MOCK_SESSION="$SID" flux submit --quantum-vendor mock -n1 \ -- sh -c 'echo QUANTUM_SESSION=$QUANTUM_SESSION_ID' 2>"$ERR") sed 's/^/ /' "$ERR" main=$(grep -oE 'held classical job [0-9]+' "$ERR" | awk '{print $NF}') diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index 0d85e61..eac411c 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -73,10 +73,20 @@ def test_plugin_registers_options_with_quantum_prefix(stub_flux_cli): plugin = shim.QuantumCLIPlugin("submit") # flux instantiates as entry(prog) assert plugin.prefix == "quantum" names = {n for n, _ in plugin.options} - # core options always present - assert {"--vendor", "--select"} <= names - # each registered vendor backend contributes its own namespaced options - assert "--ibm-resource" in names and "--braket-device" in names + # the same options whichever vendor it is + assert { + "--vendor", + "--select", + "--device", + "--hold", + "--hold-max", + "--wait", + "--dry-run", + } <= names + # and nothing vendor specific + assert not [ + n for n in names if n.startswith(("--ibm", "--braket", "--ionq", "--mock")) + ] def test_plugin_inactive_for_other_progs(stub_flux_cli): @@ -453,3 +463,88 @@ def submit(handle, jobspec_json): # both halves were described to the check assert order[0][1] is not None and order[0][2] is not None assert any(x[0] == "submitted" for x in order) + + +def test_common_options_have_defaults_and_follow_the_duration( + stub_flux_cli, monkeypatch +): + """Setting a walltime sets the hold limit, with the scout's slack.""" + monkeypatch.delenv("FLUX_QUANTUM_MOCK", raising=False) + cli = importlib.import_module("flux_quantum.cli") + + c = cli.common_options(_make_args(), {"attributes": {"system": {"duration": 600}}}) + assert c == { + "device": None, + "hold": "session", + "hold_max": 900.0, + "wait": 0.0, + "dry_run": False, + } + assert cli.common_options(_make_args(), None)["hold_max"] == 900.0 + + c = cli.common_options( + _make_args( + quantum_device="ibm_fez", + quantum_hold="probe", + quantum_hold_max="30", + quantum_wait="5", + quantum_dry_run=True, + ) + ) + assert c["device"] == "ibm_fez" and c["hold"] == "probe" + assert c["hold_max"] == 30.0 and c["wait"] == 5.0 and c["dry_run"] + + +def test_the_prefixed_name_wins_over_a_bare_one(stub_flux_cli): + """flux has a --wait of its own. The plugin's is read by its prefixed + name first.""" + cli = importlib.import_module("flux_quantum.cli") + assert cli._opt(_make_args(wait=True, quantum_wait="7"), "wait") == "7" + assert cli._opt(_make_args(device="d"), "device") == "d" + + +def test_mock_makes_every_submit_a_dry_run(stub_flux_cli, monkeypatch): + monkeypatch.setenv("FLUX_QUANTUM_MOCK", "1") + cli = importlib.import_module("flux_quantum.cli") + assert cli.common_options(_make_args())["dry_run"] + + +def _vendor(simulator=None, holds=("session",)): + from flux_quantum.backends.base import Backend + + class V(Backend): + name = "v" + + def scout_options(self, common): + self.check_hold(common.get("hold")) + return dict(common, mapped=True) + + def probe(self): + raise NotImplementedError + + V.simulator = simulator + V.holds = holds + return V() + + +def test_scout_options_apply_the_dry_run_before_the_vendor_maps( + stub_flux_cli, monkeypatch +): + monkeypatch.delenv("FLUX_QUANTUM_MOCK", raising=False) + cli = importlib.import_module("flux_quantum.cli") + opts = cli.scout_options(_vendor(simulator="sim"), _make_args(quantum_dry_run=True)) + assert opts["device"] == "sim" and opts["dry_run"] and opts["mapped"] + opts = cli.scout_options(_vendor(simulator="sim"), _make_args(quantum_device="qpu")) + assert opts["device"] == "qpu" and not opts["dry_run"] + + +def test_a_hold_the_vendor_lacks_fails_at_submit(stub_flux_cli, monkeypatch): + """Before anything is held, with the vendor's explanation.""" + from flux_quantum.backends import BackendError + + monkeypatch.delenv("FLUX_QUANTUM_MOCK", raising=False) + cli = importlib.import_module("flux_quantum.cli") + with pytest.raises(BackendError, match="no probe hold"): + cli.scout_options(_vendor(), _make_args(quantum_hold="probe")) + with pytest.raises(BackendError, match="unknown hold"): + cli.scout_options(_vendor(), _make_args(quantum_hold="grab")) diff --git a/tests/unit/test_mock_queue.py b/tests/unit/test_mock_queue.py index f818d4e..d32ba85 100644 --- a/tests/unit/test_mock_queue.py +++ b/tests/unit/test_mock_queue.py @@ -18,20 +18,12 @@ def mock(monkeypatch): return mockmod.MockBackend() -class Args: - mock_session = None - mock_latency = None - mock_queue_depth = None - mock_service_time = None - mock_base_overhead = None - mock_jitter = None - mock_seed = None - - -def test_defaults_match_what_ibm_measured(mock): +def test_defaults_match_what_ibm_measured(mock, monkeypatch): """ibm_marrakesh at depth 0 took 10 to 12 seconds to dequeue, so the fixed overhead is not invented.""" - opts = mock.scout_options(Args()) + for v in ("QUEUE", "SERVICE_TIME", "BASE_OVERHEAD", "JITTER"): + monkeypatch.delenv("FLUX_QUANTUM_MOCK_" + v, raising=False) + opts = mock.scout_options({}) assert opts["base_overhead"] == 10.0 assert opts["queue_depth"] == 0 assert opts["service_time"] == 0.1 @@ -109,3 +101,14 @@ def test_position_is_reported_while_draining(mock, capsys): ) out = capsys.readouterr().out assert "position 3" in out and "position 1" in out + + +def test_the_knobs_are_environment_variables(mock, monkeypatch): + """Experiment parameters, not submit options. A run sets them once in + the environment and every submit in it sees the same queue.""" + monkeypatch.setenv("FLUX_QUANTUM_MOCK_QUEUE", "25") + monkeypatch.setenv("FLUX_QUANTUM_MOCK_SERVICE_TIME", "0.5") + monkeypatch.setenv("FLUX_QUANTUM_MOCK_SESSION", "forced") + opts = mock.scout_options({"hold": "probe"}) + assert opts["queue_depth"] == 25 and opts["service_time"] == 0.5 + assert opts["session"] == "forced" diff --git a/flux_quantum/backends/hold/__init__.py b/tests/unit/vendors/__init__.py similarity index 100% rename from flux_quantum/backends/hold/__init__.py rename to tests/unit/vendors/__init__.py diff --git a/tests/unit/vendors/braket/__init__.py b/tests/unit/vendors/braket/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/test_braket.py b/tests/unit/vendors/braket/test_backend.py similarity index 97% rename from tests/unit/test_braket.py rename to tests/unit/vendors/braket/test_backend.py index 5bf103e..35c2ea4 100644 --- a/tests/unit/test_braket.py +++ b/tests/unit/vendors/braket/test_backend.py @@ -109,15 +109,10 @@ def test_defaults_to_the_sv1_simulator(fake_braket): """SV1 is cents per task, so it is the sane default for a probe.""" from flux_quantum.backends.braket import SV1 - class Args: - braket_device = None - braket_region = None - braket_shots = None - braket_queue_timeout = None - - opts = _backend().scout_options(Args()) + opts = _backend().scout_options({}) assert opts["device"] == SV1 assert opts["shots"] == 1 + assert opts["queue_timeout"] == 0 def test_region_comes_from_the_device_arn(monkeypatch): diff --git a/tests/unit/test_braket_hold.py b/tests/unit/vendors/braket/test_hold.py similarity index 52% rename from tests/unit/test_braket_hold.py rename to tests/unit/vendors/braket/test_hold.py index 8bc1b32..feb8c60 100644 --- a/tests/unit/test_braket_hold.py +++ b/tests/unit/vendors/braket/test_hold.py @@ -10,7 +10,7 @@ def test_entry_publishes_the_token_then_waits_for_release(): - from flux_quantum.backends.hold import entry + from flux_quantum.backends.braket import hold as entry written = {} released = {"yet": False} @@ -46,7 +46,7 @@ def sleep(_): def test_entry_gives_up_rather_than_billing_forever(): """A scout that is never released has to stop. The instance bills by the minute for as long as it sits there.""" - from flux_quantum.backends.hold import entry + from flux_quantum.backends.braket import hold as entry class FakeS3: def put_object(self, Bucket, Key, Body): @@ -82,36 +82,38 @@ def _backend(): return b -def test_scout_options_carry_the_hold_mode(): - - class Args: - braket_device = None - braket_region = None - braket_shots = None - braket_queue_timeout = None - braket_ungate_position = None - braket_hold = "job" - braket_hold_instance = None - braket_hold_max_seconds = None - - opts = _backend().scout_options(Args()) +def test_the_session_hold_is_the_hybrid_job(): + """The common session hold is Braket's hybrid job, and the hold limit + arrives in seconds, which is what the hybrid job's stopping condition + wants.""" + opts = _backend().scout_options({"hold": "session", "hold_max": 1200}) assert opts["hold"] == "job" # the cheapest instance, since it does no work assert opts["hold_instance"] == "ml.m5.large" - assert opts["hold_max_seconds"] == 900 + assert opts["hold_max_seconds"] == 1200 -def test_probe_stays_the_default(): - """The older mode keeps working for anyone not asking for a hybrid job.""" +def test_session_is_the_default_and_probe_is_the_other_hold(): + assert _backend().scout_options({})["hold"] == "job" + assert _backend().scout_options({"hold": "probe"})["hold"] == "probe" - class Args: - braket_device = None - braket_region = None - braket_shots = None - braket_queue_timeout = None - braket_ungate_position = None - assert _backend().scout_options(Args())["hold"] == "probe" +def test_tuning_comes_from_the_environment(monkeypatch): + """The instance and the probe shots are an operator's settings, not a + submitter's, so they are not submit options.""" + monkeypatch.setenv("FLUX_QUANTUM_BRAKET_HOLD_INSTANCE", "ml.m5.xlarge") + monkeypatch.setenv("FLUX_QUANTUM_BRAKET_SHOTS", "10") + opts = _backend().scout_options({}) + assert opts["hold_instance"] == "ml.m5.xlarge" and opts["shots"] == 10 + + +def test_a_dry_run_is_sv1(): + from flux_quantum.backends.braket import SV1 + + b = _backend() + common = b.dry_run({"device": "arn:aws:braket:us-east-1::device/qpu/ionq/Forte-1"}) + assert common["device"] == SV1 and common["dry_run"] + assert b.scout_options(common)["device"] == SV1 def test_the_token_replaces_the_job_arn_as_the_session(): @@ -210,88 +212,6 @@ def test_a_job_that_died_says_so_and_for_how_long(): assert not ok and "FAILED" in why and "10s" in why -def test_the_verdict_survives_a_task_that_finishes_instantly(): - """jobArn is set at creation and is still there afterwards, so the check - never has to win a race. A one qubit task on a simulator can finish before - the first poll returns.""" - import probe_hold - - class Instant: - def get_quantum_task(self, quantumTaskArn): - return { - "status": "COMPLETED", - "jobArn": "arn:job" if quantumTaskArn == "t1" else None, - } - - seen = probe_hold.settle(Instant(), ["t1", "t2"], sleep=lambda _: None) - assert seen["t1"]["job"] == "arn:job" - assert not seen["t2"].get("job") - - -def test_a_queue_glimpse_is_kept_but_not_required(): - """queueInfo describes a waiting task and is gone once it runs. Record it - if it shows up, never depend on it.""" - import probe_hold - - class Fleeting: - def __init__(self): - self.n = 0 - - def get_quantum_task(self, quantumTaskArn): - self.n += 1 - if self.n <= 2: - return { - "status": "QUEUED", - "jobArn": "arn:job", - "queueInfo": {"queuePriority": "Priority", "position": "1"}, - } - return {"status": "COMPLETED", "jobArn": "arn:job"} - - seen = probe_hold.settle(Fleeting(), ["t1"], sleep=lambda _: None) - assert seen["t1"]["priority"] == "Priority" - assert seen["t1"]["status"] == "COMPLETED" - - -def test_settle_gives_up_rather_than_hanging(): - import probe_hold - - class Stuck: - def get_quantum_task(self, quantumTaskArn): - return {"status": "QUEUED", "jobArn": "arn:job"} - - calls = {"n": 0} - - def sleep(_): - calls["n"] += 1 - if calls["n"] > 500: - raise AssertionError("did not respect the timeout") - - seen = probe_hold.settle(Stuck(), ["t1"], timeout=0, sleep=sleep) - assert seen["t1"] == {} - - -def test_verdict_handles_a_queue_that_was_never_visible(): - """The ordinary simulator case. settle only records fields it saw, so - reading priority unconditionally raised KeyError on the very run this - probe is for.""" - from probe_hold import SV1, verdict - - out = "\n".join(verdict({"job": "arn:job"}, {}, SV1)) - assert "ASSOCIATED" in out - assert "cannot be seen" in out - - -def test_verdict_reports_priority_when_both_were_seen(): - from probe_hold import verdict - - out = "\n".join( - verdict( - {"job": "arn:job", "priority": "Priority"}, {"priority": "Normal"}, "qpu" - ) - ) - assert "PRIORITISED" in out - - def test_verdict_calls_out_a_control_that_is_also_associated(): """If the control is linked too, the comparison says nothing.""" from probe_hold import verdict @@ -380,49 +300,6 @@ def queue_depth(self): assert qpus_with_a_queue(cls=Broken) == [] -def test_priority_is_read_from_which_queue_grew(): - from probe_hold import verdict - - out = "\n".join( - verdict( - {"job": "arn:job"}, - {}, - "qpu", - {"normal": 0, "priority": 1}, - {"normal": 1, "priority": 0}, - ) - ) - assert "PRIORITISED" in out and "NOT PRIORITISED" not in out - - -def test_a_token_that_queues_normally_is_called_out(): - from probe_hold import verdict - - out = "\n".join( - verdict( - {"job": "arn:job"}, - {}, - "qpu", - {"normal": 1, "priority": 0}, - {"normal": 1, "priority": 0}, - ) - ) - assert "NOT PRIORITISED" in out - - -def test_an_idle_device_gives_one_verdict_not_two(): - from probe_hold import SV1, verdict - - out = verdict( - {"job": "arn:job"}, - {}, - SV1, - {"normal": 0, "priority": 0}, - {"normal": 0, "priority": 0}, - ) - assert len([x for x in out if "PRIORITIS" in x or "cannot be seen" in x]) == 1 - - def test_a_maintenance_gap_counts_as_shut(): """Garnet runs weekdays but stops between 15:30 and 17:15. From outside that looks the same as a long queue.""" @@ -513,14 +390,6 @@ def queue_depth(self): assert "prioritized hybrid job running" in q["jobs_note"] -def test_moved_ignores_the_note(): - from probe_hold import moved - - before = {"normal": 1, "priority": 0, "jobs": 0, "jobs_note": "0 (1 running)"} - after = {"normal": 1, "priority": 1, "jobs": 0, "jobs_note": "0 (1 running)"} - assert moved(before, after) == {"normal": 0, "priority": 1, "jobs": 0} - - def _device(wins): return type( "D", @@ -909,48 +778,6 @@ def fake_create(**kw): assert seen["region"] == "us-west-1" -def test_wait_for_queue_returns_as_soon_as_the_filler_registers(): - """Submitting returns before the service counts the task, so reading the - depth straight after can miss it.""" - import probe_hold - - depths = [ - {"normal": 0, "priority": 0}, - {"normal": 0, "priority": 0}, - {"normal": 3, "priority": 0}, - ] - probe_hold.queue_depth = lambda d: depths.pop(0) if depths else {"normal": 3} - try: - assert probe_hold.wait_for_queue("arn:dev", sleep=lambda _: None)["normal"] == 3 - finally: - from flux_quantum.backends.braket import queue_depth - - probe_hold.queue_depth = queue_depth - - -def test_wait_for_queue_gives_up_on_a_device_too_fast_to_queue(): - """Cepheus ran four ten shot tasks before the hold was ready. Waiting - forever for a queue that will never form helps nobody.""" - import probe_hold - - probe_hold.queue_depth = lambda d: {"normal": 0, "priority": 0} - calls = {"n": 0} - - def sleep(_): - calls["n"] += 1 - if calls["n"] > 200: - raise AssertionError("did not respect the timeout") - - try: - assert ( - probe_hold.wait_for_queue("arn:dev", timeout=0, sleep=sleep)["normal"] == 0 - ) - finally: - from flux_quantum.backends.braket import queue_depth - - probe_hold.queue_depth = queue_depth - - def test_a_completed_task_is_not_reported_as_a_cancel_failure(): """On a fast device every filler task has already run, and four lines of ConflictException are noise.""" @@ -1065,28 +892,6 @@ def test_a_calibration_we_cannot_read_is_not_invented(): assert calibration(type("D", (), {})()) is None -def test_a_queue_that_drained_during_the_hold_says_so(): - """Seen four times: the device has a backlog when the run starts and none - by the time the hold is up. That is not the same as never having one.""" - from probe_hold import verdict - - out = "\n".join( - verdict({"job": "j"}, {}, "qpu", {"normal": -7, "priority": 0}, {"normal": 0}) - ) - assert "lost 7 tasks while the hold was provisioning" in out - assert "Nothing was waiting" not in out - - -def test_an_idle_device_still_reads_as_idle(): - from probe_hold import verdict - - out = "\n".join( - verdict({"job": "j"}, {}, "qpu", {"normal": 0, "priority": 0}, {"normal": 0}) - ) - assert "Nothing was waiting" in out - assert "provisioning" not in out - - def test_used_qubits_does_not_block_on_an_unfinished_task(): """result() polls until the task completes, five days by default. On a queued task that is a hang, and it hung a live run.""" @@ -1155,68 +960,876 @@ def test_every_device_gets_its_windows(): assert "02:00 to 12:00" in out and "18:00 to 20:59" in out -def test_a_task_sitting_in_priority_is_the_plainest_evidence(): - """Priority read zero on every survey of these devices. One task in it - after we submitted exactly one token task is the answer, and movement - between two laggy counts is not.""" +# --------------------------------------------------------------------------- +# the probe reads the queue from the task, not from the device counters + + +def test_the_poll_asks_for_queue_info(): + """queueInfo only comes back when GetQuantumTask is asked for it. Every + run before this one polled without asking, saw no queue, no position and + no priority, and fell back to device counters that lag by minutes.""" + from probe_hold import describe + + asked = {} + + class Client: + def get_quantum_task(self, **kw): + asked.update(kw) + return {"status": "QUEUED"} + + describe(Client(), "arn:t") + assert asked == { + "quantumTaskArn": "arn:t", + "additionalAttributeNames": ["QueueInfo"], + } + + +def test_a_queued_task_reports_its_queue(): + from probe_hold import describe + + class Client: + def get_quantum_task(self, **kw): + return { + "status": "QUEUED", + "jobArn": "arn:job", + "queueInfo": { + "queue": "QUANTUM_TASKS_QUEUE", + "position": "1", + "queuePriority": "Priority", + }, + } + + seen = describe(Client(), "arn:t") + assert (seen["priority"], seen["position"], seen["job"]) == ( + "Priority", + "1", + "arn:job", + ) + + +def test_a_position_of_none_is_none(): + """The service spells a missing position as the string None.""" + from probe_hold import describe + + class Client: + def get_quantum_task(self, **kw): + return {"status": "RUNNING", "queueInfo": {"position": "None"}} + + assert describe(Client(), "arn:t")["position"] is None + + +class _Ticker: + """A clock and a sleep that agree, so follow can be run without waiting.""" + + def __init__(self): + self.t = 1000.0 + + def clock(self): + return self.t + + def sleep(self, s): + self.t += s + + +def test_follow_keeps_the_first_queue_sighting(): + """The queue fields describe a waiting task and are gone once it runs. + The first sighting is the task's queue, the last position is where it + was when it left, and the timeline holds every change in between.""" + import probe_hold + + polls = [ + { + "status": "QUEUED", + "jobArn": "arn:job", + "queueInfo": {"queuePriority": "Priority", "position": "2"}, + }, + { + "status": "QUEUED", + "jobArn": "arn:job", + "queueInfo": {"queuePriority": "Priority", "position": "1"}, + }, + {"status": "RUNNING", "jobArn": "arn:job"}, + {"status": "COMPLETED", "jobArn": "arn:job", "endedAt": "later"}, + ] + + class Client: + def get_quantum_task(self, **kw): + return polls.pop(0) + + tick = _Ticker() + v = probe_hold.follow(Client(), ["t1"], sleep=tick.sleep, clock=tick.clock)["t1"] + assert v["priority"] == "Priority" + assert v["position"] == "2" + assert v["position_last"] == "1" + assert v["status"] == "COMPLETED" + assert v["ended"] == "later" + assert [row[1:] for row in v["timeline"]] == [ + ["QUEUED", "Priority", "2"], + ["QUEUED", "Priority", "1"], + ["RUNNING", None, None], + ["COMPLETED", None, None], + ] + + +def test_follow_gives_up_rather_than_hanging(): + """The old poll gave up after a fixed three minutes whatever the hold + was, and released a hold with both tasks still queued. Now the caller + sets the budget, and it is honoured.""" + import probe_hold + + class Stuck: + def get_quantum_task(self, **kw): + return {"status": "QUEUED", "jobArn": "arn:job"} + + tick = _Ticker() + v = probe_hold.follow( + Stuck(), ["t1"], timeout=60, interval=5, sleep=tick.sleep, clock=tick.clock + ) + assert v["t1"]["status"] == "QUEUED" + assert tick.t - 1000 <= 60 + + +def test_follow_stops_when_the_token_task_is_done(): + """The control can wait all afternoon. The hold only has to stay up + until the token task has finished, so that is the stop.""" + import probe_hold + + class Client: + def __init__(self): + self.n = 0 + + def get_quantum_task(self, quantumTaskArn, **kw): + self.n += 1 + if quantumTaskArn == "token": + return {"status": "COMPLETED" if self.n > 2 else "QUEUED"} + return { + "status": "QUEUED", + "queueInfo": {"queuePriority": "Normal", "position": "4"}, + } + + tick = _Ticker() + v = probe_hold.follow( + Client(), + ["control", "token"], + sleep=tick.sleep, + clock=tick.clock, + done=lambda v: v["token"]["status"] in probe_hold.TERMINAL, + ) + assert v["token"]["status"] == "COMPLETED" + assert v["control"]["status"] == "QUEUED" + assert v["control"]["position_last"] == "4" + + +def test_follow_resumes_without_re_polling_finished_tasks(): + import probe_hold + + asked = [] + + class Client: + def get_quantum_task(self, quantumTaskArn, **kw): + asked.append(quantumTaskArn) + return {"status": "COMPLETED"} + + tick = _Ticker() + views = {"done": {"status": "COMPLETED", "timeline": []}} + probe_hold.follow( + Client(), ["done", "other"], sleep=tick.sleep, clock=tick.clock, views=views + ) + assert asked == ["other"] + assert views["other"]["status"] == "COMPLETED" + + +def _at(seconds): + import datetime + + return datetime.datetime(2026, 9, 29, 0, 0, 0, tzinfo=datetime.timezone.utc) + ( + datetime.timedelta(seconds=seconds) + ) + + +def test_the_verdict_reads_the_queue_from_the_task(): + """Priority against Normal, read from the tasks themselves, and the + token task submitted after the control but finished before it.""" from probe_hold import verdict out = "\n".join( verdict( - {"job": "j"}, - {}, + { + "job": "arn:job", + "priority": "Priority", + "position": "1", + "created": _at(5), + "ended": _at(100), + }, + { + "priority": "Normal", + "position": "3", + "created": _at(0), + "ended": _at(400), + }, "qpu", - {"normal": 6, "priority": 0}, - {"normal": 0}, - after_token={"normal": 8, "priority": 1}, ) ) + assert "ASSOCIATED" in out assert "IN THE PRIORITY QUEUE" in out + assert "Priority at position 1" in out and "Normal at position 3" in out + assert "SERVED FIRST" in out + assert "5s after the control" in out and "300s before it" in out assert "NOT PRIORITISED" not in out -def test_an_empty_priority_queue_falls_back_to_movement(): +def test_a_token_in_the_normal_queue_is_not_prioritised(): from probe_hold import verdict out = "\n".join( verdict( - {"job": "j"}, - {}, + {"job": "arn:job", "priority": "Normal", "position": "4"}, + {"priority": "Normal", "position": "3"}, "qpu", - {"normal": 1, "priority": 0}, - {"normal": 1}, - after_token={"normal": 8, "priority": 0}, ) ) assert "NOT PRIORITISED" in out -def test_a_reading_is_taken_only_once_it_stops_moving(): - """The counters lag a submit. Six filler tasks read as one, then seven a - moment later, and a token task's arrival in Priority was credited to the - filler catching up.""" +def test_finishing_in_order_buys_no_place(): + from probe_hold import verdict + + out = "\n".join( + verdict( + { + "job": "arn:job", + "priority": "Priority", + "created": _at(1), + "ended": _at(500), + }, + {"priority": "Normal", "created": _at(0), "ended": _at(400)}, + "qpu", + ) + ) + assert "SERVED IN ORDER" in out and "100s after it" in out + + +def test_a_control_still_queued_when_the_token_finished_counts(): + """The control may sit for an hour. The token task finishing while the + control is still in line is the ordering, without waiting for the rest.""" + from probe_hold import verdict + + out = "\n".join( + verdict( + { + "job": "arn:job", + "priority": "Priority", + "created": _at(1), + "ended": _at(200), + }, + { + "priority": "Normal", + "status": "QUEUED", + "position_last": "2", + "created": _at(0), + }, + "qpu", + ) + ) + assert "SERVED FIRST" in out + assert "still QUEUED at position 2" in out + + +def test_a_token_task_that_did_not_finish_says_so(): + from probe_hold import verdict + + out = "\n".join( + verdict( + { + "job": "arn:job", + "priority": "Priority", + "status": "QUEUED", + "created": _at(1), + }, + {"priority": "Normal", "created": _at(0)}, + "qpu", + ) + ) + assert "did not finish" in out and "max-seconds" in out + + +def test_the_token_task_has_to_go_in_last(): + """Finishing first after being submitted first proves nothing.""" + from probe_hold import verdict + + out = "\n".join( + verdict( + { + "job": "arn:job", + "priority": "Priority", + "created": _at(0), + "ended": _at(10), + }, + {"priority": "Normal", "created": _at(5), "ended": _at(20)}, + "qpu", + ) + ) + assert "WARNING" in out and "created before the control" in out + assert "SERVED FIRST" not in out + + +def test_nothing_waiting_is_said_once(): + """The ordinary simulator case: neither task was ever seen in a queue, + so there is no queue to name and no order worth claiming.""" + from probe_hold import SV1, verdict + + out = verdict( + {"job": "arn:job", "created": _at(1), "ended": _at(2)}, + {"created": _at(0), "ended": _at(3)}, + SV1, + ) + assert len(out) == 2 + assert "cannot be seen" in out[1] + + +def test_the_estimate_covers_both_tasks_and_the_hold(): + """Read before anything is created, because the token task and the + control both bill per shot and the hold bills by the minute.""" + from probe_hold import estimate + + d = type( + "D", + (), + { + "properties": type( + "P", + (), + { + "service": type( + "S", + (), + { + "deviceCost": type( + "C", (), {"price": 0.08, "unit": "shot"} + )() + }, + )() + }, + )() + }, + )() + # two tasks at 0.30 + 100 * 0.08, plus half an hour of the hold instance + assert abs(estimate(d, 100, 0, 500, 1800) - (2 * 8.30 + 0.115 / 2)) < 1e-6 + assert estimate(type("D", (), {})(), 100, 0, 500, 1800) is None + + +def test_a_cancelled_control_is_not_a_finish(): + """--cancel-control ends the control with an end time, but it never + ran. Reading that as finished would turn a cancel into a race result.""" + from probe_hold import verdict + + out = "\n".join( + verdict( + { + "job": "arn:job", + "priority": "Priority", + "status": "COMPLETED", + "created": _at(1), + "ended": _at(200), + }, + { + "priority": "Normal", + "status": "CANCELLED", + "position_last": "2", + "created": _at(0), + "ended": _at(230), + }, + "qpu", + ) + ) + assert "SERVED FIRST" in out + assert "cancelled afterwards" in out and "at position 2" in out + assert "before it" not in out + + +def test_a_failed_token_task_is_not_a_race_result(): + from probe_hold import verdict + + out = "\n".join( + verdict( + { + "job": "arn:job", + "priority": "Priority", + "status": "FAILED", + "created": _at(1), + "ended": _at(5), + }, + { + "priority": "Normal", + "status": "COMPLETED", + "created": _at(0), + "ended": _at(400), + }, + "qpu", + ) + ) + assert "FAILED" in out and "never ran" in out + assert "SERVED" not in out + + +class _SimulatedDevice: + """A Braket client over a device that serves one task every `each` + seconds, Priority tasks first. Creating a task takes a second, as it + does for real, so the first filler is already running when the token + task arrives and the token task queues at Priority position 1.""" + + def __init__(self, clock, each=30): + self.clock = clock + self.each = each + self.tasks = {} + self.n = 0 + + def _now(self): + return self.clock["t"] + + def create_quantum_task(self, **kw): + self.n += 1 + arn = "arn:aws:braket:us-west-1:1:quantum-task/%d" % self.n + self.tasks[arn] = { + "created": self._now(), + "priority": bool(kw.get("jobToken")), + "job": "arn:job" if kw.get("jobToken") else None, + "cancelled": None, + } + self.clock["t"] += 1 + return {"quantumTaskArn": arn} + + def _schedule(self, skip=()): + """When each task starts, back to back, priority first.""" + free = min(t["created"] for t in self.tasks.values()) + left = {a: t for a, t in self.tasks.items() if a not in skip} + starts = {} + while left: + ready = [a for a, t in left.items() if t["created"] <= free] + if not ready: + free = min(t["created"] for t in left.values()) + continue + ready.sort(key=lambda a: (not left[a]["priority"], left[a]["created"])) + starts[ready[0]] = free + free += self.each + del left[ready[0]] + return starts + + def _view(self, arn): + starts = self._schedule() + gone = [ + a + for a, t in self.tasks.items() + if t["cancelled"] is not None and starts[a] >= t["cancelled"] + ] + starts = self._schedule(skip=gone) + t = self.tasks[arn] + if arn in gone: + return "CANCELLED", None, t["cancelled"] + start = starts[arn] + now = self._now() + if now >= start + self.each: + return "COMPLETED", None, start + self.each + if now >= start: + return "RUNNING", None, None + ahead = [ + a + for a, s in starts.items() + if a not in gone + and self.tasks[a]["priority"] == t["priority"] + and now < s < start + ] + return "QUEUED", str(len(ahead) + 1), None + + def get_quantum_task(self, quantumTaskArn, additionalAttributeNames=None): + import datetime + + assert additionalAttributeNames == ["QueueInfo"], "the poll must ask" + t = self.tasks[quantumTaskArn] + status, position, ended = self._view(quantumTaskArn) + stamp = lambda x: datetime.datetime.fromtimestamp(x, datetime.timezone.utc) + out = {"status": status, "jobArn": t["job"], "createdAt": stamp(t["created"])} + if ended: + out["endedAt"] = stamp(ended) + if status == "QUEUED": + out["queueInfo"] = { + "queue": "QUANTUM_TASKS_QUEUE", + "position": position, + "queuePriority": "Priority" if t["priority"] else "Normal", + } + return out + + def cancel_quantum_task(self, quantumTaskArn, clientToken): + status, _, _ = self._view(quantumTaskArn) + if status != "QUEUED": + raise RuntimeError( + "ConflictException: cannot cancel a task in the %s status" % status + ) + self.tasks[quantumTaskArn]["cancelled"] = self._now() + + +def test_the_whole_run_against_a_simulated_device(): + """main, end to end, the way it will be run on Cepheus: six filler + tasks, a control, then the token task. Nothing here touches AWS, and + nothing sleeps, because the clock is ours.""" + import json + import os + import sys + import tempfile + import types + + import braket.aws as aws_mod + import probe_hold - readings = [ - {"normal": 1, "priority": 0}, - {"normal": 7, "priority": 0}, - {"normal": 7, "priority": 1}, - {"normal": 7, "priority": 1}, + clock = {"t": 1000.0} + fake_time = types.SimpleNamespace( + time=lambda: clock["t"], + sleep=lambda s: clock.__setitem__("t", clock["t"] + s), + ) + device = _SimulatedDevice(clock) + + class FakeBackend: + credential_note = "fake credentials" + closed = 0 + + def __init__(self): + self._hold_job = types.SimpleNamespace(state=lambda: "RUNNING") + self._hold_ready_after = None + + def open_session(self, opts): + assert opts["hold"] == "job" and opts["hold_max_seconds"] == 900 + return "arn:job" + + def wait_for_priority(self, opts): + clock["t"] += 100 + self._hold_ready_after = 100 + return True, "holding the device after 100s" + + def session_id(self, arn): + return "tok-secret-0123456789abcdef" + + def close_session(self): + FakeBackend.closed += 1 + + class FakeDevice: + def __init__(self, arn): + self.properties = type( + "P", + (), + { + "service": type( + "S", + (), + { + "executionWindows": ALWAYS, + "shotsRange": (10, 50000), + "deviceCost": type( + "C", (), {"price": 0.000425, "unit": "shot"} + )(), + }, + )(), + "provider": type("V", (), {"specs": {}})(), + }, + )() + + tmp = tempfile.mkdtemp() + rec, log = os.path.join(tmp, "runs.jsonl"), os.path.join(tmp, "probe.log") + saved = ( + probe_hold.time, + probe_hold.BraketBackend, + probe_hold.boto3, + probe_hold.AwsSession, + probe_hold.queue_depth, + aws_mod.AwsDevice, + sys.argv, + sys.stdout, + sys.stderr, + os.environ.get("AWS_DEFAULT_REGION"), + ) + probe_hold.time = fake_time + probe_hold.BraketBackend = FakeBackend + probe_hold.boto3 = types.SimpleNamespace(client=lambda name: device) + probe_hold.AwsSession = lambda: types.SimpleNamespace( + default_bucket=lambda: "bucket" + ) + probe_hold.queue_depth = lambda d: { + "normal": 0, + "priority": 0, + "jobs": 0, + "jobs_note": "0", + } + aws_mod.AwsDevice = FakeDevice + sys.argv = [ + "probe_hold.py", + "--device", + "arn:aws:braket:us-west-1::device/qpu/rigetti/Cepheus-1-108Q", + "--shots", + "10", + "--make-queue", + "6", + "--filler-shots", + "500", + "--max-seconds", + "900", + "--record", + rec, + "--log", + log, ] - real, probe_hold.queue_depth = probe_hold.queue_depth, lambda d: readings.pop(0) try: - got = probe_hold.settled_depth("arn:dev", sleep=lambda _: None) + probe_hold.main() + finally: + ( + probe_hold.time, + probe_hold.BraketBackend, + probe_hold.boto3, + probe_hold.AwsSession, + probe_hold.queue_depth, + aws_mod.AwsDevice, + sys.argv, + sys.stdout, + sys.stderr, + region, + ) = saved + if region is None: + os.environ.pop("AWS_DEFAULT_REGION", None) + else: + os.environ["AWS_DEFAULT_REGION"] = region + + run = json.loads(open(rec).read().strip().splitlines()[-1]) + text = "\n".join(run["verdict"]) + assert "ASSOCIATED" in text + assert "Priority at position 1" in text + assert "Normal at position 6" in text + assert "SERVED FIRST" in text and "180s before it" in text + assert run["with_token"]["job"] == "arn:job" and not run["without"].get("job") + # the filler is part of the measurement, so its ARNs are on record too. + # Forte's went unrecorded and its finish time could not be read back + assert len(run["filler_arns"]) == 6 + assert run["with_token"]["status"] == "COMPLETED" + assert run["without"]["status"] == "COMPLETED" + # the hold is released once, as soon as the token task is done + assert FakeBackend.closed == 1 + + logged = open(log).read() + assert "hybrid job: arn:job" in logged + assert "this run costs about $3." in logged + assert "Priority queue position 1" in logged + assert "SERVED FIRST" in logged + assert "recorded to" in logged + # the token is not written down. The prefix printed is not the token + assert "tok-secret-0123456789abcdef" not in logged + assert "tok-secret-0" in logged + + +def test_log_to_tees_and_headers(): + import io + import os + import sys + import tempfile + + import probe_hold + + path = os.path.join(tempfile.mkdtemp(), "x.log") + out, err = io.StringIO(), io.StringIO() + saved = sys.stdout, sys.stderr + sys.stdout, sys.stderr = out, err + try: + probe_hold.log_to(path, argv=["probe_hold.py", "--survey"]) + print("hello") + print("oops", file=sys.stderr) + finally: + sys.stdout, sys.stderr = saved + logged = open(path).read() + assert "=== " in logged and "probe_hold.py --survey" in logged + assert "hello\n" in logged and "oops\n" in logged + assert "hello" in out.getvalue() and "oops" in err.getvalue() + + +def test_labels_without_positions_mean_nobody_waited(): + """Cepheus, 01:03 on the 29th: both tasks finished 450ms after creation + and Braket still labelled one Priority and one Normal. The labels are + real, but finishing in submission order on an idle device is not a + result, and the old wording called it one.""" + from probe_hold import verdict + + out = "\n".join( + verdict( + { + "job": "arn:job", + "status": "COMPLETED", + "priority": "Priority", + "created": _at(0.35), + "ended": _at(0.81), + }, + { + "status": "COMPLETED", + "priority": "Normal", + "created": _at(0.0), + "ended": _at(0.46), + }, + "qpu", + ) + ) + assert "IN THE PRIORITY QUEUE" in out + assert "NO CONTENTION" in out and "ran in 0.5s" in out + assert "SERVED" not in out and "bought no place" not in out + + +def test_quera_gets_an_analog_program_and_everyone_else_a_circuit(): + """Aquila runs Hamiltonians and refuses a circuit. It is also the one + device with a queue that does not drain the moment it opens.""" + import json + + import probe_hold + + real = probe_hold.ahs_action + probe_hold.ahs_action = lambda: "analog" + try: + assert ( + probe_hold.program_action( + "arn:aws:braket:us-east-1::device/qpu/quera/Aquila" + ) + == "analog" + ) + circuit = json.loads( + probe_hold.program_action( + "arn:aws:braket:us-west-1::device/qpu/rigetti/Cepheus-1-108Q" + ) + ) + finally: + probe_hold.ahs_action = real + assert circuit["braketSchemaHeader"]["name"] == "braket.ir.openqasm.program" + assert "measure" in circuit["source"] + + +def test_submit_sends_the_device_program(): + import probe_hold + + sent = {} + + class Client: + def create_quantum_task(self, **kw): + sent.update(kw) + return {"quantumTaskArn": "arn:t"} + + real = probe_hold.ahs_action + probe_hold.ahs_action = lambda: "analog" + try: + probe_hold.submit( + Client(), + "arn:aws:braket:us-east-1::device/qpu/quera/Aquila", + "b", + "tok", + 10, + ) + finally: + probe_hold.ahs_action = real + assert ( + sent["action"] == "analog" and sent["jobToken"] == "tok" and sent["shots"] == 10 + ) + + +def test_inspect_reads_old_tasks_back(): + """Free. Two tasks from the Forte run sat five minutes behind two others, + and how long they finally took is the number that says whether Forte + can hold a queue.""" + import builtins + + import probe_hold + + class Client: + def get_quantum_task(self, quantumTaskArn, additionalAttributeNames): + assert additionalAttributeNames == ["QueueInfo"] + return { + "status": "COMPLETED", + "createdAt": _at(0), + "endedAt": _at(400), + "queueInfo": { + "queuePriority": ( + "Priority" if quantumTaskArn.endswith("/a") else "Normal" + ) + }, + } + + said = [] + real = builtins.print + builtins.print = lambda *a, **k: said.append(" ".join(str(x) for x in a)) + try: + rows = probe_hold.inspect(Client(), ["arn:x/a", "arn:x/b"]) finally: - probe_hold.queue_depth = real - assert got == {"normal": 7, "priority": 1} + builtins.print = real + assert [r["priority"] for r in rows] == ["Priority", "Normal"] + assert any("400s" in line and "Priority" in line for line in said) + + +def test_pace_reports_a_device_that_serves_one_task_at_a_time(): + """Filler only builds a queue on a device like this. The simulated + device takes 30s a task, so three tasks span about 90s.""" + import probe_hold + + clock = {"t": 1000.0} + device = _SimulatedDevice(clock) + tick = lambda s: clock.__setitem__("t", clock["t"] + s) + rows = probe_hold.pace( + device, "arn:dev", "bucket", 3, 10, sleep=tick, clock=lambda: clock["t"] + ) + assert [r["status"] for r in rows] == ["COMPLETED"] * 3 + assert rows[1]["position"] == "1" and rows[2]["position"] == "2" + text = "\n".join(probe_hold.pace_report(rows)) + assert "SERIAL" in text and "filler will hold" in text + + +def test_pace_reports_a_device_that_runs_everything_at_once(): + """Cepheus: eight tasks, each done 450ms after creation, overlapping.""" + import probe_hold + + rows = [ + { + "arn": "arn:x/%d" % i, + "status": "COMPLETED", + "created": _at(i * 0.3), + "ended": _at(i * 0.3 + 0.45), + "took": 0.45, + "position": None, + } + for i in range(3) + ] + text = "\n".join(probe_hold.pace_report(rows)) + assert "OVERLAPPING" in text and "will not build a queue" in text -def test_settling_gives_up_rather_than_waiting_forever(): +def test_a_long_wait_says_it_is_still_waiting(): + """Forte prints nothing for minutes while a task runs, which reads as a + hang. A line a minute says otherwise, and what the queue looks like.""" + import builtins + import probe_hold - flip = [{"normal": i % 2, "priority": 0} for i in range(200)] - real, probe_hold.queue_depth = probe_hold.queue_depth, lambda d: flip.pop(0) + class Stuck: + def get_quantum_task(self, **kw): + return { + "status": "QUEUED", + "queueInfo": {"queuePriority": "Priority", "position": "1"}, + } + + said = [] + real = builtins.print + builtins.print = lambda *a, **k: said.append(" ".join(str(x) for x in a)) + tick = _Ticker() try: - probe_hold.settled_depth("arn:dev", timeout=0, sleep=lambda _: None) + probe_hold.follow( + Stuck(), + ["arn:x/token"], + timeout=200, + interval=5, + sleep=tick.sleep, + clock=tick.clock, + say=lambda a, s: None, + ) finally: - probe_hold.queue_depth = real + builtins.print = real + beats = [x for x in said if "still waiting" in x] + assert 2 <= len(beats) <= 4 + assert "QUEUED at 1" in beats[0] diff --git a/tests/unit/vendors/ibm/__init__.py b/tests/unit/vendors/ibm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/test_qrmi.py b/tests/unit/vendors/ibm/test_qrmi.py similarity index 92% rename from tests/unit/test_qrmi.py rename to tests/unit/vendors/ibm/test_qrmi.py index baffe64..7da4bb2 100644 --- a/tests/unit/test_qrmi.py +++ b/tests/unit/vendors/ibm/test_qrmi.py @@ -109,19 +109,29 @@ def test_close_without_open_is_harmless(fake_qrmi): def test_resource_is_inferred_when_only_one_is_configured(fake_qrmi): - class Args: - ibm_resource = None - ibm_type = None - - assert _ibm().scout_options(Args())["resource"] == RESOURCE + assert _ibm().scout_options({})["resource"] == RESOURCE def test_explicit_resource_wins(fake_qrmi): - class Args: - ibm_resource = "ibm_fez" - ibm_type = None + assert _ibm().scout_options({"device": "ibm_fez"})["resource"] == "ibm_fez" + + +def test_ibm_has_no_probe_hold(fake_qrmi): + """QRMI acquires a session or nothing, so a probe is refused at submit + time, before anything is held.""" + from flux_quantum.backends import BackendError + + with pytest.raises(BackendError, match="no probe hold"): + _ibm().scout_options({"hold": "probe"}) + - assert _ibm().scout_options(Args())["resource"] == "ibm_fez" +def test_the_warmup_wait_is_the_common_wait(fake_qrmi, monkeypatch): + monkeypatch.setenv("FLUX_QUANTUM_IBM_SKIP_WARMUP", "1") + monkeypatch.setenv("FLUX_QUANTUM_IBM_TYPE", "ibm-quantum-system") + opts = _ibm().scout_options({"wait": 600}) + assert opts["warmup_timeout"] == 600 + assert opts["skip_warmup"] is True + assert opts["type"] == "ibm-quantum-system" def test_job_environment_matches_the_qrmi_convention(fake_qrmi): @@ -314,4 +324,4 @@ def test_missing_qiskit_says_what_to_install(fake_qrmi, monkeypatch): with pytest.raises(RuntimeError) as e: b.wait_for_priority({}) assert "qrmi[ibm]" in str(e.value) - assert "skip-warmup" in str(e.value) + assert "SKIP_WARMUP" in str(e.value) From 5871cca86c86f12c82f6ac43c6e83b60da6f05e6 Mon Sep 17 00:00:00 2001 From: vsoch Date: Thu, 1 Oct 2026 17:00:17 +0000 Subject: [PATCH 19/26] feat: IonQ backend over the v0.4 REST API, with a dry run on the simulator Nothing but the standard library, the key from IONQ_API_KEY. The session hold creates a session on the backend, warms it up with a one qubit job inside it, and releases the classical job once IonQ reports the session active. The session id is what the classical job is handed. Sessions are in beta, and an account without them is told to hold with a probe. 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 the circuit still compiles the way the hardware would. The fake service in ionq/fake.py stands in for api.ionq.co so the backend and the whole pipeline run with no key at all, which CI now does, with and without sessions. Signed-off-by: vsoch --- .github/workflows/test.yml | 10 +- README.md | 72 +++- examples/ionq/workload.py | 80 ++++ flux_quantum/backends/__init__.py | 1 + flux_quantum/backends/ionq/__init__.py | 268 ++++++++++++ flux_quantum/backends/ionq/client.py | 58 +++ flux_quantum/backends/ionq/fake.py | 213 ++++++++++ tests/integration/test-ionq-e2e.sh | 90 ++++ tests/unit/test_backends.py | 23 +- tests/unit/vendors/ionq/__init__.py | 0 tests/unit/vendors/ionq/test_backend.py | 528 ++++++++++++++++++++++++ 11 files changed, 1320 insertions(+), 23 deletions(-) create mode 100755 examples/ionq/workload.py create mode 100644 flux_quantum/backends/ionq/__init__.py create mode 100644 flux_quantum/backends/ionq/client.py create mode 100755 flux_quantum/backends/ionq/fake.py create mode 100755 tests/integration/test-ionq-e2e.sh create mode 100644 tests/unit/vendors/ionq/__init__.py create mode 100644 tests/unit/vendors/ionq/test_backend.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 83d36d3..39ae700 100644 --- a/.github/workflows/test.yml +++ b/.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-resource --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 @@ -138,6 +138,12 @@ 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: | diff --git a/README.md b/README.md index 6166c03..93f75ab 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ branch of https://github.com/vsoch/flux-sched, and also loads against flux-core ## Components - **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 per vendor, IBM through QRMI, AWS Braket, and a mock. They probe queue depth and cost and open and close the session. They only ever run as the user. +- **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_ -> qpu` to the fluxion graph. @@ -30,17 +30,54 @@ flux submit --help | grep quantum Populate the graph once, before any job runs. Growing it while jobs hold resources corrupts fluxion. ```bash -flux python -m flux_quantum.populate ibm braket +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 \ +flux jobtap load $PWD/flux_quantum/jobtap/quantum.so vendors=ibm,braket,ionq \ total_cores=128 reserve_cores=32 protect_types=qpu preempt_after=30 ``` +## 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_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. @@ -53,11 +90,40 @@ FLUX_QUANTUM_MOCK=1 flux start bash tests/integration/test-admission.sh # job 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. + +## 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 flux-quantum is distributed under the terms of the MIT license. diff --git a/examples/ionq/workload.py b/examples/ionq/workload.py new file mode 100755 index 0000000..388f183 --- /dev/null +++ b/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() diff --git a/flux_quantum/backends/__init__.py b/flux_quantum/backends/__init__.py index 6c12bca..475018c 100644 --- a/flux_quantum/backends/__init__.py +++ b/flux_quantum/backends/__init__.py @@ -18,6 +18,7 @@ ) from . import ibm # noqa: F401 (registers IBMBackend) from . import braket # noqa: F401 (registers BraketBackend) +from . import ionq # noqa: F401 (registers IonQBackend) if os.environ.get("FLUX_QUANTUM_MOCK"): from . import mock # noqa: F401 (registers MockBackend, MockBusyBackend) diff --git a/flux_quantum/backends/ionq/__init__.py b/flux_quantum/backends/ionq/__init__.py new file mode 100644 index 0000000..1fe1b06 --- /dev/null +++ b/flux_quantum/backends/ionq/__init__.py @@ -0,0 +1,268 @@ +"""IonQ backend, over the v0.4 REST API with nothing but the standard library. + +Credentials come from IONQ_API_KEY, or IONQ_API_TOKEN. IONQ_API_URL overrides +the endpoint, which is how the tests point it at the fake service in fake.py. + +Two ways to hold the device, set with --quantum-hold. + +session is the default. It creates a session on the backend and submits a one +qubit warm-up job inside it. The scout is ready once IonQ reports the session +active, which is when jobs carrying the session id are being served. The +session id is what the classical job is handed. Sessions are in beta and only +some accounts have them, and a refusal says so plainly. + +probe submits the warm-up job alone and is ready once it has started. That +holds nothing, like the Braket probe, but it does say the device is serving us. + +Dry run. --quantum-dry-run, or FLUX_QUANTUM_MOCK in the environment, sends +every job to the simulator, which IonQ does not charge for, with the noise +model of the backend asked for, so the circuit compiles the way it would on +hardware. The simulator has no queue, so a dry run exercises the plumbing and +not the hold. + +Tuning, from the environment: + + FLUX_QUANTUM_IONQ_SHOTS shots for the warm-up job, 100 +""" + +import math +import os +import time + +from ..base import Backend, BackendError, Signals, register, truthy, tuning +from .client import API_URL, APIError, Client + +SIMULATOR = "simulator" +DEFAULT_BACKEND = "qpu.forte-1" + +# v0.4 says started where v0.3 said running, and a job that has already +# finished was served too +STARTED = ("started", "running", "completed") +TERMINAL = ("completed", "canceled", "failed") + + +def noise_model(backend): + """The simulator noise model for a QPU name, or None for the simulator + itself. qpu.forte-1 is modelled as forte-1.""" + if backend.startswith("qpu."): + return backend[len("qpu.") :] + return None + + +# one qubit, one gate. The smallest job the service accepts. +WARMUP = {"qubits": 1, "gateset": "qis", "circuit": [{"gate": "h", "target": 0}]} + + +@register +class IonQBackend(Backend): + name = "ionq" + + def __init__(self, client=None): + if client is None: + key = os.environ.get("IONQ_API_KEY") or os.environ.get("IONQ_API_TOKEN") + if not key: + raise BackendError( + "ionq: no IONQ_API_KEY in the environment. Create one at " + "cloud.ionq.com and export it" + ) + client = Client(key, os.environ.get("IONQ_API_URL") or API_URL) + self.credential_note = "ionq: credentials present via IONQ_API_KEY" + else: + self.credential_note = "ionq: client supplied" + self.client = client + self._session = None + self._job = None + self._hold_started = None + self._hold_ready_after = None + + simulator = SIMULATOR + holds = ("session", "probe") + + def dry_run(self, common): + """The simulator, with the noise model of the backend that was asked + for, so the circuit still compiles the way the hardware would.""" + out = super().dry_run(common) + out["noise"] = noise_model(common.get("device") or DEFAULT_BACKEND) + return out + + def scout_options(self, common): + hold = self.check_hold(common.get("hold")) + return { + # where jobs are sent, and the noise that stands in for the + # hardware when it is the simulator + "target": common.get("device") or DEFAULT_BACKEND, + "noise": common.get("noise"), + "dry_run": bool(common.get("dry_run")), + "hold": hold, + # the session's own limit is in minutes + "max_minutes": max( + 1, int(math.ceil(float(common.get("hold_max") or 900) / 60)) + ), + "shots": int(tuning("ionq_shots", 100)), + "timeout": float(common.get("wait") or 0), + } + + def job_environment(self, options): + """What the classical job needs to submit into the hold: the target + and, on the simulator, the noise model. The key stays the user's.""" + env = { + "IONQ_BACKEND": options.get("target") or DEFAULT_BACKEND, + "IONQ_API_URL": self.client.url, + } + if options.get("noise"): + env["IONQ_NOISE_MODEL"] = options["noise"] + return env + + def _warmup(self, options, session=None): + """Submit the one qubit job and return its id.""" + body = { + "type": "ionq.circuit.v1", + "name": "flux-quantum-warmup", + "metadata": {"flux_quantum": "warmup"}, + "shots": int(options.get("shots") or 100), + "backend": options.get("target") or DEFAULT_BACKEND, + "input": WARMUP, + } + if session: + body["session_id"] = session + if options.get("noise"): + body["noise"] = {"model": options["noise"]} + return self.client.post("/jobs", body)["id"] + + def open_session(self, options): + """Take the device. session returns the session id, probe the job id. + + A session the account cannot create is refused with the fix named. + In a dry run the simulator may refuse one too, and since the run is + there to exercise the plumbing, it falls back to a probe and says so. + """ + self._hold_started = time.time() + if (options.get("hold") or "session") == "session": + body = { + "backend": options.get("target") or DEFAULT_BACKEND, + "settings": { + "duration_limit_min": int(options.get("max_minutes") or 15) + }, + } + try: + self._session = self.client.post("/sessions", body)["id"] + except APIError as e: + if e.status not in (400, 401, 403, 404): + raise + if not options.get("dry_run"): + raise BackendError( + "ionq: this account cannot create a session on {} " + "(HTTP {}). Sessions are in beta. Ask IonQ, or hold " + "with --quantum-hold probe".format( + body["backend"], e.status + ) + ) + print( + "ionq: {} refused a session (HTTP {}), holding with a " + "probe job instead".format(body["backend"], e.status), + flush=True, + ) + else: + self._job = self._warmup(options, session=self._session) + return self._session + self._job = self._warmup(options) + return self._job + + def session_id(self, opened): + """A session id goes over as is, since jobs are submitted into it. A + probe job's id is prefixed, so the classical job can tell it holds + nothing it can submit into.""" + if self._session and opened == self._session: + return opened + return "job:{}".format(opened) + + def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep): + """Poll until the hold is real. + + With a session, that is IonQ saying it is active, or the warm-up job + inside it having started, whichever is reported first. With a probe, + the warm-up job having started. A warm-up that fails or is cancelled + is a failure, not readiness. + """ + options = options or {} + timeout = float(options.get("timeout") or 0) + started = self._hold_started or time.time() + deadline = started + timeout if timeout > 0 else None + while True: + job = self.client.get("/jobs/{}".format(self._job)) + status = job.get("status") + waited = time.time() - started + if status in ("failed", "canceled"): + why = (job.get("failure") or {}).get("error") or status + return False, "warm-up job {} after {:.0f}s: {}".format( + status, waited, why + ) + if self._session: + s = self.client.get("/sessions/{}".format(self._session)) + if s.get("ended_at") or s.get("status") in ("ended", "expired"): + return False, "session {} before it became active".format( + s.get("status") or "ended" + ) + if s.get("active") or s.get("started_at") or status in STARTED: + self._hold_ready_after = waited + return True, "session active after {:.0f}s".format(waited) + said = "session {}, warm-up job {}".format( + s.get("status") or "pending", status + ) + else: + if status in STARTED: + self._hold_ready_after = waited + return True, "warm-up job {} after {:.0f}s".format(status, waited) + said = "warm-up job {}".format(status) + print("ionq: {} ({:.0f}s)".format(said, waited), flush=True) + if deadline is not None and time.time() + interval >= deadline: + return False, "gave up after {:.0f}s, still {}".format(waited, said) + sleep(interval) + + def close_session(self, session=None): + """End the session, which cancels anything still queued in it, or + cancel a probe job that has not run. Safe to call twice.""" + sess, job, self._session, self._job = self._session, self._job, None, None + if sess: + try: + self.client.post("/sessions/{}/end".format(sess)) + except Exception as e: + print("ionq: could not end session {}: {}".format(sess, e)) + return + if job: + try: + if ( + self.client.get("/jobs/{}".format(job)).get("status") + not in TERMINAL + ): + self.client.put("/jobs/{}/status/cancel".format(job)) + except Exception as e: + print("ionq: could not cancel job {}: {}".format(job, e)) + + def probe(self): + """The backend's status and queue, from GET /backends. + + IonQ reports an average queue time rather than a count. It goes in + queue_depth so --quantum-select queue can rank it, lower being + better either way, and the unit is noted in detail. + """ + backend = DEFAULT_BACKEND + if truthy(os.environ.get("FLUX_QUANTUM_MOCK")): + backend = SIMULATOR + try: + b = self.client.get("/backends/{}".format(backend)) + except Exception as e: + return Signals( + available=False, detail={"backend": backend, "error": str(e)} + ) + return Signals( + available=b.get("status") == "available", + queue_depth=b.get("average_queue_time"), + detail={ + "backend": backend, + "status": b.get("status"), + "degraded": b.get("degraded"), + "qubits": b.get("qubits"), + "queue_unit": "average_queue_time", + }, + ) diff --git a/flux_quantum/backends/ionq/client.py b/flux_quantum/backends/ionq/client.py new file mode 100644 index 0000000..4a225e4 --- /dev/null +++ b/flux_quantum/backends/ionq/client.py @@ -0,0 +1,58 @@ +"""The few calls the IonQ backend makes, with the key in a header. + +Nothing but the standard library. The endpoint is overridable, which is how +the tests and a dry run without a key point it at the fake service. +""" + +import json + +from ..base import BackendError + +API_URL = "https://api.ionq.co/v0.4" + + +class APIError(BackendError): + """The service answered with an error status.""" + + def __init__(self, status, method, path, body): + self.status = status + text = body.decode(errors="replace") if isinstance(body, bytes) else body + super().__init__( + "ionq: {} {} returned {}: {}".format(method, path, status, text[:200]) + ) + + +class Client: + """The few calls the backend makes, with the key in a header and the + endpoint overridable.""" + + def __init__(self, key, url=API_URL, timeout=30): + self.url = url.rstrip("/") + self._key = key + self.timeout = timeout + + def request(self, method, path, body=None): + import urllib.error + import urllib.request + + data = None if body is None else json.dumps(body).encode() + req = urllib.request.Request(self.url + path, data=data, method=method) + req.add_header("Authorization", "apiKey " + self._key) + req.add_header("Content-Type", "application/json") + try: + with urllib.request.urlopen(req, timeout=self.timeout) as r: + raw = r.read() + except urllib.error.HTTPError as e: + raise APIError(e.code, method, path, e.read()) + except urllib.error.URLError as e: + raise BackendError("ionq: cannot reach {}: {}".format(self.url, e.reason)) + return json.loads(raw) if raw else {} + + def get(self, path): + return self.request("GET", path) + + def post(self, path, body=None): + return self.request("POST", path, body if body is not None else {}) + + def put(self, path, body=None): + return self.request("PUT", path, body) diff --git a/flux_quantum/backends/ionq/fake.py b/flux_quantum/backends/ionq/fake.py new file mode 100755 index 0000000..c76e94a --- /dev/null +++ b/flux_quantum/backends/ionq/fake.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""A stand-in for api.ionq.co, so the IonQ backend runs with no key. + +Implements the handful of v0.4 calls the backend makes: backends, jobs, +cancel, sessions and end. Jobs advance one state per GET, submitted to ready +to started to completed, so a poll loop sees every state and nothing has to +sleep. A session becomes active when a job inside it starts, and ending it +cancels whatever is still queued. + + python3 -m flux_quantum.backends.ionq.fake --port 8765 + IONQ_API_KEY=anything IONQ_API_URL=http://127.0.0.1:8765 ... + +FAKE_IONQ_NO_SESSIONS=1 refuses POST /sessions with 404, which is what an +account without the beta sees. FAKE_IONQ_STEPS is how many GETs a job spends +in each state before moving on, default 1. +""" + +import argparse +import json +import os +import threading +import uuid +from http.server import BaseHTTPRequestHandler, HTTPServer + +ORDER = ["submitted", "ready", "started", "completed"] + + +class State: + def __init__(self, no_sessions=False, steps=1): + self.jobs = {} + self.sessions = {} + self.no_sessions = no_sessions + self.steps = max(1, int(steps)) + self.requests = [] + self.lock = threading.Lock() + + def create_job(self, body): + jid = str(uuid.uuid4()) + session = body.get("session_id") + if session and session not in self.sessions: + return 404, {"error": "session not found"} + self.jobs[jid] = { + "id": jid, + "status": "submitted", + "backend": body.get("backend"), + "shots": body.get("shots"), + "session_id": session, + "noise": body.get("noise"), + "name": body.get("name"), + "polls": 0, + } + return 200, {"id": jid, "status": "submitted", "session_id": session} + + def get_job(self, jid): + job = self.jobs.get(jid) + if job is None: + return 404, {"error": "job not found"} + if job["status"] in ORDER and job["status"] != "completed": + job["polls"] += 1 + if job["polls"] >= self.steps: + job["polls"] = 0 + job["status"] = ORDER[ORDER.index(job["status"]) + 1] + sess = self.sessions.get(job["session_id"]) + if sess and job["status"] == "started" and not sess["active"]: + sess["active"] = True + sess["status"] = "active" + sess["started_at"] = "now" + return 200, {k: v for k, v in job.items() if k != "polls"} + + def cancel_job(self, jid): + job = self.jobs.get(jid) + if job is None: + return 404, {"error": "job not found"} + if job["status"] not in ("completed", "failed"): + job["status"] = "canceled" + return 200, {"id": jid, "status": "canceled"} + + def create_session(self, body): + if self.no_sessions: + return 404, {"error": "not found"} + sid = str(uuid.uuid4()) + self.sessions[sid] = { + "id": sid, + "backend": body.get("backend"), + "settings": body.get("settings") or {}, + "status": "pending", + "active": False, + "started_at": None, + "ended_at": None, + } + return 200, dict(self.sessions[sid]) + + def get_session(self, sid): + s = self.sessions.get(sid) + return (404, {"error": "session not found"}) if s is None else (200, dict(s)) + + def end_session(self, sid): + s = self.sessions.get(sid) + if s is None: + return 404, {"error": "session not found"} + s["status"], s["active"], s["ended_at"] = "ended", False, "now" + for job in self.jobs.values(): + if job["session_id"] == sid and job["status"] not in ( + "completed", + "failed", + ): + job["status"] = "canceled" + return 200, dict(s) + + def backend(self, name): + return 200, { + "backend": name, + "status": "available", + "degraded": False, + "qubits": 36, + "average_queue_time": 0 if name == "simulator" else 600, + } + + +def make_handler(state): + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def _body(self): + n = int(self.headers.get("Content-Length") or 0) + raw = self.rfile.read(n) if n else b"" + return json.loads(raw) if raw else {} + + def _reply(self, code, payload): + data = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def _route(self, method): + auth = self.headers.get("Authorization") or "" + if not auth.startswith("apiKey ") or len(auth) <= len("apiKey "): + return self._reply(401, {"error": "unauthorized"}) + parts = [p for p in self.path.split("?")[0].split("/") if p] + body = self._body() if method in ("POST", "PUT") else {} + with state.lock: + state.requests.append((method, "/" + "/".join(parts), body)) + code, payload = self._dispatch(method, parts, body) + self._reply(code, payload) + + def _dispatch(self, method, parts, body): + if parts[:1] == ["backends"] and len(parts) == 2 and method == "GET": + return state.backend(parts[1]) + if parts == ["jobs"] and method == "POST": + return state.create_job(body) + if parts[:1] == ["jobs"] and len(parts) == 2 and method == "GET": + return state.get_job(parts[1]) + if parts[:1] == ["jobs"] and parts[2:] == ["results"] and method == "GET": + return 200, {"0": 0.5, "1": 0.5} + if parts[:1] == ["jobs"] and parts[2:] == ["status", "cancel"]: + return state.cancel_job(parts[1]) + if parts == ["sessions"] and method == "POST": + return state.create_session(body) + if parts[:1] == ["sessions"] and len(parts) == 2 and method == "GET": + return state.get_session(parts[1]) + if parts[:1] == ["sessions"] and parts[2:] == ["end"] and method == "POST": + return state.end_session(parts[1]) + return 404, { + "error": "no such route: {} {}".format(method, "/".join(parts)) + } + + def do_GET(self): + self._route("GET") + + def do_POST(self): + self._route("POST") + + def do_PUT(self): + self._route("PUT") + + return Handler + + +def serve(port=0, state=None): + """Start serving in a thread. Returns (server, state). The port is + server.server_address[1], which matters when 0 was asked for.""" + state = state or State( + no_sessions=bool(os.environ.get("FAKE_IONQ_NO_SESSIONS")), + steps=os.environ.get("FAKE_IONQ_STEPS", 1), + ) + server = HTTPServer(("127.0.0.1", port), make_handler(state)) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return server, state + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--port", type=int, default=8765) + args = ap.parse_args() + server, _ = serve(args.port) + print( + "fake ionq listening on http://127.0.0.1:%d" % server.server_address[1], + flush=True, + ) + try: + threading.Event().wait() + except KeyboardInterrupt: + pass + finally: + server.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/integration/test-ionq-e2e.sh b/tests/integration/test-ionq-e2e.sh new file mode 100755 index 0000000..29ba17f --- /dev/null +++ b/tests/integration/test-ionq-e2e.sh @@ -0,0 +1,90 @@ +#!/bin/bash +# End-to-end test of the pipeline with the IonQ backend against the fake +# api.ionq.co, so no key is needed and nothing is billed. +# +# One flux submit with --quantum-vendor ionq produces the held classical and +# the scout. The scout opens a session on the fake service, the classical +# job submits a circuit into it, and the scout ends the session afterwards. +# +# flux start -s1 bash tests/integration/test-ionq-e2e.sh +# FAKE_IONQ_NO_SESSIONS=1 flux start bash tests/integration/test-ionq-e2e.sh +set -u +HERE=$(cd "$(dirname "$0")/../.." && pwd) +export FLUX_CLI_PLUGINPATH="$HERE/cli-plugins" +export FLUX_QUANTUM_MOCK=1 +export IONQ_API_KEY="fake-key-$$" +ERR=$(mktemp) +LOG=$(mktemp) +rc=0 + +echo "=== 0. start the fake service ===" +flux python -m flux_quantum.backends.ionq.fake --port 0 >"$LOG" 2>&1 & +FAKE=$! +for _ in $(seq 1 50); do + grep -q "listening on" "$LOG" && break + sleep 0.1 +done +IONQ_API_URL=$(grep -oE 'http://[0-9.]+:[0-9]+' "$LOG") +export IONQ_API_URL +[ -n "$IONQ_API_URL" ] || { echo "FAIL: fake service did not start"; cat "$LOG"; exit 1; } +echo "fake ionq at $IONQ_API_URL" +trap 'kill $FAKE 2>/dev/null; rm -f "$ERR" "$LOG"' EXIT + +echo "" +echo "=== 1. plugin discovery ===" +flux submit --help 2>&1 | grep -q -- --quantum-dry-run || { echo "FAIL: no --quantum-dry-run"; exit 1; } + +echo "" +echo "=== 2. load fluxion + qmanager ===" +flux module remove -f sched-fluxion-qmanager 2>/dev/null || true +flux module remove -f sched-fluxion-feasibility 2>/dev/null || true +flux module remove -f sched-fluxion-resource 2>/dev/null || true +flux module remove -f sched-simple 2>/dev/null || true +flux module load sched-fluxion-resource +flux module load sched-fluxion-feasibility 2>/dev/null || true +flux module load sched-fluxion-qmanager + +echo "" +echo "=== 3. one quantum submit -> held classical and scout ===" +scout_id=$(flux submit --quantum-vendor ionq --quantum-dry-run \ + --quantum-device qpu.forte-1 -n1 \ + -- flux python "$HERE/examples/ionq/workload.py" 2>"$ERR") +cat "$ERR" >&2 +main_id=$(grep -oE 'held classical job [0-9]+' "$ERR" | awk '{print $NF}') +echo "scout=$scout_id classical=$main_id" +[ -n "$main_id" ] || { echo "FAIL: no held classical job"; rc=1; } + +echo "" +echo "=== 4. the classical job submitted into the session and got a result ===" +if [ -n "$main_id" ]; then + if flux job wait-event -t 60 "$main_id" clean &1) + echo "$out" | sed 's/^/ /' + echo "$out" | grep -q "^session: " || { echo "FAIL: no session reached the classical job"; rc=1; } + echo "$out" | grep -q "^target: simulator" || { echo "FAIL: dry run did not target the simulator"; rc=1; } + echo "$out" | grep -q "^status: completed" || { echo "FAIL: the circuit did not complete"; rc=1; } + else + echo "FAIL: classical never completed"; rc=1 + flux jobs -a 2>&1 | head + fi +fi + +echo "" +echo "=== 5. the scout held the session for the classical's lifetime, then ended it ===" +if [ -n "$scout_id" ]; then + if flux job wait-event -t 60 "$scout_id" clean /dev/null 2>&1; then + sout=$(flux job attach "$scout_id" &1) + echo "$sout" | sed 's/^/ /' + echo "$sout" | grep -q "classical job .* finished" || { echo "FAIL: scout did not wait"; rc=1; } + echo "$sout" | grep -q "closed ionq session" || { echo "FAIL: scout did not close the session"; rc=1; } + else + echo "FAIL: scout never completed"; rc=1 + fi +fi + +echo "" +echo "=== 6. teardown ===" +flux module remove -f sched-fluxion-qmanager 2>/dev/null || true +flux module remove -f sched-fluxion-feasibility 2>/dev/null || true +flux module remove -f sched-fluxion-resource 2>/dev/null || true +exit "$rc" diff --git a/tests/unit/test_backends.py b/tests/unit/test_backends.py index 9d28a1a..89a7419 100644 --- a/tests/unit/test_backends.py +++ b/tests/unit/test_backends.py @@ -12,7 +12,7 @@ def test_real_backends_registered(backends_real): - assert backends_real.known_vendors() == {"ibm", "braket"} + assert backends_real.known_vendors() == {"ibm", "braket", "ionq"} def test_mock_gated_by_env(backends_real, backends_mock): @@ -75,21 +75,12 @@ def test_mock_needs_no_creds(backends_mock): assert backends_mock.get_backend("mock") is not None -def test_backend_declares_options_and_scout_options_roundtrip(fresh, monkeypatch): - """A vendor backend declares CLI options and extracts them from args.""" +def test_scout_options_map_the_common_device_to_the_resource(fresh, monkeypatch): + """A vendor backend takes the common options and speaks its own terms.""" _fake_qrmi(monkeypatch) b = fresh(mock=False, env=IBM_CREDS) - ibm = b.get_backend("ibm") - declared = [] - ibm.add_options(lambda name, **kw: declared.append(name)) - assert "--ibm-resource" in declared and "--ibm-type" in declared - - class _Args: - ibm_resource = "ibm_kingston" - ibm_type = None - - opts = ibm.scout_options(_Args()) + opts = b.get_backend("ibm").scout_options({"device": "ibm_kingston"}) assert opts["resource"] == "ibm_kingston" assert opts["type"] == "qiskit-runtime-service" @@ -99,11 +90,7 @@ def test_scout_options_infer_the_resource_from_the_environment(fresh, monkeypatc _fake_qrmi(monkeypatch) b = fresh(mock=False, env=IBM_CREDS) - class _Args: - ibm_resource = None - ibm_type = None - - assert b.get_backend("ibm").scout_options(_Args())["resource"] == "ibm_kingston" + assert b.get_backend("ibm").scout_options({})["resource"] == "ibm_kingston" def test_job_environment_follows_the_qrmi_convention(fresh, monkeypatch): diff --git a/tests/unit/vendors/ionq/__init__.py b/tests/unit/vendors/ionq/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/vendors/ionq/test_backend.py b/tests/unit/vendors/ionq/test_backend.py new file mode 100644 index 0000000..b02d032 --- /dev/null +++ b/tests/unit/vendors/ionq/test_backend.py @@ -0,0 +1,528 @@ +"""The IonQ backend, without a key. + +Most of these drive the backend with a scripted client. The last few run it +against the fake server over real HTTP, so the request shapes are checked +end to end: the header, the JSON, the paths. +""" + +import pytest + + +class Scripted: + """A client that records every call and answers from a script. + + answers maps (method, path) to a response or a list of responses that + are handed out in order, the last one repeating. + """ + + url = "https://fake" + + def __init__(self, answers): + self.answers = { + k: (v if isinstance(v, list) else [v]) for k, v in answers.items() + } + self.calls = [] + + def request(self, method, path, body=None): + self.calls.append((method, path, body)) + key = (method, path) + if key not in self.answers: + raise AssertionError("unexpected call {} {}".format(method, path)) + queue = self.answers[key] + answer = queue.pop(0) if len(queue) > 1 else queue[0] + if isinstance(answer, Exception): + raise answer + return answer + + def get(self, path): + return self.request("GET", path) + + def post(self, path, body=None): + return self.request("POST", path, body if body is not None else {}) + + def put(self, path, body=None): + return self.request("PUT", path, body) + + +def _common(**kw): + """The common options as the CLI hands them over, dry run applied.""" + from flux_quantum.backends.ionq import IonQBackend + + common = { + "device": None, + "hold": "session", + "hold_max": 900.0, + "wait": 0.0, + "dry_run": False, + } + common.update(kw) + if common["dry_run"]: + common = IonQBackend(client=Scripted({})).dry_run(common) + return common + + +def test_no_key_is_reported_by_name(monkeypatch): + from flux_quantum.backends import BackendError + from flux_quantum.backends.ionq import IonQBackend + + monkeypatch.delenv("IONQ_API_KEY", raising=False) + monkeypatch.delenv("IONQ_API_TOKEN", raising=False) + with pytest.raises(BackendError) as e: + IonQBackend() + assert "IONQ_API_KEY" in str(e.value) + + +def test_the_key_is_never_in_the_note(monkeypatch): + from flux_quantum.backends.ionq import IonQBackend + + monkeypatch.setenv("IONQ_API_KEY", "supersecret") + b = IonQBackend() + assert "supersecret" not in b.credential_note + assert "IONQ_API_KEY" in b.credential_note + + +def test_ionq_is_a_registered_vendor(backends_real): + assert "ionq" in backends_real.known_vendors() + + +def test_defaults_target_the_qpu_with_a_session(): + from flux_quantum.backends.ionq import IonQBackend + + opts = IonQBackend(client=Scripted({})).scout_options(_common()) + assert opts["target"] == "qpu.forte-1" + assert opts["hold"] == "session" + assert opts["noise"] is None and not opts["dry_run"] + # 900 seconds is 15 minutes, the session's own unit + assert opts["max_minutes"] == 15 and opts["shots"] == 100 + assert opts["timeout"] == 0 + + +def test_dry_run_goes_to_the_simulator_with_the_qpu_noise(): + """The simulator is free. Keeping the noise model means the circuit still + compiles the way the hardware would compile it.""" + from flux_quantum.backends.ionq import IonQBackend + + b = IonQBackend(client=Scripted({})) + common = b.dry_run({"device": "qpu.forte-enterprise-1"}) + assert common["device"] == "simulator" and common["noise"] == "forte-enterprise-1" + opts = b.scout_options(common) + assert opts["target"] == "simulator" + assert opts["noise"] == "forte-enterprise-1" + + +def test_the_hold_limit_rounds_up_to_whole_minutes(): + from flux_quantum.backends.ionq import IonQBackend + + b = IonQBackend(client=Scripted({})) + assert b.scout_options(_common(hold_max=61))["max_minutes"] == 2 + assert b.scout_options(_common(hold_max=1))["max_minutes"] == 1 + + +def test_shots_are_tuning(monkeypatch): + from flux_quantum.backends.ionq import IonQBackend + + monkeypatch.setenv("FLUX_QUANTUM_IONQ_SHOTS", "250") + assert IonQBackend(client=Scripted({})).scout_options(_common())["shots"] == 250 + + +def test_the_simulator_itself_gets_no_noise_model(): + from flux_quantum.backends.ionq import IonQBackend + + b = IonQBackend(client=Scripted({})) + common = b.dry_run({"device": "simulator"}) + assert common["device"] == "simulator" and common["noise"] is None + + +def test_the_classical_job_learns_the_target_not_the_key(): + from flux_quantum.backends.ionq import IonQBackend + + b = IonQBackend(client=Scripted({})) + env = b.job_environment({"target": "simulator", "noise": "forte-1"}) + assert env["IONQ_BACKEND"] == "simulator" + assert env["IONQ_NOISE_MODEL"] == "forte-1" + assert env["IONQ_API_URL"] == "https://fake" + assert not any("KEY" in k for k in env) + + +def test_a_session_is_opened_then_warmed_up_inside_it(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/sessions"): { + "id": "sess-1", + "status": "pending", + "active": False, + }, + ("POST", "/jobs"): { + "id": "job-1", + "status": "submitted", + "session_id": "sess-1", + }, + } + ) + b = IonQBackend(client=c) + opened = b.open_session( + {"target": "qpu.forte-1", "hold": "session", "max_minutes": 7, "shots": 100} + ) + assert opened == "sess-1" + assert b.session_id(opened) == "sess-1" + _, _, session = c.calls[0] + assert session == {"backend": "qpu.forte-1", "settings": {"duration_limit_min": 7}} + _, _, job = c.calls[1] + assert job["session_id"] == "sess-1" + assert job["backend"] == "qpu.forte-1" + assert job["shots"] == 100 + assert job["input"]["qubits"] == 1 + assert "noise" not in job + + +def test_a_dry_run_warmup_carries_the_noise_model(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted({("POST", "/jobs"): {"id": "job-1", "status": "submitted"}}) + IonQBackend(client=c).open_session( + {"target": "simulator", "noise": "forte-1", "hold": "probe", "shots": 10} + ) + assert c.calls[0][2]["noise"] == {"model": "forte-1"} + assert c.calls[0][2]["backend"] == "simulator" + + +def test_an_account_without_sessions_is_told_what_to_do(): + from flux_quantum.backends import BackendError + from flux_quantum.backends.ionq import APIError, IonQBackend + + c = Scripted( + {("POST", "/sessions"): APIError(403, "POST", "/sessions", b"forbidden")} + ) + with pytest.raises(BackendError) as e: + IonQBackend(client=c).open_session({"target": "qpu.forte-1", "hold": "session"}) + assert "beta" in str(e.value) and "--quantum-hold probe" in str(e.value) + assert len(c.calls) == 1 + + +def test_a_dry_run_falls_back_to_a_probe_and_says_so(capsys): + """The simulator may not take sessions. A dry run is there to exercise + the plumbing, so it carries on with a probe rather than stopping.""" + from flux_quantum.backends.ionq import APIError, IonQBackend + + c = Scripted( + { + ("POST", "/sessions"): APIError(404, "POST", "/sessions", b"nope"), + ("POST", "/jobs"): {"id": "job-1", "status": "submitted"}, + } + ) + b = IonQBackend(client=c) + opened = b.open_session({"target": "simulator", "hold": "session", "dry_run": True}) + assert opened == "job-1" + assert "session_id" not in c.calls[1][2] + assert "probe job instead" in capsys.readouterr().out + + +def test_a_server_error_on_sessions_is_not_swallowed(): + from flux_quantum.backends.ionq import APIError, IonQBackend + + c = Scripted({("POST", "/sessions"): APIError(500, "POST", "/sessions", b"boom")}) + with pytest.raises(APIError): + IonQBackend(client=c).open_session({"target": "qpu.forte-1", "dry_run": True}) + + +def test_ready_once_the_session_is_active(capsys): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/sessions"): {"id": "sess-1"}, + ("POST", "/jobs"): {"id": "job-1"}, + ("GET", "/jobs/job-1"): [{"status": "submitted"}, {"status": "ready"}], + ("GET", "/sessions/sess-1"): [ + {"status": "pending", "active": False}, + {"status": "active", "active": True, "started_at": "t"}, + ], + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1"}) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda _: None) + assert ok and "session active" in why + assert "session pending, warm-up job submitted" in capsys.readouterr().out + + +def test_a_started_warmup_counts_as_active_even_if_the_session_lags(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/sessions"): {"id": "sess-1"}, + ("POST", "/jobs"): {"id": "job-1"}, + ("GET", "/jobs/job-1"): {"status": "started"}, + ("GET", "/sessions/sess-1"): {"status": "pending", "active": False}, + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1"}) + ok, _ = b.wait_for_priority({}, interval=0, sleep=lambda _: None) + assert ok + + +def test_a_failed_warmup_is_a_failure_not_readiness(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/jobs"): {"id": "job-1"}, + ("GET", "/jobs/job-1"): { + "status": "failed", + "failure": {"error": "circuit too wide"}, + }, + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1", "hold": "probe"}) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda _: None) + assert not ok and "failed" in why and "circuit too wide" in why + + +def test_a_session_that_ends_first_is_a_failure(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/sessions"): {"id": "sess-1"}, + ("POST", "/jobs"): {"id": "job-1"}, + ("GET", "/jobs/job-1"): {"status": "submitted"}, + ("GET", "/sessions/sess-1"): {"status": "expired", "ended_at": "t"}, + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1"}) + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda _: None) + assert not ok and "expired" in why + + +def test_the_wait_gives_up_at_the_timeout(): + import time + + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/jobs"): {"id": "job-1"}, + ("GET", "/jobs/job-1"): {"status": "submitted"}, + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1", "hold": "probe"}) + b._hold_started = time.time() - 100 + ok, why = b.wait_for_priority({"timeout": 60}, interval=1, sleep=lambda _: None) + assert not ok and "gave up" in why + + +def test_probe_mode_is_ready_when_the_job_starts(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/jobs"): {"id": "job-1"}, + ("GET", "/jobs/job-1"): [{"status": "submitted"}, {"status": "started"}], + } + ) + b = IonQBackend(client=c) + assert b.open_session({"target": "qpu.forte-1", "hold": "probe"}) == "job-1" + # a job id is not something the classical job can submit into, and the + # handover says so + assert b.session_id("job-1") == "job:job-1" + ok, why = b.wait_for_priority({}, interval=0, sleep=lambda _: None) + assert ok and "started" in why + assert not any(path.startswith("/sessions") for _, path, _ in c.calls) + + +def test_closing_ends_the_session_once(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/sessions"): {"id": "sess-1"}, + ("POST", "/jobs"): {"id": "job-1"}, + ("POST", "/sessions/sess-1/end"): {"status": "ended"}, + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1"}) + b.close_session() + b.close_session() + ends = [p for m, p, _ in c.calls if p.endswith("/end")] + assert ends == ["/sessions/sess-1/end"] + + +def test_closing_a_probe_cancels_only_an_unfinished_job(): + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted( + { + ("POST", "/jobs"): {"id": "job-1"}, + ("GET", "/jobs/job-1"): {"status": "ready"}, + ("PUT", "/jobs/job-1/status/cancel"): {"status": "canceled"}, + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1", "hold": "probe"}) + b.close_session() + assert ("PUT", "/jobs/job-1/status/cancel", None) in c.calls + + c = Scripted( + { + ("POST", "/jobs"): {"id": "job-2"}, + ("GET", "/jobs/job-2"): {"status": "completed"}, + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1", "hold": "probe"}) + b.close_session() + assert not any(m == "PUT" for m, _, _ in c.calls) + + +def test_a_failed_close_is_reported_not_raised(capsys): + from flux_quantum.backends.ionq import APIError, IonQBackend + + c = Scripted( + { + ("POST", "/sessions"): {"id": "sess-1"}, + ("POST", "/jobs"): {"id": "job-1"}, + ("POST", "/sessions/sess-1/end"): APIError(500, "POST", "/x", b"down"), + } + ) + b = IonQBackend(client=c) + b.open_session({"target": "qpu.forte-1"}) + b.close_session() + assert "could not end session sess-1" in capsys.readouterr().out + + +def test_probe_reads_the_backend(monkeypatch): + from flux_quantum.backends.ionq import IonQBackend + + monkeypatch.delenv("FLUX_QUANTUM_MOCK", raising=False) + monkeypatch.delenv("IONQ_BACKEND", raising=False) + c = Scripted( + { + ("GET", "/backends/qpu.forte-1"): { + "status": "available", + "degraded": True, + "qubits": 36, + "average_queue_time": 1200, + } + } + ) + sig = IonQBackend(client=c).probe() + assert sig.available and sig.queue_depth == 1200 + assert sig.detail["degraded"] is True and sig.detail["qubits"] == 36 + + +def test_probe_follows_mock_to_the_simulator(monkeypatch): + from flux_quantum.backends.ionq import IonQBackend + + monkeypatch.setenv("FLUX_QUANTUM_MOCK", "1") + c = Scripted({("GET", "/backends/simulator"): {"status": "available"}}) + assert IonQBackend(client=c).probe().available + + +def test_an_unreachable_service_is_unavailable_not_an_exception(): + from flux_quantum.backends import BackendError + from flux_quantum.backends.ionq import IonQBackend + + c = Scripted({("GET", "/backends/qpu.forte-1"): BackendError("ionq: cannot reach")}) + sig = IonQBackend(client=c).probe() + assert not sig.available and "cannot reach" in sig.detail["error"] + + +def test_noise_model_names(): + from flux_quantum.backends.ionq import noise_model + + assert noise_model("qpu.forte-1") == "forte-1" + assert noise_model("qpu.aria-2") == "aria-2" + assert noise_model("simulator") is None + + +# --------------------------------------------------------------------------- +# over real HTTP, against the fake server + + +@pytest.fixture +def fake_ionq(): + from flux_quantum.backends.ionq import fake + + server, state = fake.serve(0) + yield "http://127.0.0.1:%d" % server.server_address[1], state + server.shutdown() + + +def test_the_client_sends_the_key_and_json(fake_ionq, monkeypatch): + from flux_quantum.backends.ionq import Client + + url, state = fake_ionq + c = Client("k-123", url) + b = c.get("/backends/simulator") + assert b["status"] == "available" + made = c.post("/jobs", {"backend": "simulator", "shots": 5}) + assert made["status"] == "submitted" + assert state.requests[-1] == ("POST", "/jobs", {"backend": "simulator", "shots": 5}) + + +def test_a_bad_key_is_an_api_error_with_the_status(fake_ionq): + from flux_quantum.backends.ionq import APIError, Client + + url, _ = fake_ionq + with pytest.raises(APIError) as e: + Client("", url).get("/backends/simulator") + assert e.value.status == 401 + + +def test_the_whole_scout_lifecycle_against_the_fake_server( + fake_ionq, monkeypatch, capsys +): + """open, wait, hand over, close. What the scout does, over HTTP.""" + from flux_quantum.backends.ionq import IonQBackend + + url, state = fake_ionq + monkeypatch.setenv("IONQ_API_KEY", "k") + monkeypatch.setenv("IONQ_API_URL", url) + monkeypatch.setenv("FLUX_QUANTUM_MOCK", "1") + b = IonQBackend() + opts = b.scout_options(b.dry_run({"device": "qpu.forte-1", "hold_max": 900})) + assert opts["target"] == "simulator" and opts["noise"] == "forte-1" + + opened = b.open_session(opts) + assert opened in state.sessions + ok, why = b.wait_for_priority(opts, interval=0, sleep=lambda _: None) + assert ok, why + assert state.sessions[opened]["active"] + session = b.session_id(opened) + assert session == opened + + # the warm-up went to the simulator, in the session, with the noise model + (job,) = state.jobs.values() + assert job["backend"] == "simulator" + assert job["session_id"] == opened + assert job["noise"] == {"model": "forte-1"} + + b.close_session(session) + assert state.sessions[opened]["status"] == "ended" + + +def test_the_lifecycle_without_sessions_falls_back_in_a_dry_run(fake_ionq, monkeypatch): + from flux_quantum.backends.ionq import IonQBackend + + url, state = fake_ionq + state.no_sessions = True + monkeypatch.setenv("IONQ_API_KEY", "k") + monkeypatch.setenv("IONQ_API_URL", url) + monkeypatch.setenv("FLUX_QUANTUM_MOCK", "1") + b = IonQBackend() + opts = b.scout_options(b.dry_run({})) + opened = b.open_session(opts) + assert opened in state.jobs + ok, why = b.wait_for_priority(opts, interval=0, sleep=lambda _: None) + assert ok and "warm-up job" in why + assert b.session_id(opened).startswith("job:") + b.close_session(opened) + assert state.sessions == {} From bf2cb9eaf2226333d4c4caf970994ceed2b8edbd Mon Sep 17 00:00:00 2001 From: vsoch Date: Sat, 3 Oct 2026 19:43:08 +0000 Subject: [PATCH 20/26] ionq: the fake service reports sessions and job times the way IonQ does Sessions go created, started, ended, which are the service's status values, and jobs carry submitted_at, started_at, completed_at and execution_duration_ms. FAKE_IONQ_SECONDS gives a job a minimum time in each state, so an experiment against the fake has timings with a shape. The nested experiments checkout is ignored. Signed-off-by: vsoch --- .gitignore | 3 ++ flux_quantum/backends/ionq/fake.py | 55 +++++++++++++++++++++++------- 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index a9b62e7..09021f6 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,6 @@ dist __pycache__/ *.pyc *.so + +# a sibling checkout the experiments live in +flux-quantum-experiments/ diff --git a/flux_quantum/backends/ionq/fake.py b/flux_quantum/backends/ionq/fake.py index c76e94a..6221634 100755 --- a/flux_quantum/backends/ionq/fake.py +++ b/flux_quantum/backends/ionq/fake.py @@ -4,33 +4,46 @@ Implements the handful of v0.4 calls the backend makes: backends, jobs, cancel, sessions and end. Jobs advance one state per GET, submitted to ready to started to completed, so a poll loop sees every state and nothing has to -sleep. A session becomes active when a job inside it starts, and ending it -cancels whatever is still queued. +sleep, and they carry the service's timestamps, submitted_at, started_at and +completed_at. A session goes created, started, ended, as the service's does, +started when a job inside it starts, and ending it cancels whatever is still +queued. python3 -m flux_quantum.backends.ionq.fake --port 8765 IONQ_API_KEY=anything IONQ_API_URL=http://127.0.0.1:8765 ... FAKE_IONQ_NO_SESSIONS=1 refuses POST /sessions with 404, which is what an account without the beta sees. FAKE_IONQ_STEPS is how many GETs a job spends -in each state before moving on, default 1. +in each state before moving on, default 1. FAKE_IONQ_SECONDS is the least +time a job spends in each state, default 0, for timings that look like a +device's. """ import argparse +import datetime import json import os import threading +import time import uuid from http.server import BaseHTTPRequestHandler, HTTPServer ORDER = ["submitted", "ready", "started", "completed"] +def _stamp(t=None): + return datetime.datetime.fromtimestamp( + time.time() if t is None else t, datetime.timezone.utc + ).isoformat() + + class State: - def __init__(self, no_sessions=False, steps=1): + def __init__(self, no_sessions=False, steps=1, seconds=0.0): self.jobs = {} self.sessions = {} self.no_sessions = no_sessions self.steps = max(1, int(steps)) + self.seconds = float(seconds) self.requests = [] self.lock = threading.Lock() @@ -47,7 +60,12 @@ def create_job(self, body): "session_id": session, "noise": body.get("noise"), "name": body.get("name"), + "submitted_at": _stamp(), + "started_at": None, + "completed_at": None, + "execution_duration_ms": None, "polls": 0, + "since": time.time(), } return 200, {"id": jid, "status": "submitted", "session_id": session} @@ -57,15 +75,24 @@ def get_job(self, jid): return 404, {"error": "job not found"} if job["status"] in ORDER and job["status"] != "completed": job["polls"] += 1 - if job["polls"] >= self.steps: + if ( + job["polls"] >= self.steps + and time.time() - job["since"] >= self.seconds + ): job["polls"] = 0 + job["since"] = time.time() job["status"] = ORDER[ORDER.index(job["status"]) + 1] - sess = self.sessions.get(job["session_id"]) - if sess and job["status"] == "started" and not sess["active"]: - sess["active"] = True - sess["status"] = "active" - sess["started_at"] = "now" - return 200, {k: v for k, v in job.items() if k != "polls"} + if job["status"] == "started": + job["started_at"] = _stamp() + sess = self.sessions.get(job["session_id"]) + if sess and not sess["active"]: + sess["active"] = True + sess["status"] = "started" + sess["started_at"] = _stamp() + elif job["status"] == "completed": + job["completed_at"] = _stamp() + job["execution_duration_ms"] = int(max(1.0, self.seconds) * 1000) + return 200, {k: v for k, v in job.items() if k not in ("polls", "since")} def cancel_job(self, jid): job = self.jobs.get(jid) @@ -83,8 +110,9 @@ def create_session(self, body): "id": sid, "backend": body.get("backend"), "settings": body.get("settings") or {}, - "status": "pending", + "status": "created", "active": False, + "created_at": _stamp(), "started_at": None, "ended_at": None, } @@ -98,7 +126,7 @@ def end_session(self, sid): s = self.sessions.get(sid) if s is None: return 404, {"error": "session not found"} - s["status"], s["active"], s["ended_at"] = "ended", False, "now" + s["status"], s["active"], s["ended_at"] = "ended", False, _stamp() for job in self.jobs.values(): if job["session_id"] == sid and job["status"] not in ( "completed", @@ -185,6 +213,7 @@ def serve(port=0, state=None): state = state or State( no_sessions=bool(os.environ.get("FAKE_IONQ_NO_SESSIONS")), steps=os.environ.get("FAKE_IONQ_STEPS", 1), + seconds=os.environ.get("FAKE_IONQ_SECONDS", 0), ) server = HTTPServer(("127.0.0.1", port), make_handler(state)) thread = threading.Thread(target=server.serve_forever, daemon=True) From 407373a6de51a59e2469e1a4d5fc60b0970d8a4f Mon Sep 17 00:00:00 2001 From: vsoch Date: Sat, 3 Oct 2026 19:46:24 +0000 Subject: [PATCH 21/26] ionq: sessions in the fake service expire at their duration limit As the service's do. A session past its limit reads as ended, a job submitted into it is refused, and anything queued in it is cancelled. The experiments' hold-max check needs that to be exercisable without a key. Signed-off-by: vsoch --- README.md | 2 +- flux_quantum/backends/ionq/__init__.py | 13 +- flux_quantum/backends/ionq/fake.py | 226 ++++++++++++++++++++++-- flux_quantum/scout.py | 67 ++++++- tests/integration/test-scout-cancel.sh | 60 +++++++ tests/unit/vendors/ionq/test_backend.py | 70 ++++++++ 6 files changed, 410 insertions(+), 28 deletions(-) create mode 100755 tests/integration/test-scout-cancel.sh diff --git a/README.md b/README.md index 93f75ab..adadf68 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ flux submit --quantum-vendor ionq --quantum-dry-run -n1 -- flux python examples/ 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. +`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 diff --git a/flux_quantum/backends/ionq/__init__.py b/flux_quantum/backends/ionq/__init__.py index 1fe1b06..545f35f 100644 --- a/flux_quantum/backends/ionq/__init__.py +++ b/flux_quantum/backends/ionq/__init__.py @@ -22,7 +22,11 @@ Tuning, from the environment: - FLUX_QUANTUM_IONQ_SHOTS shots for the warm-up job, 100 + FLUX_QUANTUM_IONQ_SHOTS shots for the warm-up job, 100 + FLUX_QUANTUM_IONQ_COST_LIMIT_USD a cost limit on every session opened, + none by default. IonQ ends the session + when its jobs reach it, so a session + nobody is watching cannot bill past it """ import math @@ -144,6 +148,9 @@ def open_session(self, options): "duration_limit_min": int(options.get("max_minutes") or 15) }, } + limit = tuning("ionq_cost_limit_usd") + if limit: + body["settings"]["cost_limit"] = {"unit": "usd", "value": float(limit)} try: self._session = self.client.post("/sessions", body)["id"] except APIError as e: @@ -153,9 +160,7 @@ def open_session(self, options): raise BackendError( "ionq: this account cannot create a session on {} " "(HTTP {}). Sessions are in beta. Ask IonQ, or hold " - "with --quantum-hold probe".format( - body["backend"], e.status - ) + "with --quantum-hold probe".format(body["backend"], e.status) ) print( "ionq: {} refused a session (HTTP {}), holding with a " diff --git a/flux_quantum/backends/ionq/fake.py b/flux_quantum/backends/ionq/fake.py index 6221634..83ebbc8 100755 --- a/flux_quantum/backends/ionq/fake.py +++ b/flux_quantum/backends/ionq/fake.py @@ -12,21 +12,36 @@ python3 -m flux_quantum.backends.ionq.fake --port 8765 IONQ_API_KEY=anything IONQ_API_URL=http://127.0.0.1:8765 ... +It listens on the loopback. On an instance of more than one node, jobs run +where the scheduler puts them, so --bind 0.0.0.0 listens on every interface +and the URL printed is the address other nodes reach this one at. + FAKE_IONQ_NO_SESSIONS=1 refuses POST /sessions with 404, which is what an account without the beta sees. FAKE_IONQ_STEPS is how many GETs a job spends in each state before moving on, default 1. FAKE_IONQ_SECONDS is the least time a job spends in each state, default 0, for timings that look like a device's. + +FAKE_IONQ_QUEUE is the device's queue: the seconds a job sits in submitted +before it is served, a number or a range like 30-90 drawn per job. A job +in a session that has started skips it, which is what a session buys: the +session's first job queues like anyone's, and once it starts the device is +the session's. /backends reports the mean as average_queue_time. The queue +can be changed while running, POST /fake/config {"queue": "30-90"}, so an +experiment can sweep it. """ import argparse import datetime import json import os +import random +import socket import threading import time import uuid from http.server import BaseHTTPRequestHandler, HTTPServer +from urllib.parse import parse_qs ORDER = ["submitted", "ready", "started", "completed"] @@ -37,21 +52,68 @@ def _stamp(t=None): ).isoformat() +def parse_queue(spec): + """A queue spec as (low, high) seconds: "30" is 30 to 30, "30-90" a range, + nothing is no queue.""" + if spec is None: + return (0.0, 0.0) + spec = str(spec).strip() + if not spec: + return (0.0, 0.0) + if "-" in spec: + lo, hi = spec.split("-", 1) + lo, hi = float(lo), float(hi) + return (min(lo, hi), max(lo, hi)) + return (float(spec), float(spec)) + + class State: - def __init__(self, no_sessions=False, steps=1, seconds=0.0): + def __init__(self, no_sessions=False, steps=1, seconds=0.0, queue=None): self.jobs = {} self.sessions = {} self.no_sessions = no_sessions self.steps = max(1, int(steps)) self.seconds = float(seconds) + self.queue = parse_queue(queue) self.requests = [] self.lock = threading.Lock() + def queue_wait(self): + lo, hi = self.queue + return random.uniform(lo, hi) if hi > lo else lo + + def config(self): + return { + "queue": "%g-%g" % self.queue, + "seconds": self.seconds, + "steps": self.steps, + "no_sessions": self.no_sessions, + } + + def configure(self, body): + """Change the service while it runs, for a sweep. Jobs already + submitted keep the wait they were given.""" + if "queue" in body: + self.queue = parse_queue(body["queue"]) + if "seconds" in body: + self.seconds = float(body["seconds"]) + if "steps" in body: + self.steps = max(1, int(body["steps"])) + if "no_sessions" in body: + self.no_sessions = bool(body["no_sessions"]) + return 200, self.config() + def create_job(self, body): jid = str(uuid.uuid4()) session = body.get("session_id") if session and session not in self.sessions: return 404, {"error": "session not found"} + if session: + self._expire(session) + if self.sessions[session]["status"] == "ended": + return 400, {"error": "session has ended"} + # the device's queue, unless a started session already has the device + served = bool(session and self.sessions[session]["active"]) self.jobs[jid] = { "id": jid, "status": "submitted", @@ -66,6 +128,7 @@ def create_job(self, body): "execution_duration_ms": None, "polls": 0, "since": time.time(), + "queue": 0.0 if served else self.queue_wait(), } return 200, {"id": jid, "status": "submitted", "session_id": session} @@ -75,10 +138,10 @@ def get_job(self, jid): return 404, {"error": "job not found"} if job["status"] in ORDER and job["status"] != "completed": job["polls"] += 1 - if ( - job["polls"] >= self.steps - and time.time() - job["since"] >= self.seconds - ): + least = self.seconds + if job["status"] == "submitted": + least = max(least, job["queue"]) + if job["polls"] >= self.steps and time.time() - job["since"] >= least: job["polls"] = 0 job["since"] = time.time() job["status"] = ORDER[ORDER.index(job["status"]) + 1] @@ -92,7 +155,36 @@ def get_job(self, jid): elif job["status"] == "completed": job["completed_at"] = _stamp() job["execution_duration_ms"] = int(max(1.0, self.seconds) * 1000) - return 200, {k: v for k, v in job.items() if k not in ("polls", "since")} + return 200, { + k: v for k, v in job.items() if k not in ("polls", "since", "queue") + } + + def job_cost(self, jid): + """GET /jobs/{id}/cost, shaped like the service's: the estimate and, + once the job ran, what was charged. Dry runs charge nothing.""" + job = self.jobs.get(jid) + if job is None: + return 404, {"error": "job not found"} + # the service answers with only dry_run for a job that cost nothing + if job.get("backend") == "simulator": + return 200, {"dry_run": False} + est = self.estimate({"shots": job.get("shots") or 1})[1]["estimated_total_cost"] + charged = est if job["status"] == "completed" else 0.0 + return 200, { + "dry_run": False, + "estimated_cost": {"value": est, "unit": "usd"}, + "cost": {"value": charged, "unit": "usd"}, + } + + def list_jobs(self, params): + """GET /jobs, newest first, with the service's paging keys.""" + limit = int(params.get("limit", 25)) + jobs = sorted(self.jobs.values(), key=lambda j: j["submitted_at"], reverse=True) + view = [ + {k: v for k, v in j.items() if k not in ("polls", "since", "queue")} + for j in jobs[:limit] + ] + return 200, {"jobs": view, "next": None} def cancel_job(self, jid): job = self.jobs.get(jid) @@ -106,21 +198,42 @@ def create_session(self, body): if self.no_sessions: return 404, {"error": "not found"} sid = str(uuid.uuid4()) + settings = body.get("settings") or {} + minutes = settings.get("duration_limit_min") + expires = time.time() + 60 * float(minutes) if minutes else None self.sessions[sid] = { "id": sid, "backend": body.get("backend"), - "settings": body.get("settings") or {}, + "settings": settings, "status": "created", "active": False, "created_at": _stamp(), "started_at": None, "ended_at": None, + "expires_at": _stamp(expires) if expires else None, + "expires": expires, } - return 200, dict(self.sessions[sid]) + return 200, self._session_view(sid) - def get_session(self, sid): + def _session_view(self, sid): + return {k: v for k, v in self.sessions[sid].items() if k != "expires"} + + def _expire(self, sid): + """A session past its duration limit ends, as the service's does.""" s = self.sessions.get(sid) - return (404, {"error": "session not found"}) if s is None else (200, dict(s)) + if ( + s + and s["expires"] + and s["status"] != "ended" + and time.time() >= s["expires"] + ): + self.end_session(sid) + + def get_session(self, sid): + if sid not in self.sessions: + return 404, {"error": "session not found"} + self._expire(sid) + return 200, self._session_view(sid) def end_session(self, sid): s = self.sessions.get(sid) @@ -133,15 +246,44 @@ def end_session(self, sid): "failed", ): job["status"] = "canceled" - return 200, dict(s) + return 200, self._session_view(sid) + + # rates for GET /jobs/estimate, shaped like the service's answer. The + # minimum is what a direct account was quoted for a small circuit, and + # a small circuit never gets above it, so that is the price of a job + RATE = {"cost_model": "2QGE_operations", "job_cost_minimum": 25.79, "fake": True} + + def estimate(self, params): + shots = int(params.get("shots", 1)) + cost = self.RATE["job_cost_minimum"] + return 200, { + "input_values": { + "backend": params.get("backend"), + "type": "ionq.circuit.v1", + "qubits": int(params.get("qubits", 1)), + "shots": shots, + "1q_gates": int(params.get("1q_gates", 0)), + "2q_gates": int(params.get("2q_gates", 0)), + "error_mitigation": params.get("error_mitigation", "false"), + }, + "estimated_at": _stamp(), + "estimated_total_cost": round(cost, 4), + "cost_unit": "usd", + "rate_information": dict(self.RATE), + } def backend(self, name): + lo, hi = self.queue + if hi > 0: + queue_time = (lo + hi) / 2 + else: + queue_time = 0 if name == "simulator" else 600 return 200, { "backend": name, "status": "available", "degraded": False, "qubits": 36, - "average_queue_time": 0 if name == "simulator" else 600, + "average_queue_time": queue_time, } @@ -167,8 +309,12 @@ def _route(self, method): auth = self.headers.get("Authorization") or "" if not auth.startswith("apiKey ") or len(auth) <= len("apiKey "): return self._reply(401, {"error": "unauthorized"}) - parts = [p for p in self.path.split("?")[0].split("/") if p] + path, _, query = self.path.partition("?") + parts = [p for p in path.split("/") if p] + params = {k: v[0] for k, v in parse_qs(query).items()} body = self._body() if method in ("POST", "PUT") else {} + if parts in (["jobs", "estimate"], ["jobs"]) and method == "GET": + body = params with state.lock: state.requests.append((method, "/" + "/".join(parts), body)) code, payload = self._dispatch(method, parts, body) @@ -179,6 +325,12 @@ def _dispatch(self, method, parts, body): return state.backend(parts[1]) if parts == ["jobs"] and method == "POST": return state.create_job(body) + if parts == ["jobs", "estimate"] and method == "GET": + return state.estimate(body) + if parts == ["jobs"] and method == "GET": + return state.list_jobs(body) + if parts[:1] == ["jobs"] and parts[2:] == ["cost"] and method == "GET": + return state.job_cost(parts[1]) if parts[:1] == ["jobs"] and len(parts) == 2 and method == "GET": return state.get_job(parts[1]) if parts[:1] == ["jobs"] and parts[2:] == ["results"] and method == "GET": @@ -191,6 +343,11 @@ def _dispatch(self, method, parts, body): return state.get_session(parts[1]) if parts[:1] == ["sessions"] and parts[2:] == ["end"] and method == "POST": return state.end_session(parts[1]) + # not IonQ's: the fake's own knobs, for a sweep + if parts == ["fake", "config"] and method == "GET": + return 200, state.config() + if parts == ["fake", "config"] and method == "POST": + return state.configure(body) return 404, { "error": "no such route: {} {}".format(method, "/".join(parts)) } @@ -207,27 +364,62 @@ def do_PUT(self): return Handler -def serve(port=0, state=None): +def serve(port=0, state=None, bind="127.0.0.1"): """Start serving in a thread. Returns (server, state). The port is server.server_address[1], which matters when 0 was asked for.""" state = state or State( no_sessions=bool(os.environ.get("FAKE_IONQ_NO_SESSIONS")), steps=os.environ.get("FAKE_IONQ_STEPS", 1), seconds=os.environ.get("FAKE_IONQ_SECONDS", 0), + queue=os.environ.get("FAKE_IONQ_QUEUE"), ) - server = HTTPServer(("127.0.0.1", port), make_handler(state)) + server = HTTPServer((bind, port), make_handler(state)) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() return server, state +def advertised(bind): + """The address to publish for a bind address. The loopback and a real + address are themselves. 0.0.0.0 is every interface, so publish the one + the default route leaves by, which is what the other nodes reach.""" + if bind not in ("0.0.0.0", ""): + return bind + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + # no packet is sent, this only picks the interface for the route + s.connect(("10.255.255.255", 1)) + return s.getsockname()[0] + except OSError: + return "127.0.0.1" + finally: + s.close() + + def main(): ap = argparse.ArgumentParser() ap.add_argument("--port", type=int, default=8765) + ap.add_argument( + "--bind", + default="127.0.0.1", + help="address to listen on. 0.0.0.0 for every interface, so other " + "nodes of the instance reach it (default: the loopback)", + ) args = ap.parse_args() - server, _ = serve(args.port) + server, state = serve(args.port, bind=args.bind) + print( + "fake ionq listening on http://%s:%d" + % (advertised(args.bind), server.server_address[1]), + flush=True, + ) print( - "fake ionq listening on http://127.0.0.1:%d" % server.server_address[1], + "fake ionq: queue %ss, %gs per state, %d poll(s) per state%s" + % ( + state.config()["queue"], + state.seconds, + state.steps, + ", no sessions" if state.no_sessions else "", + ), flush=True, ) try: diff --git a/flux_quantum/scout.py b/flux_quantum/scout.py index a83e3ea..0789ff4 100644 --- a/flux_quantum/scout.py +++ b/flux_quantum/scout.py @@ -20,10 +20,10 @@ # takes an injectable so the unit tests can run without them. try: import flux - from flux.job import JobID, event_wait + from flux.job import JobID, event_watch_async from flux.job import cancel as flux_cancel except ImportError: - flux = JobID = event_wait = flux_cancel = None + flux = JobID = event_watch_async = flux_cancel = None def post_session(handle, jobid, session, rpc=None): @@ -40,10 +40,55 @@ def post_session(handle, jobid, session, rpc=None): def wait_for_job(handle, jobid, waiter=None): - """Block until the job reaches clean.""" - waiter = waiter or event_wait - # a failed job still reaches clean, and we close the session either way - waiter(handle, jobid, "clean", raiseJobException=False) + """Block until the job reaches clean. A failed job still reaches clean, + and the session is closed either way. + + SIGTERM while waiting unwinds with SystemExit so the caller's finally + closes the session. The bindings' synchronous wait blocks inside the + reactor, where a Python signal handler never gets to run, and a + cancelled scout sat there until flux killed it with the session still + open. So the watch runs in the reactor with a signal watcher beside it. + """ + if waiter is not None: + waiter(handle, jobid, "clean", raiseJobException=False) + return + + outcome = {} + + def on_event(future, *_): + try: + event = future.get_event() + except Exception as e: # the eventlog is gone, or the job is unknown + outcome["error"] = e + handle.reactor_stop() + return + if event is None or event.name == "clean": + outcome["done"] = True + handle.reactor_stop() + + def on_signal(h, _watcher, signum, _args): + outcome["signal"] = signum + h.reactor_stop() + + future = event_watch_async(handle, jobid) + future.then(on_event) + watcher = handle.signal_watcher_create(signal.SIGTERM, on_signal) + watcher.start() + try: + handle.reactor_run() + finally: + watcher.stop() + # the reactor's handler is gone with the watcher, so the Python one + # covers the close that follows + _install_signal_handlers() + try: + future.cancel() + except Exception: + pass + if "signal" in outcome: + raise SystemExit("quantum-scout: received signal {}".format(outcome["signal"])) + if "error" in outcome: + raise outcome["error"] def abort_held(handle, jobid, why, cancel=None): @@ -96,6 +141,12 @@ def main(): ) args = ap.parse_args() + # what the scout says has to survive a kill, so no block buffering + try: + sys.stdout.reconfigure(line_buffering=True) + except (AttributeError, ValueError): + pass + if flux is None: sys.exit( "quantum-scout: the flux bindings are not importable, run under flux python" @@ -128,15 +179,19 @@ def main(): # opening is not the same as having the device, so wait until it is ours # before letting the classical job start + # say which hold is being released, so a run can check the service + # afterwards even when nothing was ever memoed to the classical try: ok, why = backend.wait_for_priority(opts) except Exception as e: backend.close_session(session) + print("quantum-scout: closed {} session {}".format(args.vendor, session)) abort_held( h, jobid, "waiting for {} priority failed: {}".format(args.vendor, e) ) if not ok: backend.close_session(session) + print("quantum-scout: closed {} session {}".format(args.vendor, session)) abort_held(h, jobid, "{} never became ours: {}".format(args.vendor, why)) print("quantum-scout: {}".format(why)) # what the classical job needs may only exist once the device is diff --git a/tests/integration/test-scout-cancel.sh b/tests/integration/test-scout-cancel.sh new file mode 100755 index 0000000..b682afe --- /dev/null +++ b/tests/integration/test-scout-cancel.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# A cancelled scout must get to its finally and close the session. The wait +# on the classical job runs in the reactor, where a SIGTERM used to be +# ignored until flux escalated to SIGKILL. This drives the wait directly, +# so it needs only a flux instance, no scheduler or plugin. +# +# flux start -s1 bash tests/integration/test-scout-cancel.sh +set -u +HERE=$(cd "$(dirname "$0")/../.." && pwd) +export PYTHONPATH="$HERE${PYTHONPATH:+:$PYTHONPATH}" +rc=0 +LOG=$(mktemp) +trap 'rm -f "$LOG"' EXIT + +waiter() { + exec flux python - "$1" <<'PY' +import sys, flux +from flux.job import JobID +from flux_quantum.scout import wait_for_job +try: + wait_for_job(flux.Flux(), int(JobID(sys.argv[1]))) +except SystemExit as e: + print("unwound: %s" % e, flush=True) + raise +print("returned", flush=True) +PY +} + +echo "=== 1. SIGTERM while waiting unwinds, promptly ===" +job=$(flux submit sleep 120) +waiter "$job" >"$LOG" 2>&1 & +pid=$! +sleep 3 # the bindings take a moment to import, the watcher is not up before +t0=$(date +%s%N) +kill -TERM "$pid" +wait "$pid"; wrc=$? +dt=$(( ($(date +%s%N) - t0) / 1000000 )) +cat "$LOG" | sed 's/^/ /' +grep -q "unwound: quantum-scout: received signal 15" "$LOG" && echo " ok unwound on SIGTERM" \ + || { echo " FAIL did not unwind on SIGTERM (rc=$wrc)"; rc=1; } +[ "$dt" -lt 3000 ] && echo " ok within ${dt}ms" || { echo " FAIL took ${dt}ms"; rc=1; } +flux cancel "$job" 2>/dev/null + +echo "=== 2. the wait returns when the job reaches clean ===" +job=$(flux submit true) +if timeout 30 bash -c "$(declare -f waiter); waiter $job" >"$LOG" 2>&1; then + grep -q "^returned" "$LOG" && echo " ok returned after clean" || { echo " FAIL no return"; cat "$LOG"; rc=1; } +else + echo " FAIL the wait did not return"; cat "$LOG"; rc=1 +fi + +echo "=== 3. a job that already finished does not block ===" +flux job wait-event -t 10 "$job" clean >/dev/null 2>&1 +if timeout 10 bash -c "$(declare -f waiter); waiter $job" >"$LOG" 2>&1 && grep -q "^returned" "$LOG"; then + echo " ok returned on a finished job" +else + echo " FAIL blocked or failed on a finished job"; cat "$LOG"; rc=1 +fi +[ "$rc" = 0 ] && echo "=== PASS ===" || echo "=== FAIL ===" +exit "$rc" diff --git a/tests/unit/vendors/ionq/test_backend.py b/tests/unit/vendors/ionq/test_backend.py index b02d032..c5b97d6 100644 --- a/tests/unit/vendors/ionq/test_backend.py +++ b/tests/unit/vendors/ionq/test_backend.py @@ -5,6 +5,8 @@ end to end: the header, the JSON, the paths. """ +import time + import pytest @@ -526,3 +528,71 @@ def test_the_lifecycle_without_sessions_falls_back_in_a_dry_run(fake_ionq, monke assert b.session_id(opened).startswith("job:") b.close_session(opened) assert state.sessions == {} + + +# --------------------------------------------------------------------------- +# the fake's device queue + + +def test_the_fake_queue_holds_a_job_in_submitted_until_its_wait_is_up(): + from flux_quantum.backends.ionq import fake + + state = fake.State(queue="0.3") + _, made = state.create_job({"backend": "simulator"}) + for _ in range(3): + assert state.get_job(made["id"])[1]["status"] == "submitted" + time.sleep(0.35) + assert state.get_job(made["id"])[1]["status"] == "ready" + + +def test_a_started_sessions_jobs_skip_the_fake_queue(): + """The session's first job queues like anyone's. Once it starts the + device is the session's, and the next job is served on arrival.""" + from flux_quantum.backends.ionq import fake + + state = fake.State(queue="0.3") + _, s = state.create_session({"backend": "qpu.forte-1"}) + _, first = state.create_job({"backend": "qpu.forte-1", "session_id": s["id"]}) + assert state.get_job(first["id"])[1]["status"] == "submitted" + time.sleep(0.35) + state.get_job(first["id"]) + state.get_job(first["id"]) + assert state.sessions[s["id"]]["active"] + _, second = state.create_job({"backend": "qpu.forte-1", "session_id": s["id"]}) + assert state.get_job(second["id"])[1]["status"] == "ready" + + +def test_the_fake_queue_is_a_range_and_backends_reports_its_mean(fake_ionq): + from flux_quantum.backends.ionq import Client, fake + + assert fake.parse_queue("30-90") == (30.0, 90.0) + assert fake.parse_queue("30") == (30.0, 30.0) + assert fake.parse_queue("") == (0.0, 0.0) + state = fake.State(queue="30-90") + assert 30 <= state.queue_wait() <= 90 + assert state.backend("qpu.forte-1")[1]["average_queue_time"] == 60 + + url, state = fake_ionq + c = Client("k", url) + assert c.get("/fake/config")["queue"] == "0-0" + assert c.post("/fake/config", {"queue": "5-10"})["queue"] == "5-10" + assert state.queue == (5.0, 10.0) + assert c.get("/backends/simulator")["average_queue_time"] == 7.5 + made = c.post("/jobs", {"backend": "simulator"}) + assert "queue" not in c.get("/jobs/%s" % made["id"]) + + +def test_a_cost_limit_from_the_environment_goes_on_the_session(fake_ionq, monkeypatch): + from flux_quantum.backends.ionq import IonQBackend + + url, state = fake_ionq + monkeypatch.setenv("IONQ_API_KEY", "k") + monkeypatch.setenv("IONQ_API_URL", url) + monkeypatch.setenv("FLUX_QUANTUM_IONQ_COST_LIMIT_USD", "25") + b = IonQBackend() + b.open_session(b.scout_options({"device": "qpu.forte-1"})) + posted = [ + body for m, path, body in state.requests if (m, path) == ("POST", "/sessions") + ] + assert posted[0]["settings"]["cost_limit"] == {"unit": "usd", "value": 25.0} + assert posted[0]["settings"]["duration_limit_min"] == 15 From 04356f43c1472811938320c9f514f8e396bb4e18 Mon Sep 17 00:00:00 2001 From: vsoch Date: Thu, 8 Oct 2026 09:36:13 -0700 Subject: [PATCH 22/26] fix: dryrun bug should not submit Signed-off-by: vsoch --- flux_quantum/backends/ionq/__init__.py | 85 +++++++++++++--- flux_quantum/backends/ionq/fake.py | 97 +++++++++++++++--- flux_quantum/cli.py | 129 +++++++++++++++--------- tests/unit/test_cli_plugin.py | 74 +++++++++++--- tests/unit/vendors/ionq/test_backend.py | 117 +++++++++++++++++++++ 5 files changed, 412 insertions(+), 90 deletions(-) diff --git a/flux_quantum/backends/ionq/__init__.py b/flux_quantum/backends/ionq/__init__.py index 545f35f..194cc46 100644 --- a/flux_quantum/backends/ionq/__init__.py +++ b/flux_quantum/backends/ionq/__init__.py @@ -6,9 +6,19 @@ Two ways to hold the device, set with --quantum-hold. session is the default. It creates a session on the backend and submits a one -qubit warm-up job inside it. The scout is ready once IonQ reports the session -active, which is when jobs carrying the session id are being served. The -session id is what the classical job is handed. Sessions are in beta and only +qubit warm-up job inside it. IonQ marks the session started when that first +job is submitted, not when it runs, so the scout is ready when the warm-up +job itself has started: that is when the device is serving the session. The +session id is what the classical job is handed. + +What IonQ has told us about sessions, which the design leans on: a session +does not start by itself, it needs a job. Its jobs go into a separate queue +with about half the device's capacity, so they wait on the order of a minute +or two rather than the public queue's hours. Any session that runs a job is +billed a one hour minimum of device time when it ends, which no cost_limit +caps, so one session per classical job is the expensive way to use them. +duration_limit_min counts from creation. Creation returns 409 when the +device's session capacity is full, and the scout retries that with backoff. Sessions are in beta and only some accounts have them, and a refusal says so plainly. probe submits the warm-up job alone and is ready once it has started. That @@ -23,6 +33,16 @@ Tuning, from the environment: FLUX_QUANTUM_IONQ_SHOTS shots for the warm-up job, 100 + FLUX_QUANTUM_IONQ_WARMUP 1, the default, sends a warm-up job into + the session and is ready when it starts. + 0 sends nothing and waits on the session + status alone, which IonQ's sessions do + not reach by themselves, so it is for a + service that queues the session itself. + A probe hold always sends one + FLUX_QUANTUM_IONQ_SESSION_RETRIES how many times to retry a 409 on session + creation, the device's session capacity + being full, 20, with backoff from 15s FLUX_QUANTUM_IONQ_COST_LIMIT_USD a cost limit on every session opened, none by default. IonQ ends the session when its jobs reach it, so a session @@ -152,7 +172,7 @@ def open_session(self, options): if limit: body["settings"]["cost_limit"] = {"unit": "usd", "value": float(limit)} try: - self._session = self.client.post("/sessions", body)["id"] + self._session = self._create_session(body)["id"] except APIError as e: if e.status not in (400, 401, 403, 404): raise @@ -168,11 +188,38 @@ def open_session(self, options): flush=True, ) else: - self._job = self._warmup(options, session=self._session) + if truthy(tuning("ionq_warmup", "1")): + self._job = self._warmup(options, session=self._session) + else: + self._job = None + print( + "ionq: session created, no warm-up, waiting on the session", + flush=True, + ) return self._session self._job = self._warmup(options) return self._job + def _create_session(self, body, sleep=time.sleep): + """POST /sessions, retrying a 409. IonQ answers 409 when the + device's session capacity is full, and a slot comes free when + someone's session ends, so wait and try again, backing off.""" + retries = int(tuning("ionq_session_retries", 20)) + wait = 15.0 + for attempt in range(retries + 1): + try: + return self.client.post("/sessions", body) + except APIError as e: + if e.status != 409 or attempt == retries: + raise + print( + "ionq: session capacity on {} is full (HTTP 409), retry {} of {} " + "in {:.0f}s".format(body["backend"], attempt + 1, retries, wait), + flush=True, + ) + sleep(wait) + wait = min(wait * 1.5, 300.0) + def session_id(self, opened): """A session id goes over as is, since jobs are submitted into it. A probe job's id is prefixed, so the classical job can tell it holds @@ -184,17 +231,19 @@ def session_id(self, opened): def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep): """Poll until the hold is real. - With a session, that is IonQ saying it is active, or the warm-up job - inside it having started, whichever is reported first. With a probe, - the warm-up job having started. A warm-up that fails or is cancelled - is a failure, not readiness. + With a session and a warm-up job, that is the warm-up job having + started: IonQ marks the session started as soon as its first job is + submitted, so the session's status says nothing about the device + serving it yet. Without a warm-up, the session's own status is all + there is. With a probe, the warm-up job having started. A warm-up + that fails or is cancelled is a failure, not readiness. """ options = options or {} timeout = float(options.get("timeout") or 0) started = self._hold_started or time.time() deadline = started + timeout if timeout > 0 else None while True: - job = self.client.get("/jobs/{}".format(self._job)) + job = self.client.get("/jobs/{}".format(self._job)) if self._job else {} status = job.get("status") waited = time.time() - started if status in ("failed", "canceled"): @@ -208,12 +257,18 @@ def wait_for_priority(self, options=None, interval=5.0, sleep=time.sleep): return False, "session {} before it became active".format( s.get("status") or "ended" ) - if s.get("active") or s.get("started_at") or status in STARTED: + if self._job: + ready = status in STARTED + else: + ready = bool(s.get("active") or s.get("started_at")) + if ready: self._hold_ready_after = waited - return True, "session active after {:.0f}s".format(waited) - said = "session {}, warm-up job {}".format( - s.get("status") or "pending", status - ) + return True, "session {} and being served after {:.0f}s".format( + s.get("status") or "started", waited + ) + said = "session {}".format(s.get("status") or "pending") + if self._job: + said += ", warm-up job {}".format(status) else: if status in STARTED: self._hold_ready_after = waited diff --git a/flux_quantum/backends/ionq/fake.py b/flux_quantum/backends/ionq/fake.py index 83ebbc8..3632fa1 100755 --- a/flux_quantum/backends/ionq/fake.py +++ b/flux_quantum/backends/ionq/fake.py @@ -22,13 +22,24 @@ time a job spends in each state, default 0, for timings that look like a device's. -FAKE_IONQ_QUEUE is the device's queue: the seconds a job sits in submitted -before it is served, a number or a range like 30-90 drawn per job. A job -in a session that has started skips it, which is what a session buys: the -session's first job queues like anyone's, and once it starts the device is -the session's. /backends reports the mean as average_queue_time. The queue +FAKE_IONQ_QUEUE is the device's public queue: the seconds a job sits in +submitted before it is served, a number or a range like 30-90 drawn per job. +A job in a session that has started waits FAKE_IONQ_SESSION_QUEUE instead, +the session queue, 0 by default as the fake campaign ran. IonQ's sessions +do not skip the queue: their jobs go into a separate queue with about half +the device's capacity and wait a minute or two where the public queue is +hours, so a faithful run sets both, say FAKE_IONQ_QUEUE=300 and +FAKE_IONQ_SESSION_QUEUE=1-2, at whatever scale fits the time available. /backends reports the mean as average_queue_time. The queue can be changed while running, POST /fake/config {"queue": "30-90"}, so an experiment can sweep it. + +FAKE_IONQ_SESSION_START says how a session reaches started. job, the default +and what the fake campaign ran with, starts it when its first job starts. +submit is what IonQ does: the session is started the moment its first job +is submitted, and the job then waits its turn, so a scout that trusts the +session's status releases too early. queue starts the session by itself, +one queue wait after it is created, with no job in it, which IonQ's do not; +it models a service whose sessions queue alone. """ import argparse @@ -68,26 +79,42 @@ def parse_queue(spec): class State: - def __init__(self, no_sessions=False, steps=1, seconds=0.0, queue=None): + def __init__( + self, + no_sessions=False, + steps=1, + seconds=0.0, + queue=None, + session_start="job", + session_queue=None, + ): self.jobs = {} self.sessions = {} self.no_sessions = no_sessions self.steps = max(1, int(steps)) self.seconds = float(seconds) self.queue = parse_queue(queue) + self.session_queue = parse_queue(session_queue) + self.session_start = session_start or "job" + if self.session_start not in ("job", "submit", "queue"): + raise ValueError( + "session_start is job, submit or queue, not %r" % self.session_start + ) self.requests = [] self.lock = threading.Lock() - def queue_wait(self): - lo, hi = self.queue + def queue_wait(self, which=None): + lo, hi = which if which is not None else self.queue return random.uniform(lo, hi) if hi > lo else lo def config(self): return { "queue": "%g-%g" % self.queue, + "session_queue": "%g-%g" % self.session_queue, "seconds": self.seconds, "steps": self.steps, "no_sessions": self.no_sessions, + "session_start": self.session_start, } def configure(self, body): @@ -95,12 +122,18 @@ def configure(self, body): submitted keep the wait they were given.""" if "queue" in body: self.queue = parse_queue(body["queue"]) + if "session_queue" in body: + self.session_queue = parse_queue(body["session_queue"]) if "seconds" in body: self.seconds = float(body["seconds"]) if "steps" in body: self.steps = max(1, int(body["steps"])) if "no_sessions" in body: self.no_sessions = bool(body["no_sessions"]) + if "session_start" in body: + if body["session_start"] not in ("job", "submit", "queue"): + return 400, {"error": "session_start is job, submit or queue"} + self.session_start = body["session_start"] return 200, self.config() def create_job(self, body): @@ -110,10 +143,17 @@ def create_job(self, body): return 404, {"error": "session not found"} if session: self._expire(session) + self._start(session) if self.sessions[session]["status"] == "ended": return 400, {"error": "session has ended"} # the device's queue, unless a started session already has the device served = bool(session and self.sessions[session]["active"]) + # IonQ marks the session started on its first submission. The job + # still waits its turn, so this does not make it served + if session and self.session_start == "submit": + sess = self.sessions[session] + if sess["status"] == "created": + sess["status"], sess["started_at"] = "started", _stamp() self.jobs[jid] = { "id": jid, "status": "submitted", @@ -128,7 +168,9 @@ def create_job(self, body): "execution_duration_ms": None, "polls": 0, "since": time.time(), - "queue": 0.0 if served else self.queue_wait(), + "queue": ( + self.queue_wait(self.session_queue) if served else self.queue_wait() + ), } return 200, {"id": jid, "status": "submitted", "session_id": session} @@ -212,11 +254,32 @@ def create_session(self, body): "ended_at": None, "expires_at": _stamp(expires) if expires else None, "expires": expires, + # when the session starts by itself, under session_start=queue + "starts": ( + time.time() + self.queue_wait() + if self.session_start == "queue" + else None + ), } return 200, self._session_view(sid) def _session_view(self, sid): - return {k: v for k, v in self.sessions[sid].items() if k != "expires"} + return { + k: v + for k, v in self.sessions[sid].items() + if k not in ("expires", "starts") + } + + def _start(self, sid): + """A session queued for the device starts once its wait is up.""" + s = self.sessions.get(sid) + if ( + s + and s["status"] == "created" + and s.get("starts") is not None + and time.time() >= s["starts"] + ): + s["active"], s["status"], s["started_at"] = True, "started", _stamp() def _expire(self, sid): """A session past its duration limit ends, as the service's does.""" @@ -233,6 +296,7 @@ def get_session(self, sid): if sid not in self.sessions: return 404, {"error": "session not found"} self._expire(sid) + self._start(sid) return 200, self._session_view(sid) def end_session(self, sid): @@ -255,7 +319,12 @@ def end_session(self, sid): def estimate(self, params): shots = int(params.get("shots", 1)) - cost = self.RATE["job_cost_minimum"] + # the simulator is free; a small circuit on a QPU is the minimum + cost = ( + 0.0 + if params.get("backend") == "simulator" + else self.RATE["job_cost_minimum"] + ) return 200, { "input_values": { "backend": params.get("backend"), @@ -372,6 +441,8 @@ def serve(port=0, state=None, bind="127.0.0.1"): steps=os.environ.get("FAKE_IONQ_STEPS", 1), seconds=os.environ.get("FAKE_IONQ_SECONDS", 0), queue=os.environ.get("FAKE_IONQ_QUEUE"), + session_queue=os.environ.get("FAKE_IONQ_SESSION_QUEUE"), + session_start=os.environ.get("FAKE_IONQ_SESSION_START") or "job", ) server = HTTPServer((bind, port), make_handler(state)) thread = threading.Thread(target=server.serve_forever, daemon=True) @@ -413,11 +484,13 @@ def main(): flush=True, ) print( - "fake ionq: queue %ss, %gs per state, %d poll(s) per state%s" + "fake ionq: queue %ss, session queue %ss, %gs per state, %d poll(s) per state, sessions start on %s%s" % ( state.config()["queue"], + state.config()["session_queue"], state.seconds, state.steps, + state.session_start, ", no sessions" if state.no_sessions else "", ), flush=True, diff --git a/flux_quantum/cli.py b/flux_quantum/cli.py index e56623b..973076f 100644 --- a/flux_quantum/cli.py +++ b/flux_quantum/cli.py @@ -42,43 +42,32 @@ def _wrap_commands(jobspec_dict, wrap_path): def _check_pair_fits(handle, classical, scout, rpc_fn=None): """Refuse a pair the graph can never hold, before either half exists. - The match_coschedule check in the add-hold branch is a satisfiability - test. Each half is checked alone against the whole graph with nothing - else on it, so this catches a missing qdevice or a half bigger than the - machine, not a busy machine. Fit against what is reachable right now is - the jobtap plugin's admission check. + Each half is checked alone against the whole graph through fluxion's + feasibility.check, which is what the job manager asks before it accepts + a job. This catches a missing qdevice or a half bigger than the machine, + not a busy machine. Fit against what is reachable right now is the + jobtap plugin's admission check. """ - payload = { - "check": True, - "jobs": [ - {"jobspec": scout, "op": "allocate"}, - {"jobspec": classical, "op": "reserve"}, - ], - } try: - if rpc_fn is None: - handle.rpc("sched-fluxion-resource.match_coschedule", payload).get() - else: - rpc_fn(handle, payload) + for half in (scout, classical): + payload = {"jobspec": half} + if rpc_fn is None: + handle.rpc("feasibility.check", payload).get() + else: + rpc_fn(handle, payload) except OSError as exc: - # ENODEV means the graph can never hold the pair. EBUSY would mean a - # fluxion that checks the live graph found no room now. Anything else - # means the question could not be asked, and that must not block a - # submit that worked before the check existed. - if exc.errno not in (errno.EBUSY, errno.ENODEV): - if exc.errno in (errno.EINVAL, errno.EPROTO): - print( - "flux quantum: WARNING fluxion did not understand the pair " - "check, skipping it: {}".format(exc), - file=sys.stderr, - ) + # ENODEV and EINVAL from feasibility.check mean the graph can never + # hold this half. Anything else means the question could not be + # asked, and that must not block a submit that worked before the + # check existed. + if exc.errno not in (errno.EBUSY, errno.ENODEV, errno.EINVAL): return raise SystemExit( - "flux quantum: this pair cannot be scheduled together on this " - "cluster. The classical half needs its cores and the scout needs " - "one more plus the vendor device, and the graph cannot hold both. " - "Submitting anyway would open a vendor session for a pair that " - "cannot run. ({})".format(exc) + "flux quantum: this pair cannot be scheduled on this cluster. The " + "classical half needs its cores and the scout needs one more plus " + "the vendor device, and the graph cannot hold both. Submitting " + "anyway would open a vendor session for a pair that cannot run. " + "({})".format(exc) ) except Exception: # no usable handle, so there is nothing to ask @@ -96,6 +85,17 @@ def _safe_cancel(cancel_fn, handle, jobid, reason): ) +def flux_dry_run(args): + """True when flux's own --dry-run is set, as opposed to --quantum-dry-run. + + Inside a plugin callback flux hands us a proxy that aliases the bare + option names to ours, so args.dry_run is --quantum-dry-run there. Flux's + flag lives on the namespace underneath the proxy. + """ + ns = getattr(args, "_ns", args) + return bool(getattr(ns, "dry_run", False)) + + def prepare_pair( handle, jobspec, @@ -109,10 +109,17 @@ def prepare_pair( cancel_fn=None, wrap_path=None, scout_path=None, + dry_run=False, ): """Submit the user work held, then rewrite the jobspec in place into the scout that releases it. Returns the id of the held job. + With dry_run nothing is submitted. The classical half is printed to + stderr, the jobspec is still rewritten into the scout with a placeholder + job id of 0, and flux prints that as it would any dry run. Without this, + flux's --dry-run skipped only flux's own submit and left a held job + parked with no scout to ever release it. + flux submits whatever is left in the jobspec, so it submits the scout. The flux calls are injectable so this is testable without a broker. """ @@ -172,17 +179,31 @@ def prepare_pair( if probe is not None: _check_pair_fits(handle, classical, probe) - # if feasibility validation is on, an unsatisfiable request is rejected here - # and we abort before spending any quantum quota - try: - main_id = int(submit_fn(handle, json.dumps(classical))) - except SystemExit: - raise - except Exception as exc: - raise SystemExit( - "flux quantum: classical request not accepted " - "(unsatisfiable, or ingest error): {}".format(exc) + if dry_run: + print( + "flux quantum: dry run, nothing submitted. The held classical job " + "would be:", + file=sys.stderr, + ) + print(json.dumps(classical, indent=2), file=sys.stderr) + print( + "flux quantum: and the scout, printed below by flux, would release " + "it. Its --job 0 stands in for the classical job id.", + file=sys.stderr, ) + main_id = 0 + else: + # if feasibility validation is on, an unsatisfiable request is rejected + # here and we abort before spending any quantum quota + try: + main_id = int(submit_fn(handle, json.dumps(classical))) + except SystemExit: + raise + except Exception as exc: + raise SystemExit( + "flux quantum: classical request not accepted " + "(unsatisfiable, or ingest error): {}".format(exc) + ) try: scout = build_scout_jobspec( @@ -197,12 +218,13 @@ def prepare_pair( except SystemExit: raise except Exception as exc: - _safe_cancel( - cancel_fn, - handle, - main_id, - "flux quantum: aborting held classical (scout build failed)", - ) + if not dry_run: + _safe_cancel( + cancel_fn, + handle, + main_id, + "flux quantum: aborting held classical (scout build failed)", + ) raise SystemExit( "flux quantum: could not build the scout jobspec: {}".format(exc) ) @@ -389,9 +411,18 @@ def modify_jobspec(self, args, jobspec): raise SystemExit("flux quantum: {}".format(e)) job_env = backend.job_environment(options) if backend else {} handle = flux.Flux() + dry_run = flux_dry_run(args) main_id = prepare_pair( - handle, jobspec, vendor, scout_cores=1, options=options, job_env=job_env + handle, + jobspec, + vendor, + scout_cores=1, + options=options, + job_env=job_env, + dry_run=dry_run, ) + if dry_run: + return print( "flux quantum: held classical job {} (vendor={}); this submit " "launches its scout".format(main_id, vendor), diff --git a/tests/unit/test_cli_plugin.py b/tests/unit/test_cli_plugin.py index eac411c..f5283d6 100644 --- a/tests/unit/test_cli_plugin.py +++ b/tests/unit/test_cli_plugin.py @@ -392,35 +392,31 @@ def test_check_pair_fits_rejects_when_the_pair_cannot_be_placed(stub_flux_cli): from flux_quantum import cli def rpc(handle, payload): - raise OSError(errno.EBUSY, "the group does not fit right now") + raise OSError(errno.ENODEV, "unsatisfiable request") try: cli._check_pair_fits(None, {"a": 1}, {"b": 2}, rpc_fn=rpc) except SystemExit as exc: - assert "cannot be scheduled together" in str(exc) + assert "cannot be scheduled" in str(exc) else: raise AssertionError("a pair that cannot be placed was allowed through") -def test_check_pair_fits_sends_both_halves_and_their_ops(stub_flux_cli): - """The question is about the pair, so both jobspecs go, the scout to be - allocated and the classical to be reserved. No jobids, since neither job +def test_check_pair_fits_asks_feasibility_for_each_half(stub_flux_cli): + """Both jobspecs are checked, each on its own, through the same + feasibility question the job manager asks. No jobids, since neither job exists yet.""" from flux_quantum import cli - seen = {} + seen = [] def rpc(handle, payload): - seen.update(payload) - return {"fits": True} + seen.append(payload) + return {} cli._check_pair_fits(None, {"classical": 1}, {"scout": 1}, rpc_fn=rpc) - assert seen["check"] is True - ops = [j["op"] for j in seen["jobs"]] - assert sorted(ops) == ["allocate", "reserve"] - assert {"scout": 1} in [j["jobspec"] for j in seen["jobs"]] - assert {"classical": 1} in [j["jobspec"] for j in seen["jobs"]] - assert all("jobid" not in j for j in seen["jobs"]) + assert [p["jobspec"] for p in seen] == [{"scout": 1}, {"classical": 1}] + assert all(set(p) == {"jobspec"} for p in seen) def test_check_pair_fits_is_permissive_when_it_cannot_ask(stub_flux_cli): @@ -465,6 +461,56 @@ def submit(handle, jobspec_json): assert any(x[0] == "submitted" for x in order) +def test_flux_dry_run_submits_nothing(stub_flux_cli, capsys): + """flux --dry-run skips only flux's own submit. The classical half must + not be submitted either, or a held job is left parked with no scout.""" + from flux_quantum import cli + + def submit(handle, jobspec_json): + raise AssertionError("submitted the classical on a dry run") + + def cancel(handle, jobid, reason): + raise AssertionError("cancelled a job that was never submitted") + + js = _fake_jobspec(["myprog"]) + main_id = cli.prepare_pair( + handle=None, + jobspec=js, + vendor="mock", + submit_fn=submit, + populate_fn=lambda h, v: None, + get_graph_fn=lambda h: _live_graph(), + cancel_fn=cancel, + dry_run=True, + ) + assert main_id == 0 + # the classical half is shown, and the jobspec is still the scout so flux + # prints that as the dry run output + err = capsys.readouterr().err + assert "nothing submitted" in err and '"hold": 1' in err + types = [r["type"] for r in js.jobspec["resources"]] + assert any(t.startswith("qdevice_") for t in types) + assert "--job" in js.jobspec["tasks"][0]["command"] + + +def test_flux_dry_run_is_read_from_under_the_proxy(stub_flux_cli): + """Inside a callback args.dry_run is aliased to --quantum-dry-run, so + flux's own flag has to come from the namespace the proxy wraps.""" + from flux_quantum import cli + + class Proxy: + def __init__(self, ns): + object.__setattr__(self, "_ns", ns) + + def __getattr__(self, name): + # the alias flux installs: dry_run -> quantum_dry_run + return getattr(self._ns, {"dry_run": "quantum_dry_run"}.get(name, name)) + + assert cli.flux_dry_run(Proxy(_make_args(dry_run=True, quantum_dry_run=False))) + assert not cli.flux_dry_run(Proxy(_make_args(dry_run=False, quantum_dry_run=True))) + assert not cli.flux_dry_run(_make_args()) + + def test_common_options_have_defaults_and_follow_the_duration( stub_flux_cli, monkeypatch ): diff --git a/tests/unit/vendors/ionq/test_backend.py b/tests/unit/vendors/ionq/test_backend.py index c5b97d6..2f5282b 100644 --- a/tests/unit/vendors/ionq/test_backend.py +++ b/tests/unit/vendors/ionq/test_backend.py @@ -596,3 +596,120 @@ def test_a_cost_limit_from_the_environment_goes_on_the_session(fake_ionq, monkey ] assert posted[0]["settings"]["cost_limit"] == {"unit": "usd", "value": 25.0} assert posted[0]["settings"]["duration_limit_min"] == 15 + + +def test_without_a_warmup_the_scout_waits_on_the_session_alone(monkeypatch): + """A session that takes its place in the queue when it is created needs + no job to start, so the scout sends none and pays for none.""" + from flux_quantum.backends.ionq import Client, IonQBackend, fake + + state = fake.State(queue="0.3", session_start="queue") + server, _ = fake.serve(0, state=state) + try: + url = "http://127.0.0.1:%d" % server.server_address[1] + monkeypatch.setenv("IONQ_API_KEY", "k") + monkeypatch.setenv("IONQ_API_URL", url) + monkeypatch.setenv("FLUX_QUANTUM_IONQ_WARMUP", "0") + b = IonQBackend() + opts = b.scout_options({"device": "qpu.forte-1"}) + sid = b.open_session(opts) + assert state.jobs == {} + ok, why = b.wait_for_priority(opts, interval=0.05) + assert ok, why + assert state.sessions[sid]["status"] == "started" and state.jobs == {} + # the session's first job is served on arrival + c = Client("k", url) + j = c.post("/jobs", {"backend": "qpu.forte-1", "session_id": sid}) + assert c.get("/jobs/%s" % j["id"])["status"] == "ready" + finally: + server.shutdown() + + +def test_the_default_session_start_is_the_first_job_as_the_campaign_ran(): + from flux_quantum.backends.ionq import fake + + state = fake.State(queue="0") + _, s = state.create_session({"backend": "qpu.forte-1"}) + assert state.get_session(s["id"])[1]["status"] == "created" + + +def test_the_scout_waits_for_the_warmup_to_run_not_for_the_session_status(monkeypatch): + """IonQ marks a session started when its first job is submitted. The + device is serving the session only when that job runs, so that is what + the scout waits for, or it would release the classical job too early.""" + from flux_quantum.backends.ionq import IonQBackend, fake + + state = fake.State(queue="0.3", session_start="submit") + server, _ = fake.serve(0, state=state) + try: + monkeypatch.setenv("IONQ_API_KEY", "k") + monkeypatch.setenv( + "IONQ_API_URL", "http://127.0.0.1:%d" % server.server_address[1] + ) + monkeypatch.delenv("FLUX_QUANTUM_IONQ_WARMUP", raising=False) + b = IonQBackend() + opts = b.scout_options({"device": "qpu.forte-1"}) + sid = b.open_session(opts) + assert state.sessions[sid]["status"] == "started" # at submission, as IonQ does + (job,) = state.jobs.values() + assert job["status"] == "submitted" + t = time.time() + ok, why = b.wait_for_priority(opts, interval=0.05) + assert ok, why + assert time.time() - t >= 0.25 # it waited for the job, not the status + assert state.jobs[job["id"]]["status"] in ("started", "completed") + finally: + server.shutdown() + + +def test_session_creation_retries_a_409_with_backoff(): + from flux_quantum.backends.ionq import APIError, IonQBackend + + class _Client: + url = "http://fake" + + def __init__(self): + self.calls = 0 + + def post(self, path, body=None): + self.calls += 1 + if self.calls < 3: + raise APIError(409, "POST", path, b"session capacity full") + return {"id": "s-1"} + + waits = [] + b = IonQBackend(client=_Client()) + got = b._create_session({"backend": "qpu.forte-1"}, sleep=waits.append) + assert got == {"id": "s-1"} and b.client.calls == 3 + assert waits == [15.0, 22.5] + + +def test_session_creation_gives_up_after_the_retries(monkeypatch): + from flux_quantum.backends.ionq import APIError, IonQBackend + + class _Client: + url = "http://fake" + + def post(self, path, body=None): + raise APIError(409, "POST", path, b"full") + + monkeypatch.setenv("FLUX_QUANTUM_IONQ_SESSION_RETRIES", "2") + with pytest.raises(APIError) as e: + IonQBackend(client=_Client())._create_session( + {"backend": "qpu.forte-1"}, sleep=lambda _: None + ) + assert e.value.status == 409 + + +def test_a_started_sessions_jobs_wait_the_session_queue_not_the_public_one(): + """IonQ's session jobs do not skip the queue, they get a shorter one.""" + from flux_quantum.backends.ionq import fake + + state = fake.State(queue="10", session_queue="0.3", session_start="submit") + _, s = state.create_session({"backend": "qpu.forte-1"}) + state.sessions[s["id"]]["active"] = True # as if its first job had started + _, j = state.create_job({"backend": "qpu.forte-1", "session_id": s["id"]}) + assert state.jobs[j["id"]]["queue"] == 0.3 + _, k = state.create_job({"backend": "qpu.forte-1"}) + assert state.jobs[k["id"]]["queue"] == 10.0 + assert fake.State().session_queue == (0.0, 0.0) # as the campaign ran From 18c14b0198b1643dc31c81c3d7d6541049582d19 Mon Sep 17 00:00:00 2001 From: vsoch Date: Thu, 8 Oct 2026 12:06:53 -0700 Subject: [PATCH 23/26] container: add demo Signed-off-by: vsoch --- .github/workflows/test.yml | 3 +- examples/qrmi/container/Dockerfile | 147 +++++++++ .../qrmi/container/Dockerfile.dockerignore | 15 + examples/qrmi/container/README.md | 281 ++++++++++++++++++ .../qrmi/container/bin/flux-quantum-check | 103 +++++++ examples/qrmi/container/bin/flux-quantum-demo | 133 +++++++++ examples/qrmi/container/bin/flux-quantum-init | 65 ++++ .../qrmi/container/bin/flux-quantum-start | 70 +++++ examples/qrmi/container/coschedule.toml | 19 ++ examples/qrmi/container/workload.py | 129 ++++++++ flux_quantum/backends/__init__.py | 1 + flux_quantum/backends/base.py | 18 ++ flux_quantum/backends/ionq/__init__.py | 6 + flux_quantum/backends/qrmi.py | 18 ++ flux_quantum/cli.py | 52 +++- flux_quantum/scout.py | 270 ++++++++++++----- tests/integration/fluxion.sh | 28 ++ tests/integration/setup-mock-registry.sh | 18 +- tests/integration/test-admission.sh | 4 + tests/integration/test-ionq-e2e.sh | 9 +- tests/integration/test-mock-e2e.sh | 15 +- tests/integration/test-preemption.sh | 4 + tests/integration/test-reload.sh | 4 + tests/integration/test-scout-cancel.sh | 23 +- tests/unit/test_cli_plugin.py | 79 +++++ tests/unit/test_handoff.py | 132 +++++++- tests/unit/vendors/braket/conftest.py | 61 ++++ tests/unit/vendors/braket/test_hold.py | 16 +- tests/unit/vendors/ibm/test_qrmi.py | 22 ++ tests/unit/vendors/ionq/test_backend.py | 25 +- 30 files changed, 1635 insertions(+), 135 deletions(-) create mode 100644 examples/qrmi/container/Dockerfile create mode 100644 examples/qrmi/container/Dockerfile.dockerignore create mode 100644 examples/qrmi/container/README.md create mode 100755 examples/qrmi/container/bin/flux-quantum-check create mode 100755 examples/qrmi/container/bin/flux-quantum-demo create mode 100755 examples/qrmi/container/bin/flux-quantum-init create mode 100755 examples/qrmi/container/bin/flux-quantum-start create mode 100644 examples/qrmi/container/coschedule.toml create mode 100644 examples/qrmi/container/workload.py create mode 100644 tests/integration/fluxion.sh create mode 100644 tests/unit/vendors/braket/conftest.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 39ae700..b976140 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -151,7 +151,8 @@ jobs: 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 diff --git a/examples/qrmi/container/Dockerfile b/examples/qrmi/container/Dockerfile new file mode 100644 index 0000000..f3f0b32 --- /dev/null +++ b/examples/qrmi/container/Dockerfile @@ -0,0 +1,147 @@ +# syntax=docker/dockerfile:1 +# The flux-quantum coscheduling stack in one container, for the QRMI demo. +# docker build -f examples/qrmi/container/Dockerfile -t flux-quantum-qrmi --build-arg FLUX_QUANTUM_REF= . + +FROM ubuntu:24.04 +ENV DEBIAN_FRONTEND=noninteractive LANG=C.UTF-8 + +ARG FLUX_SECURITY_VERSION=0.15.0 +ARG FLUX_CORE_VERSION=0.87.0 +ARG FLUX_SCHED_REPO=https://github.com/vsoch/flux-sched +ARG FLUX_SCHED_BRANCH=hold +# a shallow clone has no tags to describe, so cmake is told a version. It +# only names the build, the branch is what matters. +ARG FLUX_SCHED_VERSION=0.56.0 +ARG FLUX_QUANTUM_REPO=https://github.com/converged-computing/flux-quantum +# the branch to clone, as the cluster's start-script does. Empty means use the +# build context, this checkout, which is the only way to get work that is not +# pushed yet into the image. +ARG FLUX_QUANTUM_REF=braket-hybrid +# pip spec for qrmi. Pin here, e.g. "qrmi[ibm]==0.26.0", for a reproducible image +ARG QRMI_SPEC="qrmi[ibm]" +ARG INSTALL_IBMCLOUD=1 + +# build and runtime dependencies of flux-core and flux-sched, the list from +# the AMI build script, in ubuntu 24.04 package names +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates curl wget git jq vim less sudo apt-utils \ + build-essential pkg-config autotools-dev libtool libffi-dev \ + autoconf automake make gcc g++ cmake ninja-build flex \ + lua5.4 liblua5.4-dev lua-posix \ + python3-dev python3-pip python3-yaml \ + libsodium-dev libzmq3-dev libczmq-dev libjansson-dev \ + munge libmunge-dev libncursesw5-dev liblz4-dev libsqlite3-dev \ + uuid-dev libhwloc-dev libevent-dev libarchive-dev libpam-dev \ + libboost-dev libboost-graph-dev libboost-system-dev \ + libboost-filesystem-dev libboost-regex-dev libyaml-cpp-dev \ + libedit-dev \ + && rm -rf /var/lib/apt/lists/* + +# flux-core's python bindings need cffi and ply, and qrmi later needs a newer +# cffi than debian ships. A debian cffi has no pip RECORD, so pip cannot +# replace it and the qrmi install fails. Install these with pip from the +# start, like the CI image does, so pip owns them and can upgrade them. +RUN python3 -m pip install --break-system-packages --no-cache-dir \ + --upgrade --ignore-installed cffi ply pyyaml jsonschema + +WORKDIR /opt + +# flux-core 0.87 needs flux-security 0.13 or newer +RUN wget -q https://github.com/flux-framework/flux-security/releases/download/v${FLUX_SECURITY_VERSION}/flux-security-${FLUX_SECURITY_VERSION}.tar.gz \ + && tar -xzf flux-security-${FLUX_SECURITY_VERSION}.tar.gz \ + && cd flux-security-${FLUX_SECURITY_VERSION} \ + && ./configure --prefix=/usr --sysconfdir=/etc \ + && make -j"$(nproc)" && make install && ldconfig \ + && cd /opt && rm -rf flux-security-${FLUX_SECURITY_VERSION}* + +# prefix /usr and sysconfdir /etc, like the AMI. flux python is then +# /usr/bin/python3, the python everything below installs into. +RUN wget -q https://github.com/flux-framework/flux-core/releases/download/v${FLUX_CORE_VERSION}/flux-core-${FLUX_CORE_VERSION}.tar.gz \ + && tar -xzf flux-core-${FLUX_CORE_VERSION}.tar.gz \ + && cd flux-core-${FLUX_CORE_VERSION} \ + && ./configure --prefix=/usr --sysconfdir=/etc --runstatedir=/run/flux \ + --with-flux-security --disable-docs \ + && make -j"$(nproc)" && make install && ldconfig \ + && cd /opt && rm -rf flux-core-${FLUX_CORE_VERSION}* + +# CMAKE_INSTALL_SYSCONFDIR must be the absolute /etc. flux-sched derives its +# rc1.d path from it before adopting flux-core's prefix, and a relative "etc" +# lands the fluxion rc script in /usr/etc/flux/rc1.d where flux-core never +# looks, so fluxion silently never loads and the instance runs sched-simple. +RUN git clone --depth 1 -b "${FLUX_SCHED_BRANCH}" "${FLUX_SCHED_REPO}" /opt/flux-sched \ + && cd /opt/flux-sched && git log --oneline -1 \ + && FLUX_SCHED_VERSION=${FLUX_SCHED_VERSION} cmake -B build \ + -DCMAKE_INSTALL_PREFIX=/usr -DCMAKE_INSTALL_SYSCONFDIR=/etc \ + -DCMAKE_BUILD_TYPE=RelWithDebInfo \ + && cmake --build build -j"$(nproc)" \ + && cmake --build build -t install && ldconfig \ + && test -x /etc/flux/rc1.d/01-sched-fluxion \ + && test -f /etc/flux/modprobe/modprobe.d/fluxion.toml \ + && rm -rf /opt/flux-sched/build + +# qrmi into the python flux uses. The flux-quantum ibm backend still speaks +# the pre-0.23 names (ResourceType.IBMQiskitRuntimeService, the QRMI_IBM_QRS_ +# variables), which qrmi keeps as deprecated aliases until 2026-11-21. Fail +# the build now, rather than at the first submit, if the installed qrmi has +# dropped them. +RUN python3 -m pip install --break-system-packages --no-cache-dir "${QRMI_SPEC}" +RUN python3 - <<'PY' +import qrmi, qiskit +from qrmi import ResourceType +print("qrmi", getattr(qrmi, "__version__", "?"), "qiskit", qiskit.__version__) +assert hasattr(ResourceType, "IBMQiskitRuntimeService"), ( + "this qrmi has dropped ResourceType.IBMQiskitRuntimeService. Pin an older " + "qrmi with --build-arg QRMI_SPEC='qrmi[ibm]==' or migrate " + "flux_quantum/backends/qrmi.py to the ibm-quantum-compute-service names") +PY + +RUN if [ "${INSTALL_IBMCLOUD}" = 1 ]; then \ + curl -fsSL https://clis.cloud.ibm.com/install/linux | sh \ + && ibmcloud version; \ + fi + +# the demo runs as a user, the scout and the classical job both run as the +# submitting user and that is where the credentials live +ARG USERNAME=ubuntu +RUN echo "${USERNAME} ALL=(ALL) NOPASSWD: ALL" > /etc/sudoers.d/${USERNAME} \ + && chmod 0440 /etc/sudoers.d/${USERNAME} \ + && mkdir -p /etc/flux/system/jobtap /etc/flux/system/conf.d + +# flux-quantum, from the branch or from this checkout. Editable, so a bind +# mount of the repository over /opt/flux-quantum picks up edits without a +# rebuild. The jobtap plugin is rebuilt from source, since a checkout may +# carry a quantum.so built against another flux-core. The scripts under +# examples/qrmi/container/bin come from the same tree, so a cloned branch +# has to carry this directory already. +COPY --chown=${USERNAME}:${USERNAME} . /opt/flux-quantum-checkout +RUN if [ -n "${FLUX_QUANTUM_REF}" ]; then \ + git clone --depth 1 -b "${FLUX_QUANTUM_REF}" "${FLUX_QUANTUM_REPO}" /opt/flux-quantum \ + && rm -rf /opt/flux-quantum-checkout; \ + else \ + mv /opt/flux-quantum-checkout /opt/flux-quantum; \ + fi \ + && cd /opt/flux-quantum && (git log --oneline -1 || echo "flux-quantum: this checkout") \ + && { test -d /opt/flux-quantum/examples/qrmi/container/bin \ + || { echo "flux-quantum: ${FLUX_QUANTUM_REF:-the checkout} has no examples/qrmi/container/bin." \ + "Push the directory, or build with --build-arg FLUX_QUANTUM_REF= to use this checkout." >&2; exit 1; }; } \ + && python3 -m pip install --break-system-packages --no-cache-dir -e /opt/flux-quantum \ + && flux python -c "from flux_quantum.cli import QuantumCLIPlugin; print('flux_quantum importable')" \ + && rm -f /opt/flux-quantum/flux_quantum/jobtap/quantum.so \ + && make -C /opt/flux-quantum/flux_quantum/jobtap \ + && cp /opt/flux-quantum/flux_quantum/jobtap/quantum.so /etc/flux/system/jobtap/ \ + && chown -R ${USERNAME}:${USERNAME} /opt/flux-quantum \ + && install -m 0755 /opt/flux-quantum/examples/qrmi/container/bin/* /usr/local/bin/ \ + && chmod +x /opt/flux-quantum/examples/qrmi/ibm-run.sh \ + && ln -sf /opt/flux-quantum/examples/qrmi/ibm-run.sh /usr/local/bin/ibm-run + +# FLUX_CLI_PLUGINPATH holds only the discovery shim, never the package. +# FLUX_QUANTUM_MOCK registers the token free mock vendor. It has to be in the +# environment before flux start, because the ingest validator is a separate +# process. Real vendors are registered regardless, so it stays on. +ENV FLUX_CLI_PLUGINPATH=/opt/flux-quantum/cli-plugins \ + FLUX_QUANTUM_MOCK=1 \ + FLUX_QUANTUM_JOBTAP=/etc/flux/system/jobtap/quantum.so + +USER ${USERNAME} +WORKDIR /opt/flux-quantum/examples/qrmi/container +CMD ["flux-quantum-start"] diff --git a/examples/qrmi/container/Dockerfile.dockerignore b/examples/qrmi/container/Dockerfile.dockerignore new file mode 100644 index 0000000..d91dd13 --- /dev/null +++ b/examples/qrmi/container/Dockerfile.dockerignore @@ -0,0 +1,15 @@ +# Context is the repository root. Keep the scratch and the build trees out. +.git +.claude +.devcontainer +.vscode +flux-sched +flux_quantum.egg-info +**/__pycache__ +**/*.pyc +**/*.so +*.jsonl +*.log +*.patch +*.diff +*.png diff --git a/examples/qrmi/container/README.md b/examples/qrmi/container/README.md new file mode 100644 index 0000000..56c5eff --- /dev/null +++ b/examples/qrmi/container/README.md @@ -0,0 +1,281 @@ +# flux-quantum with QRMI + +This container builds a single node Flux instance you can start, and test a pipeline for Flux with QRMI first without a token (mock) and then run one cheap coscheduled job against an IBM QPU through QRMI. Feel free to tweak that and run a more expensive one, if you please. + +| component | version | why? | +|---------------|------------------------------------------|-----------------------------------------| +| ubuntu | 24.04, python 3.12 | qrmi needs python 3.11 or newer | +| flux-security | 0.15.0 | release tarball | +| flux-core | 0.87.0 | release tarball, prefix /usr | +| flux-sched | vsoch/flux-sched branch `hold` | git, cmake, sysconfdir /etc | +| flux-quantum | branch `braket-hybrid`, or this checkout | git clone like the cluster, editable | +| qrmi | latest `qrmi[ibm]` on PyPI, with qiskit | pip | +| ibmcloud cli | latest | clis.cloud.ibm.com | + +Since this is in a container, we have no systemd or munge. Keep in mind these would be part of a production setup, and are part of the setup where we are developing. The instance is started with `flux start`, run as the `ubuntu` user, and the scheduler is configured through `--config-path` with the same two config tables the cluster writes to `/etc/flux/system/conf.d`. Everything the scout and the classical job do is the same. + +## 1. Build + +From the repository root: + +```bash +uri=ghcr.io/converged-computing/flux-quantum:qrmi-demo-10-8-2026 +docker build -f examples/qrmi/container/Dockerfile -t $uri . +``` + +flux-quantum is cloned from the `braket-hybrid` branch, the same clone the +cluster's start-script makes, so the image runs what is pushed. To put this +checkout in instead, uncommitted changes and all, pass an empty ref. The +build context is the repository root either way, and `Dockerfile.dockerignore` +keeps the build trees out of it. + +```bash +docker build -f examples/qrmi/container/Dockerfile -t flux-quantum-qrmi \ + --build-arg FLUX_QUANTUM_REF= . +``` + +A cloned branch has to carry this directory already, since the scripts are +installed from it. Until it is pushed, build from the checkout. + +To build another flux-sched branch or flux-core, or to pin qrmi: + +```bash +docker build -f examples/qrmi/container/Dockerfile -t flux-quantum-qrmi \ + --build-arg FLUX_SCHED_BRANCH=add-hold \ + --build-arg FLUX_CORE_VERSION=0.87.0 \ + --build-arg QRMI_SPEC='qrmi[ibm]==0.26.0' . +``` + +## 2. Design + +Let's talk about the coscheduling design. We add a flux quantum client plugin, which is an interface to customize flux commands for submit, run, batch, etc. With flux-quantum, we can add flags however we like that will say "I want to run a quantum job." Our strategy is simple. We take a request for hybrid quantum work (classical and quantum resources) and split it into a hybrid pair. + +- A classical job that is "held" until specifically told to release (a form of reservation) +- A scout job or probe that will create a quantum session, monitor it, and release the classical resources when it is ready. + +With this approach, no classical resources are running until the quantum is ready. Smaller jobs can backfill into the space with lower priority that can be pre-empted. The design uses priority and ensures that we do not accept more hybrid jobs than could be supported given the size of the cluster resources, and contender pre-emptible jobs. Full experiments to demonstrate the approach are pending - we have only mocked the setup due to funding limitations. + +## 3. Run the container + +The pair runs as the user (you) and for this case you are the Ubuntu user in the container. QRMI reads the credentials from your environment, so the credentials come in as environment variables. For the token free +mock, none are needed. + +```bash +docker run -it --rm $uri +``` + +You will be in a Flux instance. Take a look. + +```bash +flux resource list +``` +```console + STATE NNODES NCORES NGPUS NODELIST + free 1 8 0 eda311ea73f7 + allocated 0 0 0 + down 0 0 0 +``` + +Take a look to see that fluxion is handling scheduling. + +```bash +flux module list +``` + +Take a look that we have the quantum plugin loaded. + +```bash +flux jobtap list +quantum.so +``` +```bash +$ flux jobtap query quantum.so | jq +{ + "name": "quantum.so", + "path": "/etc/flux/system/jobtap/quantum.so", + "total_cores": 16, + "reserve_cores": 0, + "preempt_after": 0.0, + "vendors": "ibm,mock,braket,ionq", + "protect_types": "qpu", + "free_cores": 16, + "preemptible_cores": 0, + "promised_cores": 0, + "tracked_jobs": 0 +} +``` + +The plugin above is for validation, ensuring that we can support a vendor, that the request is feasible given that size of the cluster and the plugin will also be able to restore state on reload of Flux. Note that when you log in, the entrypoint is a `flux start --config-path=... flux-quantum-init` and note the config path. The config is what loads the jobtap plugin, and configures the job manager. + +```toml +[sched-fluxion-qmanager] +queue-policy = "coschedule" + +[job-manager] +plugins = [ + { load = "/etc/flux/system/jobtap/quantum.so", + conf = { vendors = "ibm,mock,braket,ionq", protect_types = "qpu", + total_cores = 8, reserve_cores = 0 } } +] +``` + +We use a `coschedule` queue policy that lets a held job reserve. + + +## 4. Run a Mock Example + +Let's start with a mock example, and check the setup before anything else: + +```bash +flux-quantum-check # fluxion, policy, jobtap, cli plugin, graph, qrmi +flux-quantum-check --ibm # also the credentials, by name +``` + +Here is a full demo without you having to think (mocked) + +```bash +flux-quantum-demo mock +``` + +Some notes on the output: + +- **the submit** returns the scout id on stdout and says `held classical job ` on stderr. One submit, two jobs. +- **flux jobs** shows the classical in `SCHED` while the scout runs. It is held, with its cores reserved, until the scout releases it. +- **the scout's output** says it opened the session, waited for priority, released the classical, then waited for the classical to finish before closing the session. The qpu vertex is allocated to the scout the whole time. +- **the classical eventlog** carries a `memo` with `quantum_session` before its `alloc`. That memo is the handoff. No shared filesystem is involved, which is why the two halves can land on different nodes. +- **the classical's output** prints the session it was given. + +And here is doing it by hand: + +```bash +flux submit --quantum-vendor mock -n1 -- flux python /opt/flux-quantum/examples/qrmi/container/workload.py + +# Look at the jobs to see the scout and classical +flux jobs -a +flux job attach +flux job eventlog | grep memo +flux job attach +``` + +`FLUX_QUANTUM_MOCK_QUEUE=30 flux submit ...` gives the mock a queue so the scout visibly waits, and `flux-quantum-start -s 2` gives the pair two ranks to land on. + +## 5. IBM through QRMI + +### Credentials + +QRMI reads its configuration from the environment. The easiest thing is just to export them to the environment in the container. Exporting them in the shell inside the running instance is enough. flux copies your environment into both halves of the pair, and the ingest validator checks, by name, that the variables the chosen resource needs are in the job's environment. It never reads the values and has no credentials of its own. Note that you usually need to login to ibmcloud. This does not usually work for me and I will copy paste the quick credentials from the IBM Cloud interface. + +```bash +export IBM_CLOUD_TOKEN=... +ibmcloud login --apikey "$IBM_CLOUD_TOKEN" +export IBM_CLOUD_CRN=$(ibmcloud resource service-instances --service-name quantum-computing --output json | jq -r '.[0].crn') +export RESOURCE=ibm_kingston +``` + +You can check credentials are set: + +### Choose a backend + +Then choose a backend: + +```bash +flux python -c " +from qiskit_ibm_runtime import QiskitRuntimeService +import os +svc = QiskitRuntimeService(channel='ibm_quantum_platform', + token=os.environ['IBM_CLOUD_TOKEN'], instance=os.environ['IBM_CLOUD_CRN']) +for b in svc.backends(operational=True, simulator=False): + print(b.name, b.status().pending_jobs)" +``` + +Export given your choice: + +```bash +export ibm_kingston_QRMI_IBM_QRS_ENDPOINT=https://quantum.cloud.ibm.com/api/v1 +export ibm_kingston_QRMI_IBM_QRS_IAM_ENDPOINT=https://iam.cloud.ibm.com +export ibm_kingston_QRMI_IBM_QRS_IAM_APIKEY=$IBM_CLOUD_TOKEN +export ibm_kingston_QRMI_IBM_QRS_SERVICE_CRN=$IBM_CLOUD_CRN +``` + +```bash +flux-quantum-check --ibm +``` + +Your plan has to allow sessions. The `qiskit-runtime-service` type opens a real session, which IBM permits on plans that support them. On a plan that does not, IBM answers 403, the scout explains, and cancels the held +classical job. + + +### Do a test run + +`ibm-run` submits a classical job that prints the session and exits. The billed window is the one shot warm-up task plus a few seconds, and it ends with what IBM says it charged. + +```bash +RESOURCE=ibm_kingston ibm-run +SKIP_WARMUP=1 ibm-run # free: opens and closes a session, no task +``` + +### Do a session run + +`flux-quantum-demo ibm` is the same pair with this directory's `workload.py` as the classical half. The workload joins the session the scout opened, transpiles a two qubit Bell circuit for the device, runs it +through the QRMI SamplerV2 with 100 shots, and prints the counts. It asks before submitting. + +```bash +RESOURCE=ibm_kingston flux-quantum-demo ibm +SHOTS=50 WALLTIME=3m flux-quantum-demo ibm -y +``` + +The above shows a lot of how it works and metadata. If you need to debug further, in another terminal you can shell into the container to check the queue: + +```bash +docker exec -it sweet_mclean bash + +# This will be at a similar path +flux proxy local:///tmp/flux-2Ge9tP/local-0 bash +flux resource list +``` +```console +STATE NNODES NCORES NGPUS NODELIST + free 1 7 0 1dcdc3d5f315 + allocated 1 1 0 1dcdc3d5f315 + down 0 0 0 +``` +Export your access key again: + +```bash +export IBM_CLOUD_TOKEN=... +export IBM_CLOUD_CRN=$(ibmcloud resource service-instances --service-name quantum-computing --output json | jq -r '.[0].crn') +``` + +Then look at the queue: + +```bash +$ flux python -c " +import os +from qiskit_ibm_runtime import QiskitRuntimeService +svc = QiskitRuntimeService(channel='ibm_quantum_platform', token=os.environ['IBM_CLOUD_TOKEN'], instance=os.environ['IBM_CLOUD_CRN']) +print('pending on ibm_kingston:', svc.backend('ibm_kingston').status().pending_jobs) +for j in svc.jobs(limit=3): print(j.job_id(), j.status())" +``` + +For my testing, I saw previous jobs, and the queue depth. + +```bash +qiskit_runtime_service._discover_account:WARNING:2026-10-08 21:11:12,194: Loading account with the given token. A saved account will not be used. +pending on ibm_kingston: 13 +db40assvf2bc73ctt390 QUEUED +d9uj04535hes73fk0950 DONE +d9uip7343mgs73et1g4g CANCELLED +``` + +To show the command here that the script prints: + +```bash +flux submit -t 5m --quantum-vendor ibm --quantum-device ibm_kingston --quantum-wait 1800 \ + -n1 -- flux python examples/qrmi/container/workload.py --shots 100 +flux jobs -a +flux job attach # a 403 or a credential problem shows up here +flux job eventlog | grep memo +flux job attach # metadata, circuit, status, counts +``` + +`--quantum-wait` is how long the scout waits for the one shot warm-up task to leave the queue. An IBM session goes active when its first task reaches the head of the queue. If it times out the scout closes the session and cancels the classical. diff --git a/examples/qrmi/container/bin/flux-quantum-check b/examples/qrmi/container/bin/flux-quantum-check new file mode 100755 index 0000000..de67014 --- /dev/null +++ b/examples/qrmi/container/bin/flux-quantum-check @@ -0,0 +1,103 @@ +#!/bin/bash +# Verify the stack from inside the instance, before spending anything. +# Every check prints what it found. Hard failures make the exit code nonzero. +# +# flux-quantum-check inside flux-quantum-start +# flux-quantum-check --ibm also check the IBM credentials by name +set -u +rc=0 +ok() { echo " ok $*"; } +bad() { echo " FAIL $*"; rc=1; } +warn() { echo " warn $*"; } + +echo "== versions ==" +flux version | sed 's/^/ /' +flux python -c "import sys; print(' flux python', sys.version.split()[0], sys.executable)" +flux python -c "import flux_quantum; print(' flux_quantum', flux_quantum.__file__)" || bad "flux_quantum is not importable from flux python" + +echo "== instance ==" +if flux resource list >/dev/null 2>&1; then + flux resource list | sed 's/^/ /' +else + bad "no instance. Run this inside flux-quantum-start"; exit 1 +fi + +echo "== scheduler ==" +if flux module list | grep -qE "sched-fluxion-(resource|qmanager)"; then + ok "fluxion is the scheduler" +else + bad "fluxion is not loaded, sched-simple cannot hold a job" +fi +policy=$(flux config get sched-fluxion-qmanager.queue-policy 2>/dev/null || true) +if [ "$policy" = coschedule ]; then ok "queue-policy coschedule"; else bad "queue-policy is '${policy:-unset}', want coschedule"; fi + +echo "== jobtap plugin ==" +if flux jobtap list 2>/dev/null | grep -q quantum; then + ok "quantum.so loaded" + flux jobtap query quantum.so 2>/dev/null | flux python -c " +import json, sys +try: + q = json.load(sys.stdin) +except Exception: + sys.exit(0) +for k in ('vendors', 'protect_types', 'total_cores', 'reserve_cores', 'preempt_after'): + if k in q: + print(' {} = {}'.format(k, q[k])) +" +else + warn "quantum.so is not loaded, no admission control (flux dmesg | grep quantum)" +fi + +echo "== cli plugin ==" +if flux submit --help 2>&1 | grep -q -- --quantum-vendor; then + ok "flux submit has the --quantum options (FLUX_CLI_PLUGINPATH=${FLUX_CLI_PLUGINPATH:-unset})" +else + bad "no --quantum options. FLUX_CLI_PLUGINPATH=${FLUX_CLI_PLUGINPATH:-unset}" +fi + +echo "== vendor devices in the graph ==" +# the same find RPC populate uses, so this works without the flux-sched +# python commands installed +found=$(flux python -c " +import flux +from flux_quantum import graph +print(' '.join('qdevice_' + v for v in sorted(graph.vendors_present(graph.get_live_graph(flux.Flux())))))" 2>/dev/null) +if [ -n "$found" ]; then ok "$found"; else bad "no qdevice vertices, run flux python -m flux_quantum.populate ibm mock"; fi + +echo "== qrmi ==" +if flux python -c "import qrmi, qiskit; print(' qrmi', getattr(qrmi, '__version__', '?'), 'qiskit', qiskit.__version__)" 2>/dev/null; then + ok "qrmi and qiskit import" + flux python -c " +from qrmi import ResourceType +print(' legacy ibm type present:', hasattr(ResourceType, 'IBMQiskitRuntimeService'))" +else + bad "qrmi does not import. pip install 'qrmi[ibm]' needs python 3.11 or newer" +fi +if command -v ibmcloud >/dev/null; then + ok "ibmcloud $(ibmcloud version 2>/dev/null | head -1)" +else + warn "no ibmcloud cli, set IBM_CLOUD_CRN by hand for ibm-run" +fi + +echo "== mock vendor ==" +if [ -n "${FLUX_QUANTUM_MOCK:-}" ]; then ok "FLUX_QUANTUM_MOCK set, mock vendor registered"; else warn "FLUX_QUANTUM_MOCK unset, no token free rehearsal"; fi + +if [ "${1:-}" = "--ibm" ]; then + echo "== ibm credentials (names only, never values) ==" + # constructing the backend is the check, it raises naming what is missing + if out=$(flux python -c " +from flux_quantum.backends import get_backend +print(get_backend('ibm').credential_note)" 2>&1); then + ok "$out" + else + bad "$(echo "$out" | tail -1)" + fi +fi + +echo +if [ "$rc" -eq 0 ]; then + echo "flux-quantum-check: all good" +else + echo "flux-quantum-check: something is off, see FAIL above" +fi +exit "$rc" diff --git a/examples/qrmi/container/bin/flux-quantum-demo b/examples/qrmi/container/bin/flux-quantum-demo new file mode 100755 index 0000000..da3d9e3 --- /dev/null +++ b/examples/qrmi/container/bin/flux-quantum-demo @@ -0,0 +1,133 @@ +#!/bin/bash +# One coscheduled pair, end to end, with the output that shows the handoff. +# +# flux-quantum-demo mock token free rehearsal, the mock vendor +# flux-quantum-demo ibm the real thing, IBM through QRMI. Costs. +# flux-quantum-demo ibm -y skip the confirmation +# +# The ibm run needs QRMI credentials for one resource in the environment, +# either the four _QRMI_IBM_QRS_* variables (see ../env.example), +# or IBM_CLOUD_TOKEN plus IBM_CLOUD_CRN plus RESOURCE, from which the four are +# derived here the same way ibm-run derives them. +# +# RESOURCE the QRMI resource id, an IBM backend name like ibm_kingston. +# Optional when exactly one resource is configured +# SHOTS shots for the Bell circuit the classical job submits, 100 +# WALLTIME the classical job's -t, 5m. The session lives this long at most +# WAIT --quantum-wait, how long the scout waits for the warm-up task +# to reach the head of the queue, 1800. Waiting is free +# NTASKS cores for the classical job, 1 +# OUT where the transcript goes, ~/flux-quantum-demo-