feat: Maximal file size for checksums (snakemake#3368)

JZentgraf · johannaelenaschmitz · johannaschmitz · web-flow · commit b039f8aaeece · 2025-03-12T17:30:41.000+01:00
Add a new parameter `--max-checksum-file-size` to adjust the maximal file size for which a checksum is computed when the DAG is build. This is related to snakemake#2949 ### QC  * [ ] The PR contains a test case for the changes or the changes are already covered by an existing test case. * [ ] 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).  ## Summary by CodeRabbit - **New Features** - Added a new command-line option to let users specify a maximum file size for checksum processing, with a default threshold of 1,000,000 bytes. - **Refactor** - Updated file validation and checksum routines to enforce the new file size limit, ensuring more efficient processing across workflow operations. - Made method signatures more explicit by removing default values for threshold parameters in checksum-related methods.  --------- Co-authored-by: Johanna <johanna_schmitz@posteo.de> Co-authored-by: johannaschmitz <56937653+johannaschmitz@users.noreply.github.com> Co-authored-by: Johannes Köster <johannes.koester@uni-due.de>
diff --git a/snakemake/cli.py b/snakemake/cli.py
@@ -70,6 +70,12 @@
 from snakemake.utils import available_cpu_count, update_config
 
 
+def parse_size_in_bytes(value):
+    from humanfriendly import parse_size
+
+    return parse_size(value)
+
+
 def expandvars(atype):
     def inner(args):
         if isinstance(args, list):
@@ -1270,6 +1276,15 @@ def get_argument_parser(profiles=None):
         "network file systems. Hence, we do not spend more than a given amount of time and fall back "
         "to individual checks for the rest.",
     )
+    group_behavior.add_argument(
+        "--max-checksum-file-size",
+        default=1000000,
+        metavar="SIZE",
+        parse_func=parse_size_in_bytes,
+        help="Compute the checksum during DAG computation and job postprocessing "
+        "only for files that are smaller than the provided threshold (given in any valid size "
+        "unit, e.g. 1MB, which is also the default). ",
+    )
     group_behavior.add_argument(
         "--latency-wait",
         "--output-wait",
@@ -2052,6 +2067,7 @@ def args_to_api(args, parser):
                             allowed_rules=args.allowed_rules,
                             rerun_triggers=args.rerun_triggers,
                             max_inventory_wait_time=args.max_inventory_time,
+                            max_checksum_file_size=args.max_checksum_file_size,
                             strict_evaluation=args.strict_dag_evaluation,
                         ),
                     )
diff --git a/snakemake/dag.py b/snakemake/dag.py
@@ -145,6 +145,7 @@ def __init__(
         self._running = set()
         self._jobs_with_finished_queue_input = set()
         self._storage_input_jobs = defaultdict(list)
+        self.max_checksum_file_size = self.workflow.dag_settings.max_checksum_file_size
 
         self.job_factory = JobFactory()
         self.group_job_factory = GroupJobFactory()
@@ -649,7 +650,9 @@ async def is_not_same_checksum(f, checksum):
                         e,
                         rule=job.rule,
                     )
-            return not await f.is_same_checksum(checksum, force=True)
+            return not await f.is_same_checksum(
+                checksum, self.max_checksum_file_size, force=True
+            )
 
         checksum_failed_output = [
             f
@@ -1254,7 +1257,7 @@ async def is_same_checksum(f, job):
             try:
                 return is_same_checksum_cache[(f, job)]
             except KeyError:
-                if not await f.is_checksum_eligible():
+                if not await f.is_checksum_eligible(self.max_checksum_file_size):
                     # no chance to compute checksum, cannot be assumed the same
                     is_same = False
                 else:
@@ -1268,7 +1271,9 @@ async def is_same_checksum(f, job):
                         # no checksums recorded, we cannot assume them to be the same
                         is_same = False
                     else:
-                        is_same = await f.is_same_checksum(checksums.pop())
+                        is_same = await f.is_same_checksum(
+                            checksums.pop(), self.max_checksum_file_size
+                        )
 
                 is_same_checksum_cache[(f, job)] = is_same
                 return is_same
diff --git a/snakemake/io.py b/snakemake/io.py
@@ -596,19 +596,19 @@ async def size_local(self):
         await self.check_broken_symlink()
         return os.path.getsize(self.file)
 
-    async def is_checksum_eligible(self):
+    async def is_checksum_eligible(self, threshold):
         return (
             await self.exists_local()
             and not os.path.isdir(self.file)
-            and await self.size() <= 1000000
+            and await self.size() <= threshold
             and not self.is_fifo()
         )
 
-    async def checksum(self, force=False):
+    async def checksum(self, threshold, force=False):
         """Return checksum if file is small enough, else None.
         Returns None if file does not exist. If force is True,
         omit eligibility check."""
-        if force or await self.is_checksum_eligible():  # less than 1 MB
+        if force or await self.is_checksum_eligible(threshold):
             checksum = sha256()
             if await self.size() > 0:
                 # only read if file is bigger than zero
@@ -623,8 +623,8 @@ async def checksum(self, force=False):
         else:
             return None
 
-    async def is_same_checksum(self, other_checksum, force=False):
-        checksum = await self.checksum(force=force)
+    async def is_same_checksum(self, other_checksum, threshold, force=False):
+        checksum = await self.checksum(threshold, force=force)
         if checksum is None or other_checksum is None:
             # if no checksum available or files too large, not the same
             return False
diff --git a/snakemake/persistence.py b/snakemake/persistence.py
@@ -130,6 +130,9 @@ def __init__(
             self.unlock = self.noop
 
         self._read_record = self._read_record_cached
+        self.max_checksum_file_size = (
+            self.dag.workflow.dag_settings.max_checksum_file_size
+        )
 
     @property
     def path(self) -> Path:
@@ -326,8 +329,10 @@ async def finished(self, job):
                 else fallback_time
             )
 
-            checksums = ((infile, await infile.checksum()) for infile in job.input)
-
+            checksums = (
+                (infile, await infile.checksum(self.max_checksum_file_size))
+                for infile in job.input
+            )
             self._record(
                 self._metadata_path,
                 {
diff --git a/snakemake/settings/types.py b/snakemake/settings/types.py
@@ -203,6 +203,7 @@ class DAGSettings(SettingsBase):
     allowed_rules: AnySet[str] = frozenset()
     rerun_triggers: AnySet[RerunTrigger] = RerunTrigger.all()
     max_inventory_wait_time: int = 20
+    max_checksum_file_size: int = 1000000
     strict_evaluation: AnySet[StrictDagEvaluation] = frozenset()
     # strict_functions_evaluation: bool = False
     # strict_cycle_evaluation: bool = False
