Repository: Quantum-Accelerators/quacc
Pull request #1581 (2024-01-22 00:05:44)

Increase multithreading safety


Patch:
diff --git a/docs/dev/contributing.md b/docs/dev/contributing.md
index 218a96f7..cc2894aa 100644
--- a/docs/dev/contributing.md
+++ b/docs/dev/contributing.md
@@ -30,7 +30,7 @@ Please abide by the following guidelines when contributing code to quacc:
 
 - All code should include type hints and have internally consistent documentation for the inputs and outputs.
 
-- Use absolute paths where possible and avoid using `os.chdir`. This helps ensure that quacc can be run in multi-threading mode.
+- Use absolute paths where possible and avoid using `os.chdir` or `os.getcwd()`. This helps ensure that quacc can be run in multi-threading mode.
 
 - `gzip` large test files to save space.
 
diff --git a/src/quacc/runners/ase.py b/src/quacc/runners/ase.py
index e0253434..369e0d76 100644
--- a/src/quacc/runners/ase.py
+++ b/src/quacc/runners/ase.py
@@ -13,7 +13,7 @@ from monty.dev import requires
 from monty.os.path import zpath
 
 from quacc import SETTINGS
-from quacc.atoms.core import copy_atoms
+from quacc.atoms.core import copy_atoms, get_final_atoms_from_dyn
 from quacc.runners.prep import calc_cleanup, calc_setup
 from quacc.utils.dicts import recursive_dict_merge
 
