Si_2V1_2 batched bo job

TODO.md

@@ -1,13 +1,21 @@
-TODO:
+# TODO:
 
-- Clean up BO states
+- unify `BOEngine.__init__()`
-    - unify `BOEngine.__init__()`
 - BO train_x/y transfer between engines for single job
 - multi-GPU dispatcher
+    - synchronous batched BO
+    - asynchronous BO
 - template job sequences / yamls
     - mobo
 - metric() function(s) for each ptycho engine
     - FRC score
+- separate `config` into `bo_config` and `ptycho_config`; let `BOEngine` have no knowledge of ptychography and vice versa.
+- add GPU version of ExamplePtychoEngine
+- change `PtychoEngine.metric()` to accept list of names and return dict
+
+
+# TODAY:
+
 - fold_slice: load diffractions / hdf5 files; change only param per job
     - restructure `FoldSlicePtychoEngine.__init__()` to load data but not params
         - write new `prepare_data.m`

bo/base.py

@@ -7,7 +7,7 @@ class BOEngine(ABC):
         self.state = None
 
     @abstractmethod
-    def ask(self):
+    def ask(self, n: int = 1) -> list[dict]:
         pass
 
     @abstractmethod

bo/random.py

@@ -37,20 +37,23 @@ class RandomBOEngine(BOEngine):
                 self.train_y = train_y
         
 
-    def ask(self):
+    def ask(self, n = 1):
         config = self.config
-        next_config = copy.deepcopy(config)
+        next_configs = []
 
-        for param in self.params:
+        for _ in range(n):
-            radius = config['bo']['params'][param]['radius']
+            next_config = copy.deepcopy(config)
-            center = config['ptycho']['params'][param]
+            for param in self.params:
-            modulation = radius * (np.random.rand() - 0.5) * 2
+                radius = config['bo']['params'][param]['radius']
-            next_value = center + modulation
+                center = config['ptycho']['params'][param]
-            if self.param_types[param] == 'int':
+                modulation = radius * (np.random.rand() - 0.5) * 2
-                next_value = round(next_value)
+                next_value = center + modulation
-            next_config['ptycho']['params'][param] = next_value
+                if self.param_types[param] == 'int':
+                    next_value = round(next_value)
+                next_config['ptycho']['params'][param] = next_value
+            next_configs.append(next_config)
 
-        return next_config
+        return next_configs
 
 
     def tell(self, job_config, y_value):

bo/sobo.py

@@ -52,7 +52,7 @@ class SingleObjectiveBOEngine(BOEngine):
         self.acquisition = config['bo']['acquisition']
 
 
-    def ask(self):
+    def ask(self, n = 1):
 
         train_x = torch.from_numpy(self.train_x)
         train_y = torch.from_numpy(self.train_y).unsqueeze(-1) # shape: (BOiter, 1)
@@ -98,44 +98,41 @@ class SingleObjectiveBOEngine(BOEngine):
         sampler = SobolQMCNormalSampler(sample_shape=torch.Size([512]))
 
         if self.acquisition == 'ucb':
-            beta = 0.2
+            acqf = qUpperConfidenceBound(gp, beta=self.config['bo']['beta'], sampler=sampler)
-            print("Acquisition: UCB | Beta: {} (fixed)".format(beta))
-            acqf = qUpperConfidenceBound(gp, beta=beta, sampler=sampler)
         elif self.acquisition == 'ei':
-            best_f = train_y.max()
+            acqf = qLogExpectedImprovement(gp, best_f=train_y.max(), sampler=sampler)
-            print("Acquisition: LogEI  best_f: {:.6f}".format(best_f.item()))
-            acqf = qLogExpectedImprovement(gp, best_f=best_f, sampler=sampler)
         else:
             raise NotImplementedError(f"Acquisition function {self.acquisition} is not implemented. Current options: 'ucb', 'ei'")
     
