nntool 2.0.0rc0__py3-none-any.whl

nntool/utils/__init__.py

@@ -0,0 +1,12 @@
+# this is for backward compatibility
+from ..experiment.utils import (
+    get_current_time,
+    get_output_path,
+    read_toml_file,
+)
+
+__all__ = [
+    "get_current_time",
+    "get_output_path",
+    "read_toml_file",
+]

nntool/version.py

@@ -0,0 +1,11 @@
+_MAJOR = "2"
+_MINOR = "0"
+# On main and in a nightly release the patch should be one ahead of the last
+# released build.
+_PATCH = "0"
+# This is mainly for nightly builds which have the suffix ".dev$DATE". See
+# https://semver.org/#is-v123-a-semantic-version for the semantics.
+_SUFFIX = "rc"
+
+VERSION_SHORT = "{0}.{1}".format(_MAJOR, _MINOR)
+VERSION = "{0}.{1}.{2}{3}".format(_MAJOR, _MINOR, _PATCH, _SUFFIX)

nntool/wandb/__init__.py

@@ -0,0 +1,7 @@
+from .config import WandbConfig, init_wandb
+
+
+__all__ = [
+    "WandbConfig",
+    "init_wandb",
+]

nntool/wandb/config.py

@@ -0,0 +1,116 @@
+import os
+import git
+import wandb
+import tomli
+import warnings
+
+from dataclasses import dataclass, field
+from typing import Optional, Union
+from wandb.sdk.lib.disabled import RunDisabled
+from wandb.sdk.wandb_run import Run
+
+
+@dataclass
+class WandbConfig:
+    """
+    Configuration class for Weights and Biases (wandb) integration.
+
+    :param api_key_config_file: The file path to the configuration file containing the wandb API key.
+                                The file should be a toml file with a `[wandb]` section. Default is an empty string.
+
+    :param dir: The directory to save the wandb logs. Default is None.
+
+    :param project: The name of the project in wandb. Default is an empty string.
+
+    :param entity: The wandb user or team name. Default is an empty string.
+
+    :param name: The name of the wandb run. Leave blank to use the default run name. Default is an empty string.
+
+    :param notes: Notes or comments for the wandb run. Default is an empty string.
+
+    :param log_git_hash: Whether to log the current Git hash. Default is True.
+
+    :param log_code: Whether to log the current codebase. Default is True.
+
+    :param code_root: The root directory of the codebase to be logged. Default is the current directory (`.`).
+
+    :param code_ext: A list of file extensions for the code files to be logged. Default includes `.py` and `.sh`.
+    """
+
+    # wandb api key (toml file with [wandb] key field)
+    api_key_config_file: str = ""
+
+    # dir to save the wandb logs
+    dir: Optional[str] = None
+
+    # project name in wandb
+    project: str = ""
+
+    # wandb user name
+    entity: str = ""
+
+    # wandb run name (leave blacnk to use default name)
+    name: str = ""
+
+    # wandb run notes
+    notes: str = ""
+
+    # log git hash
+    log_git_hash: bool = True
+
+    # log code
+    log_code: bool = True
+
+    # code root
+    code_root: str = "."
+
+    # code file extensions
+    code_ext: list[str] = field(default_factory=lambda: [".py", ".sh"])
+
+
+def is_wandb_enabled():
+    return wandb.run is not None
+
+
+def init_wandb(args: WandbConfig, run_config: dict) -> Union[Run, RunDisabled, None]:
+    """initialize wandb and log the configuration
+
+    :param args: WandbConfig object
+    :param run_config: configuration dictionary to be logged
+    :return: wandb run object
+    """
+    project, entity = args.project, args.entity
+    if "WANDB_API_KEY" in os.environ:
+        warnings.warn("WANDB_API_KEY is found in environment variables. Using it.")
+        wandb.login(key=os.environ["WANDB_API_KEY"])
+    elif args.api_key_config_file:
+        with open(args.api_key_config_file, "rb") as config_file:
+            config_data = tomli.load(config_file)
+        wandb.login(key=config_data["wandb"]["key"])
+        project = config_data["wandb"].get("project", args.project)
+        entity = config_data["wandb"].get("entity", args.entity)
+    else:
+        warnings.warn("WANDB_API_KEY is not found in environment variables or the local key file.")
+        pass
+
+    if args.log_git_hash:
+        repo = git.Repo(search_parent_directories=True)
+        git_hash = repo.head.object.hexsha
+        args.notes = f"{args.notes + ', ' if args.notes else ''}" + f"git hash: {git_hash}"
+
+    wandb_run = wandb.init(
+        dir=args.dir,
+        project=project,
+        entity=entity,
+        name=args.name if args.name else None,
+        notes=args.notes,
+        config=run_config,
+    )
+    if not is_wandb_enabled():
+        warnings.warn("wandb is not enabled after intialization. Please check `wandb enabled`.")
+
+    wandb_run.log_code(
+        root=args.code_root,
+        include_fn=lambda path, root: any([path.endswith(ext) for ext in args.code_ext]),
+    )
+    return wandb_run

