Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| """SmartPaper R3-Free independent task executor. | |
| Launched by Streamlit as a detached subprocess. It claims a D1 task with a | |
| server-to-server service key, runs the headless generation pipeline, renews the | |
| lease only while real execution activity is recent, and atomically asks the | |
| Worker to commit/release quota at terminal completion. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import logging | |
| import os | |
| import re | |
| import sys | |
| import threading | |
| import time | |
| from pathlib import Path | |
| from urllib.parse import urlsplit | |
| import requests | |
| from config import APP_CONFIG | |
| from core.generation_runner import GenerationRunner | |
| from core.pipeline_agent import ( | |
| DeepSeekAPIError, DeepSeekAuthenticationError, DeepSeekBalanceError, DeepSeekRequestError, | |
| ) | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(message)s", | |
| handlers=[logging.StreamHandler(sys.stdout)], | |
| ) | |
| def _safe(value, limit=500): | |
| text = str(value or "") | |
| text = re.sub(r"(?i)(api[_ -]?key|authorization|bearer|token)\s*[:=]\s*[^\s,;]+", r"\1=[redacted]", text) | |
| return text[:limit] | |
| def _api_origin() -> str: | |
| value = str(os.environ.get("SMARTPAPER_AUTH_API") or "").strip().rstrip("/") | |
| parsed = urlsplit(value) | |
| if parsed.scheme != "https" or not parsed.netloc or parsed.path not in ("", "/") or parsed.query or parsed.fragment: | |
| raise RuntimeError("SMARTPAPER_AUTH_API_invalid") | |
| return f"https://{parsed.netloc}" | |
| def _service_key() -> str: | |
| value = str(os.environ.get("SMARTPAPER_SERVICE_KEY") or "").strip() | |
| if not value: | |
| raise RuntimeError("SMARTPAPER_SERVICE_KEY_missing") | |
| return value | |
| class ExecutorApiError(RuntimeError): | |
| def __init__(self, error: str, status_code: int = 0, payload: dict | None = None): | |
| super().__init__(str(error or "executor_api_error")) | |
| self.error = str(error or "executor_api_error") | |
| self.status_code = int(status_code or 0) | |
| self.payload = dict(payload or {}) | |
| def _is_recoverable_generation_failure(exc: Exception) -> bool: | |
| """Return True only for transient execution failures that can use checkpoints. | |
| Authentication/balance/request-validation failures are terminal and release | |
| quota immediately. Transport/stream failures can be resumed from the last | |
| durable task checkpoint without consuming a second quota unit. | |
| """ | |
| if isinstance(exc, ExecutorApiError): | |
| return exc.status_code == 429 or exc.status_code >= 500 | |
| if isinstance(exc, (DeepSeekAuthenticationError, DeepSeekBalanceError, DeepSeekRequestError)): | |
| return False | |
| if isinstance(exc, (requests.exceptions.Timeout, requests.exceptions.ConnectionError)): | |
| return True | |
| if isinstance(exc, DeepSeekAPIError): | |
| return True | |
| text = str(exc or "").lower() | |
| transient_markers = ( | |
| "connection", "timed out", "timeout", "temporarily unavailable", | |
| "streaming failed after", "api failed after", "stream was interrupted", | |
| "insufficient eligible references", | |
| ) | |
| return any(marker in text for marker in transient_markers) | |
| class ExecutorApi: | |
| def __init__(self): | |
| self.origin = _api_origin() | |
| self.headers = { | |
| "Authorization": f"Bearer {_service_key()}", | |
| "Accept": "application/json", | |
| "Content-Type": "application/json", | |
| "User-Agent": "SmartPaper-R3-Free-Executor/1.0", | |
| } | |
| def post(self, path: str, payload: dict, *, timeout=20) -> dict: | |
| response = requests.post( | |
| self.origin + path, | |
| headers=self.headers, | |
| json=payload, | |
| timeout=timeout, | |
| ) | |
| try: | |
| data = response.json() | |
| except Exception as exc: | |
| raise RuntimeError(f"executor_worker_non_json_http_{response.status_code}") from exc | |
| if not response.ok or data.get("ok") is False: | |
| error = str(data.get("error") or data.get("reason") or f"HTTP_{response.status_code}") | |
| raise ExecutorApiError(error, response.status_code, data) | |
| return data | |
| def claim(self, *, task_id: str, user_id: int, request_hash: str) -> dict: | |
| return self.post( | |
| "/executor/task/claim", | |
| {"task_id": task_id, "user_id": user_id, "request_hash": request_hash}, | |
| ) | |
| def heartbeat(self, payload: dict) -> dict: | |
| return self.post("/executor/task/heartbeat", payload) | |
| def finish(self, payload: dict) -> dict: | |
| return self.post("/executor/task/finish", payload, timeout=30) | |
| class ExecutionState: | |
| def __init__(self, *, task_id: str, user_id: int, epoch: int, instance_id: str, api: ExecutorApi): | |
| self.task_id = task_id | |
| self.user_id = int(user_id) | |
| self.epoch = int(epoch) | |
| self.instance_id = str(instance_id) | |
| self.api = api | |
| self._lock = threading.Lock() | |
| self._last_activity = time.monotonic() | |
| self._progress_pct = 0 | |
| self._current_node = "executor_claimed" | |
| self._checkpoint_seq = 0 | |
| self._checkpoint_stage = "" | |
| self._ds_calls = 0 | |
| self._s2_calls = 0 | |
| self._oa_calls = 0 | |
| self._cancel = False | |
| self._stop = threading.Event() | |
| self._thread = None | |
| def touch(self, *, progress_pct=None, current_node=None, checkpoint_seq=None, checkpoint_stage=None, ds_calls=None, s2_calls=None, oa_calls=None, activity_only=False): | |
| with self._lock: | |
| self._last_activity = time.monotonic() | |
| if not activity_only: | |
| if progress_pct is not None: | |
| self._progress_pct = max(self._progress_pct, max(0, min(100, int(progress_pct)))) | |
| if current_node: | |
| self._current_node = str(current_node)[:240] | |
| if checkpoint_seq is not None and int(checkpoint_seq) >= self._checkpoint_seq: | |
| self._checkpoint_seq = int(checkpoint_seq) | |
| self._checkpoint_stage = str(checkpoint_stage or self._checkpoint_stage)[:120] | |
| if ds_calls is not None: | |
| self._ds_calls = max(self._ds_calls, int(ds_calls)) | |
| if s2_calls is not None: | |
| self._s2_calls = max(self._s2_calls, int(s2_calls)) | |
| if oa_calls is not None: | |
| self._oa_calls = max(self._oa_calls, int(oa_calls)) | |
| def cancelled(self) -> bool: | |
| with self._lock: | |
| return bool(self._cancel) | |
| def start_heartbeat(self): | |
| self._thread = threading.Thread(target=self._heartbeat_loop, name="smartpaper-executor-heartbeat", daemon=True) | |
| self._thread.start() | |
| def stop_heartbeat(self): | |
| self._stop.set() | |
| if self._thread and self._thread.is_alive(): | |
| self._thread.join(timeout=2) | |
| def _payload(self): | |
| with self._lock: | |
| return { | |
| "task_id": self.task_id, | |
| "user_id": self.user_id, | |
| "execution_epoch": self.epoch, | |
| "executor_instance_id": self.instance_id, | |
| "current_node": self._current_node, | |
| "progress_pct": self._progress_pct, | |
| "checkpoint_seq": self._checkpoint_seq, | |
| "checkpoint_stage": self._checkpoint_stage, | |
| "ds_calls": self._ds_calls, | |
| "s2_calls": self._s2_calls, | |
| "oa_calls": self._oa_calls, | |
| }, time.monotonic() - self._last_activity | |
| def heartbeat_once(self): | |
| payload, _ = self._payload() | |
| result = self.api.heartbeat(payload) | |
| if result.get("cancel_requested"): | |
| with self._lock: | |
| self._cancel = True | |
| return result | |
| def _heartbeat_loop(self): | |
| interval = int(APP_CONFIG.get("executor_heartbeat_interval_seconds", 15)) | |
| grace = int(APP_CONFIG.get("executor_activity_grace_seconds", 45)) | |
| # Claim itself already set the first heartbeat; delay before the next one. | |
| while not self._stop.wait(interval): | |
| payload, idle_seconds = self._payload() | |
| if idle_seconds > grace: | |
| logging.warning( | |
| "[SP_EXECUTOR] heartbeat paused task_id=%s idle_s=%.1f node=%s", | |
| self.task_id, idle_seconds, payload.get("current_node"), | |
| ) | |
| continue | |
| try: | |
| result = self.api.heartbeat(payload) | |
| if result.get("cancel_requested"): | |
| with self._lock: | |
| self._cancel = True | |
| return | |
| except ExecutorApiError as exc: | |
| if exc.error in { | |
| "executor_superseded", "generation_lock_lost", | |
| "task_already_finalized", "task_cancel_requested", | |
| } or exc.status_code in {404, 409, 410}: | |
| with self._lock: | |
| self._cancel = True | |
| logging.warning( | |
| "[SP_EXECUTOR] executor_no_longer_authoritative task_id=%s error=%s", | |
| self.task_id, _safe(exc.error, 180), | |
| ) | |
| return | |
| logging.warning( | |
| "[SP_EXECUTOR] heartbeat_error task_id=%s error=%s", | |
| self.task_id, _safe(exc, 180), | |
| ) | |
| except Exception as exc: | |
| # A transport outage alone does not authorize another local | |
| # writer. D1 lease/epoch ownership remains authoritative. | |
| logging.warning( | |
| "[SP_EXECUTOR] heartbeat_error task_id=%s error=%s", | |
| self.task_id, _safe(exc, 180), | |
| ) | |
| def _load_request(task_dir: Path) -> dict: | |
| path = task_dir / "request.json" | |
| if not path.is_file(): | |
| raise RuntimeError("request_manifest_missing") | |
| value = json.loads(path.read_text(encoding="utf-8")) | |
| if not isinstance(value, dict): | |
| raise RuntimeError("request_manifest_invalid") | |
| return value | |
| def main() -> int: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--task-dir", required=True) | |
| parser.add_argument("--task-id", required=True) | |
| parser.add_argument("--user-id", required=True, type=int) | |
| args = parser.parse_args() | |
| task_dir = Path(args.task_dir).resolve() | |
| request_data = _load_request(task_dir) | |
| if str(request_data.get("task_id") or "") != args.task_id: | |
| raise RuntimeError("request_task_id_mismatch") | |
| if int(request_data.get("user_id") or 0) != int(args.user_id): | |
| raise RuntimeError("request_user_id_mismatch") | |
| request_hash = str(request_data.get("request_hash") or "").strip() | |
| if not request_hash: | |
| raise RuntimeError("request_hash_missing") | |
| api = ExecutorApi() | |
| claim = api.claim(task_id=args.task_id, user_id=args.user_id, request_hash=request_hash) | |
| task = dict(claim.get("task") or {}) | |
| execution = dict(claim.get("execution") or {}) | |
| epoch = int(execution.get("execution_epoch") or task.get("execution_epoch") or 0) | |
| instance_id = str(execution.get("executor_instance_id") or "") | |
| if epoch <= 0 or not instance_id: | |
| raise RuntimeError("executor_claim_identity_missing") | |
| if int(task.get("target_words") or 0) != int(request_data.get("target_words") or 0): | |
| raise RuntimeError("target_words_claim_mismatch") | |
| state = ExecutionState( | |
| task_id=args.task_id, | |
| user_id=args.user_id, | |
| epoch=epoch, | |
| instance_id=instance_id, | |
| api=api, | |
| ) | |
| state.start_heartbeat() | |
| logging.info( | |
| "[SP_EXECUTOR] claimed task_id=%s epoch=%s instance=%s", | |
| args.task_id, epoch, instance_id[:12], | |
| ) | |
| # If D1 says a prior checkpoint exists but the free Space's local disk no | |
| # longer has it (container rebuild), do not pretend resume is possible. | |
| remote_checkpoint_seq = int(task.get("checkpoint_seq") or 0) | |
| local_checkpoints = list((task_dir / "checkpoints").glob("checkpoint_*.json")) if (task_dir / "checkpoints").is_dir() else [] | |
| if remote_checkpoint_seq > 0 and not local_checkpoints: | |
| try: | |
| api.finish( | |
| { | |
| "task_id": args.task_id, | |
| "user_id": args.user_id, | |
| "execution_epoch": epoch, | |
| "executor_instance_id": instance_id, | |
| "status": "failed", | |
| "safe_error_code": "E_CHECKPOINT_MISSING", | |
| "safe_error_detail": "Checkpoint metadata exists but local checkpoint files are unavailable after runtime replacement; quota released.", | |
| "current_node": "checkpoint_missing", | |
| } | |
| ) | |
| finally: | |
| state.stop_heartbeat() | |
| return 2 | |
| try: | |
| runner = GenerationRunner( | |
| task_dir=task_dir, | |
| request_data=request_data, | |
| execution_epoch=epoch, | |
| progress_callback=state.touch, | |
| cancel_checker=state.cancelled, | |
| ) | |
| result = runner.run() | |
| state.touch(progress_pct=100, current_node="package_validated") | |
| # One last accepted heartbeat publishes the final checkpoint metadata | |
| # before the atomic terminal transaction. | |
| try: | |
| state.heartbeat_once() | |
| except Exception as exc: | |
| logging.warning("[SP_EXECUTOR] final_heartbeat_error task_id=%s error=%s", args.task_id, _safe(exc, 180)) | |
| if state.cancelled(): | |
| raise RuntimeError("task_cancel_requested") | |
| api.finish( | |
| { | |
| "task_id": args.task_id, | |
| "user_id": args.user_id, | |
| "execution_epoch": epoch, | |
| "executor_instance_id": instance_id, | |
| "status": "completed", | |
| "current_node": "package_validated", | |
| "output_filename": str(result.get("output_filename") or "")[:240], | |
| } | |
| ) | |
| logging.info("[SP_EXECUTOR] completed task_id=%s epoch=%s", args.task_id, epoch) | |
| return 0 | |
| except Exception as exc: | |
| if state.cancelled(): | |
| logging.warning( | |
| "[SP_EXECUTOR] stopped_non_authoritative task_id=%s epoch=%s", | |
| args.task_id, epoch, | |
| ) | |
| return 3 | |
| if _is_recoverable_generation_failure(exc): | |
| # Do not terminally fail or release quota. Stopping the real-execution | |
| # heartbeat lets the Worker move this epoch to `resumable`; the next | |
| # epoch continues from the latest immutable local checkpoint. | |
| logging.exception( | |
| "[SP_EXECUTOR] recoverable_interruption task_id=%s epoch=%s error=%s", | |
| args.task_id, epoch, _safe(exc, 240), | |
| ) | |
| return 75 | |
| logging.exception("[SP_EXECUTOR] failed task_id=%s epoch=%s error=%s", args.task_id, epoch, _safe(exc, 240)) | |
| try: | |
| api.finish( | |
| { | |
| "task_id": args.task_id, | |
| "user_id": args.user_id, | |
| "execution_epoch": epoch, | |
| "executor_instance_id": instance_id, | |
| "status": "failed", | |
| "safe_error_code": "E_GENERATION_FAILED", | |
| "safe_error_detail": _safe(exc, 800), | |
| "current_node": "generation_failed", | |
| } | |
| ) | |
| except Exception as finish_exc: | |
| # A newer epoch or user cancellation may already own/finalize the | |
| # task. Never override it from a stale executor. | |
| logging.warning("[SP_EXECUTOR] failure_finalize_rejected task_id=%s error=%s", args.task_id, _safe(finish_exc, 220)) | |
| return 1 | |
| finally: | |
| state.stop_heartbeat() | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |