nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/util/environment.py

@@ -0,0 +1,119 @@
+import os
+from contextlib import contextmanager
+from logging import getLogger
+
+log = getLogger(__name__)
+
+
+@contextmanager
+def remove_slurm_environment_variables():
+    """
+    SLURM_CPU_BIND_* environment variables are set by SLURM in the current environment.
+    We need to remove all of these environment variables during the codepath in which we create the new SLURM runs, so that the new SLURM runs do not inherit the environment variables from the current environment.
+    To make things easier, we will patch the environment to remove all "SLURM_" environment variables.
+    Otherwise, the runs will faill with an error like shown below:
+        srun: error: CPU binding outside of job step allocation, allocated CPUs are: 0x01F000000001F0000000.
+        srun: error: Task launch for StepId=5216715.0 failed on node learnfair0537: Unable to satisfy cpu bind request
+        srun: error: Application launch failed: Unable to satisfy cpu bind request
+        srun: Job step aborted
+
+    See https://www.mail-archive.com/slurm-users@lists.schedmd.com/msg09157.html for more details.
+    """
+
+    removed_env_vars = {}
+    for key in list(os.environ.keys()):
+        if not key.startswith("SLURM_"):
+            continue
+        removed_env_vars[key] = os.environ.pop(key)
+
+    log.debug(
+        f"Removed environment variables before launching new SLURM job: {list(removed_env_vars.keys())}"
+    )
+    try:
+        yield
+    finally:
+        os.environ.update(removed_env_vars)
+        log.debug(
+            f"Restored environment variables after launching new SLURM job: {list(removed_env_vars.keys())}"
+        )
+
+
+@contextmanager
+def remove_lsf_environment_variables():
+    """Same as above, but for the LSF task scheduler."""
+    removed_env_vars = {}
+    for key in list(os.environ.keys()):
+        if (
+            not key.startswith("LS_")
+            and not key.startswith("LSF_")
+            and not key.startswith("LSB_")
+            and not key.startswith("BSUB_")
+        ):
+            continue
+        removed_env_vars[key] = os.environ.pop(key)
+
+    log.debug(
+        f"Removed environment variables before launching new LSF job: {list(removed_env_vars.keys())}"
+    )
+    try:
+        yield
+    finally:
+        os.environ.update(removed_env_vars)
+        log.debug(
+            f"Restored environment variables after launching new LSF job: {list(removed_env_vars.keys())}"
+        )
+
+
+@contextmanager
+def remove_wandb_environment_variables():
+    """
+    Similar to above, but removes all "WANDB_" environment variables.
+    """
+
+    removed_env_vars = {}
+    for key in list(os.environ.keys()):
+        if not key.startswith("WANDB_"):
+            continue
+        removed_env_vars[key] = os.environ.pop(key)
+
+    log.debug(
+        f"Removed environment variables before launching new SLURM job: {list(removed_env_vars.keys())}"
+    )
+    try:
+        yield
+    finally:
+        os.environ.update(removed_env_vars)
+        log.debug(
+            f"Restored environment variables after launching new SLURM job: {list(removed_env_vars.keys())}"
+        )
+
+
+@contextmanager
+def set_additional_env_vars(additional_env_vars: dict[str, str] | None = None):
+    """
+    Set additional environment variables for the run.
+    Newly set environment variables will be removed after the run is finished.
+    Existing environment variables will be restored to their original values after the run is finished.
+    """
+    if additional_env_vars is None:
+        additional_env_vars = {}
+
+    removed_env_vars = {}
+    for key, value in additional_env_vars.items():
+        removed_env_vars[key] = os.environ.pop(key, None)
+        os.environ[key] = value
+
+    log.debug(
+        f"Set additional environment variables for the run: {list(additional_env_vars.keys())}"
+    )
+    try:
+        yield
+    finally:
+        for key in additional_env_vars.keys():
+            if removed_env_vars[key] is None:
+                del os.environ[key]
+            else:
+                os.environ[key] = removed_env_vars[key]
+        log.debug(
+            f"Restored environment variables after launching new SLURM job: {list(additional_env_vars.keys())}"
+        )

