From 9c66bc84ad0731291ac8467b33e4f2856b5c6878 Mon Sep 17 00:00:00 2001 From: Sooyoung Cheong <64125280+c-sooyoung@users.noreply.github.com> Date: Wed, 12 Aug 2026 22:06:27 +0900 Subject: [PATCH] sandbox pipeline for testing --- pipelines/test.py | 198 +++++++++++++++++++++++++++++++++++++++------- 1 file changed, 169 insertions(+), 29 deletions(-) diff --git a/pipelines/test.py b/pipelines/test.py index a944d78..481bf23 100644 --- a/pipelines/test.py +++ b/pipelines/test.py @@ -1,49 +1,189 @@ 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 test_pipeline(config): - result_dir = config['io']['result_dir'] + result_dir = config["io"]["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') - METRIC = config['bo']['metric'] - PTYCHOENGINE = ptycho.engines[config['ptycho']['engine']] + 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) - # 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_configs = randombo.ask(n=4) - for job_config in job_configs: - 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") + 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): - print(f"SOBO sampling; iteration {RANDOM_ITERS + j}") - job_configs = sobo.ask(n=4) - for job_config in job_configs: - 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") + 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, + )