-        # Full [0,1]^d search (trust region disabled)
         acqf_bounds = torch.stack([
             torch.zeros(train_x.shape[1], dtype=torch.double),
             torch.ones(train_x.shape[1], dtype=torch.double),
         ])
 
-        candidate, _ = optimize_acqf(
+        candidates, _ = optimize_acqf(
             acq_function=acqf,
             bounds=acqf_bounds,
-            q=1,
+            q=n,
             num_restarts=20,
             raw_samples=1024,
             post_processing_func=self._pr_post_processing,  # PR applied here
             sequential=True,
         )
 
-        new_x = candidate.detach() * (bounds[1] - bounds[0]) + bounds[0]
+        new_xs = candidates.detach() * (bounds[1] - bounds[0]) + bounds[0]
 
         # Hard-round integer dims (final guarantee)
         for i in self.integer_indices:
-            new_x[:, i] = torch.round(new_x[:, i])
+            new_xs[:, i] = torch.round(new_xs[:, i])
 
-        next_config = copy.deepcopy(self.config)
+        next_configs = []
-
+        for i in range(n):
-        for i, param in enumerate(self.params):
+            next_config = copy.deepcopy(self.config)
-            next_config['ptycho']['params'][param] = new_x[0,i].item()
+            for j, param in enumerate(self.params):
+                next_config['ptycho']['params'][param] = new_xs[i,j].item()
+            next_configs.append(next_config)
         
-        return next_config
+        return next_configs
 
 
     def _pr_post_processing(self, X):

config.yaml

@@ -1,16 +1,16 @@
 job:
   type: "random+sobo"
-  random_iters: 3
+  random_iters: 8
-  sobo_iters: 30
+  sobo_iters: 128
 
 io:
-  input_data_path: '/home/swim/Si_project/data/Si2V1_1.mat'
+  input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/test'
+  result_dir: '/home/swim/bo-ptycho/results/260812-Si/Si2V1_2'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
+  path: '/home/swim/fold_slice-stable'
   params:
     voltage: 200
     alpha_max: 30
@@ -25,8 +25,10 @@ ptycho:
     tilt_x: 2
     tilt_y: 0 
     scan_step_size: 0.36
+
     Niter: 100
     Niter_save_results: 100
+
     CBED_size: 192
     ADU: 1
     extra_print_info: ''
@@ -36,14 +38,15 @@ ptycho:
     diff_pattern_blur: 1
     probe_change_start: 1
     object_change_start: 1
-    grouping: 64
+    grouping: 512
-    probe_posiiton_search: 1
+    probe_position_search: 1
     regularize_layers: 0.2
     variable_probe: 'false'
 
 bo:
-  mode: 'sobo'
+  batch: 4
   acquisition: 'ucb'
+  beta: 0.1 # for ucb only
   metric: 'log_fourier'
   params:
     alpha_max:

examples/260809-Si/Si_02V1_2.sub

@@ -0,0 +1,26 @@
+#!/bin/bash
+#SBATCH --job-name=Si2V
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=1
+#SBATCH --gres=gpu:rtx-6000ada:1
+#SBATCH --time=100:00:00
+#SBATCH --output=/home/swim/slurm-logs/job_%j.log
+
+echo "2V"
+pwd
+hostname
+date
+
+export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
+export PATH=/usr/local/cuda-11.4/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
+
+source /home/swim/bo-ptycho/venv/bin/activate
+
+YAML=examples/260809-Si/Si_02V1_2.yaml
+cat $YAML
+python -u main.py $YAML
+
+date
+echo "SLURM JOB FINISHED"

examples/260809-Si/Si_02V1_2.yaml

@@ -0,0 +1,59 @@
+job:
+  type: "random+sobo"
+  random_iters: 100
+  sobo_iters: 1000
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si2V1_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  params:
+    voltage: 200
+    alpha_max: 30
+    defocus: -250
+    rot_ang: 0.3
+    Nlayers: 20
+    thickness: 250
+    rbf: 37
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64
+    tilt_x: 2
+    tilt_y: 0 
+    scan_step_size: 0.36
+    Niter: 100
+    Niter_save_results: 100
+    CBED_size: 192
+    ADU: 1
+    extra_print_info: 'FIB'
+    scan_number: 1
+    gpu_id: 1
+    roi_label: '0_Ndp64'
+    diff_pattern_blur: 1
+    probe_change_start: 1
+    object_change_start: 1
+    grouping: 64
+    probe_posiiton_search: 1
+    regularize_layers: 0.2
+    variable_probe: false
+
+bo:
+  mode: 'sobo'
+  acquisition: 'ucb'
+  metric: 'log_fourier'
+  params:
+    alpha_max:
+    defocus:
+      radius: 100
+    rot_ang:
+    Nlayers:
+      radius: 5
+      type: int
+    thickness:
+      radius: 100
+  train_x:
+  train_y:

examples/260809-Si/Si_05V1_2.sub

@@ -0,0 +1,26 @@
+#!/bin/bash
+#SBATCH --job-name=Si5V
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=1
+#SBATCH --gres=gpu:rtx-6000ada:1
+#SBATCH --time=100:00:00
+#SBATCH --output=/home/swim/slurm-logs/job_%j.log
+
+echo "5V"
+pwd
+hostname
+date
+
+export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
+export PATH=/usr/local/cuda-11.4/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
+
+source /home/swim/bo-ptycho/venv/bin/activate
+
+YAML=examples/260809-Si/Si_05V1_2.yaml
+cat $YAML
+python -u main.py $YAML
+
+date
+echo "SLURM JOB FINISHED"

examples/260809-Si/Si_05V1_2.yaml

@@ -0,0 +1,59 @@
+job:
+  type: "random+sobo"
+  random_iters: 100
+  sobo_iters: 1000
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si5V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si5V1_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  params:
+    voltage: 200
+    alpha_max: 30
+    defocus: -200
+    rot_ang: 1
+    Nlayers: 20
+    thickness: 250
+    rbf: 37
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64
+    tilt_x: 3
+    tilt_y: -1 
+    scan_step_size: 0.36
+    Niter: 100
+    Niter_save_results: 100
+    CBED_size: 192
+    ADU: 1
+    extra_print_info: 'FIB'
+    scan_number: 1
+    gpu_id: 1
+    roi_label: '0_Ndp64'
+    diff_pattern_blur: 1
+    probe_change_start: 1
+    object_change_start: 1
+    grouping: 64
+    probe_posiiton_search: 1
+    regularize_layers: 0.2
+    variable_probe: false
+
+bo:
+  mode: 'sobo'
+  acquisition: 'ucb'
+  metric: 'log_fourier'
+  params:
+    alpha_max:
+    defocus:
+      radius: 100
+    rot_ang:
+    Nlayers:
+      radius: 5
+      type: int
+    thickness:
+      radius: 150
+  train_x:
+  train_y:

examples/260809-Si/Si_08V1_2.sub

@@ -0,0 +1,26 @@
+#!/bin/bash
+#SBATCH --job-name=Si8V
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=1
+#SBATCH --gres=gpu:rtx-6000ada:1
+#SBATCH --time=100:00:00
+#SBATCH --output=/home/swim/slurm-logs/job_%j.log
+
+echo "8V"
+pwd
+hostname
+date
+
+export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
+export PATH=/usr/local/cuda-11.4/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
+
+source /home/swim/bo-ptycho/venv/bin/activate
+
+YAML=examples/260809-Si/Si_08V1_2.yaml
+cat $YAML
+python -u main.py $YAML
+
+date
+echo "SLURM JOB FINISHED"

examples/260809-Si/Si_08V1_2.yaml

@@ -0,0 +1,59 @@
+job:
+  type: "random+sobo"
+  random_iters: 100
+  sobo_iters: 1000
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si8V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si8V1_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  params:
+    voltage: 200
+    alpha_max: 30
+    defocus: -200
+    rot_ang: 1.3
+    Nlayers: 25
+    thickness: 350
+    rbf: 37
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64
+    tilt_x: 4
+    tilt_y: 0 
+    scan_step_size: 0.36
+    Niter: 100
+    Niter_save_results: 100
+    CBED_size: 192
+    ADU: 1
+    extra_print_info: 'FIB'
+    scan_number: 1
+    gpu_id: 1
+    roi_label: '0_Ndp64'
+    diff_pattern_blur: 1
+    probe_change_start: 1
+    object_change_start: 1
+    grouping: 64
+    probe_posiiton_search: 1
+    regularize_layers: 0.2
+    variable_probe: false
+
+bo:
+  mode: 'sobo'
+  acquisition: 'ucb'
+  metric: 'log_fourier'
+  params:
+    alpha_max:
+    defocus:
+      radius: 100
+    rot_ang:
+    Nlayers:
+      radius: 5
+      type: int
+    thickness:
+      radius: 150
+  train_x:
+  train_y:

examples/260809-Si/Si_30V3_2.sub

@@ -0,0 +1,26 @@
+#!/bin/bash
+#SBATCH --job-name=Si30V
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=1
+#SBATCH --gres=gpu:rtx-6000ada:1
+#SBATCH --time=100:00:00
+#SBATCH --output=/home/swim/slurm-logs/job_%j.log
+
+echo "30V"
+pwd
+hostname
+date
+
+export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
+export PATH=/usr/local/cuda-11.4/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
+
+source /home/swim/bo-ptycho/venv/bin/activate
+
+YAML=examples/260809-Si/Si_30V3_2.yaml
+cat $YAML
+python -u main.py $YAML
+
+date
+echo "SLURM JOB FINISHED"

examples/260809-Si/Si_30V3_2.yaml

@@ -0,0 +1,59 @@
+job:
+  type: "random+sobo"
+  random_iters: 100
+  sobo_iters: 1000
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si30V3_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si30V3_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  params:
+    voltage: 200
+    alpha_max: 30
+    defocus: 175
+    rot_ang: 0.1
+    Nlayers: 30
+    thickness: 680
+    rbf: 37
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64
+    tilt_x: 5
+    tilt_y: 3 
+    scan_step_size: 0.35
+    Niter: 100
+    Niter_save_results: 100
+    CBED_size: 192
+    ADU: 1
+    extra_print_info: 'FIB'
+    scan_number: 1
+    gpu_id: 1
+    roi_label: '0_Ndp64'
+    diff_pattern_blur: 1
+    probe_change_start: 1
+    object_change_start: 1
+    grouping: 64
+    probe_posiiton_search: 1
+    regularize_layers: 0.2
+    variable_probe: false
+
+bo:
+  mode: 'sobo'
+  acquisition: 'ucb'
+  metric: 'log_fourier'
+  params:
+    alpha_max:
+    defocus:
+      radius: 50
+    rot_ang:
+    Nlayers:
+      radius: 5
+      type: int
+    thickness:
+      radius: 100
+  train_x:
+  train_y:

examples/260812-Si/2V.yaml

@@ -0,0 +1,60 @@
+job:
+  type: "random+sobo"
+  random_iters: 8
+  sobo_iters: 128
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260812-Si/Si2V1_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  params:
+    voltage: 200
+    alpha_max: 30
+    defocus: -250
+    rot_ang: 0.3
+    Nlayers: 20
+    thickness: 250
+    rbf: 37
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64
+    tilt_x: 2
+    tilt_y: 0 
+    scan_step_size: 0.36
+    Niter: 100
+    Niter_save_results: 100
+    CBED_size: 192
+    ADU: 1
+    extra_print_info: 'FIB'
+    scan_number: 1
+    gpu_id: 1
+    roi_label: '0_Ndp64'
+    diff_pattern_blur: 1
+    probe_change_start: 1
+    object_change_start: 1
+    grouping: 512
+    probe_posiiton_search: 1
+    regularize_layers: 0.2
+    variable_probe: false
+
+bo:
+  batch: 4
+  acquisition: 'ucb'
+  beta: 0.1
+  metric: 'log_fourier'
+  params:
+    alpha_max:
+    defocus:
+      radius: 100
+    rot_ang:
+    Nlayers:
+      radius: 5
+      type: int
+    thickness:
+      radius: 100
+  train_x:
+  train_y:

examples/260812-Si/job.sub

@@ -0,0 +1,26 @@
+#!/bin/bash
+#SBATCH --job-name=Si2V
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=1
+#SBATCH --gres=gpu:rtx-6000ada:4
+#SBATCH --time=100:00:00
+#SBATCH --output=/home/swim/slurm-logs/job_%j.log
+
+echo "2V"
+pwd
+hostname
+date
+
+export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
+export PATH=/usr/local/cuda-11.4/bin:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
+
+source /home/swim/bo-ptycho/venv/bin/activate
+
+YAML=examples/260812-Si/2V.yaml
+cat $YAML
+python -u main.py $YAML
+
+date
+echo "SLURM JOB FINISHED"

job.sub

@@ -3,7 +3,7 @@
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
 #SBATCH --cpus-per-task=1
-#SBATCH --gres=gpu:rtx-6000ada:1
+#SBATCH --gres=gpu:rtx-6000ada:4
 #SBATCH --time=100:00:00
 #SBATCH --output=/home/swim/slurm-logs/job_%j.log
 

main.py

@@ -8,11 +8,7 @@ def main(config_yaml):
     with open(config_yaml, 'r') as f:
         config = yaml.safe_load(f)
 
-    job_types = {
+    pipelines.job_types[config['job']['type']](config)
-        'random+sobo': pipelines.sobo_pipeline
-    }
-
-    job_types[config['job']['type']](config)
     
 
 if __name__ == "__main__":

pipelines/__init__.py

@@ -1 +1,8 @@
-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,
+    'test': test_pipeline,
+}

