nntool 1.6.2__cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl

nntool/__init__.py

@@ -0,0 +1,7 @@
+from .version import VERSION, VERSION_SHORT
+
+__all__ = [
+    "VERSION",
+    "VERSION_SHORT",
+]
+__version__ = VERSION

nntool/experiment/__init__.py

@@ -0,0 +1,13 @@
+from .config import BaseExperimentConfig
+from .utils import (
+    get_current_time,
+    get_output_path,
+    read_toml_file,
+)
+
+__all__ = [
+    "BaseExperimentConfig",
+    "get_current_time",
+    "get_output_path",
+    "read_toml_file",
+]

nntool/experiment/config.py

@@ -0,0 +1,108 @@
+import os
+
+from typing import Any, Dict
+from pathlib import Path
+from dataclasses import dataclass
+from .utils import get_output_path, read_toml_file
+
+
+@dataclass
+class BaseExperimentConfig:
+    """
+    Configuration class for setting up an experiment.
+
+    :param config_name: The name of the configuration.
+    :param output_folder: The folder path where the outputs will be saved.
+    :param experiment_name_key: Key for experiment name in the environment variable, default is 'EXP_NAME'.
+    :param env_toml_path: Path to the `env.toml` file, default is 'env.toml'.
+    :param append_date_to_path: If True, the current date and time will be appended to the output path, default is True.
+    :param existing_output_path_ok: If True, the existing output path is ok to be reused, default is False.
+    """
+
+    # config name
+    config_name: str
+
+    # the output folder for the outputs
+    output_folder: str
+
+    # key for experiment name in the environment variable
+    experiment_name_key: str = "EXP_NAME"
+
+    # the path to the env.toml file
+    env_toml_path: str = "env.toml"
+
+    # append date time to the output path
+    append_date_to_path: bool = True
+
+    # exisiting output path is ok
+    existing_output_path_ok: bool = False
+
+    def __post_init__(self):
+        # annotations
+        self.experiment_name: str
+        self.project_path: str
+        self.output_path: str
+        self.current_time: str
+        self.env_toml: Dict[str, Any] = self.__prepare_env_toml_dict()
+
+        self.experiment_name = self.__prepare_experiment_name()
+        self.project_path, self.output_path, self.current_time = self.__prepare_experiment_paths()
+
+        # custom post update for the derived class
+        self.set_up_stateful_fields()
+
+    def __prepare_env_toml_dict(self):
+        env_toml_path = Path(self.env_toml_path)
+        if not env_toml_path.exists():
+            raise FileNotFoundError(f"{env_toml_path} does not exist")
+
+        config = read_toml_file(env_toml_path)
+        return config
+
+    def __prepare_experiment_name(self):
+        return os.environ.get(self.experiment_name_key, "default")
+
+    def __prepare_experiment_paths(self):
+        project_path = self.env_toml["project"]["path"]
+
+        output_path, current_time = get_output_path(
+            output_path=os.path.join(self.output_folder, self.config_name, self.experiment_name),
+            append_date=self.append_date_to_path,
+            cache_into_env=False,
+        )
+        output_path = f"{project_path}/{output_path}"
+        return project_path, output_path, current_time
+
+    def get_output_path(self) -> str:
+        """Return the output path prepared for the experiment.
+
+        :return: output path for the experiment
+        """
+        return self.output_path
+
+    def get_current_time(self) -> str:
+        """Return the current time for the experiment.
+
+        :return: current time for the experiment
+        """
+        return self.current_time
+
+    def set_up_stateful_fields(self):
+        """
+        Post configuration steps for stateful fields such as `output_path` in the derived class.
+        This method should be overridden in the derived class.
+        """
+        pass
+
+    def start(self):
+        """
+        Start the experimen. This will
+        - cache `NNTOOL_OUTPUT_PATH` and `NNTOOL_OUTPUT_PATH_DATE` into environment variables, which means the later launched processes would inherit these variables.
+        - create the output path if it does not exist.
+        """
+        os.environ["NNTOOL_OUTPUT_PATH"] = self.get_output_path()
+        os.environ["NNTOOL_OUTPUT_PATH_DATE"] = self.get_current_time()
+
+        # create the output path
+        output_path = Path(self.get_output_path())
+        output_path.mkdir(parents=True, exist_ok=self.existing_output_path_ok)