nntool-2.0.0rc0.dist-info/METADATA

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: nntool
+Version: 2.0.0rc0
+Summary: NNTool is a package built on top of submitit designed to provide simple abstractions to conduct experiments on Slurm for machine learning research.
+Author-email: Junhao Liu <junhaoliu17@gmail.com>
+Project-URL: Homepage, https://github.com/jhliu17/nntool
+Keywords: deep learning,neural network,research
+Requires-Python: >=3.10
+Requires-Dist: tyro
+Requires-Dist: submitit>=1.5.4
+Requires-Dist: wandb>=0.15.0
+Requires-Dist: tomli>=2.0.1

nntool-2.0.0rc0.dist-info/RECORD

@@ -0,0 +1,25 @@
+nntool/__init__.py,sha256=ISEQYjns0E37EbFA__hp6bJJES5t_3Y_ZgtwZVKP8ko,117
+nntool/build_backend.py,sha256=kGAtEbEyjxq8YHXNun95Dl6hQ0dI6o4VVdP3RAmL24E,759
+nntool/version.py,sha256=GeySxzv0TjgCbYsd3GB6Xo-twF0ZgnVguEaHdn9wgmI,409
+nntool/experiment/__init__.py,sha256=GIr3ehBQSTTiQyxn7WDiqna7SQVO5YvAG7WbMl7gg8s,239
+nntool/experiment/config.py,sha256=6cjk9pBdKjs9-Syawbjw0RtgfQybnoJkRpxM-g993O0,3807
+nntool/experiment/utils.py,sha256=ajrjYNltyNwNfmFciygCvbcRBtn8Gc7-dVR07Q8xbKo,2119
+nntool/slurm/__init__.py,sha256=YHFIpTTVDR-xmj7xDmfsRfa25UNThO4qISVS01NtiPc,429
+nntool/slurm/config.py,sha256=wDsdHhcc5tNR9nyLWG5SwcvKJoqo84dedz_3v0WcqdM,8260
+nntool/slurm/function.py,sha256=_awI1g8hEaOQxBbRMmO7jJwqBXtYHPTaaZ9g1XK7JsM,8865
+nntool/slurm/task.py,sha256=3rEgEXKL9q3VnDRtxRH6LzqaKjXDyiCOxe3aV6mgSeQ,11132
+nntool/slurm/wrap.py,sha256=ZtHDqJD6Uk2rwumdiewz6Yl6tVtau5tZ5ZCo9VF_JJw,5291
+nntool/slurm/accelerator/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nntool/slurm/accelerator/utils.py,sha256=VuGfPGjcanDXmDHtZh7U_jEx0k45aBtAJAsUJs_LoRc,1276
+nntool/slurm/core/__init__.py,sha256=920Q5lxUFOpngHN0ZAXRjgb4GqPXSyJAhlnQk_hbKYM,79
+nntool/slurm/core/_slurm.py,sha256=1tZ_TDFFsolsjTXbFsBTJUd7pBRSXfVbJrBrwUukfdE,24660
+nntool/slurm/core/_slurm_context.py,sha256=oW9R9fwa5IAnA8LtbLkyn_wPf-54aDMRoF8vhyhUQN8,1912
+nntool/slurm/parser/__init__.py,sha256=yxIckzha7nSlsdjPYKrP9ozneeQFtBdziTPJeVPXAFY,72
+nntool/slurm/parser/parse.py,sha256=IqlcqQfd7GoeTGJkMOrb2_-x0s49Zz_2fA3Fsk_6ysQ,552
+nntool/utils/__init__.py,sha256=YnGpWfnDUhxOR1-wx0-tjKOdHr9cP3_B242xoIT5Kww,219
+nntool/wandb/__init__.py,sha256=lExtgvdu-E60vOyqlNIc7cFF565U4fGFDHFKrZyMxtg,97
+nntool/wandb/config.py,sha256=TzHZmatV8bGl_kixWgJx0MI_DiXjscEESXlbYsokIos,3712
+nntool-2.0.0rc0.dist-info/METADATA,sha256=czd8bC5BJGXcmEqouT27A9H5Vjqpnlvb0ECmsuvafn8,496
+nntool-2.0.0rc0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nntool-2.0.0rc0.dist-info/top_level.txt,sha256=hXO7zSAI-o9plYeUkCiKbkDpQCfTDIfjivyoSEqT28Y,7
+nntool-2.0.0rc0.dist-info/RECORD,,

nntool-2.0.0rc0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nntool-2.0.0rc0.dist-info/top_level.txt

@@ -0,0 +1 @@
+nntool