Blame for task_executor.py

# -*- 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())