Skip to content

Commit

Permalink
fix: various bug fixes for resource parsing (#2711)
Browse files Browse the repository at this point in the history
### Description

<!--Add a description of your PR here-->

### QC
<!-- Make sure that you can tick the boxes below. -->

* [x] The PR contains a test case for the changes or the changes are
already covered by an existing test case.
* [x] The documentation (`docs/`) is updated to reflect the changes or
this is not necessary (e.g. if the change does neither modify the
language nor the behavior or functionalities of Snakemake).
  • Loading branch information
johanneskoester committed Feb 20, 2024
1 parent 476ae37 commit d1daf0b
Show file tree
Hide file tree
Showing 6 changed files with 53 additions and 8 deletions.
5 changes: 4 additions & 1 deletion snakemake/cli.py
Expand Up @@ -35,6 +35,7 @@
from snakemake.io import flag
from snakemake.resources import (
DefaultResources,
ParsedResource,
ResourceScopes,
eval_resource_expression,
parse_resources,
Expand Down Expand Up @@ -96,7 +97,9 @@ def parse_set_resources(args):
value = eval_resource_expression(orig_value)
if isinstance(value, int) and value < 0:
raise ValueError(errmsg)
assignments[rule][resource] = flag(value, "orig_arg", orig_value)
assignments[rule][resource] = ParsedResource(
value=value, orig_arg=orig_value
)
return assignments


Expand Down
11 changes: 10 additions & 1 deletion snakemake/resources.py
@@ -1,8 +1,10 @@
from collections import UserDict, defaultdict
from dataclasses import dataclass
import itertools as it
import operator as op
import re
import tempfile
from typing import Any

from snakemake.exceptions import (
ResourceScopesException,
Expand All @@ -12,6 +14,12 @@
from snakemake.common.tbdstring import TBDString


@dataclass
class ParsedResource:
orig_arg: str
value: Any


class DefaultResources:
defaults = {
"mem_mb": "min(max(2*input.size_mb, 1000), 8000)",
Expand Down Expand Up @@ -536,7 +544,8 @@ def generic_callable(val, threads_arg, **kwargs):
f"'{val}'.\n"
" String arguments may need additional "
"quoting. E.g.: --default-resources "
"\"tmpdir='/home/user/tmp'\".",
"\"tmpdir='/home/user/tmp'\" or "
"--set-resources \"compute1:slurm_extra='--nice=100'\".",
e,
)
raise e
Expand Down
2 changes: 1 addition & 1 deletion snakemake/settings.py
Expand Up @@ -20,7 +20,7 @@

from snakemake.common import dict_to_key_value_args, get_container_image
from snakemake.common.configfile import load_configfile
from snakemake.resources import DefaultResources
from snakemake.resources import DefaultResources, ParsedResource
from snakemake.utils import update_config
from snakemake.exceptions import WorkflowError

Expand Down
11 changes: 8 additions & 3 deletions snakemake/spawn_jobs.py
Expand Up @@ -4,6 +4,7 @@
import sys
from typing import Mapping, TypeVar, TYPE_CHECKING, Any
from snakemake_interface_executor_plugins.utils import format_cli_arg, join_cli_args
from snakemake.resources import ParsedResource
from snakemake_interface_storage_plugins.registry import StoragePluginRegistry

from snakemake import common
Expand Down Expand Up @@ -73,19 +74,23 @@ def get_storage_provider_args(self):
for plugin, field, field_settings in self._get_storage_provider_setting_items():
if not field.metadata.get("env_var", False):
cli_arg = plugin.get_cli_arg(field.name)
if isinstance(field_settings[0], bool) and field.default is not True:
# so far no tagged settings for flags with default value False
assert len(field_settings) == 1
field_settings = field_settings[0]
yield format_cli_arg(cli_arg, field_settings)

def get_storage_provider_envvars(self):
return {
plugin.get_envvar(field.name): " ".join(map(str, field_settings))
for plugin, field, field_settings in self._get_storage_provider_setting_items()
if "env_var" in field.metadata
if field.metadata.get("env_var")
}

def get_set_resources_args(self):
def get_orig_arg(value):
if is_flagged(value, "orig_arg"):
return get_flag_value(value, "orig_arg")
if isinstance(value, ParsedResource):
return value.orig_arg
else:
return value

Expand Down
16 changes: 14 additions & 2 deletions snakemake/workflow.py
Expand Up @@ -109,7 +109,7 @@
)
from snakemake.utils import simplify_path
from snakemake.checkpoints import Checkpoints
from snakemake.resources import ResourceScopes
from snakemake.resources import ParsedResource, ResourceScopes
from snakemake.caching.local import OutputFileCache as LocalOutputFileCache
from snakemake.caching.storage import OutputFileCache as StorageOutputFileCache
from snakemake.modules import ModuleInfo, WorkflowModifier, get_name_modifier_func
Expand Down Expand Up @@ -1628,7 +1628,19 @@ def decorate(ruleinfo):
)
rule.resources.update(resources)
if name in self.resource_settings.overwrite_resources:
rule.resources.update(self.resource_settings.overwrite_resources[name])

def get_value(value):
if isinstance(value, ParsedResource):
return value.value
else:
return value

rule.resources.update(
(resource, get_value(value))
for resource, value in self.resource_settings.overwrite_resources[
name
].items()
)

if ruleinfo.priority:
if not isinstance(ruleinfo.priority, int) and not isinstance(
Expand Down
16 changes: 16 additions & 0 deletions tests/tests.py
Expand Up @@ -2002,3 +2002,19 @@ def test_queue_input_forceall():
@skip_on_windows # OS independent
def test_issue2685():
run(dpath("test_issue2685"))


@skip_on_windows
def test_set_resources_complex():
run(
dpath("test05"),
shellcmd="snakemake -c1 --set-resources \"compute1:slurm_extra='--nice=10'\"",
)


@skip_on_windows
def test_set_resources_human_readable():
run(
dpath("test05"),
shellcmd="snakemake -c1 --set-resources \"compute1:runtime='50h'\"",
)

0 comments on commit d1daf0b

Please sign in to comment.