151 lines
6.5 KiB
Python
151 lines
6.5 KiB
Python
import time
|
|
|
|
from ..utils.config import load_config
|
|
from ..utils.constants import HEARTBEAT_TIMEOUT
|
|
from ..utils.logging import debug_log, log
|
|
from ..utils.network import build_worker_url, probe_worker
|
|
from .job_models import BaseJobState
|
|
from .job_store import ensure_tile_jobs_initialized
|
|
|
|
|
|
def _find_worker_record(worker_id):
|
|
"""Return worker config entry by id, or None when missing."""
|
|
workers = load_config().get("workers", [])
|
|
return next((w for w in workers if str(w.get("id")) == str(worker_id)), None)
|
|
|
|
|
|
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
|
|
"""Check timed out workers and requeue their tasks. Returns requeued count."""
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
current_time = time.time()
|
|
|
|
# Allow override via config setting 'worker_timeout_seconds'
|
|
cfg = load_config()
|
|
hb_timeout = int(cfg.get("settings", {}).get("worker_timeout_seconds", HEARTBEAT_TIMEOUT))
|
|
|
|
# Snapshot timed-out workers and job details under lock.
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
if not isinstance(job_data, BaseJobState):
|
|
return 0
|
|
|
|
completed_tasks_snapshot = set(job_data.completed_tasks.keys())
|
|
batched_static_snapshot = bool(job_data.batched_static)
|
|
num_tiles_per_image_snapshot = int(job_data.num_tiles_per_image or 1)
|
|
batch_size_snapshot = int(job_data.batch_size or 1)
|
|
|
|
timed_out_workers = []
|
|
for worker, last_heartbeat in list(job_data.worker_status.items()):
|
|
age = current_time - float(last_heartbeat)
|
|
debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
|
|
if age > hb_timeout:
|
|
timed_out_workers.append(
|
|
{
|
|
"worker_id": worker,
|
|
"last_heartbeat": float(last_heartbeat),
|
|
"assigned_tasks": list(job_data.assigned_to_workers.get(worker, [])),
|
|
}
|
|
)
|
|
|
|
if not timed_out_workers:
|
|
return 0
|
|
|
|
# Probe outside lock to avoid lock contention on network latency.
|
|
workers_to_requeue = []
|
|
workers_graced = []
|
|
for worker_info in timed_out_workers:
|
|
worker = worker_info["worker_id"]
|
|
assigned = worker_info["assigned_tasks"]
|
|
age = current_time - worker_info["last_heartbeat"]
|
|
|
|
incomplete_assigned = 0
|
|
try:
|
|
if assigned:
|
|
if batched_static_snapshot:
|
|
for task_id in assigned:
|
|
for b in range(batch_size_snapshot):
|
|
gidx = b * num_tiles_per_image_snapshot + task_id
|
|
if gidx not in completed_tasks_snapshot:
|
|
incomplete_assigned += 1
|
|
break
|
|
else:
|
|
for task_id in assigned:
|
|
if task_id not in completed_tasks_snapshot:
|
|
incomplete_assigned += 1
|
|
debug_log(
|
|
f"Assigned diagnostics: total_assigned={len(assigned)} "
|
|
f"incomplete_assigned={incomplete_assigned}"
|
|
)
|
|
except Exception as e:
|
|
debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
|
|
|
|
busy = False
|
|
probe_queue = None
|
|
try:
|
|
worker_record = _find_worker_record(worker)
|
|
if worker_record:
|
|
worker_url = build_worker_url(worker_record)
|
|
debug_log(f"Probing worker {worker} at {worker_url}/prompt")
|
|
payload = await probe_worker(worker_url, timeout=2.0)
|
|
if payload is not None:
|
|
probe_queue = int(payload.get("exec_info", {}).get("queue_remaining", 0))
|
|
busy = probe_queue is not None and probe_queue > 0
|
|
else:
|
|
debug_log(f"Probe skipped; worker {worker} not found in config")
|
|
except Exception as e:
|
|
debug_log(f"Probe failed for worker {worker}: {e}")
|
|
finally:
|
|
debug_log(
|
|
f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}"
|
|
)
|
|
|
|
if busy:
|
|
workers_graced.append(worker)
|
|
debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
|
|
continue
|
|
|
|
log(f"Worker {worker} heartbeat timed out after {age:.1f}s")
|
|
workers_to_requeue.append((worker, assigned))
|
|
|
|
# Re-acquire lock and apply requeue/cleanup decisions.
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
if not isinstance(job_data, BaseJobState):
|
|
return 0
|
|
|
|
# Refresh heartbeat for workers that we proved are still busy.
|
|
for worker in workers_graced:
|
|
if worker in job_data.worker_status:
|
|
job_data.worker_status[worker] = current_time
|
|
|
|
requeued_count = 0
|
|
completed_tasks = job_data.completed_tasks
|
|
batched_static = bool(job_data.batched_static)
|
|
num_tiles_per_image = int(job_data.num_tiles_per_image or 1)
|
|
batch_size = int(job_data.batch_size or 1)
|
|
for worker, assigned_snapshot in workers_to_requeue:
|
|
# Use current assignments if present, falling back to the snapshot.
|
|
assigned_tasks = list(job_data.assigned_to_workers.get(worker, assigned_snapshot))
|
|
for task_id in assigned_tasks:
|
|
# If batched_static, task_id is a tile_idx; consider it complete only if
|
|
# all corresponding global_idx entries are present in completed_tasks.
|
|
if batched_static:
|
|
all_done = True
|
|
for b in range(batch_size):
|
|
gidx = b * num_tiles_per_image + task_id
|
|
if gidx not in completed_tasks:
|
|
all_done = False
|
|
break
|
|
if not all_done:
|
|
await job_data.pending_tasks.put(task_id)
|
|
requeued_count += 1
|
|
else:
|
|
if task_id not in completed_tasks:
|
|
await job_data.pending_tasks.put(task_id)
|
|
requeued_count += 1
|
|
job_data.worker_status.pop(worker, None)
|
|
if worker in job_data.assigned_to_workers:
|
|
job_data.assigned_to_workers[worker] = []
|
|
|
|
return requeued_count
|