Mlflow fix

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 24.4.0
     hooks:
       - id: black
         exclude: ^apax/utils/jax_md_reduced/

apax/nodes/md.py

@@ -66,10 +66,11 @@ def _handle_parameter_file(self):
     def run(self):
         """Primary method to run which executes all steps of the model training"""
 
-        atoms = self.data[self.data_id]
-        if self.repeat is not None:
-            atoms = atoms.repeat(self.repeat)
-        ase.io.write(self.init_struc_dir.as_posix(), atoms)
+        if not self.state.restarted:
+            atoms = self.data[self.data_id]
+            if self.repeat is not None:
+                atoms = atoms.repeat(self.repeat)
+            ase.io.write(self.init_struc_dir.as_posix(), atoms)
 
         run_md(self.model._parameter, self._parameter)
 

apax/nodes/model.py

@@ -61,11 +61,11 @@ def _post_load_(self) -> None:
         self._handle_parameter_file()
 
     def _handle_parameter_file(self):
-        with self.state.use_tmp_path():
-            self._parameter = yaml.safe_load(pathlib.Path(self.config).read_text())
+        self._parameter = yaml.safe_load(self.state.fs.read_text(self.config))
 
+        with self.state.use_tmp_path():
             custom_parameters = {
-                "directory": self.model_directory.resolve().as_posix(),
+                "directory": self.model_directory.as_posix(),
                 "experiment": "",
                 "train_data_path": self.train_data_file.as_posix(),
                 "val_data_path": self.validation_data_file.as_posix(),