nshtrainer 0.8.5__py3-none-any.whl → 0.8.7__py3-none-any.whl

nshtrainer/trainer/signal_connector.py

@@ -3,6 +3,7 @@ import os
 import re
 import signal
 import subprocess
+import sys
 import threading
 from collections import defaultdict
 from collections.abc import Callable
@@ -150,7 +151,7 @@ class _SignalConnector(_LightningSignalConnector):
         cmd = ["scontrol", "requeue", job_id]
 
         # requeue job
-        log.info(f"requeing job {job_id}...")
+        log.info(f"Requeuing job {job_id}...")
         try:
             result = subprocess.call(cmd)
         except FileNotFoundError:
@@ -202,18 +203,63 @@ class _SignalConnector(_LightningSignalConnector):
             exe = str((Path(bin_dir) / exe).resolve().absolute())
 
         log.info(f"Using LSF requeue executable: {exe}")
-        cmd = [exe, job_id]
-
-        # Requeue job
-        log.info(f"Requeuing job {job_id}...")
-        try:
-            result = subprocess.call(cmd)
-        except FileNotFoundError:
-            # Retry with shell context if subprocess call fails
-            result = subprocess.call(" ".join(cmd), shell=True)
 
-        # Print result text
-        if result == 0:
-            log.info(f"Requeued LSF job: {job_id}")
+        # If NSHRUNNER_LSF_EXIT_SCRIPT_DIR exists, we should emit a bash script in that directory
+        # rather than calling the requeue command directly. This is because the requeue command
+        # is only available outside of the `jsrun` context, and the exit script is called within
+        # the `jsrun` context.
+        if not (exit_script_dir := os.getenv("NSHRUNNER_LSF_EXIT_SCRIPT_DIR")):
+            cmd = [exe, job_id]
+
+            # Requeue job
+            log.info(f"Requeuing job {job_id}...")
+            try:
+                result = subprocess.call(cmd)
+            except FileNotFoundError:
+                # Retry with shell context if subprocess call fails
+                result = subprocess.call(" ".join(cmd), shell=True)
+
+            # Print result text
+            if result == 0:
+                log.info(f"Requeued LSF job: {job_id}")
+            else:
+                log.warning(
+                    f"Requeuing LSF job {job_id} failed with error code {result}"
+                )
         else:
-            log.warning(f"Requeuing LSF job {job_id} failed with error code {result}")
+            log.critical(
+                "Environment variable NSHRUNNER_LSF_EXIT_SCRIPT_DIR found.\n"
+                "Writing requeue script to exit script directory."
+            )
+            exit_script_dir = Path(exit_script_dir)
+            assert (
+                exit_script_dir.is_dir()
+            ), f"Exit script directory {exit_script_dir} does not exist"
+
+            exit_script_path = exit_script_dir / f"requeue_{job_id}.sh"
+            log.info(f"Writing requeue script to {exit_script_path}")
+
+            with exit_script_path.open("w") as f:
+                f.write(f"#!/bin/bash\n{exe} {job_id}\n")
+
+            # Make the script executable
+            os.chmod(exit_script_path, 0o755)
+
+            log.info(f"Requeue script written to {exit_script_path}")
+
+            # Kill the current session to trigger the exit script
+            log.info("Killing current session to trigger exit script")
+            self._kill_current_session()
+
+    def _kill_current_session(self):
+        from lightning.pytorch.trainer.call import _interrupt
+
+        _interrupt(self.trainer, KeyboardInterrupt())
+        self.trainer._teardown()
+        if (launcher := self.trainer.strategy.launcher) is not None:
+            launcher.kill(_get_sigkill_signal())
+        exit(1)
+
+
+def _get_sigkill_signal() -> _SIGNUM:
+    return signal.SIGTERM if sys.platform == "win32" else signal.SIGKILL

{nshtrainer-0.8.5.dist-info → nshtrainer-0.8.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.8.5
+Version: 0.8.7
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.8.5.dist-info → nshtrainer-0.8.7.dist-info}/RECORD

@@ -65,13 +65,13 @@ nshtrainer/runner.py,sha256=7EumpnBkdNWjSNT9Gm-pkxAJ3W6-iMC-yae-WNeZcLw,3771
 nshtrainer/scripts/check_env.py,sha256=IMl6dSqsLYppI0XuCsVq8lK4bYqXwY9KHJkzsShz4Kg,806
 nshtrainer/scripts/find_packages.py,sha256=FbdlfmAefttFSMfaT0A46a-oHLP_ioaQKihwBfBeWeA,1467
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
-nshtrainer/trainer/signal_connector.py,sha256=QAoPM_C5JJOVQebcrJOimUUD3GHyoeZUqCEAvzZlT4U,8710
+nshtrainer/trainer/signal_connector.py,sha256=JSP8W2PSdzwO3iWX1WOL1l8dufh2dKgUWeJ2gEWCppg,10626
 nshtrainer/trainer/trainer.py,sha256=eYEYfY9v70MuorHcSf8nqM7f2CkmUHhpPcjCk4FJD7k,14034
 nshtrainer/util/environment.py,sha256=_SEtiQ_s5bL5pllUlf96AOUv15kNvCPvocVC13S7mIk,4166
 nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.8.5.dist-info/METADATA,sha256=swroihBvsK71kkd51ekyULjRaBGr3ujP9fZ4HDQqhN8,647
-nshtrainer-0.8.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.8.5.dist-info/RECORD,,
+nshtrainer-0.8.7.dist-info/METADATA,sha256=O1kFYWXIuVK1EU0TpwbpbADX1lJQmPM1-9xLTuNaNB8,647
+nshtrainer-0.8.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.8.7.dist-info/RECORD,,