nshtrainer/util/seed.py

@@ -0,0 +1,11 @@
+from logging import getLogger
+
+import lightning.fabric.utilities.seed as LS
+
+log = getLogger(__name__)
+
+
+def seed_everything(seed: int | None, *, workers: bool = False):
+    seed = LS.seed_everything(seed, workers=workers)
+    log.critical(f"Set global seed to {seed}.")
+    return seed

nshtrainer/util/singleton.py

@@ -0,0 +1,89 @@
+from logging import getLogger
+from typing import Any
+
+from typing_extensions import Self, TypeVar, override
+
+log = getLogger(__name__)
+
+
+class Singleton:
+    singleton_key = "_singleton_instance"
+
+    @classmethod
+    def get(cls) -> Self | None:
+        return getattr(cls, cls.singleton_key, None)
+
+    @classmethod
+    def set(cls, instance: Self) -> None:
+        if cls.get() is not None:
+            log.warning(f"{cls.__qualname__} instance is already set")
+
+        setattr(cls, cls.singleton_key, instance)
+
+    @classmethod
+    def reset(cls) -> None:
+        if cls.get() is not None:
+            delattr(cls, cls.singleton_key)
+
+    @classmethod
+    def register(cls, instance: Self) -> None:
+        cls.set(instance)
+
+    def register_self(self):
+        self.register(self)
+
+    @classmethod
+    def instance(cls) -> Self:
+        instance = cls.get()
+        if instance is None:
+            raise RuntimeError(f"{cls.__qualname__} instance is not set")
+
+        return instance
+
+    @override
+    def __init_subclass__(cls, *args, **kwargs) -> None:
+        super().__init_subclass__(*args, **kwargs)
+
+        cls.reset()
+
+
+T = TypeVar("T", infer_variance=True)
+
+
+class Registry:
+    _registry: dict[type, Any] = {}
+
+    @staticmethod
+    def register(cls_: type[T], instance: T):
+        if not isinstance(instance, cls_):
+            raise ValueError(f"{instance} is not an instance of {cls_.__qualname__}")
+
+        if cls_ in Registry._registry:
+            raise ValueError(f"{cls_.__qualname__} is already registered")
+
+        Registry._registry[cls_] = instance
+
+    @staticmethod
+    def try_get(cls_: type[T]) -> T | None:
+        return Registry._registry.get(cls_)
+
+    @staticmethod
+    def get(cls_: type[T]) -> T:
+        instance = Registry.try_get(cls_)
+        if instance is None:
+            raise ValueError(f"{cls_.__qualname__} is not registered")
+
+        return instance
+
+    @staticmethod
+    def instance(cls_: type[T]) -> T:
+        return Registry.get(cls_)
+
+    @staticmethod
+    def reset(cls_: type[T]):
+        if cls_ in Registry._registry:
+            del Registry._registry[cls_]
+
+    @staticmethod
+    def reset_all():
+        Registry._registry.clear()

nshtrainer/util/slurm.py

