Skip to content

Split pending-worker checks out of task status polling #120

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yoonspark merged 1 commit into from
Aug 18, 2026
Merged

Split pending-worker checks out of task status polling #120

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 53 additions & 51 deletions src/tigerflow/pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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)}")

Expand All @@ -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.
Expand All @@ -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
Expand Down
5 changes: 2 additions & 3 deletions tests/integration/pipeline/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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] = []

Expand Down
102 changes: 59 additions & 43 deletions tests/integration/pipeline/test_task_supervision.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()

Expand All @@ -103,15 +103,15 @@ 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

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.
"""
Expand Down Expand Up @@ -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]
Expand All @@ -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 == []

Expand All @@ -227,14 +215,42 @@ 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
)

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",
Expand All @@ -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:
Expand Down
Loading