Dev (#6)

* Installing all the submititnow subpackages. * Fixed rich file content viewing using StringIO. * Added support for all major slurm arguments. * Implemented custom action for --nodelist * Removed `slurm` prefix from the option strings * Changed `exp_name` to `exp-name` * Added gpu_matmul.py in examples for GPU testing * Bumping the version to 0.9.2 --------- Authored-by: Maharshi Gor <[email protected]>
maharshi95 · Aug 1, 2023 · e1cffe3 · e1cffe3
1 parent 2e77fa2
commit e1cffe3
Show file tree

Hide file tree

Showing 4 changed files with 132 additions and 16 deletions.
diff --git a/examples/gpu_matmul.py b/examples/gpu_matmul.py
@@ -0,0 +1,34 @@
+import argparse
+from typing import Optional
+import torch
+from tqdm import trange
+
+def add_arguments(
+    parser: Optional[argparse.ArgumentParser] = None,
+) -> argparse.ArgumentParser:
+    if parser is None:
+        parser = argparse.ArgumentParser('Perform a matrix multiplication on a GPU')
+
+    parser.add_argument("--matrix-size", type=int, default=1000)
+    parser.add_argument("--n-iter", type=int, default=10)
+
+    return parser
+
+
+def main(args: argparse.Namespace):
+    # Set torch seed to 
+    torch.manual_seed(42)
+
+    for i in trange(args.n_iter):
+        M1 = torch.randn(args.matrix_size, args.matrix_size).cuda()
+        M2 = torch.randn(args.matrix_size, args.matrix_size).cuda()
+
+    result = M1 @ M2
+    norm = torch.norm(result, p="fro")
+    print(f"Norm of result: {norm}")
+
+
+if __name__ == "__main__":
+    parser = add_arguments()
+    args = parser.parse_args()
+    main(args)
diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setuptools.setup(
     name="submititnow",
-    version="0.9.1",
+    version="0.9.2",
     author="Maharshi Gor",
     author_email="[email protected]",
     description="A package to make submitit easier to use",
@@ -26,6 +26,7 @@
         "typer[all]>=0.7.0",
         "rich-cli>=1.8.0",
         "rich>=12.6.0",
+        "tqdm>=4.0.0",
     ],
     python_requires=">=3.8",
 )
diff --git a/submititnow/cli.py b/submititnow/cli.py
@@ -1,5 +1,5 @@
 from __future__ import annotations
-
+import io
 import time
 
 from rich import print as rich_print
@@ -13,12 +13,15 @@
 
 
 def show_file_content(filepath: str):
-    with open(filepath) as fp:
-        text = fp.read().replace("]\n ", "]\r")  # Handle tqdm progress bars
+    rich_print("[bold bright_yellow]Reading file:[/bold bright_yellow] [bold cyan]{}[/bold cyan]\n".format(filepath))
+    with open(filepath, "r", newline='') as fp:
+        text = fp.read()
         for line in text.split("\n"):
-            if "\r" in line:
-                line = line[line.rindex("\r") + 1 :]
-            rich_print(line)
+            line_buffer = io.StringIO()
+            for chunks in line.split("\r"):
+                line_buffer.seek(0)
+                line_buffer.write(chunks)
+            rich_print(line_buffer.getvalue())
 
 
 def _generate_console_table(exp: Experiment):

diff --git a/submititnow/options.py b/submititnow/options.py
@@ -2,31 +2,109 @@
 from typing import Dict, Any
 
 
+class SlurmAdditionalArgAction(argparse.Action):
+    def __init__(self, check_func, *args, **kwargs):
+        """
+        argparse custom action.
+        :param check_func: callable to do the real check.
+        """
+        self._check_func = check_func
+        super(SlurmAdditionalArgAction, self).__init__(*args, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string):
+        if isinstance(values, list):
+            values = [self._check_func(parser, v) for v in values]
+        else:
+            values = self._check_func(parser, values)
+        if option_string.startswith("--"):
+            option_string = option_string[2:]
+        setattr(namespace, self.dest, {option_string: values})
+
+
 def add_slurm_arguments(parser: argparse.ArgumentParser):
     slurm_group = parser.add_argument_group("SLURM parameters")
     slurm_group.add_argument(
-        "--slurm_profile", default=None, help="SubmititNow profile for SLURM."
+        "--profile",
+        default=None,
+        help="SubmititNow profile for SLURM.",
+        dest="slurm_profile",
+    )
+    slurm_group.add_argument(
+        "--account", default=None, help="SLURM account", dest="slurm_account"
+    )
+    slurm_group.add_argument(
+        "--partition", default=None, help="SLURM partition", dest="slurm_partition"
+    )
+    slurm_group.add_argument("--qos", default=None, help="SLURM qos", dest="slurm_qos")
+    slurm_group.add_argument(
+        "--mem", default=None, help="SLURM memory requirement", dest="slurm_mem"
+    )
+    slurm_group.add_argument(
+        "--gres", default=None, help="SLURM GPU Resource requirement", dest="slurm_gres"
     )
-    slurm_group.add_argument("--slurm_account", default=None, help="SLURM account")
-    slurm_group.add_argument("--slurm_partition", default=None, help="SLURM partition")
-    slurm_group.add_argument("--slurm_qos", default=None, help="SLURM qos")
     slurm_group.add_argument(
-        "--slurm_mem", default=None, help="SLURM memory requirement"
+        "--time", default=None, help="SLURM time requirement", dest="slurm_time"
     )
     slurm_group.add_argument(
-        "--slurm_gres", default=None, help="SLURM GPU Resource requirement"
+        "--nodes",
+        default=1,
+        help="SLURM nodes requirement",
+        dest="slurm_nodes",
+        type=int,
     )
     slurm_group.add_argument(
-        "--slurm_time", default=None, help="SLURM time requirement"
+        "--ntasks-per-node",
+        default=None,
+        help="SLURM ntasks per node",
+        dest="slurm_ntasks_per_node",
+        type=int,
+    )
+    slurm_group.add_argument(
+        "--cpus-per-task",
+        default=None,
+        help="SLURM cpus per task",
+        dest="slurm_cpus_per_task",
+        type=int,
+    )
+    slurm_group.add_argument(
+        "--cpus-per-gpu",
+        default=None,
+        help="SLURM cpus per gpu",
+        dest="slurm_cpus_per_gpu",
+        type=int,
+    )
+    slurm_group.add_argument(
+        "--gpus-per-node",
+        default=None,
+        help="SLURM gpus per node",
+        dest="slurm_gpus_per_node",
+        type=int,
+    )
+    slurm_group.add_argument(
+        "--gpus-per-task",
+        default=None,
+        help="SLURM gpus per task",
+        dest="slurm_gpus_per_task",
+        type=int,
+    )
+
+    # Additional arguments
+    slurm_group.add_argument(
+        "--nodelist",
+        default=None,
+        help="SLURM nodelist",
+        action=SlurmAdditionalArgAction,
+        check_func=lambda parser, value: value,
+        dest="slurm_additional_parameters",
     )
     return parser
 
 
 def add_submititnow_arguments(parser: argparse.ArgumentParser):
     submititnow_group = parser.add_argument_group("SubmititNow parameters")
-    submititnow_group.add_argument("--exp_name", default=None, help="Experiment Name.")
+    submititnow_group.add_argument("--exp-name", default=None, help="Experiment Name.")
     submititnow_group.add_argument(
-        "--submititnow_dir", default=None, help="Root directory for submititnow."
+        "--submititnow-dir", default=None, help="Root directory for submititnow."
     )
     return parser