@@ -0,0 +1,49 @@
+class SlurmParseException(Exception):
+    pass
+
+
+def _expand_id_suffix(suffix_parts: str) -> list[str]:
+    """Parse the a suffix formatted like "1-3,5,8" into
+    the list of numeric values 1,2,3,5,8.
+    """
+    suffixes = []
+    for suffix_part in suffix_parts.split(","):
+        if "-" in suffix_part:
+            low, high = suffix_part.split("-")
+            int_length = len(low)
+            for num in range(int(low), int(high) + 1):
+                suffixes.append(f"{num:0{int_length}}")
+        else:
+            suffixes.append(suffix_part)
+    return suffixes
+
+
+def _parse_node_group(node_list: str, pos: int, parsed: list[str]) -> int:
+    """Parse a node group of the form PREFIX[1-3,5,8] and return
+    the position in the string at which the parsing stopped
+    """
+    prefixes = [""]
+    while pos < len(node_list):
+        c = node_list[pos]
+        if c == ",":
+            parsed.extend(prefixes)
+            return pos + 1
+        if c == "[":
+            last_pos = node_list.index("]", pos)
+            suffixes = _expand_id_suffix(node_list[pos + 1 : last_pos])
+            prefixes = [prefix + suffix for prefix in prefixes for suffix in suffixes]
+            pos = last_pos + 1
+        else:
+            for i, prefix in enumerate(prefixes):
+                prefixes[i] = prefix + c
+            pos += 1
+    parsed.extend(prefixes)
+    return pos
+
+
+def parse_slurm_node_list(node_list: str):
+    pos = 0
+    parsed: list[str] = []
+    while pos < len(node_list):
+        pos = _parse_node_group(node_list, pos, parsed)
+    return parsed

nshtrainer/util/typed.py

@@ -0,0 +1,2 @@
+from ..nn.module_dict import TypedModuleDict as TypedModuleDict
+from ..nn.module_list import TypedModuleList as TypedModuleList

nshtrainer/util/typing_utils.py

@@ -0,0 +1,19 @@
+from typing import TYPE_CHECKING
+
+from typing_extensions import TypeVar
+
+TBase = TypeVar("TBase")
+
+
+def mixin_base_type(base_class: type[TBase]) -> type[TBase]:
+    """
+    Useful function to make mixins with baseclass typehint
+
+    ```
+    class ReadonlyMixin(mixin_base_type(BaseAdmin))):
+        ...
+    ```
+    """
+    if TYPE_CHECKING:
+        return base_class
+    return object

nshtrainer-0.1.0.dist-info/METADATA

@@ -0,0 +1,18 @@
+Metadata-Version: 2.1
+Name: nshtrainer
+Version: 0.1.0
+Summary: 
+Author: Nima Shoghi
+Author-email: nimashoghi@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: nshconfig (>=0.2.0,<0.3.0)
+Requires-Dist: nshrunner (>=0.1.0,<0.2.0)
+Requires-Dist: torch
+Requires-Dist: typing-extensions
+Description-Content-Type: text/markdown
+
+

nshtrainer-0.1.0.dist-info/RECORD

