nshtrainer 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

nshtrainer/_directory.py

@@ -19,20 +19,8 @@ class DirectoryConfig(C.Config):
     This isn't specific to the run; it is the parent directory of all runs.
     """
 
-    log: Path | None = None
-    """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 nshtrainer/{id}/stdio/."""
-
-    checkpoint: Path | None = None
-    """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 nshtrainer/{id}/activation/."""
-
-    profile: Path | None = None
-    """Directory to save profiling information to. If None, will use nshtrainer/{id}/profile/."""
+    logdir_basename: str = "nshtrainer"
+    """Base name for the log directory."""
 
     setup_callback: DirectorySetupCallbackConfig = DirectorySetupCallbackConfig()
     """Configuration for the directory setup PyTorch Lightning callback."""
@@ -41,11 +29,11 @@ class DirectoryConfig(C.Config):
         if (project_root_dir := self.project_root) is None:
             project_root_dir = Path.cwd()
 
-        # The default base dir is $CWD/nshtrainer/{id}/
-        base_dir = project_root_dir / "nshtrainer"
+        # The default base dir is $CWD/{logdir_basename}/{id}/
+        base_dir = project_root_dir / self.logdir_basename
         base_dir.mkdir(exist_ok=True)
 
-        # Add a .gitignore file to the nshtrainer directory
+        # Add a .gitignore file to the {logdir_basename} directory
         #   which will ignore all files except for the .gitignore file itself
         gitignore_path = base_dir / ".gitignore"
         if not gitignore_path.exists():
@@ -57,13 +45,8 @@ class DirectoryConfig(C.Config):
 
         return base_dir
 
-    def resolve_subdirectory(
-        self,
-        run_id: str,
-        # subdirectory: Literal["log", "stdio", "checkpoint", "activation", "profile"],
-        subdirectory: str,
-    ) -> Path:
-        # The subdir will be $CWD/nshtrainer/{id}/{log, stdio, checkpoint, activation}/
+    def resolve_subdirectory(self, run_id: str, subdirectory: str) -> Path:
+        # The subdir will be $CWD/{logdir_basename}/{id}/{log, stdio, checkpoint, activation}/
         if (subdir := getattr(self, subdirectory, None)) is not None:
             assert isinstance(subdir, Path), (
                 f"Expected a Path for {subdirectory}, got {type(subdir)}"
@@ -79,7 +62,7 @@ class DirectoryConfig(C.Config):
         if (log_dir := logger.log_dir) is not None:
             return log_dir
 
-        # Save to nshtrainer/{id}/log/{logger name}
+        # Save to {logdir_basename}/{id}/log/{logger name}
         log_dir = self.resolve_subdirectory(run_id, "log")
         log_dir = log_dir / logger.resolve_logger_dirname()
         # ^ NOTE: Logger must have a `name` attribute, as this is

nshtrainer/callbacks/directory_setup.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import os
 from pathlib import Path
 from typing import Literal
 
@@ -27,6 +26,7 @@ class DirectorySetupCallbackConfig(CallbackConfigBase):
     def __bool__(self):
         return self.enabled
 
+    @override
     def create_callbacks(self, trainer_config):
         if not self:
             return
@@ -35,21 +35,28 @@ class DirectorySetupCallbackConfig(CallbackConfigBase):
 
 
 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.")
+    try:
+        import nshrunner as nr
+    except ImportError:
+        log.info("nshrunner is not installed. Skipping symlink creation to nshrunner.")
+        return
+
+    # Check if we are in a nshrunner session
+    if (session := nr.Session.from_current_session()) is None:
+        log.info("No current nshrunner session found. Skipping symlink creation.")
         return
-    session_dir = Path(session_dir)
+
+    session_dir = session.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}. "
+            f"nshrunner's 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 symlink_path.exists(follow_symlinks=False):
         # If it already points to the correct directory, we're done
         if symlink_path.resolve() == session_dir.resolve():
             return
@@ -61,7 +68,7 @@ def _create_symlink_to_nshrunner(base_dir: Path):
         )
         symlink_path.unlink()
 
-    symlink_path.symlink_to(session_dir)
+    symlink_path.symlink_to(session_dir, target_is_directory=True)
 
 
 class DirectorySetupCallback(NTCallbackBase):

nshtrainer/trainer/_config.py

@@ -761,6 +761,7 @@ class TrainerConfig(C.Config):
         )
 
     def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
+        yield self.directory.setup_callback
         yield self.early_stopping
         yield self.checkpoint_saving
         yield self.lr_monitor

{nshtrainer-1.2.0.dist-info → nshtrainer-1.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.2.0
+Version: 1.2.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.2.0.dist-info → nshtrainer-1.2.1.dist-info}/RECORD

@@ -3,7 +3,7 @@ nshtrainer/__init__.py,sha256=VcqBfL8RgCcZDaY645nxeDmOspqerx4x46wggCMnS0E,692
 nshtrainer/_callback.py,sha256=ZDppiJ4d65tRXTEWYPZLH_F1xFizdz1pkWJe_sQ5uII,12564
 nshtrainer/_checkpoint/metadata.py,sha256=Hh5a7OkdknUEbkEwX6vS88-XLEeuVDoR6a3en2uLzQE,5597
 nshtrainer/_checkpoint/saver.py,sha256=utcrYKSosd04N9m2GIylufO5DO05D90qVU3mvadfApU,1658
-nshtrainer/_directory.py,sha256=SuXJe9xJXZkDXWWfeOS9rEDz6vZUA6mpnEdkAW0ZQnY,3193
+nshtrainer/_directory.py,sha256=RAG8e0y3VZwGIyy_D-GXgDMK5OvitQU6qEWxHTpWEeY,2490
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
 nshtrainer/callbacks/__init__.py,sha256=m6eJuprZfBELuKpngKXre33B9yPXkG7jlKVmI-0yXRQ,4000
@@ -15,7 +15,7 @@ nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=aCs3E1eucfDlUeW2Iq_Ke7
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
 nshtrainer/callbacks/debug_flag.py,sha256=96fuP0C7C6dSs1GiMeUYzzs0X3Q4Pjt9JVWg3b75fU4,1748
-nshtrainer/callbacks/directory_setup.py,sha256=wPas_Ren8ANejogmIdKhqqgj4ulxz9AS_8xVIAfRXa0,2565
+nshtrainer/callbacks/directory_setup.py,sha256=Ln6f0tCgoBscHeigIAWtCCoAmuWB-kPyaf7SylU7MYo,2773
 nshtrainer/callbacks/distributed_prediction_writer.py,sha256=OSh2C6XF7Nki4eFByNVhwlt69izkxnlmfPx54w4rvBo,5274
 nshtrainer/callbacks/early_stopping.py,sha256=rC_qYKCQWjRQJFo0ky46uG0aDJdYP8vsSlKunk0bUVI,4765
 nshtrainer/callbacks/ema.py,sha256=dBFiUXG0xmyCw8-ayuSzJMKqSbepl6Ii5VIbhFlT5ug,12255
@@ -135,7 +135,7 @@ nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5N
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=fQ7gQRlGWX-90TYT0rttkQyvXDCzo7DAvJgr-jX1zsY,316
-nshtrainer/trainer/_config.py,sha256=tdWAYh-KGXBpgdY8fwvOejjRZN-AS2Ze0f_9s2VEuZ0,33556
+nshtrainer/trainer/_config.py,sha256=EaAZCajuvDoVAn79zGebUgg0ijz-i3eLhwxYq8oTNe8,33600
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
 nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
@@ -159,6 +159,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.2.0.dist-info/METADATA,sha256=HkNLruaJJuf3ijnGe7NqNd9emBR6QHMRh2-taC5wTrU,960
-nshtrainer-1.2.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.2.0.dist-info/RECORD,,
+nshtrainer-1.2.1.dist-info/METADATA,sha256=MrF6xvpgRjy2HsFwCIU9YZXMYcOb9ItAXY8m66P2CoQ,960
+nshtrainer-1.2.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.2.1.dist-info/RECORD,,