nntool/experiment/utils.py

@@ -0,0 +1,63 @@
+import os
+import datetime
+import tomli
+
+
+def get_current_time() -> str:
+    """get current time in this format: MMDDYYYY/HHMMSS
+
+    :return: time in the format MMDDYYYY/HHMMSS
+    """
+    # Get the current time
+    current_time = datetime.datetime.now()
+
+    # Format the time (MDY/HMS)
+    formatted_time = current_time.strftime("%m%d%Y/%H%M%S")
+
+    return formatted_time
+
+
+def read_toml_file(file_path: str) -> dict:
+    """Read a toml file and return the content as a dictionary
+
+    :param file_path: path to the toml file
+    :return: content of the toml file as a dictionary
+    """
+    with open(file_path, "rb") as f:
+        content = tomli.load(f)
+
+    return content
+
+
+def get_output_path(
+    output_path: str = "./",
+    append_date: bool = True,
+    cache_into_env: bool = True,
+) -> tuple[str, str]:
+    """Get output path based on environment variable OUTPUT_PATH and NNTOOL_OUTPUT_PATH.
+    The output path is appended with the current time if append_date is True (e.g. /OUTPUT_PATH/xxx/MMDDYYYY/HHMMSS).
+
+    :param append_date: append a children folder with the date time, defaults to True
+    :param cache_into_env: whether cache the newly created path into env, defaults to True
+    :return: (output path, current time)
+    """
+    if "OUTPUT_PATH" in os.environ:
+        output_path = os.environ["OUTPUT_PATH"]
+        current_time = "" if not append_date else get_current_time()
+    elif "NNTOOL_OUTPUT_PATH" in os.environ:
+        # reuse the NNTOOL_OUTPUT_PATH if it is set
+        output_path = os.environ["NNTOOL_OUTPUT_PATH"]
+        current_time = "" if not append_date else os.environ["NNTOOL_OUTPUT_PATH_DATE"]
+    else:
+        current_time = get_current_time()
+        if append_date:
+            output_path = os.path.join(output_path, current_time)
+        print(
+            f"OUTPUT_PATH is not found in environment variables. NNTOOL_OUTPUT_PATH is set using path: {output_path}"
+        )
+
+        if cache_into_env:
+            os.environ["NNTOOL_OUTPUT_PATH"] = output_path
+            os.environ["NNTOOL_OUTPUT_PATH_DATE"] = current_time
+
+    return output_path, current_time

nntool/plot/__init__.py

@@ -0,0 +1,10 @@
+from .csrc import latexify, savefig, is_latexify_enabled
+from .context import latexify_plot, enable_latexify
+
+__all__ = [
+    "latexify",
+    "savefig",
+    "is_latexify_enabled",
+    "latexify_plot",
+    "enable_latexify",
+]

nntool/plot/context.py

@@ -0,0 +1,46 @@
+import os
+import matplotlib
+import seaborn as sns
+
+from typing import Union
+from dataclasses import dataclass
+from .csrc.latexify import SIZE_SMALL, latexify, savefig
+
+
+@dataclass
+class latexify_plot:
+    enable: bool = True
+    width_scale_factor: float = 1
+    height_scale_factor: float = 1
+    fig_width: Union[float, None] = None
+    fig_height: Union[float, None] = None
+    font_size: int = SIZE_SMALL
+
+    def __post_init__(self):
+        self.legend_size = 7 if self.enable else None
+
+    def __enter__(self):
+        if self.enable:
+            os.environ["LATEXIFY"] = "1"
+            latexify(
+                width_scale_factor=self.width_scale_factor,
+                height_scale_factor=self.height_scale_factor,
+                fig_width=self.fig_width,
+                fig_height=self.fig_height,
+                font_size=self.font_size,
+            )
+        return self
+
+    def __exit__(self, *args):
+        if self.enable:
+            os.environ.pop("LATEXIFY")
+            matplotlib.rcParams.update(matplotlib.rcParamsDefault)
+
+    def savefig(self, filename, despine: bool = True, fig_dir: str = "tests/plot", **kwargs):
+        if despine:
+            sns.despine()
+        savefig(filename, fig_dir=fig_dir, **kwargs)
+
+
+# This is for backward compatibility
+enable_latexify = latexify_plot

