diff --git a/examples/qrmi-workload/sampler.py b/examples/qrmi-workload/sampler.py index 6f290b7..9c834f9 100644 --- a/examples/qrmi-workload/sampler.py +++ b/examples/qrmi-workload/sampler.py @@ -1,19 +1,10 @@ #!/usr/bin/env python3 """A QRMI workload for a pod that fluence placed. -fluence picked a backend for this pod and the webhook put its name in the -environment as FLUXION_BACKEND, with the backend attributes as FLUXION_. -This script takes it from there with QRMI. It acquires a session on that -backend with the credentials mounted on the pod, runs a Bell circuit through -the QRMI sampler, prints the counts, and releases the session. fluence holds -the graph allocation and this pod holds the session. A worker in a shared gang -is handed the scout session instead and submits into it without acquiring. - -For an IBM backend the credentials are the four _QRMI_IBM_QRS_* -variables, or IBM_CLOUD_TOKEN and IBM_CLOUD_CRN from which they are derived. -Without credentials this is a rehearsal. It prints what it would need and -exits cleanly, so the same pod proves the placement and the handoff on a -cluster with no account. +fluence chose the backend and the webhook put its name in FLUXION_BACKEND. +This acquires a session on it, runs a Bell circuit through the QRMI sampler, +prints the counts, and releases the session. Without credentials it prints +what it would need and exits, so the same pod works as a rehearsal. """ import argparse @@ -24,132 +15,83 @@ from qiskit import QuantumCircuit from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager from qrmi import QuantumResource, ResourceType -from qrmi.primitives.ibm import SamplerV2, get_target +from qrmi.primitives.ibm import SamplerV2, get_backend -# QRMI type strings to ResourceType names RESOURCE_TYPES = { - "qiskit-runtime-service": "IBMQiskitRuntimeService", - "ibm-quantum-system": "IBMQuantumSystem", - "pasqal-cloud": "PasqalCloud", - "pasqal-local": "PasqalLocal", - "iqm-server": "IQMServer", - "alice-bob-felis": "AliceBobFelis", -} - -# what QRMI reads for each type, 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"), + "qiskit-runtime-service": ("IBMQiskitRuntimeService", "_QRMI_IBM_QRS"), + "ibm-quantum-system": ("IBMQuantumSystem", "_QRMI_IBM_QS"), } - -# the IBM endpoints are fixed, so the credentials reduce to a key and a CRN -IBM_QRS_DERIVED = { - "_QRMI_IBM_QRS_ENDPOINT": "https://quantum.cloud.ibm.com/api/v1", - "_QRMI_IBM_QRS_IAM_ENDPOINT": "https://iam.cloud.ibm.com", +REQUIRED = ("_ENDPOINT", "_IAM_ENDPOINT", "_IAM_APIKEY", "_SERVICE_CRN") +ENDPOINTS = { + "_ENDPOINT": "https://quantum.cloud.ibm.com/api/v1", + "_IAM_ENDPOINT": "https://iam.cloud.ibm.com", } def section(title): - print() - print("== {} ==".format(title), flush=True) + print("\n== {} ==".format(title)) def print_environment(): - """What fluence injected, and which credentials are set, by name only.""" section("environment from fluence") for name in sorted(os.environ): - if name.startswith("FLUXION_") or name.startswith("FLUENCE_"): - print("{:<40} {}".format(name, os.environ[name])) - section("QRMI environment") - shown = False - for name in sorted(os.environ): - if "QRMI" in name or name in ("IBM_CLOUD_TOKEN", "IBM_CLOUD_CRN"): - print("{:<40} (set, value not shown)".format(name)) - shown = True - if not shown: - print("(none)") + if name.startswith(("FLUXION_", "FLUENCE_")): + print("{:<32} {}".format(name, os.environ[name])) + section("credentials, names only") + names = [n for n in sorted(os.environ) if "QRMI" in n or n.startswith("IBM_CLOUD_")] + print("\n".join(names) or "(none)") def backend_from_fluence(): - """The backend fluence chose and its QRMI type, from the injected env.""" backend = os.environ.get("FLUXION_BACKEND") if not backend: - sys.exit( - "no FLUXION_BACKEND: this pod was not placed by fluence, or the " - "webhook did not inject it (does the pod request a qpu?)" - ) + sys.exit("no FLUXION_BACKEND, this pod was not placed by fluence") rtype = os.environ.get("FLUXION_QRMI_TYPE", "qiskit-runtime-service") if rtype not in RESOURCE_TYPES: - sys.exit("unknown QRMI type {!r} for backend {}".format(rtype, backend)) + sys.exit("unsupported QRMI type {} for {}".format(rtype, backend)) section("backend") print("backend: ", backend) print("type: ", rtype) return backend, rtype -def credentials(backend, rtype): - """Derive the QRMI variables from the IBM Secret and return what is still missing.""" - if rtype == "qiskit-runtime-service": - token = os.environ.get("IBM_CLOUD_TOKEN") - crn = os.environ.get("IBM_CLOUD_CRN") - derived = dict(IBM_QRS_DERIVED) - if token: - derived["_QRMI_IBM_QRS_IAM_APIKEY"] = token - if crn: - derived["_QRMI_IBM_QRS_SERVICE_CRN"] = crn - for suffix, value in derived.items(): - os.environ.setdefault(backend + suffix, value) - missing = [ - backend + s - for s in REQUIRED_ENV.get(rtype, ()) - if not os.environ.get(backend + s) - ] - if missing: - print("missing: ", ", ".join(missing)) - else: - print("credentials present for", backend) - return missing - - -def acquire(backend, rtype, ready_timeout): - """Acquire a QRMI session on the backend and wait until it accepts work.""" - qr = QuantumResource(backend, getattr(ResourceType, RESOURCE_TYPES[rtype])) - lock = qr.acquire() - print("session: ", lock) - # so anything else in this pod that constructs the resource joins this session - os.environ[backend + "_QRMI_JOB_ACQUISITION_TOKEN"] = str(lock) - deadline = time.time() + ready_timeout - while ready_timeout > 0: - try: - if qr.is_accessible(): - print("backend: ", backend, "is accepting work") - break - why = "not accepting work" - except Exception as e: # keep polling until the deadline - why = str(e) +def missing_credentials(backend, rtype): + """Derive the QRMI variables from IBM_CLOUD_TOKEN and IBM_CLOUD_CRN, and + return the names still unset.""" + prefix = backend + RESOURCE_TYPES[rtype][1] + derived = dict(ENDPOINTS) + if os.environ.get("IBM_CLOUD_TOKEN"): + derived["_IAM_APIKEY"] = os.environ["IBM_CLOUD_TOKEN"] + if os.environ.get("IBM_CLOUD_CRN"): + derived["_SERVICE_CRN"] = os.environ["IBM_CLOUD_CRN"] + for suffix, value in derived.items(): + os.environ.setdefault(prefix + suffix, value) + return [prefix + s for s in REQUIRED if not os.environ.get(prefix + s)] + + +def resource(backend, rtype): + return QuantumResource(backend, getattr(ResourceType, RESOURCE_TYPES[rtype][0])) + + +def wait_operational(qr, backend, timeout): + """The backend object, once it reports itself operational.""" + deadline = time.time() + timeout + while True: + dev = get_backend(qr) + status = dev.status() + if status.operational: + print("backend: ", backend, status.status_msg) + return dev if time.time() >= deadline: - qr.release(lock) sys.exit( - "{} did not become usable within {}s: {}".format( - backend, ready_timeout, why + "{} not operational after {}s: {}".format( + backend, timeout, status.status_msg ) ) time.sleep(5) - return qr, lock -def build_circuit(qubits): +def bell_circuit(qubits): qc = QuantumCircuit(qubits) qc.h(0) for q in range(1, qubits): @@ -158,42 +100,24 @@ def build_circuit(qubits): return qc -def transpile_for_device(qc, qr): - target = get_target(qr) - isa = generate_preset_pass_manager(optimization_level=1, target=target).run(qc) +def run(qr, dev, qubits, shots): + isa = generate_preset_pass_manager(optimization_level=1, target=dev.target).run( + bell_circuit(qubits) + ) print("circuit: ", dict(isa.count_ops())) - return isa - - -def run_sampler(qr, isa, shots): - sampler = SamplerV2(qr, options={"default_shots": shots}) - job = sampler.run([isa]) - print("task: ", job.job_id(), "submitted into the session", flush=True) - return job - - -def wait_for_task(job, interval=2.0): + job = SamplerV2(qr, options={"default_shots": shots}).run([isa]) + print("task: ", job.job_id()) started = time.time() last = None while not job.in_final_state(): - status = job.status() - if status != last: - print( - "status: ", - status, - "at {:.0f}s".format(time.time() - started), - flush=True, - ) - last = status - time.sleep(interval) + if job.status() != last: + last = job.status() + print("status: ", last, "at {:.0f}s".format(time.time() - started)) + time.sleep(2) print("status: ", job.status(), "after {:.0f}s".format(time.time() - started)) if job.errored(): - print("logs:") print(job.logs()) sys.exit("the task failed") - - -def report_counts(job, qubits, shots): counts = job.result()[0].data.meas.get_counts() print("counts: ", dict(sorted(counts.items(), key=lambda kv: -kv[1]))) ideal = {"0" * qubits, "1" * qubits} @@ -201,86 +125,65 @@ def report_counts(job, qubits, shots): print("fidelity:", "{}/{} shots in {}".format(good, shots, sorted(ideal))) -def release(qr, lock): - qr.release(lock) - print("released:", lock) - - -def linger(seconds, step=5): - if seconds <= 0: - return - section("lingering for {}s so the pod can be watched".format(seconds)) - left = seconds - while left > 0: - print("working, {}s left".format(left), flush=True) - time.sleep(min(step, left)) - left -= step - - def main(): + sys.stdout.reconfigure(line_buffering=True) ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) ap.add_argument("--shots", type=int, default=100) - ap.add_argument( - "--qubits", type=int, default=2, help="width of the GHZ/Bell circuit, 2" - ) - ap.add_argument( - "--linger", - type=int, - default=int(os.environ.get("LINGER", "0")), - help="seconds to stay alive after the work, 0", - ) + ap.add_argument("--qubits", type=int, default=2, help="width of the circuit") ap.add_argument( "--ready-timeout", type=float, - default=float(os.environ.get("READY_TIMEOUT", "120")), - help="seconds to wait for the backend to accept work after acquiring, 120", + default=120, + help="seconds to wait for the backend", ) args = ap.parse_args() print_environment() backend, rtype = backend_from_fluence() - - section("credentials, names only") - if credentials(backend, rtype): - print() - print("rehearsal: no credentials, so nothing is acquired or submitted. The") - print("placement and the handoff from fluence worked. A Secret with the") - print("variables above would make this a real run.") - linger(args.linger) + missing = missing_credentials(backend, rtype) + if missing: + print("missing: ", ", ".join(missing)) + print( + "\nrehearsal, nothing acquired or submitted. The placement and the handoff worked." + ) return - role = os.environ.get("FLUENCE_COORDINATION_ROLE") session = os.environ.get("FLUENCE_QUANTUM_SESSION_ID") - if role == "worker" and session: - # a worker joins the scout session rather than acquiring its own - section("circuit in the scout session, through QRMI") - print("session: ", session, "from the scout") + if os.environ.get("FLUENCE_COORDINATION_ROLE") == "worker": + # a worker submits into the scout's session instead of acquiring one + if not session: + print( + "worker with no session, task", + os.environ.get("FLUENCE_QUANTUM_JOB_ID", "?"), + ) + return + section("circuit in the scout's session") + print("session: ", session) os.environ[backend + "_QRMI_JOB_ACQUISITION_TOKEN"] = session - qr = QuantumResource(backend, getattr(ResourceType, RESOURCE_TYPES[rtype])) - isa = transpile_for_device(build_circuit(args.qubits), qr) - job = run_sampler(qr, isa, args.shots) - wait_for_task(job) - report_counts(job, args.qubits, args.shots) - linger(args.linger) - return - if role == "worker": - print( - "role: worker with no session, task", - os.environ.get("FLUENCE_QUANTUM_JOB_ID", "?"), + qr = resource(backend, rtype) + run( + qr, + wait_operational(qr, backend, args.ready_timeout), + args.qubits, + args.shots, ) - linger(args.linger) return section("session and circuit, through QRMI") - qr, lock = acquire(backend, rtype, args.ready_timeout) + qr = resource(backend, rtype) + lock = qr.acquire() + print("session: ", lock) + os.environ[backend + "_QRMI_JOB_ACQUISITION_TOKEN"] = str(lock) try: - isa = transpile_for_device(build_circuit(args.qubits), qr) - job = run_sampler(qr, isa, args.shots) - wait_for_task(job) - report_counts(job, args.qubits, args.shots) + run( + qr, + wait_operational(qr, backend, args.ready_timeout), + args.qubits, + args.shots, + ) finally: - release(qr, lock) - linger(args.linger) + qr.release(lock) + print("released:", lock) if __name__ == "__main__":