nshtrainer 0.8.5__tar.gz → 0.8.7__tar.gz

{nshtrainer-0.8.5 → nshtrainer-0.8.7}/PKG-INFO

@@ -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 → nshtrainer-0.8.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.8.5"
+version = "0.8.7"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -36,7 +36,6 @@ strictListInference = true
 strictDictionaryInference = true
 strictSetInference = true
 reportPrivateImportUsage = false
-ignore = ["./build/"]
 
 [tool.ruff.lint]
 ignore = ["F722", "F821", "E731", "E741"]

{nshtrainer-0.8.5 → nshtrainer-0.8.7}/src/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