multi-gpu batched bo

2026-09-17 20:29:07 +09:00 · 2026-08-12 22:07:45 +09:00
parent 9c66bc84ad
commit 1c85c324c1
6 changed files with 253 additions and 56 deletions
@@ -1,5 +1,6 @@
-from .sobo import sobo_pipeline
+# from .sobo import sobo_pipeline # depreacated
 from .test import test_pipeline
+from .batched_sobo import sobo_pipeline

 job_types = {
    'random+sobo': sobo_pipeline,
@@ -0,0 +1,190 @@
+import os
+import traceback
+import multiprocessing as mp
+import shutil
+
+import bo
+import ptycho
+
+def run_ptycho_worker(
+    worker_id,
+    gpu_token,
+    job_config,
+    metric,
+    run_id,
+    result_queue,
+):
+    # This process, and MATLAB launched from it,
+    # can see exactly one GPU.
+    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_token
+
+    try:
+        PTYCHOENGINE = ptycho.engines[job_config["ptycho"]["engine"]]
+
+        ptycho_engine = PTYCHOENGINE(job_config)
+        ptycho_engine.run(run_id=run_id)
+        y_value = ptycho_engine.metric(metric)
+
+        result_queue.put(
+            (worker_id, y_value, None)
+        )
+
+    except Exception:
+        result_queue.put(
+            (worker_id, None, traceback.format_exc())
+        )
+
+
+def run_batch(
+    ctx,
+    gpu_tokens,
+    job_configs,
+    metric,
+    iteration,
+):
+    result_queue = ctx.Queue()
+    processes = []
+
+    for i, job_config in enumerate(job_configs):
+        # Important: unique run_id for simultaneous jobs.
+        run_id = f"bo-{iteration:03d}-{i:02d}"
+
+        p = ctx.Process(
+            target=run_ptycho_worker,
+            args=(
+                i,
+                gpu_tokens[i],
+                job_config,
+                metric,
+                run_id,
+                result_queue,
+            ),
+        )
+
+        p.start()
+        processes.append(p)
+
+    # Four jobs are now running concurrently.
+
+    results = [
+        result_queue.get()
+        for _ in processes
+    ]
+
+    # Synchronization barrier.
+    for p in processes:
+        p.join()
+
+    # Completion order is arbitrary.
+    results.sort(key=lambda x: x[0])
+
+    for worker_id, _, error in results:
+        if error is not None:
+            raise RuntimeError(
+                f"Ptycho worker {worker_id} failed:\n{error}"
+            )
+
+    return [y_value for _, y_value, _ in results]
+
+
+def sobo_pipeline(config):
+
+    result_dir = config["io"]["result_dir"]
+    if os.path.exists(result_dir):
+        shutil.rmtree(result_dir)
+    os.makedirs(result_dir, exist_ok=True)
+
+    RANDOM_ITERS = config["job"].get("random_iters", 0)
+    SOBO_ITERS = config["job"].get("sobo_iters", 0)
+    METRIC = config["bo"]["metric"]
+    BO_BATCH = config["bo"]["batch"]
+
+    # SLURM should expose the four GPUs allocated to this job.
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
+
+    if visible is None:
+        raise RuntimeError(
+            "CUDA_VISIBLE_DEVICES is not set"
+        )
+
+    gpu_tokens = [
+        token.strip()
+        for token in visible.split(",")
+        if token.strip()
+    ]
+
+    if len(gpu_tokens) < BO_BATCH:
+        raise RuntimeError(
+            f"Expected {BO_BATCH} allocated GPUs, got {len(gpu_tokens)}"
+        )
+
+    # Explicitly use spawn for CUDA / MATLAB isolation.
+    ctx = mp.get_context("spawn")
+
+    # ---------------------------------------------------------
+    # Random sampling
+    # ---------------------------------------------------------
+
+    randombo = bo.RandomBOEngine(config)
+
+    for j in range(RANDOM_ITERS):
+        print(
+            f"RANDOM sampling; iteration {j}",
+            flush=True,
+        )
+
+        job_configs = randombo.ask(n=BO_BATCH)
+
+        y_values = run_batch(
+            ctx=ctx,
+            gpu_tokens=gpu_tokens,
+            job_configs=job_configs,
+            metric=METRIC,
+            iteration=j,
+        )
+
+        # Only the parent touches BO state / train_x / train_y.
+        for job_config, y_value in zip(
+            job_configs,
+            y_values,
+        ):
+            randombo.tell(
+                job_config,
+                y_value,
+            )
+
+    # ---------------------------------------------------------
+    # SOBO
+    # ---------------------------------------------------------
+
+    sobo = bo.SingleObjectiveBOEngine(config)
+
+    sobo.train_x = randombo.train_x
+    sobo.train_y = randombo.train_y
+
+    for j in range(SOBO_ITERS):
+        iteration = RANDOM_ITERS + j
+
+        print(
+            f"SOBO sampling; iteration {iteration}",
+            flush=True,
+        )
+
+        job_configs = sobo.ask(n=BO_BATCH)
+
+        y_values = run_batch(
+            ctx=ctx,
+            gpu_tokens=gpu_tokens,
+            job_configs=job_configs,
+            metric=METRIC,
+            iteration=iteration,
+        )
+
+        for job_config, y_value in zip(
+            job_configs,
+            y_values,
+        ):
+            sobo.tell(
+                job_config,
+                y_value,
+            )
@@ -1,47 +1,49 @@
-import os
-import bo
-import ptycho
+# depreacated, use pipelines.batched_sobo.sobo_pipeline()

-def sobo_pipeline(config):
+# import os
+# import bo
+# import ptycho

-    result_dir = config['io']['result_dir']
-    os.makedirs(result_dir, exist_ok=True)
+# def sobo_pipeline(config):

-    RANDOM_ITERS = config['job'].get('random_iters', 0)
-    SOBO_ITERS   = config['job'].get('sobo_iters')
-    METRIC = config['bo']['metric']
-    PTYCHOENGINE = ptycho.engines[config['ptycho']['engine']]
+#     result_dir = config['io']['result_dir']
+#     os.makedirs(result_dir, exist_ok=True)
+
+#     RANDOM_ITERS = config['job'].get('random_iters', 0)
+#     SOBO_ITERS   = config['job'].get('sobo_iters')
+#     METRIC = config['bo']['metric']
+#     PTYCHOENGINE = ptycho.engines[config['ptycho']['engine']]


-    randombo = bo.RandomBOEngine(config)
+#     randombo = bo.RandomBOEngine(config)

-    bo_txt = os.path.join(result_dir, "bo.txt")
-    with open(bo_txt, "w") as f:
-        f.write(f"    iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")
+#     bo_txt = os.path.join(result_dir, "bo.txt")
+#     with open(bo_txt, "w") as f:
+#         f.write(f"    iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")

-    for j in range(RANDOM_ITERS):
-        print(f"RANDOM sampling; iteration {j}")
-        job_config = randombo.ask()
-        ptycho_engine = PTYCHOENGINE(job_config)
-        ptycho_engine.run(run_id=f"bo-{j:03d}")
-        y_value = ptycho_engine.metric(METRIC)
-        randombo.tell(job_config, y_value)
-        with open(bo_txt, "a") as f:
-            p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in randombo.params]
-            f.write(f"{j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
+#     for j in range(RANDOM_ITERS):
+#         print(f"RANDOM sampling; iteration {j}")
+#         job_config = randombo.ask()
+#         ptycho_engine = PTYCHOENGINE(job_config)
+#         ptycho_engine.run(run_id=f"bo-{j:03d}")
+#         y_value = ptycho_engine.metric(METRIC)
+#         randombo.tell(job_config, y_value)
+#         with open(bo_txt, "a") as f:
+#             p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in randombo.params]
+#             f.write(f"{j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")


-    sobo = bo.SingleObjectiveBOEngine(config)
-    sobo.train_x = randombo.train_x
-    sobo.train_y = randombo.train_y
+#     sobo = bo.SingleObjectiveBOEngine(config)
+#     sobo.train_x = randombo.train_x
+#     sobo.train_y = randombo.train_y

-    for j in range(SOBO_ITERS):
-        print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
-        job_config = sobo.ask()
-        ptycho_engine = PTYCHOENGINE(job_config)
-        ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
-        y_value = ptycho_engine.metric(METRIC)
-        sobo.tell(job_config, y_value)
-        with open(bo_txt, "a") as f:
-            p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in sobo.params]
-            f.write(f"{RANDOM_ITERS + j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
+#     for j in range(SOBO_ITERS):
+#         print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
+#         job_config = sobo.ask()
+#         ptycho_engine = PTYCHOENGINE(job_config)
+#         ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
+#         y_value = ptycho_engine.metric(METRIC)
+#         sobo.tell(job_config, y_value)
+#         with open(bo_txt, "a") as f:
+#             p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in sobo.params]
+#             f.write(f"{RANDOM_ITERS + j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")