fix(eval): serialize synthetic CUDA batch jobs

src/sampleworks/eval/generate_synthetic_density.py

@@ -11,6 +11,7 @@
 from sampleworks.core.forward_models.xray.real_space_density import XMap_torch
 from sampleworks.eval.synthetic_utils import (
     load_structure_for_synthetic_reward,
+    resolve_parallel_jobs,
     validate_occupancy_values,
 )
 from sampleworks.utils.atom_array_utils import save_structure_to_cif
@@ -284,12 +285,10 @@ def process_batch(
     from joblib import delayed, Parallel
 
     rows = load_batch_csv(csv_path)
-    logger.info(f"Processing {len(rows)} structures from {csv_path} using {n_jobs} jobs")
+    effective_n_jobs = resolve_parallel_jobs(device, n_jobs)
+    logger.info(f"Processing {len(rows)} structures from {csv_path} using {effective_n_jobs} jobs")
 
-    # TODO(`#242`): When device is CUDA and n_jobs > 1, each loky worker gets its own
-    # CUDA context, risking GPU memory contention or OOM errors. Consider explicit
-    # per-worker device assignment and/or n_jobs capping. Same issue in SF script.
-    Parallel(n_jobs=n_jobs, backend="loky")(
+    Parallel(n_jobs=effective_n_jobs, backend="loky")(
         delayed(_process_single_row)(
             row=row,
             occupancy_mode=occupancy_mode,

src/sampleworks/eval/generate_synthetic_sf.py

@@ -23,6 +23,7 @@
 from loguru import logger
 from sampleworks.eval.synthetic_utils import (
     load_structure_for_synthetic_reward,
+    resolve_parallel_jobs,
     validate_occupancy_values,
 )
 from sampleworks.utils.atom_array_utils import BLANK_ALTLOC_IDS
@@ -472,12 +473,10 @@ def process_batch(
     from joblib import delayed, Parallel
 
     rows = load_batch_csv(csv_path)
-    logger.info(f"Processing {len(rows)} structures from {csv_path} using {n_jobs} jobs")
+    effective_n_jobs = resolve_parallel_jobs(device, n_jobs)
+    logger.info(f"Processing {len(rows)} structures from {csv_path} using {effective_n_jobs} jobs")
 
-    # TODO(`#242`): When device is CUDA and n_jobs > 1, each loky worker gets its own
-    # CUDA context, risking GPU memory contention or OOM errors. Consider explicit
-    # per-worker device assignment and/or n_jobs capping. Same issue in density script.
-    Parallel(n_jobs=n_jobs, backend="loky")(
+    Parallel(n_jobs=effective_n_jobs, backend="loky")(
         delayed(_process_single_row)(
             row=row,
             base_dir=base_dir,

src/sampleworks/eval/synthetic_utils.py

@@ -4,6 +4,7 @@
 import traceback
 from pathlib import Path
 
+import torch
 from atomworks.io.transforms.atom_array import remove_waters
 from biotite.structure import AtomArray
 from loguru import logger
@@ -18,6 +19,37 @@
 )
 
 
+def resolve_parallel_jobs(device: torch.device | str, n_jobs: int) -> int:
+    """Choose a safe job count for synthetic calculations on a device.
+
+    Process-based joblib workers each create an independent CUDA context. To
+    avoid GPU-memory contention and out-of-memory failures, CUDA work is kept
+    in the current process while CPU work retains the caller's requested
+    parallelism.
+
+    Parameters
+    ----------
+    device
+        Device used for the synthetic calculation.
+    n_jobs
+        Requested joblib worker count. Negative values request multiple workers.
+
+    Returns
+    -------
+    int
+        ``1`` for CUDA requests that would use multiple workers; otherwise the
+        requested value.
+    """
+    resolved_device = torch.device(device)
+    if resolved_device.type == "cuda" and (n_jobs < 0 or n_jobs > 1):
+        logger.warning(
+            f"CUDA device {resolved_device} with n_jobs={n_jobs} would create "
+            "multiple CUDA contexts and risk GPU memory exhaustion; using n_jobs=1"
+        )
+        return 1
+    return n_jobs
+
+
 def validate_occupancy_values(occupancy_values: list[float]) -> None:
     """Raise ValueError if any value is outside [0.0, 1.0] or values don't sum to 1.0."""
     for i, v in enumerate(occupancy_values):

tests/eval/test_synthetic_utils.py

@@ -0,0 +1,22 @@
+"""Tests for shared synthetic-data generation utilities."""
+
+import pytest
+import torch
+from sampleworks.eval.synthetic_utils import resolve_parallel_jobs
+
+
+@pytest.mark.parametrize("n_jobs", [-2, -1, 2, 8])
+def test_resolve_parallel_jobs_serializes_cuda_work(n_jobs: int) -> None:
+    """CUDA requests that imply multiple processes are clamped to one worker."""
+    assert resolve_parallel_jobs(torch.device("cuda:3"), n_jobs) == 1
+
+
+@pytest.mark.parametrize("n_jobs", [-2, -1, 1, 2, 8])
+def test_resolve_parallel_jobs_preserves_cpu_parallelism(n_jobs: int) -> None:
+    """CPU calculations retain the requested joblib parallelism."""
+    assert resolve_parallel_jobs(torch.device("cpu"), n_jobs) == n_jobs
+
+
+def test_resolve_parallel_jobs_preserves_single_cuda_worker() -> None:
+    """An explicitly serial CUDA request does not need adjustment."""
+    assert resolve_parallel_jobs("cuda", 1) == 1