From 3a7f36edac5fd2f6f5f4ce4f6a71317b5eb0bce8 Mon Sep 17 00:00:00 2001 From: Sangyoon Park Date: Mon, 17 Aug 2026 16:01:18 -0400 Subject: [PATCH] Split pending-worker checks out of task status polling `_check_pending_workers` ran inside the Slurm branch of `_check_task_status`, coupling queue-warning bookkeeping to status polling. It now owns its task loop and runs as its own tracking-cycle step. That also stops the warnings during shutdown: `_drain_tasks` polls status in a loop, so it used to query the queue every second and could log pending warnings while tasks were still exiting. Rename `_handle_task_timeout` to `_resubmit_if_timed_out` and move its ordering constraint into its own docstring, next to the code that depends on it. Drop the status-polling mocks the pending-worker tests only needed to reach the check through `_check_task_status`, and add a cycle-level test so the new step cannot be dropped from `_run_tracking_cycle` unnoticed. --- src/tigerflow/pipeline.py | 104 +++++++++--------- tests/integration/pipeline/conftest.py | 5 +- .../pipeline/test_task_supervision.py | 102 +++++++++-------- 3 files changed, 114 insertions(+), 97 deletions(-) diff --git a/src/tigerflow/pipeline.py b/src/tigerflow/pipeline.py index 5cd4693..eb2ce26 100644 --- a/src/tigerflow/pipeline.py +++ b/src/tigerflow/pipeline.py @@ -216,15 +216,15 @@ def run(self): def _run_tracking_cycle(self): """Run one iteration of the pipeline tracking loop. - `_handle_task_timeout` reads the task status that `_check_task_status` - refreshes, so a stale status would resubmit an already replaced Slurm job. - `_stage_new_files` follows `_report_failed_files` and `_handle_processed_files` - so that slots those free up can be filled in the same cycle. `_check_inactivity` - runs last because staging adds to `_filenames`, and a cycle that just started - work must not be counted as idle. + `_resubmit_if_timed_out` acts on the status `_check_task_status` writes, + so it must follow it. `_stage_new_files` follows `_report_failed_files` and + `_handle_processed_files` so that slots those free up can be filled in the + same cycle. `_check_inactivity` runs last because staging adds to + `_filenames`, and a cycle that just started work must not be counted as idle. """ self._check_task_status() - self._handle_task_timeout() + self._check_pending_workers() + self._resubmit_if_timed_out() self._report_failed_files() self._handle_processed_files() self._stage_new_files() @@ -305,7 +305,6 @@ def _check_task_status(self): elif isinstance(task, SlurmTaskConfig): job_id = self._slurm_task_ids[task.name] status = get_slurm_task_status(job_id, task.worker_job_name) - self._check_pending_workers(task) else: raise ValueError(f"Unsupported task kind: {type(task)}") @@ -322,34 +321,53 @@ def _check_task_status(self): f" ({status.detail})" if status.detail else "", ) - def _check_pending_workers(self, task: SlurmTaskConfig): + def _check_pending_workers(self): """Warn every 10 minutes a worker job spends stuck in the Slurm queue.""" - pending_ids = get_pending_worker_ids(task.worker_job_name) - since = self._worker_pending_since[task.name] - alerted = self._worker_pending_alerted[task.name] - - for job_id in list(since): - if job_id not in pending_ids: - since.pop(job_id, None) - alerted.pop(job_id, None) - - now = time.time() - newly_crossed: dict[int, list[int]] = defaultdict(list) - warning_interval = settings.slurm_task_worker_warning_interval - for job_id in pending_ids: - pending_minutes = (now - since.setdefault(job_id, now)) / 60 - threshold = int(pending_minutes // warning_interval) * warning_interval - if threshold >= warning_interval and threshold > alerted.get(job_id, 0): - alerted[job_id] = threshold - newly_crossed[threshold].append(job_id) - - for threshold, job_ids in sorted(newly_crossed.items()): - logger.warning( - "[{}] Workers pending more than {} minutes: {}", - task.name, - threshold, - ", ".join(str(job_id) for job_id in sorted(job_ids)), - ) + for task in self._config.tasks: + if not isinstance(task, SlurmTaskConfig): + continue + pending_ids = get_pending_worker_ids(task.worker_job_name) + since = self._worker_pending_since[task.name] + alerted = self._worker_pending_alerted[task.name] + + for job_id in list(since): + if job_id not in pending_ids: + since.pop(job_id, None) + alerted.pop(job_id, None) + + now = time.time() + newly_crossed: dict[int, list[int]] = defaultdict(list) + warning_interval = settings.slurm_task_worker_warning_interval + for job_id in pending_ids: + pending_minutes = (now - since.setdefault(job_id, now)) / 60 + threshold = int(pending_minutes // warning_interval) * warning_interval + if threshold >= warning_interval and threshold > alerted.get(job_id, 0): + alerted[job_id] = threshold + newly_crossed[threshold].append(job_id) + + for threshold, job_ids in sorted(newly_crossed.items()): + logger.warning( + "[{}] Workers pending more than {} minutes: {}", + task.name, + threshold, + ", ".join(str(job_id) for job_id in sorted(job_ids)), + ) + + def _resubmit_if_timed_out(self): + """Resubmit Slurm tasks that Slurm killed for hitting their time limit. + + Callers must refresh `_task_status` first; acting on a stale status + replaces a job that is already running. + """ + for task in self._config.tasks: + if not isinstance(task, SlurmTaskConfig): + continue + status = self._task_status[task.name] + if not status.is_alive and status.detail and "TIMEOUT" in status.detail: + script = task.to_script() + job_id = submit_to_slurm(script) + self._slurm_task_ids[task.name] = job_id + logger.info("[{}] Re-submitted with Slurm job ID {}", task.name, job_id) def _drain_tasks(self): """Wait for terminated tasks to actually exit, both local and Slurm. @@ -374,22 +392,6 @@ def _drain_tasks(self): return time.sleep(1) - def _handle_task_timeout(self): - for task in self._config.tasks: - if isinstance(task, SlurmTaskConfig): - task_status = self._task_status[task.name] - if ( - not task_status.is_alive - and task_status.detail - and "TIMEOUT" in task_status.detail - ): - script = task.to_script() - job_id = submit_to_slurm(script) - self._slurm_task_ids[task.name] = job_id - logger.info( - "[{}] Re-submitted with Slurm job ID {}", task.name, job_id - ) - def _report_failed_files(self): for task in self._config.tasks: n_files = 0 diff --git a/tests/integration/pipeline/conftest.py b/tests/integration/pipeline/conftest.py index 7f098de..1170441 100644 --- a/tests/integration/pipeline/conftest.py +++ b/tests/integration/pipeline/conftest.py @@ -161,9 +161,8 @@ def idle_start_logs() -> Iterator[list[str]]: def pending_worker_logs() -> Iterator[list[str]]: """Yield a list collecting the message of every "Workers pending" record logged. - Filtered by message content as well as level: `_check_task_status` also logs - an INFO status-change record on the same call, which a level-only sink would - mix in and throw off exact-count assertions. + Filtered by message content as well as level so that exact-count assertions + stay valid for tests that drive the full tracking cycle. """ records: list[str] = [] diff --git a/tests/integration/pipeline/test_task_supervision.py b/tests/integration/pipeline/test_task_supervision.py index dfeccc1..ed128b3 100644 --- a/tests/integration/pipeline/test_task_supervision.py +++ b/tests/integration/pipeline/test_task_supervision.py @@ -76,7 +76,7 @@ def test_resubmits_after_timeout(self, pipeline_factory: PipelineFactory): pipeline._task_status["gpu"] = slurm_status(TaskStatusKind.INACTIVE, "TIMEOUT") with patch("tigerflow.pipeline.submit_to_slurm", return_value=222) as resubmit: - pipeline._handle_task_timeout() + pipeline._resubmit_if_timed_out() resubmit.assert_called_once() assert pipeline._slurm_task_ids["gpu"] == 222 @@ -87,7 +87,7 @@ def test_no_resubmit_while_running(self, pipeline_factory: PipelineFactory): pipeline._task_status["gpu"] = slurm_status(TaskStatusKind.ACTIVE) with patch("tigerflow.pipeline.submit_to_slurm") as resubmit: - pipeline._handle_task_timeout() + pipeline._resubmit_if_timed_out() resubmit.assert_not_called() @@ -103,7 +103,7 @@ def test_no_resubmit_on_other_failures( pipeline._task_status["gpu"] = slurm_status(TaskStatusKind.INACTIVE, detail) with patch("tigerflow.pipeline.submit_to_slurm") as resubmit: - pipeline._handle_task_timeout() + pipeline._resubmit_if_timed_out() resubmit.assert_not_called() assert pipeline._slurm_task_ids["gpu"] == 111 @@ -111,7 +111,7 @@ def test_no_resubmit_on_other_failures( def test_resubmits_once_per_timeout(self, pipeline_factory: PipelineFactory): """The replacement job ID is what the next status poll observes. - `_handle_task_timeout` records the new ID so the following cycle polls + `_resubmit_if_timed_out` records the new ID so the following cycle polls the replacement; polling the dead job would keep reporting TIMEOUT and resubmit on every cycle. """ @@ -150,30 +150,23 @@ def test_logs_when_worker_crosses_every_ten_minutes( pending_worker_logs: list[str], ): pipeline = pipeline_factory([SLURM_TASK]) - pipeline._slurm_task_ids["gpu"] = 111 clock = 1000.0 monkeypatch.setattr("tigerflow.pipeline.time.time", lambda: clock) - with ( - patch( - "tigerflow.pipeline.get_slurm_task_status", - return_value=slurm_status(TaskStatusKind.ACTIVE), - ), - patch( - "tigerflow.pipeline.get_pending_worker_ids", - return_value=[847645, 847649], - ), + with patch( + "tigerflow.pipeline.get_pending_worker_ids", + return_value=[847645, 847649], ): - pipeline._check_task_status() + pipeline._check_pending_workers() assert pending_worker_logs == [], "First observation only starts the clock" clock += 601 - pipeline._check_task_status() + pipeline._check_pending_workers() clock += 600 - pipeline._check_task_status() + pipeline._check_pending_workers() clock += 60 # Same threshold must not be logged twice - pipeline._check_task_status() + pipeline._check_pending_workers() assert len(pending_worker_logs) == 2 assert "[gpu]" in pending_worker_logs[0] @@ -189,34 +182,29 @@ def test_stops_tracking_once_worker_leaves_the_queue( pending_worker_logs: list[str], ): pipeline = pipeline_factory([SLURM_TASK]) - pipeline._slurm_task_ids["gpu"] = 111 clock = 1000.0 monkeypatch.setattr("tigerflow.pipeline.time.time", lambda: clock) with patch( - "tigerflow.pipeline.get_slurm_task_status", - return_value=slurm_status(TaskStatusKind.ACTIVE), + "tigerflow.pipeline.get_pending_worker_ids", + return_value=[847645], ): - with patch( - "tigerflow.pipeline.get_pending_worker_ids", - return_value=[847645], - ): - pipeline._check_task_status() + pipeline._check_pending_workers() - clock += 601 - with patch( - "tigerflow.pipeline.get_pending_worker_ids", - return_value=[], # Job started running before the 10-minute mark - ): - pipeline._check_task_status() + clock += 601 + with patch( + "tigerflow.pipeline.get_pending_worker_ids", + return_value=[], # Job started running before the 10-minute mark + ): + pipeline._check_pending_workers() - clock += 601 - with patch( - "tigerflow.pipeline.get_pending_worker_ids", - return_value=[847645], # Re-queued later, clock restarts - ): - pipeline._check_task_status() + clock += 601 + with patch( + "tigerflow.pipeline.get_pending_worker_ids", + return_value=[847645], # Re-queued later, clock restarts + ): + pipeline._check_pending_workers() assert pending_worker_logs == [] @@ -227,7 +215,6 @@ def test_warning_interval_is_configurable( pending_worker_logs: list[str], ): pipeline = pipeline_factory([SLURM_TASK]) - pipeline._slurm_task_ids["gpu"] = 111 monkeypatch.setattr( "tigerflow.pipeline.settings.slurm_task_worker_warning_interval", 5 ) @@ -235,6 +222,35 @@ def test_warning_interval_is_configurable( clock = 1000.0 monkeypatch.setattr("tigerflow.pipeline.time.time", lambda: clock) + with patch( + "tigerflow.pipeline.get_pending_worker_ids", + return_value=[847645], + ): + pipeline._check_pending_workers() + clock += 301 + pipeline._check_pending_workers() + + assert len(pending_worker_logs) == 1 + assert "5 minutes" in pending_worker_logs[0] + + def test_tracking_cycle_checks_pending_workers( + self, + pipeline_factory: PipelineFactory, + monkeypatch: pytest.MonkeyPatch, + pending_worker_logs: list[str], + ): + """The check runs as a tracking cycle step, not only when called directly. + + Every other test in this class calls `_check_pending_workers` itself, so + dropping it from `_run_tracking_cycle` would leave them all passing while + the warning never fires in a running pipeline. + """ + pipeline = pipeline_factory([SLURM_TASK]) + pipeline._slurm_task_ids["gpu"] = 111 + + clock = 1000.0 + monkeypatch.setattr("tigerflow.pipeline.time.time", lambda: clock) + with ( patch( "tigerflow.pipeline.get_slurm_task_status", @@ -245,12 +261,12 @@ def test_warning_interval_is_configurable( return_value=[847645], ), ): - pipeline._check_task_status() - clock += 301 - pipeline._check_task_status() + pipeline._run_tracking_cycle() + clock += 601 + pipeline._run_tracking_cycle() assert len(pending_worker_logs) == 1 - assert "5 minutes" in pending_worker_logs[0] + assert "[gpu]" in pending_worker_logs[0] class TestSlurmShutdown: