diff --git a/src/tigerflow/pipeline.py b/src/tigerflow/pipeline.py index 1c17118..5a6a554 100644 --- a/src/tigerflow/pipeline.py +++ b/src/tigerflow/pipeline.py @@ -249,43 +249,38 @@ def _start_tasks(self): else: raise ValueError(f"Unsupported task kind: {type(task)}") - def _build_staging_context(self) -> StagingContext: - """Build the current context for staging middleware.""" - failed_stems = self._failed_stems() + def _prepare_staging_inputs(self) -> tuple[list[Path], StagingContext]: + """Collect the candidates and context for staging middleware. - n_finished = sum(1 for f in self._finished_dir.iterdir() if f.is_file()) + Both are built here so the input directory is scanned once: `waiting` + is derived from the same list the middleware chain receives. + """ + failed_stems = self._failed_stems() n_staged = sum( 1 for f in self._symlinks_dir.iterdir() if f.is_file() and f.name.removesuffix(self._config.root_input_ext) not in failed_stems ) - n_waiting = sum( - 1 + candidates = [ + f for f in self._input_dir.iterdir() if f.is_file() and f.name.endswith(self._config.root_input_ext) and f.name not in self._filenames - ) - - return StagingContext( - waiting=n_waiting, + ] + context = StagingContext( + waiting=len(candidates), staged=n_staged, - completed=n_finished, + completed=self._count_finished(), failed=len(failed_stems), input_dir=self._input_dir, output_dir=self._output_dir, ) + return candidates, context def _stage_new_files(self): - context = self._build_staging_context() - candidates = [ - f - for f in self._input_dir.iterdir() - if f.is_file() - and f.name.endswith(self._config.root_input_ext) - and f.name not in self._filenames - ] + candidates, context = self._prepare_staging_inputs() to_stage = self._config.staging.process(candidates, context) for file in to_stage: self._symlinks_dir.joinpath(file.name).symlink_to(file) @@ -406,6 +401,19 @@ def _report_failed_files(self): if n_files > 0: logger.error("[{}] {} failed files", task.name, n_files) + def _count_finished(self) -> int: + return sum(1 for file in self._finished_dir.iterdir() if file.is_file()) + + def _all_tracked_files_settled(self) -> bool: + """Whether every tracked file has either finished or failed. + + `_filenames` is never pruned, so it covers every file the pipeline has + ever tracked. + """ + return self._count_finished() + len(self._failed_stems()) >= len( + self._filenames + ) + def _failed_stems(self) -> set[str]: """Stems of input files that failed in at least one task. @@ -465,13 +473,11 @@ def _handle_processed_files(self): # Log progress if completed_file_ids: logger.info("Completed processing {} files", len(completed_file_ids)) - n_finished = sum(1 for f in self._finished_dir.iterdir() if f.is_file()) - if (n_finished + len(self._failed_stems())) >= len(self._filenames): + if self._all_tracked_files_settled(): logger.info("No more files to process, starting idle time count") def _check_inactivity(self): - n_finished = sum(1 for file in self._finished_dir.iterdir() if file.is_file()) - if (n_finished + len(self._failed_stems())) < len(self._filenames): + if not self._all_tracked_files_settled(): self._last_active = datetime.now() inactivity = datetime.now() - self._last_active diff --git a/tests/integration/pipeline/test_file_staging.py b/tests/integration/pipeline/test_file_staging.py index 9916481..4ce1406 100644 --- a/tests/integration/pipeline/test_file_staging.py +++ b/tests/integration/pipeline/test_file_staging.py @@ -143,7 +143,7 @@ def test_counts_waiting_files( for i in range(3): (input_dir / f"f{i}.txt").write_text("x") - context = pipeline._build_staging_context() + _, context = pipeline._prepare_staging_inputs() assert context.waiting == 3 assert context.staged == 0 @@ -157,7 +157,7 @@ def test_counts_shift_after_staging( (input_dir / "a.txt").write_text("x") pipeline._stage_new_files() - context = pipeline._build_staging_context() + _, context = pipeline._prepare_staging_inputs() assert context.waiting == 0 assert context.staged == 1 @@ -174,7 +174,7 @@ def test_counts_completed_files( (task.output_dir / "a.txt").write_text("done") pipeline._handle_processed_files() - context = pipeline._build_staging_context() + _, context = pipeline._prepare_staging_inputs() assert context.completed == 1 assert context.staged == 0 @@ -191,7 +191,7 @@ def test_staged_count_excludes_failures( (task.output_dir / "a.err").write_text("boom") pipeline._report_failed_files() - context = pipeline._build_staging_context() + _, context = pipeline._prepare_staging_inputs() assert context.failed == 1 assert context.staged == 0 @@ -240,7 +240,7 @@ def test_counts_stay_consistent_in_mixed_state( pipeline._report_failed_files() pipeline._handle_processed_files() - context = pipeline._build_staging_context() + _, context = pipeline._prepare_staging_inputs() counts = (context.waiting, context.staged, context.completed, context.failed) assert counts == (0, 2, 1, 1) @@ -286,7 +286,7 @@ def test_fan_out_failure_counts_one_file( "Failure must not remove the symlink the remaining task reads from" ) - context = pipeline._build_staging_context() + _, context = pipeline._prepare_staging_inputs() assert (context.staged, context.failed) == (1, 1) def test_failed_file_excluded_under_multi_part_extension( @@ -312,7 +312,7 @@ def test_failed_file_excluded_under_multi_part_extension( (task.output_dir / "sample.err").write_text("boom") pipeline._report_failed_files() - context = pipeline._build_staging_context() + _, context = pipeline._prepare_staging_inputs() assert (context.staged, context.failed) == (0, 1)