@@ -76,7 +76,7 @@ def run_calc(
     atoms = copy_atoms(atoms)
 
     # Perform staging operations
-    tmpdir, job_results_dir = calc_setup(copy_files=copy_files)
+    tmpdir, job_results_dir = calc_setup(atoms, copy_files=copy_files)
 
     # Run calculation via get_potential_energy()
     atoms.get_potential_energy()
@@ -105,7 +105,7 @@ def run_calc(
         atoms.cell = atoms_new.cell
 
     # Perform cleanup operations
-    calc_cleanup(tmpdir, job_results_dir)
+    calc_cleanup(atoms, tmpdir, job_results_dir)
 
     return atoms
 
@@ -159,7 +159,7 @@ def run_opt(
     atoms = copy_atoms(atoms)
 
     # Perform staging operations
-    tmpdir, job_results_dir = calc_setup(copy_files=copy_files)
+    tmpdir, job_results_dir = calc_setup(atoms, copy_files=copy_files)
 
     # Set defaults
     optimizer_kwargs = recursive_dict_merge(
@@ -199,7 +199,7 @@ def run_opt(
     dyn.traj_atoms = read(traj_filename, index=":")
 
     # Perform cleanup operations
-    calc_cleanup(tmpdir, job_results_dir)
+    calc_cleanup(get_final_atoms_from_dyn(dyn), tmpdir, job_results_dir)
 
     return dyn
 
@@ -239,7 +239,7 @@ def run_vib(
     vib_kwargs = vib_kwargs or {}
 
     # Perform staging operations
-    tmpdir, job_results_dir = calc_setup(copy_files=copy_files)
+    tmpdir, job_results_dir = calc_setup(atoms, copy_files=copy_files)
 
     # Run calculation
     vib = Vibrations(atoms, name=str(tmpdir / "vib"), **vib_kwargs)
@@ -249,7 +249,7 @@ def run_vib(
     vib.summary(log=sys.stdout if SETTINGS.DEBUG else str(tmpdir / "vib_summary.log"))
 
     # Perform cleanup operations
-    calc_cleanup(tmpdir, job_results_dir)
+    calc_cleanup(vib.atoms, tmpdir, job_results_dir)
 
     return vib
 
diff --git a/src/quacc/runners/prep.py b/src/quacc/runners/prep.py
index 1ec46a6a..b98f0eca 100644
--- a/src/quacc/runners/prep.py
+++ b/src/quacc/runners/prep.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 import os
 from pathlib import Path
 from shutil import move, rmtree
+from typing import TYPE_CHECKING
 
 from monty.shutil import gzip_dir, remove
 
@@ -14,9 +15,12 @@ from quacc.utils.files import (
     make_unique_dir,
 )
 
+if TYPE_CHECKING:
+    from ase.atoms import Atoms
+
 
 def calc_setup(
-    copy_files: str | Path | list[str | Path] | None = None,
+    atoms: Atoms, copy_files: str | Path | list[str | Path] | None = None
 ) -> tuple[Path, Path]:
     """
     Perform staging operations for a calculation, including copying files to the scratch
@@ -25,6 +29,9 @@ def calc_setup(
 
     Parameters
     ----------
+    atoms
+        The Atoms object to run the calculation on. Must have a calculator
+        attached.
     copy_files
         Filenames to copy from source to scratch directory.
 
@@ -47,6 +54,9 @@ def calc_setup(
     tmpdir_base = SETTINGS.SCRATCH_DIR or SETTINGS.RESULTS_DIR
     tmpdir = make_unique_dir(base_path=tmpdir_base, prefix="tmp-quacc-")
 
+    # Set the calculator's directory
+    atoms.calc.directory = tmpdir
+
     # Define the results directory
     job_results_dir = SETTINGS.RESULTS_DIR
     if SETTINGS.CREATE_UNIQUE_DIR:
@@ -66,18 +76,25 @@ def calc_setup(
     elif isinstance(copy_files, (str, Path)):
         copy_decompress_files([copy_files], tmpdir)
 
+    # NOTE: Technically, this breaks thread-safety since it will change the cwd
+    # for all threads in the current process. However, elsewhere in the code,
+    # we use absolute paths to avoid issues. We keep this here for now because some
+    # old ASE calculators do not support the `directory` keyword argument.
     os.chdir(tmpdir)
 
     return tmpdir, job_results_dir
 
 
-def calc_cleanup(tmpdir: Path, job_results_dir: Path) -> None:
+def calc_cleanup(atoms: Atoms, tmpdir: Path | str, job_results_dir: Path | str) -> None:
     """
     Perform cleanup operations for a calculation, including gzipping files, copying
     files back to the original directory, and removing the tmpdir.
 
     Parameters
     ----------
+    atoms
+        The Atoms object after the calculation. Must have a calculator
+        attached.
     tmpdir
         The path to the tmpdir, where the calculation will be run. It will be
         deleted after the calculation is complete.
@@ -91,8 +108,18 @@ def calc_cleanup(tmpdir: Path, job_results_dir: Path) -> None:
     None
     """
 
-    # Make the results directory and cd into it
-    Path(job_results_dir).mkdir(parents=True, exist_ok=True)
+    job_results_dir, tmpdir = Path(job_results_dir), Path(tmpdir)
+
+    # Reset the calculator's directory
+    atoms.calc.directory = job_results_dir
+
+    # Make the results directory
+    job_results_dir.mkdir(parents=True, exist_ok=True)
+
+    # NOTE: Technically, this breaks thread-safety since it will change the cwd
+    # for all threads in the current process. However, elsewhere in the code,
+    # we use absolute paths to avoid issues. We keep this here for now because some
+    # old ASE calculators do not support the `directory` keyword argument.
     os.chdir(job_results_dir)
 
     # Gzip files in tmpdir
diff --git a/src/quacc/schemas/_aliases/ase.py b/src/quacc/schemas/_aliases/ase.py
index 63800c8a..7c190f37 100644
--- a/src/quacc/schemas/_aliases/ase.py
+++ b/src/quacc/schemas/_aliases/ase.py
@@ -53,8 +53,6 @@ class VibResults(TypedDict):
 class VibSchema(AtomsSchema):
     """Schema for `quacc.schemas.ase.summarize_vib_run`"""
 
-    nid: str
-    dir_name: str
     parameters: parameters | None
     parameters_vib: ParametersVib | None
     results: VibResults
@@ -85,8 +83,6 @@ class ThermoResults(TypedDict):
 
 
 class ThermoSchema(AtomsSchema):
-    nid: str
-    dir_name: str
     parameters_thermo: ParametersThermo
     results: ThermoResults
 
diff --git a/src/quacc/schemas/ase.py b/src/quacc/schemas/ase.py
index eab70f19..ffaf018e 100644
--- a/src/quacc/schemas/ase.py
+++ b/src/quacc/schemas/ase.py
@@ -1,7 +1,6 @@
 """Schemas for storing ASE-based data."""
 from __future__ import annotations
 
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -82,7 +81,8 @@ def summarize_run(
         msg = "ASE Atoms object's calculator has no results."
         raise ValueError(msg)
 
-    uri = get_uri(Path.cwd())
+    directory = final_atoms.calc.directory
+    uri = get_uri(directory)
     input_atoms_metadata = (
         atoms_to_metadata(input_atoms, charge_and_multiplicity=charge_and_multiplicity)
         if input_atoms
@@ -167,12 +167,6 @@ def summarize_opt_run(
     additional_fields = additional_fields or {}
     store = SETTINGS.PRIMARY_STORE if store is None else store
 
-    # Check convergence
-    is_converged = dyn.converged()
-    if check_convergence and not is_converged:
-        msg = f"Optimization did not converge. Refer to {Path.cwd()}"
-        raise RuntimeError(msg)
-
     # Get trajectory
     if not trajectory:
         trajectory = (
@@ -183,6 +177,13 @@ def summarize_opt_run(
 
     initial_atoms = trajectory[0]
     final_atoms = get_final_atoms_from_dyn(dyn)
+    directory = final_atoms.calc.directory
+
+    # Check convergence
+    is_converged = dyn.converged()
+    if check_convergence and not is_converged:
+        msg = f"Optimization did not converge. Refer to {directory}"
+        raise RuntimeError(msg)
 
     # Base task doc
     base_task_doc = summarize_run(
@@ -218,7 +219,7 @@ def summarize_opt_run(
 
 
 def summarize_vib_run(
-    vib: Vibrations,
+    vib: Vibrations | VibrationsData,
     charge_and_multiplicity: tuple[int, int] | None = None,
     additional_fields: dict[str, Any] | None = None,
     store: Store | bool | None = None,
@@ -260,25 +261,26 @@ def summarize_vib_run(
             vib_freqs_raw[i] = np.abs(f)
             vib_energies_raw[i] = np.abs(vib_energies_raw[i])
 
-    atoms = vib._atoms if isinstance(vib, VibrationsData) else vib.atoms
-
-    uri = get_uri(Path.cwd())
-    inputs = {
-        "parameters": None
-        if isinstance(vib, VibrationsData)
-        else atoms.calc.parameters,
-        "parameters_vib": None
-        if isinstance(vib, VibrationsData)
-        else {
-            "delta": vib.delta,
-            "direction": vib.direction,
-            "method": vib.method,
-            "ndof": vib.ndof,
-            "nfree": vib.nfree,
-        },
-        "nid": uri.split(":")[0],
-        "dir_name": ":".join(uri.split(":")[1:]),
-    }
+    if isinstance(vib, VibrationsData):
+        atoms = vib._atoms
+        inputs = {}
+    else:
+        atoms = vib.atoms
+        directory = atoms.calc.directory
+        uri = get_uri(directory)
+
+        inputs = {
+            "parameters": atoms.calc.parameters,
+            "parameters_vib": {
+                "delta": vib.delta,
+                "direction": vib.direction,
+                "method": vib.method,
+                "ndof": vib.ndof,
+                "nfree": vib.nfree,
+            },
+            "nid": uri.split(":")[0],
+            "dir_name": ":".join(uri.split(":")[1:]),
+        }
 
     atoms_metadata = atoms_to_metadata(
         atoms, charge_and_multiplicity=charge_and_multiplicity
@@ -368,7 +370,6 @@ def summarize_ideal_gas_thermo(
     additional_fields = additional_fields or {}
     store = SETTINGS.PRIMARY_STORE if store is None else store
 
-    uri = get_uri(Path.cwd())
     spin_multiplicity = round(2 * igt.spin + 1)
 
     inputs = {
@@ -380,9 +381,7 @@ def summarize_ideal_gas_thermo(
             "vib_freqs": [e / units.invcm for e in igt.vib_energies],
             "vib_energies": igt.vib_energies.tolist(),
             "n_imag": igt.n_imag,
-        },
-        "nid": uri.split(":")[0],
-        "dir_name": ":".join(uri.split(":")[1:]),
+        }
     }
 
     results = {
diff --git a/src/quacc/schemas/cclib.py b/src/quacc/schemas/cclib.py
index 5ffa957d..e45ef46e 100644
--- a/src/quacc/schemas/cclib.py
+++ b/src/quacc/schemas/cclib.py
@@ -36,7 +36,7 @@ logger = logging.getLogger(__name__)
 def cclib_summarize_run(
     final_atoms: Atoms,
     logfile_extensions: str | list[str],
-    dir_path: str | None = None,
+    dir_path: Path | str | None = None,
     pop_analyses: list[
         Literal[
             "cpsa",
@@ -71,7 +71,7 @@ def cclib_summarize_run(
         used. For an exact match only, put in the full file name.
     dir_path
         The path to the folder containing the calculation outputs. A value of
-        None specifies the current working directory.
+        None specifies the calculator directory.
     pop_analyses
         The name(s) of any cclib post-processing analysis to run. Note that for
         bader, ddec6, and hirshfeld, a cube file (.cube, .cub) must reside in
@@ -92,7 +92,7 @@ def cclib_summarize_run(
         Dictionary representation of the task document
     """
 
-    dir_path = dir_path or Path.cwd()
+    dir_path = Path(dir_path or final_atoms.calc.directory)
     check_convergence = (
         SETTINGS.CHECK_CONVERGENCE if check_convergence is None else check_convergence
     )
diff --git a/src/quacc/schemas/phonons.py b/src/quacc/schemas/phonons.py
index 971d5b9b..2e5663b8 100644
--- a/src/quacc/schemas/phonons.py
+++ b/src/quacc/schemas/phonons.py
@@ -1,13 +1,11 @@
 """Summarizer for phonopy."""
 from __future__ import annotations
 
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 from quacc import SETTINGS
 from quacc.schemas.atoms import atoms_to_metadata
 from quacc.utils.dicts import clean_task_doc, recursive_dict_merge
-from quacc.utils.files import get_uri
 from quacc.wflow_tools.db import results_to_db
 
 try:
@@ -58,14 +56,7 @@ def summarize_phonopy(
     additional_fields = additional_fields or {}
     store = SETTINGS.PRIMARY_STORE if store is None else store
 
-    uri = get_uri(Path.cwd())
-
-    inputs = {
-        "parameters": parameters,
-        "phonopy_metadata": {"version": phonon.version},
-        "nid": uri.split(":")[0],
-        "dir_name": ":".join(uri.split(":")[1:]),
-    }
+    inputs = {"parameters": parameters, "phonopy_metadata": {"version": phonon.version}}
 
     results = {
         "results": {
diff --git a/src/quacc/schemas/vasp.py b/src/quacc/schemas/vasp.py
index 9ed2bd4d..c2be88e8 100644
--- a/src/quacc/schemas/vasp.py
+++ b/src/quacc/schemas/vasp.py
@@ -46,8 +46,7 @@ def vasp_summarize_run(
     final_atoms
         ASE Atoms object following a calculation.
     dir_path
-        Path to VASP outputs. A value of None specifies the current working
-        directory
+        Path to VASP outputs. A value of None specifies the calculator directory.
     move_magmoms
         Whether to move the final magmoms of the original Atoms object to the
         initial magmoms of the returned Atoms object.
@@ -80,7 +79,7 @@ def vasp_summarize_run(
     check_convergence = (
         SETTINGS.CHECK_CONVERGENCE if check_convergence is None else check_convergence
     )
-    dir_path = dir_path or Path.cwd()
+    dir_path = Path(dir_path or final_atoms.calc.directory)
     store = SETTINGS.PRIMARY_STORE if store is None else store
 
     # Fetch all tabulated results from VASP outputs files. Fortunately, emmet
@@ -131,7 +130,7 @@ def vasp_summarize_run(
     return task_doc
 
 
-def _bader_runner(path: str | None = None) -> BaderSchema:
+def _bader_runner(path: Path | str) -> BaderSchema:
     """
     Runs a Bader partial charge and spin moment analysis using the VASP output
     files in the given path. This function requires that `bader` is located in
@@ -152,7 +151,6 @@ def _bader_runner(path: str | None = None) -> BaderSchema:
     BaderSchema
         Dictionary containing the Bader analysis summary
     """
-    path = path or Path.cwd()
 
     # Make sure files are present
     relevant_files = ["AECCAR0", "AECCAR2", "CHGCAR", "POTCAR"]
@@ -179,7 +177,7 @@ def _bader_runner(path: str | None = None) -> BaderSchema:
 
 
 def _chargemol_runner(
-    path: str | None = None, atomic_densities_path: str | None = None
+    path: str, atomic_densities_path: str | None = None
 ) -> ChargemolSchema:
     """
     Runs a Chargemol (i.e. DDEC6 + CM5) analysis using the VASP output files in
@@ -206,7 +204,6 @@ def _chargemol_runner(
     ChargemolSchema
         Dictionary containing the Chargemol analysis summary
     """
-    path = path or Path.cwd()
 
     # Make sure files are present
     relevant_files = ["AECCAR0", "AECCAR2", "CHGCAR", "POTCAR"]
diff --git a/src/quacc/utils/files.py b/src/quacc/utils/files.py
index 81f5a9cf..24dd507c 100644
--- a/src/quacc/utils/files.py
+++ b/src/quacc/utils/files.py
@@ -148,7 +148,7 @@ def make_unique_dir(
     base_path: Path | str | None = None, prefix: str | None = None
 ) -> Path:
     """
-    Make a directory with a unique name. Uses the same format as Jobflow.
+    Make a directory with a unique name.
 
     Parameters
     ----------


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index ec2afe0c..924a2919 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -174,7 +174,7 @@ jobs:
           name: ${{ github.job }} ${{ matrix.wflow_engine }} coverage report
           path: "coverage.xml"
           retention-days: 1
-  
+
   tests-psi4:
     strategy:
       fail-fast: true
diff --git a/tests/core/runners/test_prep.py b/tests/core/runners/test_prep.py
index c5d35b88..386b335d 100644
--- a/tests/core/runners/test_prep.py
+++ b/tests/core/runners/test_prep.py
@@ -1,6 +1,9 @@
 import os
 from pathlib import Path
 
+from ase.build import bulk
+from ase.calculators.emt import EMT
+
 from quacc.runners.prep import calc_cleanup, calc_setup
 
 
@@ -24,8 +27,10 @@ def test_calc_setup(tmp_path, monkeypatch):
 
     make_files()
     SETTINGS.SCRATCH_DIR = tmp_path
+    atoms = bulk("Cu")
+    atoms.calc = EMT()
 
-    tmpdir, results_dir = calc_setup()
+    tmpdir, results_dir = calc_setup(atoms)
 
     assert tmpdir.is_dir()
     assert "tmp" in str(tmpdir)
@@ -38,6 +43,7 @@ def test_calc_setup(tmp_path, monkeypatch):
     assert "file1.txt" not in os.listdir(tmpdir)
     assert "file2.txt" not in os.listdir(tmpdir)
     assert Path.cwd() == tmpdir
+    assert Path(atoms.calc.directory) == tmpdir
 
     SETTINGS.RESULTS_DIR = DEFAULT_SETTINGS.RESULTS_DIR
     SETTINGS.SCRATCH_DIR = DEFAULT_SETTINGS.SCRATCH_DIR
@@ -51,8 +57,10 @@ def test_calc_setup_v2(tmp_path, monkeypatch):
 
     make_files()
     SETTINGS.SCRATCH_DIR = None
+    atoms = bulk("Cu")
+    atoms.calc = EMT()
 
-    tmpdir, results_dir = calc_setup(copy_files=["file1.txt", "tmp_dir"])
+    tmpdir, results_dir = calc_setup(atoms, copy_files=["file1.txt", "tmp_dir"])
 
     assert tmpdir.is_dir()
     assert "tmp" in str(tmpdir)
@@ -72,14 +80,17 @@ def test_calc_cleanup(tmp_path, monkeypatch):
 
     DEFAULT_SETTINGS = SETTINGS.model_copy()
     monkeypatch.chdir(tmp_path)
+    atoms = bulk("Cu")
+    atoms.calc = EMT()
 
     make_files2()
     SETTINGS.SCRATCH_DIR = tmp_path
 
     p = Path(Path.cwd(), "quacc-tmp-1234").resolve()
     assert p.is_dir()
-    calc_cleanup(p, SETTINGS.RESULTS_DIR)
+    calc_cleanup(atoms, p, SETTINGS.RESULTS_DIR)
     assert not p.exists()
     assert Path.cwd() == SETTINGS.RESULTS_DIR
+    assert Path(atoms.calc.directory) == SETTINGS.RESULTS_DIR
 
     SETTINGS.SCRATCH_DIR = DEFAULT_SETTINGS.SCRATCH_DIR
diff --git a/tests/core/schemas/test_ase.py b/tests/core/schemas/test_ase.py
index 6e478230..a16fb9d9 100644
--- a/tests/core/schemas/test_ase.py
+++ b/tests/core/schemas/test_ase.py
@@ -332,8 +332,6 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
     assert results["parameters_thermo"]["pressure"] == 20.0
     assert results["parameters_thermo"]["sigma"] == 6
     assert results["parameters_thermo"]["spin_multiplicity"] == 2
-    assert "nid" in results
-    assert "dir_name" in results
 
     # test document can be jsanitized and decoded
     d = jsanitize(results, strict=True, enum_values=True)
diff --git a/tests/core/schemas/test_vasp_pop_analysis.py b/tests/core/schemas/test_vasp_pop_analysis.py
index 870fddcd..b3f58a6d 100644
--- a/tests/core/schemas/test_vasp_pop_analysis.py
+++ b/tests/core/schemas/test_vasp_pop_analysis.py
@@ -44,9 +44,6 @@ def patch_pop_analyses(monkeypatch):
 
 
 def prep_files(monkeypatch):
-    if not os.path.exists("rundir"):
-        os.mkdir("rundir")
-    monkeypatch.chdir("rundir")
     for f in ["CHGCAR", "POTCAR"]:
         with open(f, "w") as w:
             w.write("test")
@@ -59,7 +56,7 @@ def test_run_bader(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
     prep_files(monkeypatch)
 
-    bader_stats = _bader_runner()
+    bader_stats = _bader_runner(tmp_path)
     assert bader_stats["min_dist"] == [1.0]
     assert bader_stats["partial_charges"] == [1.0]
     assert bader_stats["spin_moments"] == [0.0]
@@ -73,7 +70,7 @@ def test_bader_erorr(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
 
     with pytest.raises(FileNotFoundError):
-        _bader_runner()
+        _bader_runner(tmp_path)
     with open("CHGCAR", "w") as w:
         w.write("test")
 
@@ -82,7 +79,7 @@ def test_run_chargemol(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
     prep_files(monkeypatch)
 
-    chargemol_stats = _chargemol_runner(path=".", atomic_densities_path=".")
+    chargemol_stats = _chargemol_runner(path=tmp_path, atomic_densities_path=tmp_path)
     assert chargemol_stats["ddec"]["partial_charges"] == [1.0]
     assert chargemol_stats["ddec"]["spin_moments"] == [0.0]
 
@@ -92,10 +89,10 @@ def test_chargemol_erorr(tmp_path, monkeypatch):
     prep_files(monkeypatch)
 
     with pytest.raises(OSError):
-        _chargemol_runner()
+        _chargemol_runner(tmp_path)
 
     os.remove("CHGCAR")
     with pytest.raises(FileNotFoundError):
-        _chargemol_runner(atomic_densities_path=".")
+        _chargemol_runner(tmp_path, atomic_densities_path=tmp_path)
     with open("CHGCAR", "w") as w:
         w.write("test")