Merge pull request #567 from ikondov/recovery-rerun

[BUG] Recovery rerun

Author: computron

fireworks/core/launchpad.py

@@ -410,7 +410,7 @@ def bulk_add_wfs(self, wfs) -> None:
 
         """
         # Make all fireworks workflows
-        wfs = [Workflow.from_Firework(wf) if isinstance(wf, Firework) else wf for wf in wfs]
+        wfs = [Workflow.from_firework(wf) if isinstance(wf, Firework) else wf for wf in wfs]
 
         # Initialize new firework counter, starting from the next fw id
         total_num_fws = sum(len(wf) for wf in wfs)
@@ -1700,8 +1700,11 @@ def rerun_fw(self, fw_id, rerun_duplicates=True, recover_launch=None, recover_mo
 
         Returns:
             [int]: list of firework ids that were rerun
+
+        Raises:
+            ValueError: raised in case of firework, recover_launch or recovery info not found
         """
-        m_fw = self.fireworks.find_one({"fw_id": fw_id}, {"state": 1})
+        m_fw = self.fireworks.find_one({"fw_id": fw_id}, {"state": True, "launches": True})
 
         if not m_fw:
             raise ValueError(f"FW with id: {fw_id} not found!")
@@ -1722,14 +1725,25 @@ def rerun_fw(self, fw_id, rerun_duplicates=True, recover_launch=None, recover_mo
 
         # Launch recovery
         if recover_launch is not None:
-            recovery = self.get_recovery(fw_id, recover_launch)
+            if not m_fw["launches"]:
+                raise ValueError(f"FW with id: {fw_id} has no active launches")
+            if recover_launch == "last":
+                rec_launch_id = m_fw["launches"][-1]
+            else:
+                if recover_launch not in m_fw["launches"]:
+                    raise ValueError(f"launch_id: {recover_launch} is no launch of fw_id: {fw_id}")
+                rec_launch_id = recover_launch
+            recovery = self.get_recovery(rec_launch_id)
+            if not recovery:
+                raise ValueError(f"No recovery info found in launch {rec_launch_id}")
             recovery.update(_mode=recover_mode)
             set_spec = recursive_dict({"$set": {"spec._recovery": recovery}})
             if recover_mode == "prev_dir":
-                prev_dir = self.get_launch_by_id(recovery.get("_launch_id")).launch_dir
+                launch_f = {"launch_id": recovery.get("_launch_id")}
+                launch_p = {"launch_dir": True}
+                prev_dir = self.launches.find_one(launch_f, launch_p)["launch_dir"]
                 set_spec["$set"]["spec._launch_dir"] = prev_dir
             self.fireworks.find_one_and_update({"fw_id": fw_id}, set_spec)
-
         # If no launch recovery specified, unset the firework recovery spec
         else:
             set_spec = {"$unset": {"spec._recovery": ""}}
@@ -1756,17 +1770,28 @@ def rerun_fw(self, fw_id, rerun_duplicates=True, recover_launch=None, recover_mo
 
         return reruns
 
-    def get_recovery(self, fw_id, launch_id="last"):
-        """Function to get recovery data for a given fw and launch
+    def get_recovery(self, launch_id):
+        """Function to get recovery data for a given launch.
+
         Args:
-            fw_id (int): fw id to get recovery data for
-            launch_id (int or 'last'): launch_id to get recovery data for, if 'last'
-                recovery data is generated from last launch.
-        """
-        m_fw = self.get_fw_by_id(fw_id)
-        launch = m_fw.launches[-1] if launch_id == "last" else self.get_launch_by_id(launch_id)
-        recovery = launch.state_history[-1].get("checkpoint")
-        recovery.update(_prev_dir=launch.launch_dir, _launch_id=launch.launch_id)
+            launch_id (int): launch_id to get recovery data for
+
+        Returns:
+            recovery (dict): recovery metadata, None when no recovery retrieved
+
+        Raises:
+            ValueError: raised when no launch under the the launch_id is found
+        """
+        launch_f = {"launch_id": launch_id}
+        launch_p = {"launch_dir": True, "state_history": True}
+        launch_dct = self.launches.find_one(launch_f, launch_p)
+        if launch_dct is None:
+            raise ValueError(f"launch_id: {launch_id} does not exist")
+        if not launch_dct["state_history"]:
+            return None
+        recovery = launch_dct["state_history"][-1].get("checkpoint")
+        if recovery:
+            recovery.update(_prev_dir=launch_dct["launch_dir"], _launch_id=launch_id)
         return recovery
 
     def _refresh_wf(self, fw_id) -> None:

fireworks/core/rocket.py

@@ -54,8 +54,9 @@ def do_ping(launchpad: LaunchPad, launch_id: int) -> None:
 
 def ping_launch(launchpad: LaunchPad, launch_id: int, stop_event: Event, master_thread: Thread) -> None:
     while not stop_event.is_set() and master_thread.is_alive():
-        do_ping(launchpad, launch_id)
+        # wait before pinging to avoid a race condition with main thread
         stop_event.wait(PING_TIME_SECS)
+        do_ping(launchpad, launch_id)
 
 
 def start_ping_launch(launchpad: LaunchPad, launch_id: int) -> Event | None:

fireworks/core/tests/test_launchpad.py

@@ -1102,6 +1102,38 @@ def test_task_level_rerun_prev_dir(self) -> None:
         assert ExecutionCounterTask.exec_counter == 1
         assert ExceptionTestTask.exec_counter == 2
 
+    def test_task_level_rerun_wrong_fw_id(self) -> None:
+        with pytest.raises(ValueError, match="FW with id: 999 not found!"):
+            self.lp.rerun_fw(999)
+
+    def test_task_level_rerun_recover_launch_id(self) -> None:
+        rapidfire(self.lp, self.fworker, m_dir=MODULE_DIR)
+        self.lp.rerun_fw(fw_id=1, recover_launch=1)
+        fw = self.lp.get_fw_by_id(1)
+        assert "_recovery" in fw.spec
+        assert isinstance(fw.spec["_recovery"], dict)
+        assert fw.spec["_recovery"]["_task_n"] == 2
+
+    def test_task_level_rerun_wrong_launch_id(self) -> None:
+        rapidfire(self.lp, self.fworker, m_dir=MODULE_DIR)
+        assert os.getcwd() == MODULE_DIR
+        with pytest.raises(ValueError, match="launch_id: 999 is no launch of fw_id: 1"):
+            self.lp.rerun_fw(1, recover_launch=999)
+
+    def test_task_level_rerun_wrong_state(self) -> None:
+        with pytest.raises(ValueError, match="FW with id: 1 has no active launches"):
+            self.lp.rerun_fw(1, recover_launch="last")
+
+    def test_task_level_rerun_no_recovery_info(self) -> None:
+        self.lp.add_wf(Firework(ScriptTask.from_str('echo')))
+        launch_rocket(self.lp, self.fworker, fw_id=2)
+        with pytest.raises(ValueError, match="No recovery info found in launch 1"):
+            self.lp.rerun_fw(2, recover_launch="last")
+
+    def test_get_recovery_wrong_launch_id(self) -> None:
+        with pytest.raises(ValueError, match="launch_id: 999 does not exist"):
+            self.lp.get_recovery(launch_id=999)
+
 
 class WFLockTest(unittest.TestCase):
     @classmethod

fireworks/user_objects/queue_adapters/pbs_newt_adapter.py

@@ -6,12 +6,14 @@
 import os
 
 from fireworks.queue.queue_adapter import QueueAdapterBase
+from fireworks.utilities.fw_utilities import get_fw_logger
 
 try:
     from requests import Session
 except ImportError:
     Session = None
-    print("pip install requests to use PBSAdapterNEWT")
+    msg = "install the fireworks[newt] extra or the requests package to use PBSAdapterNEWT"
+    get_fw_logger(__name__).warning(msg)
 
 __author__ = "Shreyas Cholia, Anubhav Jain"
 __copyright__ = "Copyright 2013, The Materials Project"
@@ -41,7 +43,7 @@ def submit_to_queue(self, script_file: str) -> int:
 
     def get_njobs_in_queue(self, username: str | None = None) -> int:
         if username is None:
-            username = getpass.getus
+            username = getpass.getuser()
 
         resp = Session().get(f"https://newt.nersc.gov/newt/queue/{self.resource}/?user={username}")
         return len(resp.json())