pipelines/batched_sobo.py

@@ -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,
+            )

pipelines/sobo.py

@@ -1,47 +1,49 @@
-import os
+# depreacated, use pipelines.batched_sobo.sobo_pipeline()
-import bo
-import ptycho
 
-def sobo_pipeline(config):
+# import os
+# import bo
+# import ptycho
 
-    result_dir = config['io']['result_dir']
+# def sobo_pipeline(config):
-    os.makedirs(result_dir, exist_ok=True)
 
-    RANDOM_ITERS = config['job'].get('random_iters', 0)
+#     result_dir = config['io']['result_dir']
-    SOBO_ITERS   = config['job'].get('sobo_iters')
+#     os.makedirs(result_dir, exist_ok=True)
-    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')
+#     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")
+#     bo_txt = os.path.join(result_dir, "bo.txt")
-    with open(bo_txt, "w") as f:
+#     with open(bo_txt, "w") as f:
-        f.write(f"    iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")
+#         f.write(f"    iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")
 
-    for j in range(RANDOM_ITERS):
+#     for j in range(RANDOM_ITERS):
-        print(f"RANDOM sampling; iteration {j}")
+#         print(f"RANDOM sampling; iteration {j}")
-        job_config = randombo.ask()
+#         job_config = randombo.ask()
-        ptycho_engine = PTYCHOENGINE(job_config)
+#         ptycho_engine = PTYCHOENGINE(job_config)
-        ptycho_engine.run(run_id=f"bo-{j:03d}")
+#         ptycho_engine.run(run_id=f"bo-{j:03d}")
-        y_value = ptycho_engine.metric(METRIC)
+#         y_value = ptycho_engine.metric(METRIC)
-        randombo.tell(job_config, y_value)
+#         randombo.tell(job_config, y_value)
-        with open(bo_txt, "a") as f:
+#         with open(bo_txt, "a") as f:
-            p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in randombo.params]
+#             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")
+#             f.write(f"{j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
 
 
-    sobo = bo.SingleObjectiveBOEngine(config)
+#     sobo = bo.SingleObjectiveBOEngine(config)
-    sobo.train_x = randombo.train_x
+#     sobo.train_x = randombo.train_x
-    sobo.train_y = randombo.train_y
+#     sobo.train_y = randombo.train_y
 
