mirror of
https://github.com/c-sooyoung/bo-ptycho.git
synced 2026-09-17 21:29:07 +09:00
95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
import os
|
|
import traceback
|
|
import multiprocessing as mp
|
|
|
|
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:
|
|
ptycho_engine = ptycho.engines[job_config["ptycho"]["engine"]](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)
|
|
|
|
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):
|
|
|
|
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"]
|
|
|
|
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}")
|
|
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)
|
|
for job_config, y_value in zip(job_configs, y_values):
|
|
randombo.tell(job_config, y_value)
|
|
|
|
############################ SOBO SAMPLING ###############################
|
|
sobo = bo.SingleObjectiveBOEngine(config)
|
|
sobo.train_x = randombo.train_x
|
|
sobo.train_y = randombo.train_y
|
|
|
|
for j in range(RANDOM_ITERS, SOBO_ITERS+RANDOM_ITERS):
|
|
print(f"SOBO sampling iteration {j}")
|
|
job_configs = sobo.ask(n=BO_BATCH)
|
|
y_values = run_batch(ctx=ctx, gpu_tokens=gpu_tokens, job_configs=job_configs, metric=METRIC, iteration=j)
|
|
|
|
for job_config, y_value in zip(job_configs, y_values):
|
|
sobo.tell(job_config, y_value)
|