nshtrainer 0.8.6__tar.gz → 0.9.1__tar.gz

{nshtrainer-0.8.6 → nshtrainer-0.9.1}/PKG-INFO

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

{nshtrainer-0.8.6 → nshtrainer-0.9.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.8.6"
+version = "0.9.1"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.8.6 → nshtrainer-0.9.1}/src/nshtrainer/model/config.py

@@ -805,6 +805,36 @@ class CheckpointLoadingConfig(C.Config):
     """
 
 
+def _create_symlink_to_nshrunner(base_dir: Path):
+    # Resolve the current nshrunner session directory
+    if not (session_dir := os.environ.get("NSHRUNNER_SESSION_DIR")):
+        log.warning("NSHRUNNER_SESSION_DIR is not set. Skipping symlink creation.")
+        return
+    session_dir = Path(session_dir)
+    if not session_dir.exists() or not session_dir.is_dir():
+        log.warning(
+            f"NSHRUNNER_SESSION_DIR is not a valid directory: {session_dir}. "
+            "Skipping symlink creation."
+        )
+        return
+
+    # Create the symlink
+    symlink_path = base_dir / "nshrunner"
+    if symlink_path.exists():
+        # If it already points to the correct directory, we're done
+        if symlink_path.resolve() == session_dir.resolve():
+            return
+
+        # Otherwise, we should log a warning and remove the existing symlink
+        log.warning(
+            f"A symlink pointing to {symlink_path.resolve()} already exists at {symlink_path}. "
+            "Removing the existing symlink."
+        )
+        symlink_path.unlink()
+
+    symlink_path.symlink_to(session_dir)
+
+
 class DirectoryConfig(C.Config):
     project_root: Path | None = None
     """
@@ -813,30 +843,33 @@ class DirectoryConfig(C.Config):
     This isn't specific to the run; it is the parent directory of all runs.
     """
 
+    create_symlink_to_nshrunner_root: bool = True
+    """Should we create a symlink to the root folder for the Runner (if we're in one)?"""
+
     log: Path | None = None
-    """Base directory for all experiment tracking (e.g., WandB, Tensorboard, etc.) files. If None, will use lltrainer/{id}/log/."""
+    """Base directory for all experiment tracking (e.g., WandB, Tensorboard, etc.) files. If None, will use nshtrainer/{id}/log/."""
 
     stdio: Path | None = None
-    """stdout/stderr log directory to use for the trainer. If None, will use lltrainer/{id}/stdio/."""
+    """stdout/stderr log directory to use for the trainer. If None, will use nshtrainer/{id}/stdio/."""
 
     checkpoint: Path | None = None
-    """Checkpoint directory to use for the trainer. If None, will use lltrainer/{id}/checkpoint/."""
+    """Checkpoint directory to use for the trainer. If None, will use nshtrainer/{id}/checkpoint/."""
 
     activation: Path | None = None
-    """Activation directory to use for the trainer. If None, will use lltrainer/{id}/activation/."""
+    """Activation directory to use for the trainer. If None, will use nshtrainer/{id}/activation/."""
 
     profile: Path | None = None
-    """Directory to save profiling information to. If None, will use lltrainer/{id}/profile/."""
+    """Directory to save profiling information to. If None, will use nshtrainer/{id}/profile/."""
 
     def resolve_run_root_directory(self, run_id: str) -> Path:
         if (project_root_dir := self.project_root) is None:
             project_root_dir = Path.cwd()
 
-        # The default base dir is $CWD/lltrainer/{id}/
-        base_dir = project_root_dir / "lltrainer"
+        # The default base dir is $CWD/nshtrainer/{id}/
+        base_dir = project_root_dir / "nshtrainer"
         base_dir.mkdir(exist_ok=True)
 
-        # Add a .gitignore file to the lltrainer directory
+        # Add a .gitignore file to the nshtrainer directory
         #   which will ignore all files except for the .gitignore file itself
         gitignore_path = base_dir / ".gitignore"
         if not gitignore_path.exists():
@@ -846,6 +879,10 @@ class DirectoryConfig(C.Config):
         base_dir = base_dir / run_id
         base_dir.mkdir(exist_ok=True)
 
+        # Create a symlink to the root folder for the Runner
+        if self.create_symlink_to_nshrunner_root:
+            _create_symlink_to_nshrunner(base_dir)
+
         return base_dir
 
     def resolve_subdirectory(
@@ -854,7 +891,7 @@ class DirectoryConfig(C.Config):
         # subdirectory: Literal["log", "stdio", "checkpoint", "activation", "profile"],
         subdirectory: str,
     ) -> Path:
-        # The subdir will be $CWD/lltrainer/{id}/{log, stdio, checkpoint, activation}/
+        # The subdir will be $CWD/nshtrainer/{id}/{log, stdio, checkpoint, activation}/
         if (subdir := getattr(self, subdirectory, None)) is not None:
             assert isinstance(
                 subdir, Path
@@ -874,7 +911,7 @@ class DirectoryConfig(C.Config):
         if (log_dir := logger.log_dir) is not None:
             return log_dir
 
-        # Save to lltrainer/{id}/log/{logger kind}/{id}/
+        # Save to nshtrainer/{id}/log/{logger kind}/{id}/
         log_dir = self.resolve_subdirectory(run_id, "log")
         log_dir = log_dir / logger.kind
 

{nshtrainer-0.8.6 → nshtrainer-0.9.1}/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
@@ -245,3 +246,20 @@ class _SignalConnector(_LightningSignalConnector):
             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