Coverage for bzfs_main/util/parallel_tasktree_policy.py: 100%
44 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-24 13:02 +0000
1# Copyright 2024 Wolfgang Hoschek AT mac DOT com
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#
15"""Policy layer for the generic parallel task tree scheduling algorithm.
17Purpose: Provide bzfs-specific behavior on top of the policy-free generic ``parallel_tasktree`` scheduling algorithm:
18retries, skip-on-error modes (fail/dataset/tree), and logging.
20Assumptions: Callers provide a thread-safe ``process_dataset(dataset, tid, Retry) -> bool`` callback. Dataset list is sorted
21and contains no duplicate entries (enforced by tests).
23Design rationale: Keep scheduling generic and reusable while concentrating error handling and side-effects here. This module
24exposes a stable API for callers like ``bzfs`` and ``bzfs_jobrunner``.
25"""
27from __future__ import (
28 annotations,
29)
30import logging
31import os
32import subprocess
33from concurrent.futures import (
34 Future,
35)
36from typing import (
37 Callable,
38)
40from bzfs_main.util.parallel_tasktree import (
41 CompletionCallback,
42 CompletionCallbackResult,
43 ParallelTaskTree,
44)
45from bzfs_main.util.retry import (
46 Retry,
47 RetryPolicy,
48 RetryTemplate,
49 RetryTerminationError,
50)
51from bzfs_main.util.utils import (
52 Comparable,
53 TaskTiming,
54 dry,
55 human_readable_duration,
56)
59def process_datasets_in_parallel_and_fault_tolerant(
60 *,
61 log: logging.Logger,
62 datasets: list[str], # (sorted) list of datasets to process
63 process_dataset: Callable[
64 [str, str, Retry], bool # lambda: dataset, tid, Retry; return False to skip subtree; must be thread-safe
65 ],
66 priority: Callable[[str], Comparable] = lambda dataset: dataset, # lexicographical order by default
67 skip_tree_on_error: Callable[[str], bool], # lambda: dataset # called on error; return True to skip subtree on error
68 skip_on_error: str = "fail",
69 max_workers: int = os.cpu_count() or 1,
70 interval_nanos: Callable[
71 [int, str, int], int
72 ] = lambda last_update_nanos, dataset, submit_count: 0, # optionally spread tasks out over time; e.g. for jitter
73 timing: TaskTiming = TaskTiming(), # noqa: B008
74 termination_handler: Callable[[], None] = lambda: None,
75 task_name: str = "Task",
76 enable_barriers: bool | None = None, # for testing only; None means 'auto-detect'
77 append_exception: Callable[[BaseException, str, str], None] = lambda ex, task, dataset: None, # called on nonfatal error
78 retry_template: RetryTemplate[bool] = RetryTemplate[bool]().copy(policy=RetryPolicy.no_retries()), # noqa: B008
79 dry_run: bool = False,
80 is_test_mode: bool = False,
81) -> bool: # returns True if any dataset processing failed, False if all succeeded
82 """Runs datasets in parallel with retries and skip policy.
84 Purpose: Adapt the generic engine to bzfs needs by wrapping the worker function with retries and determining skip/fail
85 behavior on completion.
87 Assumptions: ``skip_on_error`` is one of {"fail","dataset","tree"}. ``skip_tree_on_error(dataset)`` returns True if
88 subtree should be skipped.
90 Design rationale: The completion callback runs in the main thread, enabling safe cancellation of in-flight futures for
91 fail-fast while keeping worker threads free of policy decisions.
92 """
93 assert callable(process_dataset)
94 assert callable(skip_tree_on_error)
95 assert "%" not in task_name
96 assert callable(append_exception)
97 len_datasets: int = len(datasets)
98 is_debug: bool = log.isEnabledFor(logging.DEBUG)
100 def _process_dataset(dataset: str, submit_count: int) -> CompletionCallback:
101 """Wrapper around process_dataset(); adds retries plus a callback determining if to fail or skip subtree on error."""
102 tid: str = f"{submit_count}/{len_datasets}"
103 start_time_nanos: int = timing.monotonic_ns()
104 exception = None
105 no_skip: bool = False
106 try:
107 no_skip = retry_template.call_with_retries(
108 fn=lambda retry: process_dataset(dataset, tid, retry),
109 log=log,
110 )
111 except (
112 subprocess.CalledProcessError,
113 subprocess.TimeoutExpired,
114 SystemExit,
115 UnicodeDecodeError,
116 RetryTerminationError,
117 ) as e:
118 exception = e # may be reraised later
119 finally:
120 if is_debug:
121 elapsed_duration: str = human_readable_duration(timing.monotonic_ns() - start_time_nanos)
122 log.debug(dry(f"{tid} {task_name} done: %s took %s", dry_run), dataset, elapsed_duration)
124 def _completion_callback(todo_futures: set[Future[CompletionCallback]]) -> CompletionCallbackResult:
125 """CompletionCallback determining if to fail or skip subtree on error; Runs in the (single) main thread as part
126 of the coordination loop."""
127 nonlocal no_skip
128 fail: bool = False
129 if exception is not None:
130 fail = True
131 if skip_on_error == "fail" or timing.is_terminated():
132 for todo_future in todo_futures:
133 todo_future.cancel()
134 termination_handler()
135 raise exception
136 no_skip = not (skip_on_error == "tree" or skip_tree_on_error(dataset))
137 log.error("%s", exception)
138 append_exception(exception, task_name, dataset)
139 return CompletionCallbackResult(no_skip=no_skip, fail=fail)
141 return _completion_callback
143 tasktree: ParallelTaskTree = ParallelTaskTree(
144 log=log,
145 datasets=datasets,
146 process_dataset=_process_dataset,
147 priority=priority,
148 max_workers=max_workers,
149 interval_nanos=interval_nanos,
150 timing=timing,
151 enable_barriers=enable_barriers,
152 is_test_mode=is_test_mode,
153 )
154 return tasktree.process_datasets_in_parallel()