-    for j in range(SOBO_ITERS):
+#     for j in range(SOBO_ITERS):
-        print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
+#         print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
-        job_config = sobo.ask()
+#         job_config = sobo.ask()
-        ptycho_engine = PTYCHOENGINE(job_config)
+#         ptycho_engine = PTYCHOENGINE(job_config)
-        ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
+#         ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
-        y_value = ptycho_engine.metric(METRIC)
+#         y_value = ptycho_engine.metric(METRIC)
-        sobo.tell(job_config, y_value)
+#         sobo.tell(job_config, y_value)
-        with open(bo_txt, "a") as f:
+#         with open(bo_txt, "a") as f:
-            p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in sobo.params]
+#             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")
+#             f.write(f"{RANDOM_ITERS + j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")

pipelines/test.py

@@ -0,0 +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"]
+    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,
+            )

ptycho/__init__.py

@@ -1,5 +1,4 @@
-from .base import PtychoEngine
+from .example import ExamplePtychoEngine
-from .ptycho_example import ExamplePtychoEngine
 from .fold_slice import FoldSlicePtychoEngine
 
 engines = {

ptycho/example.py

@@ -8,7 +8,6 @@ class ExamplePtychoEngine(PtychoEngine):
     def __init__(self, config):
         super().__init__(config)
 
-    # run single ptychography job based on `config`
     def run(self, run_id="") -> None:
         print(f"[{run_id}] [ExamplePtychoEngine] Sleeping for 0.1 second.")
         time.sleep(0.1)

ptycho/fold_slice.py

@@ -11,9 +11,7 @@ class FoldSlicePtychoEngine(PtychoEngine):
 
     def __init__(self, config):
         super().__init__(config)
-        self._output_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
         self._fold_slice_path = self.config['ptycho']['path']
-        self._setup_txt_path = os.path.join(self._output_dir, 'setup.txt')
         self.metric_methods = {
             'log_fourier': self._log_fourier_metric,
         }
@@ -21,6 +19,9 @@ class FoldSlicePtychoEngine(PtychoEngine):
 
     def run(self, run_id="") -> None:
 
+        self._output_dir = os.path.join(self.config['io']['result_dir'], f'fold_slice-{run_id}')
+        self._setup_txt_path = os.path.join(self._output_dir, 'setup.txt')
+
         # generate setup.txt for fold_slice input
         fold_slice_dict = {}
         fold_slice_dict['raw_data'] = self.config['io']['input_data_path']
@@ -82,9 +83,11 @@ class FoldSlicePtychoEngine(PtychoEngine):
 
         log_fourier_error = self._log_fourier_metric()
         # os.makedirs(os.path.join(self.config['io']['result_dir'], "mat"), exist_ok=True)
-        os.makedirs(os.path.join(self.config['io']['result_dir'], "tiff"), exist_ok=True)
+        # os.makedirs(os.path.join(self.config['io']['result_dir'], "tiff"), exist_ok=True)
         # shutil.copy(mat_path, os.path.join(self.config['io']['result_dir'], "mat", f"{log_fourier_error:.4f}_{run_id}.mat")) # saving .mat files takes a lot of space (expect 20+ GB for 64*64 scan size, 300 iterations)
-        shutil.copy(image_path, os.path.join(self.config['io']['result_dir'], "tiff", f"{log_fourier_error:.4f}_{run_id}.tiff"))
+        # shutil.copy(image_path, os.path.join(self.config['io']['result_dir'], "tiff", f"{log_fourier_error:.4f}_{run_id}.tiff"))
+
+        shutil.rmtree(self._output_dir)
 
 
     def metric(self, names):