sampler: update to not wait at end

Signed-off-by: vsoch <vsoch@users.noreply.github.com>

Author: vsoch

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_<KEY>.
-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 <backend>_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 <backend><suffix>
-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,129 +100,90 @@ 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}
     good = sum(v for k, v in counts.items() if k in ideal)
     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__":