diff --git a/src/sampleworks/eval/generate_synthetic_density.py b/src/sampleworks/eval/generate_synthetic_density.py index 8b8521f8..9ccddd46 100644 --- a/src/sampleworks/eval/generate_synthetic_density.py +++ b/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, diff --git a/src/sampleworks/eval/generate_synthetic_sf.py b/src/sampleworks/eval/generate_synthetic_sf.py index 6b64d367..473077bf 100644 --- a/src/sampleworks/eval/generate_synthetic_sf.py +++ b/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, diff --git a/src/sampleworks/eval/synthetic_utils.py b/src/sampleworks/eval/synthetic_utils.py index 7765bbe3..d68bf689 100644 --- a/src/sampleworks/eval/synthetic_utils.py +++ b/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): diff --git a/tests/eval/test_synthetic_utils.py b/tests/eval/test_synthetic_utils.py new file mode 100644 index 00000000..77eba7d7 --- /dev/null +++ b/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