diff --git a/runner/__init__.py b/runner/__init__.py index bb34b8a..c324208 100644 --- a/runner/__init__.py +++ b/runner/__init__.py @@ -1,5 +1,7 @@ """Pull-based runner coordination layer. -Registration, heartbeat, HTTP client and active-job tracking used by the -runner process to talk to the backend runner API. No docker, no Flask. +Registration, heartbeat, HTTP client, active-job tracking, job polling and +result sending used by the runner process to talk to the backend runner API. +Job prep reuses the dispatcher's file/testdata helpers; still no docker and +no Flask in this package. """ diff --git a/runner/client.py b/runner/client.py index 117cbbc..abeb45e 100644 --- a/runner/client.py +++ b/runner/client.py @@ -25,6 +25,33 @@ class BackendAuthError(BackendAPIError): """The backend rejected our credentials (HTTP 401).""" +@dataclass +class JobPayload: + """A job claimed from GET /runners//next-job (spec §7.3 wire contract).""" + job_id: str + submission_id: str + problem_id: int + language: int + code_url: str + checker: object + tasks: list + + @classmethod + def from_dict(cls, body): + return cls( + job_id=body['job_id'], + submission_id=body['submission_id'], + # The backend job hash stores fields as strings; normalize here + # so consumers never coerce. + problem_id=int(body['problem_id']), + language=body['language'], + code_url=body['code_url'], + # The runner does not use the checker yet; keep it optional. + checker=body.get('checker'), + tasks=body['tasks'], + ) + + @dataclass class RunnerConfig: heartbeat_interval_sec: int @@ -101,6 +128,47 @@ def heartbeat(self, identity, active_job_ids): return self._raise_for_status(resp, 'heartbeat') + def next_job(self, identity): + """GET /runners//next-job. + + 200 -> JobPayload; 204 -> None (no work available); other non-2xx -> + BackendAPIError / BackendAuthError. Network errors propagate. + """ + resp = self.session.get( + f'{self.base_url}/runners/{identity.runner_id}/next-job', + headers={'Authorization': f'Bearer {identity.token}'}, + timeout=self.timeout, + ) + if resp.status_code == 200: + return JobPayload.from_dict(resp.json()) + if resp.status_code == 204: + return None + self._raise_for_status(resp, 'next-job') + + def complete(self, identity, job_id, tasks): + """PUT /runners//jobs//complete. Expect 204.""" + resp = self.session.put( + f'{self.base_url}/runners/{identity.runner_id}/jobs/{job_id}/complete', + json={'tasks': tasks}, + headers={'Authorization': f'Bearer {identity.token}'}, + timeout=self.timeout, + ) + if resp.status_code == 204: + return + self._raise_for_status(resp, 'complete') + + def abort(self, identity, job_id, reason): + """PUT /runners//jobs//abort. Expect 202.""" + resp = self.session.put( + f'{self.base_url}/runners/{identity.runner_id}/jobs/{job_id}/abort', + json={'reason': reason}, + headers={'Authorization': f'Bearer {identity.token}'}, + timeout=self.timeout, + ) + if resp.status_code == 202: + return + self._raise_for_status(resp, 'abort') + @staticmethod def _raise_for_status(resp, action): # Never include tokens in the message: only the status code is logged. diff --git a/runner/config.py b/runner/config.py index ed7736d..273be9b 100644 --- a/runner/config.py +++ b/runner/config.py @@ -27,3 +27,10 @@ DEFAULT_HEARTBEAT_INTERVAL_SEC = 15 DEFAULT_POLL_INTERVAL_SEC = 3 DEFAULT_MAX_CONCURRENT_JOBS = 8 + +# Local prep attempts before giving up and aborting with prep_failed (spec §10). +PREP_MAX_ATTEMPTS = 3 +# Backoff (seconds) between prep attempts; len == PREP_MAX_ATTEMPTS - 1. +PREP_BACKOFF_SCHEDULE = (1, 2) +# Backoff (seconds) between complete/abort resends; len == max retries (spec §7). +SEND_RETRY_BACKOFF_SCHEDULE = (1, 2, 4, 8, 16) diff --git a/runner/poller.py b/runner/poller.py new file mode 100644 index 0000000..9171996 --- /dev/null +++ b/runner/poller.py @@ -0,0 +1,142 @@ +import io +import logging +import shutil +import threading +import time + +import requests + +from dispatcher import file_manager, testdata +from dispatcher.config import SUBMISSION_DIR +from dispatcher.meta import Meta +from .client import BackendAPIError +from .config import REQUEST_TIMEOUT, PREP_MAX_ATTEMPTS, PREP_BACKOFF_SCHEDULE +from .result_sender import AbortRequest + +logger = logging.getLogger(__name__) + + +def prepare_job(payload): + """Fetch testdata + source and lay out the local job dir (spec §10). + + Talks to the legacy testdata channel (backend + redis) for now; that + coexists until the keystone slice. + """ + testdata.ensure_testdata(payload.problem_id) + resp = requests.get(payload.code_url, timeout=REQUEST_TIMEOUT) + resp.raise_for_status() + job_dir = SUBMISSION_DIR / payload.job_id + if job_dir.exists(): + shutil.rmtree(job_dir) # leftover from a failed prior attempt + file_manager.extract( + root_dir=SUBMISSION_DIR, + job_id=payload.job_id, + meta=Meta.parse_obj({ + 'language': payload.language, + 'tasks': payload.tasks, + }), + source=io.BytesIO(resp.content), + testdata=testdata.get_problem_root(payload.problem_id), + ) + + +class PollerThread(threading.Thread): + """Claims jobs from the backend and preps them for dispatch (spec §10). + + Only polls when there is spare capacity, adds the claimed job to the + tracker before prep so the heartbeat renews the lease while downloading, + and retries prep locally before giving up with a prep_failed abort. + """ + + def __init__( + self, + client, + identity, + tracker, + result_queue, + dispatch, + *, + prepare=None, + poll_interval_sec=None, + sleep=time.sleep, + ): + super().__init__(daemon=True) + self._client = client + self._identity = identity + self._tracker = tracker + self._result_queue = result_queue + self._dispatch = dispatch + self._prepare = prepare if prepare is not None else prepare_job + self._poll_interval_sec = (poll_interval_sec + if poll_interval_sec is not None else + identity.config.poll_interval_sec) + self._sleep = sleep + self._stop_event = threading.Event() + + def run(self): + while not self._stop_event.is_set(): + try: + idle = self._poll_once() + except Exception: + # A malformed payload (or any bug) must not kill the poller: + # the heartbeat would keep the runner looking alive while it + # never claims work again. + logger.exception('poller iteration failed') + idle = True + if idle: + # Interruptible wait so stop() takes effect immediately. + self._stop_event.wait(self._poll_interval_sec) + + def _poll_once(self): + """Claim and prep one job. Returns True when the loop should idle.""" + # Capacity gate: only GET next-job when there is room (spec §10). + if len(self._tracker) >= self._identity.config.max_concurrent_jobs: + return True + try: + payload = self._client.next_job(self._identity) + except (BackendAPIError, requests.RequestException) as err: + # 401 also just logs; heartbeat owns fail-fast. + logger.warning('next-job failed: %s', err) + return True + if payload is None: + return True + + # Add BEFORE prep so the heartbeat renews the lease while we download; + # removal is the sender's job once the outcome resolves. + self._tracker.add(payload.job_id) + + for i in range(PREP_MAX_ATTEMPTS): + try: + self._prepare(payload) + break + except Exception as err: + # Heterogeneous causes: network errors, extract ValueError. + logger.warning('prep for %s failed (attempt %d/%d): %s', + payload.job_id, i + 1, PREP_MAX_ATTEMPTS, err) + if i < len(PREP_BACKOFF_SCHEDULE): + self._sleep(PREP_BACKOFF_SCHEDULE[i]) + else: + logger.error('prep for %s exhausted all attempts; aborting', + payload.job_id) + self._abort(payload) + return False + + try: + self._dispatch(payload.job_id, payload.submission_id) + except Exception as err: + # One shot only: a failed handle() may have partially enqueued + # task entries, and calling it again for the same job_id would + # revive them (duplicate execution). Requeue via abort instead; + # slice 4 makes handle() atomic and closes this for good. + logger.warning('dispatch for %s failed: %s', payload.job_id, err) + self._abort(payload) + return False + + def _abort(self, payload): + self._result_queue.put( + AbortRequest(payload.job_id, payload.submission_id, 'prep_failed')) + # Do NOT remove from tracker; the sender does that as part of + # finalizing the abort. + + def stop(self): + self._stop_event.set() diff --git a/runner/result_sender.py b/runner/result_sender.py new file mode 100644 index 0000000..2a5a9df --- /dev/null +++ b/runner/result_sender.py @@ -0,0 +1,183 @@ +import logging +import queue +import threading +import time +from dataclasses import dataclass + +import requests + +from dispatcher import file_manager +from .client import BackendAPIError +from .config import SEND_RETRY_BACKOFF_SCHEDULE + +logger = logging.getLogger(__name__) + + +@dataclass +class CompleteRequest: + job_id: str + submission_id: str + tasks: list + + +@dataclass +class AbortRequest: + job_id: str + submission_id: str + reason: str # 'drain' | 'prep_failed' | 'rejected' + + +def _default_cleanup(job_id): + """Remove the local job dir, tolerating a job that was never prepped.""" + try: + file_manager.clean_data(job_id) + except FileNotFoundError: + logger.warning('cleanup: job dir for %s already absent', job_id) + + +def _default_backup(job_id): + """Move the local job dir aside, tolerating a missing dir.""" + try: + file_manager.backup_data(job_id) + except FileNotFoundError: + logger.warning('backup: job dir for %s already absent', job_id) + + +class ResultSenderThread(threading.Thread): + """Reports job outcomes (complete/abort) back to the backend (spec §7). + + Consumes CompleteRequest / AbortRequest off ``result_queue`` and retries + each report with exponential backoff. After ``stop()`` the loop keeps + draining whatever is already queued, then exits (the slice-4 drain + semantic). + """ + + def __init__( + self, + client, + identity, + tracker, + result_queue, + *, + cleanup=None, + backup=None, + sleep=time.sleep, + queue_poll_sec=0.5, + ): + super().__init__(daemon=True) + self._client = client + self._identity = identity + self._tracker = tracker + self._result_queue = result_queue + self._cleanup = cleanup if cleanup is not None else _default_cleanup + self._backup = backup if backup is not None else _default_backup + self._sleep = sleep + self._queue_poll_sec = queue_poll_sec + self._stop_event = threading.Event() + + def run(self): + while True: + try: + item = self._result_queue.get(timeout=self._queue_poll_sec) + except queue.Empty: + # Drain semantic: only exit once stop() was requested AND the + # queue is empty; otherwise keep waiting for more work. + if self._stop_event.is_set(): + return + continue + try: + self._process(item) + except Exception: + # One bad item must not kill the sender. + logger.exception('result sender failed to process %r', item) + + def _process(self, item): + if isinstance(item, CompleteRequest): + self._process_complete(item) + elif isinstance(item, AbortRequest): + self._process_abort(item.job_id, item.reason) + else: + logger.error('result sender got unknown item %r', item) + + def _process_complete(self, item): + job_id = item.job_id + outcome, status = self._send_with_retry( + 'complete', + lambda: self._client.complete(self._identity, job_id, item.tasks), + terminal_statuses=(409, 404, 400), + ) + if outcome == 'terminal' and status == 400: + # Backend rejected the payload: requeue via abort so the job can + # converge (spec §7.5, INV5). Keep the job dir as evidence of + # what the backend refused. + logger.warning( + 'complete for %s rejected with 400; aborting as rejected', + job_id, + ) + self._process_abort(job_id, 'rejected', preserve_evidence=True) + return + try: + if outcome == 'exhausted': + # Keep local evidence; lease-expiry reclaim is the safety net. + self._backup(job_id) + else: + self._cleanup(job_id) + finally: + # Even a failing cleanup must not leave the job in the tracker: + # a lingering entry would keep the lease renewed and consume + # capacity forever. + self._tracker.remove(job_id) + + def _process_abort(self, job_id, reason, preserve_evidence=False): + # Finalize local state BEFORE the send: the moment the backend + # accepts an abort it requeues the job, and this same runner may + # claim it again immediately -- no stale dir or tracker entry may + # survive to that point (ABA race). + try: + if preserve_evidence: + self._backup(job_id) + else: + self._cleanup(job_id) + except Exception: + # A failed finalize must not block the abort: an unsent abort + # leaves the job leased until its lease expires. + logger.exception('finalize before abort of %s failed', job_id) + self._tracker.remove(job_id) + self._send_abort(job_id, reason) + + def _send_abort(self, job_id, reason): + outcome, _ = self._send_with_retry( + 'abort', + lambda: self._client.abort(self._identity, job_id, reason), + terminal_statuses=(409, 404), + ) + return outcome + + def _send_with_retry(self, action, send, terminal_statuses): + """Retry ``send`` until success, a terminal status, or exhaustion. + + Retries everything except 2xx (success) and ``terminal_statuses``, + with exponential backoff, at most len(SEND_RETRY_BACKOFF_SCHEDULE) + retries (spec §7 "Runner retry 規則"). Note 401 is deliberately NOT + terminal here; heartbeat owns 401 fail-fast. + """ + retries = 0 + while True: + try: + send() + return ('sent', None) + except BackendAPIError as err: + if err.status_code in terminal_statuses: + logger.info('%s terminal with status %d', action, + err.status_code) + return ('terminal', err.status_code) + except requests.RequestException: + pass # network error -> retry + if retries == len(SEND_RETRY_BACKOFF_SCHEDULE): + logger.error('%s exhausted all retries', action) + return ('exhausted', None) + self._sleep(SEND_RETRY_BACKOFF_SCHEDULE[retries]) + retries += 1 + + def stop(self): + self._stop_event.set() diff --git a/tests/test_runner_client.py b/tests/test_runner_client.py index 021dece..2216996 100644 --- a/tests/test_runner_client.py +++ b/tests/test_runner_client.py @@ -5,6 +5,7 @@ BackendClient, BackendAPIError, BackendAuthError, + JobPayload, RunnerIdentity, RunnerConfig, ) @@ -30,6 +31,31 @@ def __init__(self, response=None, exc=None): def post(self, url, json=None, headers=None, timeout=None): self.calls.append({ + 'method': 'post', + 'url': url, + 'json': json, + 'headers': headers, + 'timeout': timeout, + }) + if self._exc is not None: + raise self._exc + return self._response + + def get(self, url, json=None, headers=None, timeout=None): + self.calls.append({ + 'method': 'get', + 'url': url, + 'json': json, + 'headers': headers, + 'timeout': timeout, + }) + if self._exc is not None: + raise self._exc + return self._response + + def put(self, url, json=None, headers=None, timeout=None): + self.calls.append({ + 'method': 'put', 'url': url, 'json': json, 'headers': headers, @@ -184,3 +210,149 @@ def test_token_absent_from_identity_repr(): RunnerConfig(15, 3, 8)) assert 'rk_super_secret' not in repr(identity) assert 'rn_abc' in repr(identity) + + +def make_identity(): + return RunnerIdentity('rn_abc', 'rk_secret', RunnerConfig(15, 3, 8)) + + +NEXT_JOB_BODY = { + 'job_id': + 'jb_1', + 'submission_id': + 'sub_1', + 'problem_id': + 42, + 'language': + 2, + 'code_url': + 'http://minio/code.zip', + 'checker': + 'diff', + 'tasks': [{ + 'taskScore': 100, + 'memoryLimit': 65536, + 'timeLimit': 1000, + 'caseCount': 3, + }], +} + + +def test_next_job_200_parses_payload(): + session = RecordingSession(StubResponse(200, NEXT_JOB_BODY)) + client = BackendClient('http://web:8080', session=session, timeout=10) + identity = make_identity() + + payload = client.next_job(identity) + + assert isinstance(payload, JobPayload) + assert payload.job_id == 'jb_1' + assert payload.submission_id == 'sub_1' + assert payload.problem_id == 42 + assert payload.language == 2 + assert payload.code_url == 'http://minio/code.zip' + assert payload.checker == 'diff' + assert payload.tasks == NEXT_JOB_BODY['tasks'] + + call = session.calls[0] + assert call['method'] == 'get' + assert call['url'] == 'http://web:8080/runners/rn_abc/next-job' + assert call['headers'] == {'Authorization': 'Bearer rk_secret'} + assert call['timeout'] == 10 + + +def test_next_job_normalizes_string_problem_id(): + # The backend job hash stores fields as strings. + body = dict(NEXT_JOB_BODY) + body['problem_id'] = '42' + session = RecordingSession(StubResponse(200, body)) + client = BackendClient('http://web:8080', session=session) + + payload = client.next_job(make_identity()) + + assert payload.problem_id == 42 + + +def test_next_job_defaults_checker_to_none_when_absent(): + body = dict(NEXT_JOB_BODY) + del body['checker'] + session = RecordingSession(StubResponse(200, body)) + client = BackendClient('http://web:8080', session=session) + + payload = client.next_job(make_identity()) + + assert payload.checker is None + + +def test_next_job_204_returns_none(): + session = RecordingSession(StubResponse(204)) + client = BackendClient('http://web:8080', session=session) + + assert client.next_job(make_identity()) is None + + +def test_next_job_401_raises_auth_error(): + session = RecordingSession(StubResponse(401)) + client = BackendClient('http://web:8080', session=session) + + with pytest.raises(BackendAuthError) as excinfo: + client.next_job(make_identity()) + assert excinfo.value.status_code == 401 + + +def test_next_job_500_raises_api_error(): + session = RecordingSession(StubResponse(500)) + client = BackendClient('http://web:8080', session=session) + + with pytest.raises(BackendAPIError) as excinfo: + client.next_job(make_identity()) + assert not isinstance(excinfo.value, BackendAuthError) + assert excinfo.value.status_code == 500 + + +def test_complete_204_sends_url_body_and_auth(): + session = RecordingSession(StubResponse(204)) + client = BackendClient('http://web:8080', session=session, timeout=10) + tasks = [{'status': 0}] + + result = client.complete(make_identity(), 'jb_1', tasks) + + assert result is None + call = session.calls[0] + assert call['method'] == 'put' + assert call['url'] == 'http://web:8080/runners/rn_abc/jobs/jb_1/complete' + assert call['json'] == {'tasks': tasks} + assert call['headers'] == {'Authorization': 'Bearer rk_secret'} + assert call['timeout'] == 10 + + +def test_complete_409_raises_api_error(): + session = RecordingSession(StubResponse(409)) + client = BackendClient('http://web:8080', session=session) + + with pytest.raises(BackendAPIError) as excinfo: + client.complete(make_identity(), 'jb_1', []) + assert excinfo.value.status_code == 409 + + +def test_abort_202_sends_url_and_reason(): + session = RecordingSession(StubResponse(202)) + client = BackendClient('http://web:8080', session=session) + + result = client.abort(make_identity(), 'jb_1', 'prep_failed') + + assert result is None + call = session.calls[0] + assert call['method'] == 'put' + assert call['url'] == 'http://web:8080/runners/rn_abc/jobs/jb_1/abort' + assert call['json'] == {'reason': 'prep_failed'} + assert call['headers'] == {'Authorization': 'Bearer rk_secret'} + + +def test_abort_404_raises_api_error(): + session = RecordingSession(StubResponse(404)) + client = BackendClient('http://web:8080', session=session) + + with pytest.raises(BackendAPIError) as excinfo: + client.abort(make_identity(), 'jb_1', 'drain') + assert excinfo.value.status_code == 404 diff --git a/tests/test_runner_poller.py b/tests/test_runner_poller.py new file mode 100644 index 0000000..67c58c2 --- /dev/null +++ b/tests/test_runner_poller.py @@ -0,0 +1,397 @@ +import io +import queue +import time + +import pytest +import requests + +from runner.client import ( + BackendAPIError, + BackendAuthError, + JobPayload, + RunnerConfig, + RunnerIdentity, +) +from runner.active_jobs import ActiveJobTracker +from runner.result_sender import AbortRequest +from runner.poller import PollerThread, prepare_job +import runner.poller as poller_mod + + +def make_identity(max_concurrent=8, poll_interval=3): + return RunnerIdentity('rn_x', 'rk_tok', + RunnerConfig(15, poll_interval, max_concurrent)) + + +def make_payload(job_id='jb_1', submission_id='sub_1', problem_id=42): + return JobPayload( + job_id=job_id, + submission_id=submission_id, + problem_id=problem_id, + language=2, + code_url='http://minio/code.zip', + checker=None, + tasks=[{ + 'taskScore': 100, + 'memoryLimit': 65536, + 'timeLimit': 1000, + 'caseCount': 1, + }], + ) + + +class ScriptedClient: + """next_job() pops scripted outcomes (JobPayload / None / Exception).""" + + def __init__(self, outcomes): + self._outcomes = list(outcomes) + self.call_count = 0 + + def next_job(self, identity): + self.call_count += 1 + outcome = self._outcomes.pop(0) if self._outcomes else None + if isinstance(outcome, Exception): + raise outcome + return outcome + + +class DispatchRecorder: + + def __init__(self, fail_times=0, exc=None): + self.calls = [] + self._fail_times = fail_times + self._exc = exc or RuntimeError('dispatch failed') + + def __call__(self, job_id, submission_id): + self.calls.append((job_id, submission_id)) + if self._fail_times > 0: + self._fail_times -= 1 + raise self._exc + + +class PrepareRecorder: + + def __init__(self, tracker, fail_times=0, exc=None): + self._tracker = tracker + self.calls = [] + self.tracker_at_call = [] + self._fail_times = fail_times + self._exc = exc or RuntimeError('prep failed') + + def __call__(self, payload): + self.calls.append(payload) + self.tracker_at_call.append(self._tracker.snapshot()) + if self._fail_times > 0: + self._fail_times -= 1 + raise self._exc + + +class SleepRecorder: + + def __init__(self): + self.slept = [] + + def __call__(self, seconds): + self.slept.append(seconds) + + +def build_poller(outcomes, + *, + tracker=None, + identity=None, + prepare=None, + dispatch=None): + tracker = tracker if tracker is not None else ActiveJobTracker() + identity = identity if identity is not None else make_identity() + client = ScriptedClient(outcomes) + dispatch = dispatch if dispatch is not None else DispatchRecorder() + q = queue.Queue() + sleep = SleepRecorder() + poller = PollerThread( + client, + identity, + tracker, + q, + dispatch, + prepare=prepare if prepare is not None else PrepareRecorder(tracker), + sleep=sleep, + ) + return poller, client, tracker, q, sleep, dispatch + + +def test_capacity_gate_returns_true_no_next_job_call(): + tracker = ActiveJobTracker() + tracker.add('a') + tracker.add('b') + identity = make_identity(max_concurrent=2) + poller, client, _, q, _, _ = build_poller([], + tracker=tracker, + identity=identity) + + assert poller._poll_once() is True + assert client.call_count == 0 + assert q.empty() + + +def test_204_returns_true_tracker_untouched(): + poller, client, tracker, q, _, _ = build_poller([None]) + + assert poller._poll_once() is True + assert client.call_count == 1 + assert len(tracker) == 0 + assert q.empty() + + +def test_api_error_returns_true_loop_survives(): + poller, _, tracker, q, _, _ = build_poller( + [BackendAPIError('boom', status_code=500)]) + + assert poller._poll_once() is True + assert len(tracker) == 0 + assert q.empty() + + +def test_connection_error_returns_true(): + poller, _, _, q, _, _ = build_poller([requests.ConnectionError('down')]) + assert poller._poll_once() is True + assert q.empty() + + +def test_auth_error_returns_true(): + poller, _, _, q, _, _ = build_poller( + [BackendAuthError('nope', status_code=401)]) + assert poller._poll_once() is True + assert q.empty() + + +def test_success_adds_to_tracker_before_prepare(): + tracker = ActiveJobTracker() + prepare = PrepareRecorder(tracker) + dispatch = DispatchRecorder() + poller, client, tracker, q, _, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + + result = poller._poll_once() + + assert result is False + # tracker.add happened BEFORE prepare was invoked. + assert prepare.tracker_at_call == [['jb_1']] + assert len(prepare.calls) == 1 + assert dispatch.calls == [('jb_1', 'sub_1')] + assert q.empty() + assert tracker.snapshot() == ['jb_1'] + + +def test_prepare_fails_once_then_succeeds(): + tracker = ActiveJobTracker() + prepare = PrepareRecorder(tracker, fail_times=1) + dispatch = DispatchRecorder() + poller, _, tracker, q, sleep, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + + result = poller._poll_once() + + assert result is False + assert len(prepare.calls) == 2 + assert dispatch.calls == [('jb_1', 'sub_1')] + assert q.empty() + assert sleep.slept == [1] # PREP_BACKOFF_SCHEDULE[0] + + +def test_prepare_fails_all_attempts_queues_abort(): + tracker = ActiveJobTracker() + prepare = PrepareRecorder(tracker, fail_times=3) + dispatch = DispatchRecorder() + poller, _, tracker, q, sleep, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + + result = poller._poll_once() + + assert result is False + assert len(prepare.calls) == 3 + assert dispatch.calls == [] # never reached a successful dispatch + item = q.get_nowait() + assert isinstance(item, AbortRequest) + assert item.job_id == 'jb_1' + assert item.submission_id == 'sub_1' + assert item.reason == 'prep_failed' + assert sleep.slept == [1, 2] # PREP_BACKOFF_SCHEDULE + # Still in tracker: the sender removes it after the abort resolves. + assert tracker.snapshot() == ['jb_1'] + + +def test_dispatch_failure_aborts_without_retry(): + # A failed handle() may have partially enqueued task entries; calling + # it again for the same job_id would revive them (duplicate execution), + # so dispatch gets exactly one shot. + tracker = ActiveJobTracker() + prepare = PrepareRecorder(tracker) + dispatch = DispatchRecorder(fail_times=1, exc=queue.Full()) + poller, _, tracker, q, sleep, _ = build_poller([make_payload()], + tracker=tracker, + prepare=prepare, + dispatch=dispatch) + + result = poller._poll_once() + + assert result is False + assert len(prepare.calls) == 1 + assert len(dispatch.calls) == 1 + assert sleep.slept == [] + item = q.get_nowait() + assert isinstance(item, AbortRequest) + assert item.reason == 'prep_failed' + # Still in tracker: the sender removes it as it finalizes the abort. + assert tracker.snapshot() == ['jb_1'] + + +def test_real_thread_polls_and_stops_promptly(): + tracker = ActiveJobTracker() + client = ScriptedClient([None] * 100) + poller = PollerThread( + client, + make_identity(poll_interval=0.01), + tracker, + queue.Queue(), + DispatchRecorder(), + prepare=PrepareRecorder(tracker), + poll_interval_sec=0.01, + ) + + poller.start() + time.sleep(0.1) + t0 = time.time() + poller.stop() + poller.join(timeout=1.0) + elapsed = time.time() - t0 + + assert not poller.is_alive() + assert elapsed < 0.5 + assert client.call_count >= 2 + + +def test_run_survives_unexpected_error_from_poll(): + # A malformed 200 body raises KeyError inside next_job; the run loop + # must log it and keep polling, not die silently. + tracker = ActiveJobTracker() + dispatch = DispatchRecorder() + client = ScriptedClient([KeyError('job_id'), make_payload()]) + poller = PollerThread( + client, + make_identity(), + tracker, + queue.Queue(), + dispatch, + prepare=PrepareRecorder(tracker), + poll_interval_sec=0.01, + ) + + poller.start() + deadline = time.time() + 1.0 + while not dispatch.calls and time.time() < deadline: + time.sleep(0.01) + poller.stop() + poller.join(timeout=1.0) + + assert not poller.is_alive() + assert dispatch.calls == [('jb_1', 'sub_1')] + + +# --- prepare_job unit tests --- + + +class StubDownloadResponse: + + def __init__(self, content=b'zipbytes', raises=None): + self.content = content + self._raises = raises + + def raise_for_status(self): + if self._raises is not None: + raise self._raises + + +def test_prepare_job_happy_path(tmp_path, monkeypatch): + ensure_calls = [] + monkeypatch.setattr(poller_mod.testdata, 'ensure_testdata', + lambda pid: ensure_calls.append(pid)) + monkeypatch.setattr(poller_mod.testdata, 'get_problem_root', + lambda pid: tmp_path / f'root_{pid}') + + get_calls = [] + + def fake_get(url, timeout=None): + get_calls.append((url, timeout)) + return StubDownloadResponse(content=b'the-zip') + + monkeypatch.setattr(poller_mod.requests, 'get', fake_get) + monkeypatch.setattr(poller_mod, 'SUBMISSION_DIR', tmp_path) + + extract_kwargs = {} + monkeypatch.setattr(poller_mod.file_manager, 'extract', + lambda **kw: extract_kwargs.update(kw)) + + payload = make_payload(problem_id=42) + prepare_job(payload) + + assert ensure_calls == [42] + assert get_calls[0][0] == 'http://minio/code.zip' + assert get_calls[0][1] == poller_mod.REQUEST_TIMEOUT + assert extract_kwargs['root_dir'] == tmp_path + assert extract_kwargs['job_id'] == 'jb_1' + meta = extract_kwargs['meta'] + assert int(meta.language) == 2 + assert len(meta.tasks) == 1 + assert isinstance(extract_kwargs['source'], io.BytesIO) + assert extract_kwargs['source'].getvalue() == b'the-zip' + assert extract_kwargs['testdata'] == tmp_path / 'root_42' + + +def test_prepare_job_download_error_raises(tmp_path, monkeypatch): + monkeypatch.setattr(poller_mod.testdata, 'ensure_testdata', + lambda pid: None) + monkeypatch.setattr(poller_mod.testdata, 'get_problem_root', + lambda pid: tmp_path) + monkeypatch.setattr(poller_mod.requests, + 'get', + lambda url, timeout=None: StubDownloadResponse( + raises=requests.HTTPError('404'))) + monkeypatch.setattr(poller_mod, 'SUBMISSION_DIR', tmp_path) + monkeypatch.setattr(poller_mod.file_manager, 'extract', lambda **kw: None) + + with pytest.raises(requests.HTTPError): + prepare_job(make_payload()) + + +def test_prepare_job_removes_leftover_job_dir(tmp_path, monkeypatch): + monkeypatch.setattr(poller_mod.testdata, 'ensure_testdata', + lambda pid: None) + monkeypatch.setattr(poller_mod.testdata, 'get_problem_root', + lambda pid: tmp_path / 'root') + monkeypatch.setattr( + poller_mod.requests, + 'get', + lambda url, timeout=None: StubDownloadResponse(content=b'z')) + monkeypatch.setattr(poller_mod, 'SUBMISSION_DIR', tmp_path) + + seen = {} + + def fake_extract(**kw): + # By the time extract runs, the leftover dir must be gone. + seen['job_dir_exists'] = (tmp_path / 'jb_1').exists() + + monkeypatch.setattr(poller_mod.file_manager, 'extract', fake_extract) + + leftover = tmp_path / 'jb_1' + leftover.mkdir() + (leftover / 'stale.txt').write_text('old') + + prepare_job(make_payload()) + + assert seen['job_dir_exists'] is False diff --git a/tests/test_runner_result_sender.py b/tests/test_runner_result_sender.py new file mode 100644 index 0000000..2aa2206 --- /dev/null +++ b/tests/test_runner_result_sender.py @@ -0,0 +1,397 @@ +import queue +import time + +import requests + +from runner.client import BackendAPIError, BackendAuthError +from runner.active_jobs import ActiveJobTracker +from runner.result_sender import ( + CompleteRequest, + AbortRequest, + ResultSenderThread, + _default_cleanup, + _default_backup, +) +import runner.result_sender as result_sender + + +class ScriptedClient: + """complete()/abort() replay scripted outcomes and record calls.""" + + def __init__(self, complete_outcomes=None, abort_outcomes=None): + self._complete = list(complete_outcomes or []) + self._abort = list(abort_outcomes or []) + self.complete_calls = [] + self.abort_calls = [] + + def complete(self, identity, job_id, tasks): + self.complete_calls.append((job_id, tasks)) + if self._complete: + outcome = self._complete.pop(0) + if isinstance(outcome, Exception): + raise outcome + + def abort(self, identity, job_id, reason): + self.abort_calls.append((job_id, reason)) + if self._abort: + outcome = self._abort.pop(0) + if isinstance(outcome, Exception): + raise outcome + + +class Recorder: + + def __init__(self, raise_once=None): + self.calls = [] + self._raise_once = raise_once + + def __call__(self, job_id): + self.calls.append(job_id) + if self._raise_once is not None: + exc = self._raise_once + self._raise_once = None + raise exc + + +class SleepRecorder: + + def __init__(self): + self.slept = [] + + def __call__(self, seconds): + self.slept.append(seconds) + + +def build_sender(client, tracker=None, job_id='jb_1'): + tracker = tracker if tracker is not None else ActiveJobTracker() + tracker.add(job_id) + cleanup = Recorder() + backup = Recorder() + sleep = SleepRecorder() + sender = ResultSenderThread( + client, + identity=object(), + tracker=tracker, + result_queue=queue.Queue(), + cleanup=cleanup, + backup=backup, + sleep=sleep, + ) + return sender, tracker, cleanup, backup, sleep + + +def test_complete_204_cleanup_and_tracker_emptied(): + client = ScriptedClient(complete_outcomes=[None]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [{'status': 0}])) + + assert client.complete_calls == [('jb_1', [{'status': 0}])] + assert client.abort_calls == [] + assert cleanup.calls == ['jb_1'] + assert backup.calls == [] + assert sleep.slept == [] + assert len(tracker) == 0 + + +def test_complete_409_terminal_no_retry_cleanup(): + client = ScriptedClient( + complete_outcomes=[BackendAPIError('conflict', status_code=409)]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert client.abort_calls == [] + assert cleanup.calls == ['jb_1'] + assert backup.calls == [] + assert sleep.slept == [] + assert len(tracker) == 0 + + +def test_complete_404_terminal_cleanup(): + client = ScriptedClient( + complete_outcomes=[BackendAPIError('gone', status_code=404)]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert client.abort_calls == [] + assert cleanup.calls == ['jb_1'] + assert len(tracker) == 0 + + +def test_complete_400_triggers_abort_rejected_with_backup(): + # rejected keeps the job dir as evidence of what the backend refused. + client = ScriptedClient( + complete_outcomes=[BackendAPIError('bad', status_code=400)], + abort_outcomes=[None], + ) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert client.abort_calls == [('jb_1', 'rejected')] + assert backup.calls == ['jb_1'] + assert cleanup.calls == [] + assert len(tracker) == 0 + + +def test_complete_400_then_abort_exhausts_backup(): + client = ScriptedClient( + complete_outcomes=[BackendAPIError('bad', status_code=400)], + abort_outcomes=[requests.ConnectionError('x')] * 6, + ) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert client.abort_calls == [('jb_1', 'rejected')] * 6 + assert backup.calls == ['jb_1'] + assert cleanup.calls == [] + assert sleep.slept == [1, 2, 4, 8, 16] + assert len(tracker) == 0 + + +def test_complete_400_then_abort_409_still_backed_up(): + client = ScriptedClient( + complete_outcomes=[BackendAPIError('bad', status_code=400)], + abort_outcomes=[BackendAPIError('conflict', status_code=409)], + ) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert client.abort_calls == [('jb_1', 'rejected')] + assert backup.calls == ['jb_1'] + assert cleanup.calls == [] + assert sleep.slept == [] + assert len(tracker) == 0 + + +def test_abort_finalizes_local_state_before_send(): + # The 202 response means the backend has already requeued the job and + # this runner may re-claim it at once: by the time the abort request + # goes out, the dir must be gone and the tracker entry removed. + events = [] + tracker = ActiveJobTracker() + tracker.add('jb_1') + + class OrderClient: + + def abort(self, identity, job_id, reason): + events.append(('abort', reason, len(tracker))) + + sender = ResultSenderThread( + OrderClient(), + identity=object(), + tracker=tracker, + result_queue=queue.Queue(), + cleanup=lambda job_id: events.append(('cleanup', job_id)), + backup=lambda job_id: events.append(('backup', job_id)), + sleep=SleepRecorder(), + ) + + sender._process(AbortRequest('jb_1', 'sub_1', 'prep_failed')) + + # cleanup first, then the send observes an already-empty tracker. + assert events == [('cleanup', 'jb_1'), ('abort', 'prep_failed', 0)] + + +def test_abort_exhaustion_does_not_backup(): + # The dir was already finalized before the first send attempt, so + # exhaustion has nothing left to back up. + client = ScriptedClient(abort_outcomes=[requests.ConnectionError('x')] * 6) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(AbortRequest('jb_1', 'sub_1', 'prep_failed')) + + assert cleanup.calls == ['jb_1'] + assert backup.calls == [] + assert sleep.slept == [1, 2, 4, 8, 16] + assert len(tracker) == 0 + + +def test_finalize_failure_still_sends_abort(): + # An unsent abort would leave the job leased until lease expiry, so a + # failing cleanup must not block the send (or the tracker removal). + client = ScriptedClient(abort_outcomes=[None]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + cleanup._raise_once = PermissionError('denied') + + sender._process(AbortRequest('jb_1', 'sub_1', 'prep_failed')) + + assert client.abort_calls == [('jb_1', 'prep_failed')] + assert len(tracker) == 0 + + +def test_complete_500_twice_then_success(): + client = ScriptedClient(complete_outcomes=[ + BackendAPIError('err', status_code=500), + BackendAPIError('err', status_code=500), + None, + ]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert len(client.complete_calls) == 3 + assert sleep.slept == [1, 2] + assert cleanup.calls == ['jb_1'] + assert backup.calls == [] + + +def test_complete_connection_error_exhausts_backup(): + client = ScriptedClient( + complete_outcomes=[requests.ConnectionError('down')] * 6) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert len(client.complete_calls) == 6 # initial + 5 retries + assert sleep.slept == [1, 2, 4, 8, 16] + assert backup.calls == ['jb_1'] + assert cleanup.calls == [] + assert len(tracker) == 0 + + +def test_complete_401_not_terminal_then_success(): + client = ScriptedClient(complete_outcomes=[ + BackendAuthError('nope', status_code=401), + None, + ]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + + assert len(client.complete_calls) == 2 # 401 retried, not terminal + assert sleep.slept == [1] + assert cleanup.calls == ['jb_1'] + assert backup.calls == [] + + +def test_abort_request_prep_failed_202_cleanup(): + client = ScriptedClient(abort_outcomes=[None]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(AbortRequest('jb_1', 'sub_1', 'prep_failed')) + + assert client.abort_calls == [('jb_1', 'prep_failed')] + assert cleanup.calls == ['jb_1'] + assert backup.calls == [] + assert len(tracker) == 0 + + +def test_abort_409_terminal_cleanup(): + client = ScriptedClient( + abort_outcomes=[BackendAPIError('conflict', status_code=409)]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + + sender._process(AbortRequest('jb_1', 'sub_1', 'drain')) + + assert cleanup.calls == ['jb_1'] + assert sleep.slept == [] + assert len(tracker) == 0 + + +def test_raising_cleanup_still_removes_job_from_tracker(): + # A lingering tracker entry would keep the lease renewed and consume + # capacity forever, so removal must survive a failing cleanup. + client = ScriptedClient(complete_outcomes=[None]) + sender, tracker, cleanup, backup, sleep = build_sender(client) + cleanup._raise_once = PermissionError('denied') + + try: + sender._process(CompleteRequest('jb_1', 'sub_1', [])) + except PermissionError: + pass + + assert len(tracker) == 0 + + +def test_default_cleanup_tolerates_missing_dir(tmp_path, monkeypatch): + monkeypatch.setattr('dispatcher.config.SUBMISSION_DIR', tmp_path) + # Must not raise even though the job dir does not exist. + _default_cleanup('does_not_exist') + + +def test_default_backup_tolerates_missing_dir(tmp_path, monkeypatch): + monkeypatch.setattr('dispatcher.config.SUBMISSION_DIR', tmp_path) + monkeypatch.setattr('dispatcher.config.SUBMISSION_BACKUP_DIR', tmp_path) + _default_backup('does_not_exist') + + +def test_real_thread_drains_queued_items_on_stop(): + client = ScriptedClient(complete_outcomes=[None, None]) + tracker = ActiveJobTracker() + tracker.add('jb_1') + tracker.add('jb_2') + cleanup = Recorder() + q = queue.Queue() + q.put(CompleteRequest('jb_1', 'sub_1', [])) + q.put(CompleteRequest('jb_2', 'sub_2', [])) + sender = ResultSenderThread( + client, + identity=object(), + tracker=tracker, + result_queue=q, + cleanup=cleanup, + backup=Recorder(), + queue_poll_sec=0.01, + ) + + sender.start() + sender.stop() + sender.join(timeout=2.0) + + assert not sender.is_alive() + assert sorted(cleanup.calls) == ['jb_1', 'jb_2'] + assert len(tracker) == 0 + + +def test_real_thread_empty_queue_stops_promptly(): + sender = ResultSenderThread( + ScriptedClient(), + identity=object(), + tracker=ActiveJobTracker(), + result_queue=queue.Queue(), + cleanup=Recorder(), + backup=Recorder(), + queue_poll_sec=0.01, + ) + + sender.start() + time.sleep(0.05) + t0 = time.time() + sender.stop() + sender.join(timeout=1.0) + assert not sender.is_alive() + assert time.time() - t0 < 0.5 + + +def test_bad_item_does_not_kill_thread(): + # cleanup raises once; the second queued item must still be processed. + client = ScriptedClient(complete_outcomes=[None, None]) + tracker = ActiveJobTracker() + tracker.add('jb_1') + tracker.add('jb_2') + cleanup = Recorder(raise_once=RuntimeError('boom')) + q = queue.Queue() + q.put(CompleteRequest('jb_1', 'sub_1', [])) + q.put(CompleteRequest('jb_2', 'sub_2', [])) + sender = ResultSenderThread( + client, + identity=object(), + tracker=tracker, + result_queue=q, + cleanup=cleanup, + backup=Recorder(), + queue_poll_sec=0.01, + ) + + sender.start() + sender.stop() + sender.join(timeout=2.0) + + assert not sender.is_alive() + # jb_1's cleanup raised, jb_2's succeeded. + assert 'jb_2' in cleanup.calls