Skip to content

Commit

Permalink
Merge pull request #5 from znamlab/dev
Browse files Browse the repository at this point in the history
Dev
  • Loading branch information
ablot authored Apr 4, 2024
2 parents a1a61c1 + 2a74256 commit 3d0d04c
Show file tree
Hide file tree
Showing 4 changed files with 13 additions and 3 deletions.
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@

## Changelog

### [v0.6] - 2024-04-04

- Change default cpu partition to `ncpu`

### [v0.5] - 2024-03-27

- Option to print job ID in log when starting.
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

setup(
name="znamutils",
version="v0.5",
version="v0.6",
packages=find_packages(exclude=["tests"]),
url="https://github.com/znamlab/znamutils",
license="MIT",
Expand Down
7 changes: 6 additions & 1 deletion znamutils/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ def slurm_it(
slurm_options=None,
imports=None,
from_imports=None,
print_job_id=False,
):
"""
Decorator to run a function on slurm.
Expand All @@ -25,7 +26,7 @@ def slurm_it(
Running the decorated function with use_slurm=True will create a slurm script and
a python script, submit the slurm script and return the job id of the slurm job.
The decorated function will have 4 new keyword arguments:
The decorated function will have 6 new keyword arguments:
use_slurm (bool): whether to use slurm or not
dependency_type (str, optional): Type of dependence on previous jobs.
Defaults to "afterok" which only runs the next job if all previous
Expand Down Expand Up @@ -55,6 +56,8 @@ def slurm_it(
instance {'numpy': 'mean'} results in `from numpy import mean`. If `None`,
the decorated function will be imported from its parent module. Defaults to
None.
print_job_id (bool, optional): Whether to print the job id of the slurm job in
the log file. Defaults to False.
Returns:
function: decorated function
Expand Down Expand Up @@ -139,6 +142,7 @@ def new_func(*args, **kwargs):
conda_env=conda_env,
slurm_options=slurm_options,
module_list=module_list,
print_job_id=print_job_id,
)

# make sure that the function does not use slurm once running on slurm
Expand All @@ -149,6 +153,7 @@ def new_func(*args, **kwargs):
arguments=kwargs,
imports=imports,
from_imports=from_imports,
print_job_id=print_job_id,
)

return slurm_helper.run_slurm_batch(
Expand Down
3 changes: 2 additions & 1 deletion znamutils/slurm_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def create_slurm_sbatch(


def python_script_single_func(
target_file, function_name, arguments=None, imports=None, from_imports=None
target_file, function_name, arguments=None, imports=None, from_imports=None, print_job_id=False,
):
"""Create a python script that will call a function
Expand All @@ -118,6 +118,7 @@ def python_script_single_func(
from_imports (dict, optional): Dictionary of imports to add to the script. Keys
are the module names, values are the functions to import. For instance
{'numpy': 'mean'} results in `from numpy import mean`. Defaults to None.
print_job_id (bool, optional): Whether to print the job ID. Defaults to False.
"""

target_file = Path(target_file)
Expand Down

0 comments on commit 3d0d04c

Please sign in to comment.