nntool/plot/csrc/__init__.py

@@ -0,0 +1,7 @@
+from .latexify import latexify, savefig, is_latexify_enabled
+
+__all__ = [
+    "latexify",
+    "savefig",
+    "is_latexify_enabled",
+]

nntool/plot/csrc/latexify.py

@@ -0,0 +1,129 @@
+"""
+This code is borrowed from https://github.com/probml/probml-utils/blob/main/probml_utils/plotting.py
+"""
+
+import os
+import matplotlib.pyplot as plt
+import warnings
+
+DEFAULT_WIDTH = 6.0
+DEFAULT_HEIGHT = 1.5
+SIZE_SMALL = 9  # Caption size in the pml book
+
+
+def latexify(
+    width_scale_factor=1,
+    height_scale_factor=1,
+    fig_width=None,
+    fig_height=None,
+    font_size=SIZE_SMALL,
+):
+    f"""
+    width_scale_factor: float, DEFAULT_WIDTH will be divided by this number, DEFAULT_WIDTH is page width: {DEFAULT_WIDTH} inches.
+    height_scale_factor: float, DEFAULT_HEIGHT will be divided by this number, DEFAULT_HEIGHT is {DEFAULT_HEIGHT} inches.
+    fig_width: float, width of the figure in inches (if this is specified, width_scale_factor is ignored)
+    fig_height: float, height of the figure in inches (if this is specified, height_scale_factor is ignored)
+    font_size: float, font size
+    """
+    if "LATEXIFY" not in os.environ:
+        warnings.warn("LATEXIFY environment variable not set, not latexifying")
+        return
+    if fig_width is None:
+        fig_width = DEFAULT_WIDTH / width_scale_factor
+    if fig_height is None:
+        fig_height = DEFAULT_HEIGHT / height_scale_factor
+
+    # use TrueType fonts so they are embedded
+    # https://stackoverflow.com/questions/9054884/how-to-embed-fonts-in-pdfs-produced-by-matplotlib
+    # https://jdhao.github.io/2018/01/18/mpl-plotting-notes-201801/
+    plt.rcParams["pdf.fonttype"] = 42
+
+    # Font sizes
+    # SIZE_MEDIUM = 14
+    # SIZE_LARGE = 24
+    # https://stackoverflow.com/a/39566040
+    plt.rc("font", size=font_size)  # controls default text sizes
+    plt.rc("axes", titlesize=font_size)  # fontsize of the axes title
+    plt.rc("axes", labelsize=font_size)  # fontsize of the x and y labels
+    plt.rc("xtick", labelsize=font_size)  # fontsize of the tick labels
+    plt.rc("ytick", labelsize=font_size)  # fontsize of the tick labels
+    plt.rc("legend", fontsize=font_size)  # legend fontsize
+    plt.rc("figure", titlesize=font_size)  # fontsize of the figure title
+
+    # latexify: https://nipunbatra.github.io/blog/visualisation/2014/06/02/latexify.html
+    plt.rcParams["backend"] = "ps"
+    plt.rc("text", usetex=True)
+    plt.rc("font", family="serif")
+    plt.rc("figure", figsize=(fig_width, fig_height))
+    plt.rcParams["text.latex.preamble"] = r"""
+    \usepackage{amsmath}
+    \usepackage{bm}
+    """
+
+
+def is_latexify_enabled():
+    """
+    returns true if LATEXIFY environment variable is set
+    """
+    return "LATEXIFY" in os.environ
+
+
+def _get_fig_name(fname_full):
+    fname_full = fname_full.replace("_latexified", "")
+    LATEXIFY = "LATEXIFY" in os.environ
+    # extention = "_latexified.pdf" if LATEXIFY else ".png"
+    extention = "_latexified.pdf" if LATEXIFY else ".pdf"
+    if fname_full[-4:] in [".png", ".pdf", ".jpg"]:
+        fname = fname_full[:-4]
+        print(
+            f"renaming {fname_full} to {fname}{extention} because LATEXIFY is {LATEXIFY}",
+        )
+    else:
+        fname = fname_full
+    return fname + extention
+
+
+def savefig(
+    f_name,
+    tight_layout=True,
+    tight_bbox=False,
+    pad_inches=0.0,
+    fig_dir=None,
+    *args,
+    **kwargs,
+):
+    if len(f_name) == 0:
+        return
+    if "FIG_DIR" not in os.environ and fig_dir is None:
+        warnings.warn("set FIG_DIR environment variable or pass fig_dir argument to save figures")
+        return
+
+    fig_dir = fig_dir if fig_dir is not None else os.environ["FIG_DIR"]
+    # Auto create the directory if it doesn't exist
+    if not os.path.exists(fig_dir):
+        os.makedirs(fig_dir)
+
+    fname_full = os.path.join(fig_dir, f_name)
+    fname_full = _get_fig_name(fname_full)
+
+    print("saving image to {}".format(fname_full))
+    if tight_layout:
+        plt.tight_layout(pad=pad_inches)
+    print("Figure size:", plt.gcf().get_size_inches())
+
+    if tight_bbox:
+        # This changes the size of the figure
+        plt.savefig(fname_full, pad_inches=pad_inches, bbox_inches="tight", *args, **kwargs)
+    else:
+        plt.savefig(fname_full, pad_inches=pad_inches, *args, **kwargs)
+
+    if "DUAL_SAVE" in os.environ:
+        if fname_full.endswith(".pdf"):
+            fname_full = fname_full[:-4] + ".png"
+        else:
+            fname_full = fname_full[:-4] + ".pdf"
+        if tight_bbox:
+            # This changes the size of the figure
+            plt.savefig(fname_full, pad_inches=pad_inches, bbox_inches="tight", *args, **kwargs)
+        else:
+            plt.savefig(fname_full, pad_inches=pad_inches, *args, **kwargs)

nntool/slurm/__init__.py

@@ -0,0 +1,21 @@
+from .config import SlurmConfig, SlurmArgs
+from .wrap import (
+    slurm_function,
+    slurm_fn,
+    slurm_launcher,
+)
+from .function import SlurmFunction
+from .task import Task, DistributedTaskConfig, PyTorchDistributedTask
+
+
+__all__ = [
+    "SlurmConfig",
+    "SlurmArgs",
+    "SlurmFunction",
+    "slurm_fn",
+    "slurm_function",
+    "slurm_launcher",
+    "Task",
+    "DistributedTaskConfig",
+    "PyTorchDistributedTask",
+]

nntool/slurm/accelerator/utils.py

@@ -0,0 +1,37 @@
+import subprocess
+
+
+def nvidia_smi_gpu_memory_stats() -> dict:
+    """
+    Parse the nvidia-smi output and extract the memory used stats.
+    """
+    out_dict = {}
+    try:
+        sp = subprocess.Popen(
+            ["nvidia-smi", "--query-gpu=index,memory.used", "--format=csv,noheader"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            close_fds=True,
+        )
+        out_str = sp.communicate()
+        out_list = out_str[0].decode("utf-8").split("\n")
+        out_dict = {}
+        for item in out_list:
+            if " MiB" in item:
+                gpu_idx, mem_used = item.split(",")
+                gpu_key = f"gpu_{gpu_idx}_mem_used_gb"
+                out_dict[gpu_key] = int(mem_used.strip().split(" ")[0]) / 1024
+    except FileNotFoundError:
+        raise Exception("Failed to find the 'nvidia-smi' executable for printing GPU stats")
+    except subprocess.CalledProcessError as e:
+        raise Exception(f"nvidia-smi returned non zero error code: {e.returncode}")
+
+    return out_dict
+
+
+def nvidia_smi_gpu_memory_stats_str() -> str:
+    """
+    Parse the nvidia-smi output and extract the memory used stats.
+    """
+    stats = nvidia_smi_gpu_memory_stats()
+    return ", ".join([f"{k}: {v:.4f}" for k, v in stats.items()])

nntool/slurm/config.py

@@ -0,0 +1,208 @@
+import os
+import sys
+
+from dataclasses import dataclass, field, replace
+from typing import List, Literal, Dict, Optional
+
+
+@dataclass
+class SlurmConfig:
+    """
+    Configuration class for SLURM job submission and execution.
+
+    Args:
+        mode (Literal["debug", "exec", "local", "slurm"]): Running mode for the job. Options include:
+            "debug" (default, run debugging which will involve pdb), "exec" (alias of local), "local" (run the job locally), or "slurm" (run the job on a SLURM cluster).
+
+        job_name (str): The name of the SLURM job. Default is 'Job'.
+
+        partition (str): The name of the SLURM partition to use. Default is ''.
+
+        output_parent_path (str): The parent directory path for saving the slurm folder. Default is './'.
+
+        output_folder (str): The folder name where SLURM output files will be stored. Default is 'slurm'.
+
+        node_list (str): A string specifying the nodes to use. Leave blank to use all available nodes. Default is an empty string.
+
+        node_list_exclude (str): A string specifying the nodes to exclude. Leave blank to use all nodes in the node list. Default is an empty string.
+
+        num_of_node (int): The number of nodes to request. Default is 1.
+
+        tasks_per_node (int): The number of tasks to run per node. Default is 1.
+
+        gpus_per_task (int): The number of GPUs to request per task. Default is 0.
+
+        cpus_per_task (int): The number of CPUs to request per task. Default is 1.
+
+        gpus_per_node (int): The number of GPUs to request per node. If this is set, `gpus_per_task` will be ignored. Default is None.
+
+        mem (str): The amount of memory (GB) to request. Leave blank to use the default memory configuration of the node. Default is an empty string.
+
+        timeout_min (int): The time limit for the job in minutes. Default is `sys.maxsize` for effectively no limit.
+
+        stderr_to_stdout (bool): Whether to redirect stderr to stdout. Default is False.
+
+        setup (List[str]): A list of environment variable setup commands. Default is an empty list.
+
+        pack_code (bool): Whether to pack the codebase before submission. Default is False.
+
+        use_packed_code (bool): Whether to use the packed code for execution. Default is False.
+
+        code_root (str): The root directory of the codebase, which will be used by the code packing. Default is the current directory (``.``).
+
+        code_file_suffixes (List[str]): A list of file extensions for code files to be included when packing. Default includes ``.py``, ``.sh``, ``.yaml``, and ``.toml``.
+
+        exclude_code_folders (List[str]): A list of folder names relative to `code_root` that will be excluded from packing. Default excludes 'wandb', 'outputs', and 'datasets'.
+
+        use_distributed_env (bool): Whether to use a distributed environment for the job. Default is False.
+
+        distributed_env_task (Literal["torch"]): The type of distributed environment task to use. Currently, only "torch" is supported. Default is "torch".
+
+        processes_per_task (int): The number of processes to run per task. This value is not used by SLURM but is relevant for correctly set up distributed environments. Default is 1.
+
+        distributed_launch_command (str): The command to launch distributed environment setup, using environment variables like ``{num_processes}``, ``{num_machines}``, ``{machine_rank}``, ``{main_process_ip}``, ``{main_process_port}``. Default is an empty string.
+
+        extra_params_kwargs (Dict[str, str]): Additional parameters for the SLURM job as a dictionary of key-value pairs. Default is an empty dictionary.
+
+        extra_submit_kwargs (Dict[str, str]): Additional submit parameters for the SLURM job as a dictionary of key-value pairs. Default is an empty dictionary.
+
+        extra_task_kwargs (Dict[str, str]): Additional task parameters for the SLURM job as a dictionary of key-value pairs. Default is an empty dictionary.
+    """
+
+    # running mode
+    mode: Literal["debug", "exec", "local", "slurm"] = "debug"
+
+    # slurm job name
+    job_name: str = "Job"
+
+    # slurm partition name
+    partition: str = ""
+
+    # slurm output parent path
+    output_parent_path: str = "./"
+
+    # slurm output folder name
+    output_folder: str = "slurm"
+
+    # node list string (leave blank to use all nodes)
+    node_list: str = ""
+
+    # node list string to be excluded (leave blank to use all nodes in the node list)
+    node_list_exclude: str = ""
+
+    # number of nodes to request
+    num_of_node: int = 1
+
+    # tasks per node
+    tasks_per_node: int = 1
+
+    # number of gpus per task to request
+    gpus_per_task: int = 0
+
+    # number of cpus per task to request
+    cpus_per_task: int = 1
+
+    # number of gpus per node to request (if this is set, gpus_per_task will be ignored)
+    gpus_per_node: Optional[int] = None
+
+    # memory (GB) to request (leave black to use default memory configurations in the node)
+    mem: str = ""
+
+    # time out min
+    timeout_min: int = sys.maxsize
+
+    # whether to redirect stderr to stdout
+    stderr_to_stdout: bool = False
+
+    # environment variables setup command
+    setup: List[str] = field(default_factory=list)
+
+    # whether to pack code
+    pack_code: bool = False
+
+    # use packed code to run
+    use_packed_code: bool = False
+
+    # code root
+    code_root: str = "."
+
+    # code file extensions
+    code_file_suffixes: list[str] = field(default_factory=lambda: [".py", ".sh", ".yaml", ".toml"])
+
+    # exclude folders (relative to the code root)
+    exclude_code_folders: list[str] = field(
+        default_factory=lambda: ["wandb", "outputs", "datasets"]
+    )
+
+    # whether to use distributed environment
+    use_distributed_env: bool = False
+
+    # distributed enviroment task
+    distributed_env_task: Literal["torch"] = "torch"
+
+    # processes per task (this value is not used by slurm, but in the distributed environment)
+    processes_per_task: int = 1
+
+    # distributed launch command (this will be called after the distributed enviroment is set up)
+    # the following environment variables are available:
+    #   num_processes: int
+    #   num_machines: int
+    #   machine_rank: int
+    #   main_process_ip: str
+    #   main_process_port: int
+    # use braces to access the environment variables, e.g. {num_processes}
+    distributed_launch_command: str = ""
+
+    # extra slurm job parameters
+    extra_params_kwargs: Dict[str, str] = field(default_factory=dict)
+
+    # extra slurm submit parameters
+    extra_submit_kwargs: Dict[str, str] = field(default_factory=dict)
+
+    # extra slurm task parameters
+    extra_task_kwargs: Dict[str, str] = field(default_factory=dict)
+
+    def _configuration_check(self):
+        # check partition
+        if self.partition == "":
+            raise ValueError("partition must be set")
+
+        # check distributed enviroment task
+        if self.use_distributed_env and self.distributed_launch_command == "":
+            raise ValueError(
+                "distributed_launch_command must be set when use_distributed_env is True"
+            )
+
+    def __post_init__(self):
+        # check configuration
+        self._configuration_check()
+
+        # normalize the output folder
+        output_folder_suffix = ""
+        if self.mode != "slurm":
+            output_folder_suffix = f"_{self.mode}"
+        if self.output_folder.endswith("slurm"):
+            self.output_folder = f"{self.output_folder}{output_folder_suffix}"
+        else:
+            self.output_folder = os.path.join(self.output_folder, f"slurm{output_folder_suffix}")
+
+        # output path
+        self.output_path: str = os.path.join(self.output_parent_path, self.output_folder)
+
+    def set_output_path(self, output_parent_path: str) -> "SlurmConfig":
+        """Set output path and date for the slurm job.
+
+        Args:
+            output_parent_path (str): The parent path for the output.
+
+        Returns:
+            SlurmConfig: The updated SlurmConfig instance.
+        """
+        new_config = replace(
+            self,
+            output_parent_path=output_parent_path,
+        )
+        return new_config
+
+
+SlurmArgs = SlurmConfig

nntool/slurm/csrc/__init__.py

@@ -0,0 +1,4 @@
+from ._slurm import SlurmFunction as _SlurmFunction
+
+
+__all__ = ["_SlurmFunction"]