Merge pull request #1 from znamlab/dev

v0.2

.vscode/settings.json

@@ -4,5 +4,10 @@
         "--line-length",
         "88"
     ],
-    "python.linting.enabled": true
+    "python.linting.enabled": true,
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
 }

CHANGELOG.md

@@ -0,0 +1,10 @@
+
+## Changelog
+
+### [v0.2] - 2023-12-05
+
+- Add option to change slurm options when calling the decorated function
+
+### [v0.1] - 2023-08-13
+
+- Initial release

ReadMe.md

@@ -90,6 +90,8 @@ use_slurm (bool): whether to use slurm or not
 job_dependency (str): job id to depend on
 slurm_folder (str): where to write the slurm script and logs
 scripts_name (str): name of the slurm script and python file
+slurm_options (dict): options to pass to sbatch, will update the default options 
+  provided in the decorator.
 ```
 
 When `use_slurm = True`, `slurm_folder` must be provided. 

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="znamutils",
-    version="v0.1",
+    version="v0.2",
     packages=find_packages(exclude=["tests"]),
     url="https://github.com/znamlab/znamutils",
     license="MIT",

tests/test-results/pytest_in_tests.xml

@@ -0,0 +1 @@
+<?xml version="1.0" encoding="utf-8"?><testsuites><testsuite name="pytest" errors="0" failures="0" skipped="0" tests="3" time="3.478" timestamp="2023-12-05T17:31:48.510906" hostname="int000"><testcase classname="tests.test_decorators" name="test_slurm_my_func" time="0.223" /><testcase classname="tests.test_decorators" name="test_update_slurm_options" time="0.052" /><testcase classname="tests.test_slurm_helpers" name="test_create_slurm_sbatch" time="0.000" /></testsuite></testsuites>

tests/test_decorators.py

@@ -85,5 +85,38 @@ def test_func(a, b):
     assert "from pandas import Dataframe" in txt
 
 
+def test_update_slurm_options():
+    slurm_folder = (
+        Path(flz.PARAMETERS["data_root"]["processed"]) / "test" / "test_slurm_it"
+    )
+    slurm_folder.mkdir(exist_ok=True)
+
+    @slurm_it(conda_env="cottage_analysis", slurm_options={"time": "00:01:00"})
+    def test_func(a, b):
+        from datetime import datetime
+
+        print("inner test_func")
+        print(a, b)
+        print(datetime.now())
+        return a + b
+
+    test_func(1, 2, use_slurm=True, slurm_folder=slurm_folder)
+    sbatch_file = slurm_folder / "test_func.sh"
+    assert sbatch_file.exists()
+    with open(sbatch_file, "r") as f:
+        txt = f.read()
+        assert "#SBATCH --time=00:01:00" in txt
+    test_func(
+        1,
+        2,
+        use_slurm=True,
+        slurm_folder=slurm_folder,
+        slurm_options={"time": "00:02:00"},
+    )
+    with open(sbatch_file, "r") as f:
+        txt = f.read()
+        assert "#SBATCH --time=00:02:00" in txt
+
+
 if __name__ == "__main__":
     test_slurm_my_func()

znamutils/__init__.py

@@ -1 +1 @@
-from .decorators import slurm_it
+from .decorators import slurm_it

znamutils/decorators.py

@@ -30,6 +30,7 @@ def slurm_it(
         job_dependency (str): job id to depend on
         slurm_folder (str): where to write the slurm script and logs
         scripts_name (str): name of the slurm script and python file
+        slurm_options (dict): options to pass to sbatch
 
     The default slurm options are:
         ntasks=1
@@ -54,6 +55,9 @@ def slurm_it(
     Returns:
         function: decorated function
     """
+    # make a copy of default slurm options to avoid modifying the original
+    default_slurm_options = slurm_options.copy()
+
     # add parameters to the wrapped function signature
     func_sig = signature(func)
     use_slurm = Parameter(
@@ -68,9 +72,13 @@ def slurm_it(
     script_name = Parameter(
         "scripts_name", kind=Parameter.POSITIONAL_OR_KEYWORD, default=None
     )
+    slurm_options = Parameter(
+        "slurm_options", kind=Parameter.POSITIONAL_OR_KEYWORD, default=None
+    )
 
     new_sig = add_signature_parameters(
-        func_sig, last=(use_slurm, job_dependency, slurm_folder, script_name)
+        func_sig,
+        last=(use_slurm, job_dependency, slurm_folder, script_name, slurm_options),
     )
     from_imports = from_imports or {func.__module__: func.__name__}
     # create the new function with modified signature
@@ -82,6 +90,10 @@ def new_func(*args, **kwargs):
         job_dependency = kwargs.pop("job_dependency")
         slurm_folder = kwargs.pop("slurm_folder")
         scripts_name = kwargs.pop("scripts_name")
+        slurm_options = kwargs.pop("slurm_options")
+        if slurm_options is None:
+            slurm_options = {}
+        slurm_options = dict(default_slurm_options, **slurm_options)
 
         if not use_slurm:
             if job_dependency is not None: