diff --git a/.vscode/settings.json b/.vscode/settings.json index f83bf4f..a7b9e2e 100644 --- a/.vscode/settings.json +++ b/.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 } \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..d76d5e6 --- /dev/null +++ b/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 diff --git a/ReadMe.md b/ReadMe.md index 6656690..a1af98f 100644 --- a/ReadMe.md +++ b/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. diff --git a/setup.py b/setup.py index 532f4b3..c839bc0 100644 --- a/setup.py +++ b/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", diff --git a/tests/test-results/pytest_in_tests.xml b/tests/test-results/pytest_in_tests.xml new file mode 100644 index 0000000..f94ae52 --- /dev/null +++ b/tests/test-results/pytest_in_tests.xml @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/tests/test_decorators.py b/tests/test_decorators.py index e4ec9f4..8ae4288 100644 --- a/tests/test_decorators.py +++ b/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() diff --git a/znamutils/__init__.py b/znamutils/__init__.py index df4521a..d0561b0 100644 --- a/znamutils/__init__.py +++ b/znamutils/__init__.py @@ -1 +1 @@ -from .decorators import slurm_it \ No newline at end of file +from .decorators import slurm_it diff --git a/znamutils/decorators.py b/znamutils/decorators.py index 6121b71..48cd2bc 100644 --- a/znamutils/decorators.py +++ b/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: