Fix batch distribution for other worker types (#82)

* fix batch distribution for other worker types * update changelog
openforcefield · Feb 9, 2024 · e1e96bf · e1e96bf
1 parent 869993d
commit e1e96bf
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@ The rules for this file:
 
 ### Fixed
 <!-- Bug fixes -->
+- Fixed batch distribution hardcoding (Issue #80, PR #82)
 - Fixed node and edge typing, adding DGL 2.0 compatibility (Issue #78, PR #79)
 
 ### Changed

diff --git a/openff/nagl/utils/_parallelization.py b/openff/nagl/utils/_parallelization.py
@@ -118,9 +118,11 @@ def batch_distributed(
         f"Setting n_workers={n_workers} for {n_batches} batches"
     )
 
-    env_extra = dask.config.get("jobqueue.lsf.job-script-prologue", default=[])
-    if not env_extra:
-        env_extra = []
+    env_extra = []
+    if worker_type != "local":
+        env_extra.extend(
+            dask.config.get(f"jobqueue.{worker_type}.job-script-prologue", default=[])
+        )
     env_extra.append(f"{package_manager} activate {conda_environment}")