@@ -0,0 +1,72 @@
+nshtrainer/__init__.py,sha256=o39TbnjwUYzE4POcncUiDx02Ey-Hzx8UGuwJDjMcKZU,2971
+nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ4b-PoWM,120
+nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
+nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
+nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
+nshtrainer/_snoop.py,sha256=Rofv1Rd92E0LY40G3A-o9Hu0ZI73RR59wJD5l4Q3PDM,7022
+nshtrainer/_submit/print_environment_info.py,sha256=enbJGl_iHIlhKN8avzKnoZSb0zUQ_fUdnsQ8a_9tbYk,963
+nshtrainer/_submit/session/_output.py,sha256=CNGH5W6_XxAC5-TRvMAMxOHd3fjGpJhK-7RGTDyvMu4,245
+nshtrainer/_submit/session/_script.py,sha256=0AeBgBduDsoIEBrY9kebARiBUEGc50JAD9oE_IDiLnA,3775
+nshtrainer/_submit/session/lsf.py,sha256=p19EP6OhROZxqfRhzeTD7GDmfYaREIKMXMOI8G933FE,14307
+nshtrainer/_submit/session/slurm.py,sha256=JpAjQvck4LjGN8o8fOvIeMuFqrg1cioANoVsX5hU-3g,17594
+nshtrainer/_submit/session/unified.py,sha256=gfh-AtnMyFHzcQOUlhlAR__vaWDk1r9XCivz_t_lHKk,11695
+nshtrainer/actsave/__init__.py,sha256=G1T-fELuGWkVqdhdyoePtj2dTOUtcIOW4VgsXv9JNTA,338
+nshtrainer/actsave/_callback.py,sha256=QoTa60F70f1RxB41VKixN9l5_htfFQxXDPHHSNFreuk,2770
+nshtrainer/actsave/_loader.py,sha256=fAhD32DrJa4onkYfcwc21YIeGEYzOSXCK_HVo9SZLgQ,4604
+nshtrainer/actsave/_saver.py,sha256=0EHmQDhqVxQWRWWSyt03eP1K9ETiACMQYmsZkDMt6HY,9451
+nshtrainer/callbacks/__init__.py,sha256=ohE_MO_kX1o4SZwcipIXUA9m7XYcijEKJtGcoU8dTkY,1667
+nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
+nshtrainer/callbacks/base.py,sha256=WESZz1VSTl1xSGVXBmxFqWwbLxXcJp97jpg9zrE0EsY,3560
+nshtrainer/callbacks/early_stopping.py,sha256=jriSU761wf_qTJ9Bos0D3h5aDvZHYpRqK62Ne8aWp5I,3768
+nshtrainer/callbacks/ema.py,sha256=zKCtvzZFo0ORlwNZHjaMk-sJoxrlTtFWOzR-yGy95W0,12134
+nshtrainer/callbacks/finite_checks.py,sha256=kX3TIJsxyqx0GuLJfYsqVgKU27zwjG9Z8324lyCFtwM,2087
+nshtrainer/callbacks/gradient_skipping.py,sha256=ModaIXpb69LbA8TpEXKRLdr4Sq7-l0CWnN6fvpaV188,3477
+nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
+nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=ZT0bn7X0BZbQXbk6fos47NsbbhD4Z9c9YmFqdcUEqus,1503
+nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
+nshtrainer/callbacks/norm_logging.py,sha256=IMrK0WiVSDFyspwyPpwELMK4mmd5Jpx4enAW_GsWbi4,6284
+nshtrainer/callbacks/on_exception_checkpoint.py,sha256=eDyB7qkpPdAaKjAY2uFMMY8Nht6TGeuDnsgHuKtp8eA,1615
+nshtrainer/callbacks/print_table.py,sha256=FcA-CBWwMf9c1NNRinvYpZC400RNQxuP28bJfgniT3Q,2840
+nshtrainer/callbacks/throughput_monitor.py,sha256=YQLdpX3LGybIiD814yT9yCCVSEXRWf8WwsvVaN5aDBE,1848
+nshtrainer/callbacks/timer.py,sha256=sDXPPcdDKu5xnuK_bjr8plIq9MBuluNJ42Mt9LvPZzc,4610
+nshtrainer/callbacks/wandb_watch.py,sha256=pUpMsNxd03ex1rzOmFw2HzGOXjnQGaH84m8cc2dXo4g,2937
+nshtrainer/config.py,sha256=0Fj5w-ry0BRl2_zJI6jwCnmMWE3p_eD8_Wn-NyFkTqU,10442
+nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
+nshtrainer/data/balanced_batch_sampler.py,sha256=bcJBcQjh1hB1yKF_xSlT9AtEWv0BJjYc1CuH2BF-ea8,4392
+nshtrainer/data/transform.py,sha256=JeGxvytQly8hougrsdMmKG8gJ6qvFPDglJCO4Tp6STk,1795
+nshtrainer/lr_scheduler/__init__.py,sha256=GNGmkcJD3jgCMk7pfaanAYrKz9957qkx6_Q0rssiHK0,738
+nshtrainer/lr_scheduler/_base.py,sha256=1tWMABevKZAuGhJN8Me2E9eqEyqoLtsG0bADPjED7a4,3752
+nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=VhsxZJ_Mw9zjkAGunFQ1KRub5_QM5NRqaEFWtmedFp8,5212
+nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=Ct-uLo8Q4t7lJ_HwoLRhNmudnCw4cSnblpBEg22aVTI,2691
+nshtrainer/model/__init__.py,sha256=PdvZkpAVkqvCLipGJvEHFU3WxnSMxYpvtuOkvLIenxg,2078
+nshtrainer/model/base.py,sha256=bhngGHxr0suQB9Ezi_3d5JgDWYqS_yPgGJZrGmc1TnI,23571
+nshtrainer/model/config.py,sha256=RMDdrbtvwm5vTFPxQ2x1hqiBIEEE-OAknhF6KTWfkkk,70293
+nshtrainer/model/modules/callback.py,sha256=JF59U9-CjJsAIspEhTJbVaGN0wGctZG7UquE3IS7R8A,6408
+nshtrainer/model/modules/debug.py,sha256=DTVty8cKnzj1GCULRyGx_sWTTsq9NLi30dzqjRTnuCU,1127
+nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
+nshtrainer/model/modules/logger.py,sha256=XEeo3QrplTNKZqfl6iWZf3fze3R4YOeOvs-RKVHFoQs,5527
+nshtrainer/model/modules/profiler.py,sha256=rQ_jRMcM1Z2AIROZlRnBRHM5rkTpq67afZPD6CIRfXs,825
+nshtrainer/model/modules/rlp_sanity_checks.py,sha256=o6gUceFwsuDHmL8eLOYuT3JGXFzq_qc4awl2RWaBygU,8900
+nshtrainer/model/modules/shared_parameters.py,sha256=mD5wrlBE3c025vzVdTpnSyC8yxzuI-aUWMmPhqPT0a0,2694
+nshtrainer/nn/__init__.py,sha256=57LPaP3G-BBGD2eGxbBUABNgYl3s_oASwrtOSS4bzTs,1339
+nshtrainer/nn/mlp.py,sha256=i-dHk0tomO_XlU6cKN4CC4HxTaYb-ukBCAgY1ySXl4I,3963
+nshtrainer/nn/module_dict.py,sha256=NOY0B6WDTnktyWH4GthsprMQo0bpehC-hCq9SfD8paE,2329
+nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,1719
+nshtrainer/nn/nonlinearity.py,sha256=IhIR8NCTY3Np9dMDnUouERR8ZhWpK3S0hTbT0i8HezU,3645
+nshtrainer/optimizer.py,sha256=JiLNRtcfYxyhAab1Z1QcEzmrX9S_JyrBS67TXy12kXI,1557
+nshtrainer/runner.py,sha256=9HsYB58aasY9RVvya_gPECDs_MBhM1fl4cbM3iJYTDc,600
+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=aGg6kRiHiqtAdGlEvEvGLmOy7AvRHTSkXdTmZpRXbjU,8435
+nshtrainer/trainer/trainer.py,sha256=oi8KdHF1AdZ54KFbCFAEI7W-C7qRtRe-KtOjNwBuS3M,14033
+nshtrainer/typecheck.py,sha256=CFkmPIxCU24CHk_7_pykb-Y1PRNhpLgsVZw1zuuOS_U,4614
+nshtrainer/util/environment.py,sha256=_SEtiQ_s5bL5pllUlf96AOUv15kNvCPvocVC13S7mIk,4166
+nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
+nshtrainer/util/singleton.py,sha256=nLhpuMZxl0zdNsnvS97o4ASUnKzCWYEKLzR_j9oP_xs,2208
+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.1.0.dist-info/METADATA,sha256=3zdNPxyB-I6Gudq2gTaU0crdgmDCcGCp6Zudef0DtuM,529
+nshtrainer-0.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.1.0.dist-info/RECORD,,

nshtrainer-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.9.0
+Root-Is-Purelib: true
+Tag: py3-none-any