adding fw configs to individual jobs

hrushikesh-s · Apr 30, 2024 · 60e9a7b · 60e9a7b
1 parent c18fe93
commit 60e9a7b
Show file tree

Hide file tree

Showing 6 changed files with 647 additions and 333 deletions.
diff --git a/src/atomate2/common/flows/elastic.py b/src/atomate2/common/flows/elastic.py
@@ -109,6 +109,7 @@ def make(
             if self.prev_calc_dir_argname is not None:
                 bulk_kwargs[self.prev_calc_dir_argname] = prev_dir
             bulk = self.bulk_relax_maker.make(structure, **bulk_kwargs)
+            bulk.update_config({"manager_config": {"_fworker": "gpu_fworker"}}) #TODO remove this line
             jobs.append(bulk)
             structure = bulk.output.structure
             prev_dir = bulk.output.dir_name

diff --git a/src/atomate2/common/flows/hiphive.py b/src/atomate2/common/flows/hiphive.py
diff --git a/src/atomate2/common/jobs/elastic.py b/src/atomate2/common/jobs/elastic.py
@@ -146,6 +146,7 @@ def run_elastic_deformations(
             elastic_job_kwargs[prev_dir_argname] = prev_dir
         # create the job
         relax_job = elastic_relax_maker.make(deformed_structure, **elastic_job_kwargs)
+        relax_job.update_config({"manager_config": {"_fworker": "gpu_reg_fworker"}}) #TODO remove this line
         relax_job.append_name(f" {idx + 1}/{len(deformations)}")
         relaxations.append(relax_job)