ls-algorithm-plugin-sdk 0.2.5__py3-none-any.whl

algorithm_plugin_sdk/service.py

@@ -0,0 +1,664 @@
+from __future__ import annotations
+
+import hmac
+import json
+import logging
+import random
+import threading
+import time
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field, replace
+from typing import Any, Callable
+
+from .context import ExecutionContext, ProgressSnapshot, utc_now
+from .errors import (ExecutionCancelled, ExecutionNotFinished,
+                     ExecutionNotFound, IdempotencyConflict)
+from .models import AlgorithmRequest, AlgorithmResult
+from .release import ReleaseManifest
+from .runner import AlgorithmRunner
+from .webui_app import announce_webui, mount_webui
+
+TERMINAL_STATES = {"succeeded", "partial", "failed", "cancelled"}
+MAX_ATTEMPTS = 10
+RETRY_SECONDS = 1.0
+RETRY_JITTER_RATIO = 0.25
+
+RunnerFactory = Callable[[], AlgorithmRunner]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class ExecutionRecord:
+    execution_id: str
+    request: AlgorithmRequest
+    accepted_at: str
+    state: str = "accepted"
+    started_at: str | None = None
+    finished_at: str | None = None
+    result: AlgorithmResult | None = None
+    error: str | None = None
+    context: ExecutionContext | None = None
+    future: Future[None] | None = None
+    attempt: int = 0
+    max_attempts: int = 1
+    events: list[dict[str, Any]] = field(default_factory=list)
+    condition: threading.Condition = field(
+        default_factory=lambda: threading.Condition(threading.RLock())
+    )
+
+    def add_progress(self, snapshot: ProgressSnapshot) -> None:
+        with self.condition:
+            self.events.append(snapshot.to_dict())
+            self.condition.notify_all()
+
+    def set_state(
+        self,
+        state: str,
+        *,
+        result: AlgorithmResult | None = None,
+        error: str | None = None,
+    ) -> None:
+        with self.condition:
+            self.state = state
+            self.result = result
+            self.error = error
+            if state == "running":
+                self.started_at = utc_now()
+            if state in TERMINAL_STATES:
+                self.finished_at = utc_now()
+            self.condition.notify_all()
+
+    def status_payload(self) -> dict[str, Any]:
+        with self.condition:
+            progress = self.context.snapshot().to_dict() if self.context else None
+            return {
+                "executionId": self.execution_id,
+                "state": self.state,
+                "acceptedAt": self.accepted_at,
+                "startedAt": self.started_at,
+                "finishedAt": self.finished_at,
+                "error": self.error,
+                "attempt": self.attempt,
+                "maxAttempts": self.max_attempts,
+                "progress": progress,
+            }
+
+    def events_after(self, sequence: int, wait_seconds: float) -> dict[str, Any]:
+        deadline = time.monotonic() + wait_seconds
+        with self.condition:
+            while (
+                not any(event["sequence"] > sequence for event in self.events)
+                and self.state not in TERMINAL_STATES
+                and time.monotonic() < deadline
+            ):
+                self.condition.wait(timeout=max(0.0, deadline - time.monotonic()))
+            events = [
+                event for event in self.events if event["sequence"] > sequence
+            ]
+            next_sequence = events[-1]["sequence"] if events else sequence
+            return {
+                "events": events,
+                "nextSequence": next_sequence,
+                "terminal": self.state in TERMINAL_STATES,
+            }
+
+
+class ExecutionManager:
+    """Protocol-independent background execution service."""
+
+    def __init__(
+        self,
+        runner: AlgorithmRunner,
+        *,
+        max_concurrent_executions: int = 1,
+        scratch_dir: str | None = None,
+        gpu_ids: list[int] | None = None,
+        runner_factory: RunnerFactory | None = None,
+        max_attempts: int = MAX_ATTEMPTS,
+        retry_seconds: float = RETRY_SECONDS,
+        retry_jitter_ratio: float = RETRY_JITTER_RATIO,
+        release_manifest: ReleaseManifest | None = None,
+    ) -> None:
+        if max_concurrent_executions < 1:
+            raise ValueError("max_concurrent_executions must be positive")
+        if max_attempts < 1:
+            raise ValueError("max_attempts must be positive")
+        if retry_seconds <= 0:
+            raise ValueError("retry_seconds must be positive")
+        if not 0 <= retry_jitter_ratio <= 1:
+            raise ValueError("retry_jitter_ratio must be in [0, 1]")
+        configured_gpu_ids = None if gpu_ids is None else list(gpu_ids)
+        if configured_gpu_ids is not None and (
+            any(
+                isinstance(gpu_id, bool) or not isinstance(gpu_id, int) or gpu_id < 0
+                for gpu_id in configured_gpu_ids
+            )
+            or len(configured_gpu_ids) != len(set(configured_gpu_ids))
+        ):
+            raise ValueError("gpu_ids must be unique non-negative integers")
+        self.runner = runner
+        algorithm_class = type(runner.algorithm)
+        self._runner_factory = runner_factory or (
+            lambda: AlgorithmRunner(algorithm_class())
+        )
+        self.scratch_dir = scratch_dir
+        self.gpu_ids = configured_gpu_ids
+        self.max_concurrent_executions = max_concurrent_executions
+        self.max_attempts = max_attempts
+        self.retry_seconds = retry_seconds
+        self.retry_jitter_ratio = retry_jitter_ratio
+        self.release_manifest = release_manifest
+        self._executor = ThreadPoolExecutor(
+            max_workers=max_concurrent_executions,
+            thread_name_prefix="algorithm-execution",
+        )
+        self._records: dict[str, ExecutionRecord] = {}
+        self._idempotency: dict[str, tuple[str, str]] = {}
+        self._lock = threading.RLock()
+        self._started_at = time.monotonic()
+        self._started = False
+        self._closed = False
+        self._runner_condition = threading.Condition(threading.RLock())
+        self._runner_generation = 0
+        self._runner_users = 0
+        self._restart_pending = False
+
+    def start(self) -> None:
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("execution manager is closed")
+            if self._started:
+                return
+            self.runner.start()
+            self._started = True
+
+    def submit(
+        self,
+        request: AlgorithmRequest,
+        *,
+        idempotency_key: str | None = None,
+    ) -> dict[str, Any]:
+        if self.gpu_ids is not None:
+            request = replace(request, gpu_ids=self.gpu_ids)
+        self.start()
+        canonical_request = json.dumps(
+            request.to_dict(), ensure_ascii=False, sort_keys=True, separators=(",", ":")
+        )
+        with self._lock:
+            if idempotency_key:
+                existing = self._idempotency.get(idempotency_key)
+                if existing is not None:
+                    existing_request, execution_id = existing
+                    if existing_request != canonical_request:
+                        raise IdempotencyConflict(idempotency_key)
+                    record = self._records[execution_id]
+                    return {
+                        "executionId": execution_id,
+                        "state": record.state,
+                        "acceptedAt": record.accepted_at,
+                    }
+        execution_id = f"exec-{uuid.uuid4().hex}"
+        record = ExecutionRecord(
+            execution_id=execution_id,
+            request=request,
+            accepted_at=utc_now(),
+            max_attempts=self.max_attempts,
+        )
+        record.context = ExecutionContext(
+            execution_id,
+            [item.dataset for item in request.inputs],
+            scratch_dir=(
+                record.request.workspace["scratchRoot"]
+                if record.request.workspace is not None
+                else self.scratch_dir
+            ),
+            progress_callback=record.add_progress,
+        )
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("execution manager is closed")
+            self._records[execution_id] = record
+            if idempotency_key:
+                self._idempotency[idempotency_key] = (
+                    canonical_request,
+                    execution_id,
+                )
+            record.future = self._executor.submit(self._run, record)
+        return {
+            "executionId": execution_id,
+            "state": record.state,
+            "acceptedAt": record.accepted_at,
+        }
+
+    def manifest(self) -> dict[str, Any]:
+        metadata = self.runner.algorithm.metadata()
+        capabilities = {
+            "cli": True,
+            "webui": True,
+            "httpService": True,
+            "cancel": True,
+            "idempotency": True,
+        }
+        if self.release_manifest is not None:
+            manifest = self.release_manifest.plugin_manifest(
+                max_concurrency=self.max_concurrent_executions,
+                capabilities=capabilities,
+            )
+            manifest["metadata"] = metadata.to_dict()
+            return manifest
+        return {
+            "apiVersion": "ldp.algorithm/v1",
+            "algorithmType": metadata.name.replace("-", "_"),
+            "name": metadata.name,
+            "version": metadata.version,
+            "inputSchema": "algorithm-plugin-sdk.request/v1",
+            "outputSchema": "algorithm-plugin-sdk.result/v1",
+            "maxConcurrency": self.max_concurrent_executions,
+            "capabilities": capabilities,
+            "metadata": metadata.to_dict(),
+        }
+
+    def heartbeat(self) -> dict[str, Any]:
+        with self._lock:
+            active = [
+                record.status_payload()
+                for record in self._records.values()
+                if record.state not in TERMINAL_STATES
+            ]
+            started = self._started and not self._closed
+        return {
+            "status": "busy" if active else ("ready" if started else "starting"),
+            "uptimeSeconds": int(time.monotonic() - self._started_at),
+            "modelLoaded": started,
+            "activeExecutions": len(active),
+            "maxConcurrency": self.max_concurrent_executions,
+            "executions": active,
+        }
+
+    def status(self, execution_id: str) -> dict[str, Any]:
+        return self._record(execution_id).status_payload()
+
+    def events(
+        self,
+        execution_id: str,
+        *,
+        after_sequence: int = 0,
+        wait_seconds: float = 0,
+    ) -> dict[str, Any]:
+        if after_sequence < 0:
+            raise ValueError("after_sequence cannot be negative")
+        if not 0 <= wait_seconds <= 30:
+            raise ValueError("wait_seconds must be in [0, 30]")
+        return self._record(execution_id).events_after(
+            after_sequence,
+            wait_seconds,
+        )
+
+    def result(self, execution_id: str) -> dict[str, Any]:
+        record = self._record(execution_id)
+        with record.condition:
+            if record.state not in TERMINAL_STATES:
+                raise ExecutionNotFinished(execution_id)
+            if record.result is None:
+                return {
+                    "executionId": execution_id,
+                    "state": record.state,
+                    "error": record.error,
+                }
+            return {
+                "executionId": execution_id,
+                "state": record.state,
+                "result": record.result.to_dict(),
+            }
+
+    def cancel(self, execution_id: str, reason: str = "cancelled by user") -> dict[str, Any]:
+        record = self._record(execution_id)
+        with record.condition:
+            if record.state in TERMINAL_STATES:
+                return {
+                    "executionId": execution_id,
+                    "state": record.state,
+                    "cancelRequested": False,
+                }
+            assert record.context is not None
+            record.context.cancel(reason)
+            cancelled_before_start = bool(record.future and record.future.cancel())
+            if cancelled_before_start:
+                record.set_state("cancelled", error=reason)
+            return {
+                "executionId": execution_id,
+                "state": record.state,
+                "cancelRequested": True,
+            }
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            active = [
+                record
+                for record in self._records.values()
+                if record.state not in TERMINAL_STATES
+            ]
+        for record in active:
+            assert record.context is not None
+            record.context.cancel("service shutdown")
+            if record.future and record.future.cancel():
+                record.set_state("cancelled", error="service shutdown")
+        self._executor.shutdown(wait=True, cancel_futures=True)
+        self.runner.close()
+
+    def _record(self, execution_id: str) -> ExecutionRecord:
+        with self._lock:
+            record = self._records.get(execution_id)
+        if record is None:
+            raise ExecutionNotFound(execution_id)
+        return record
+
+    def _run(self, record: ExecutionRecord) -> None:
+        record.set_state("running")
+        assert record.context is not None
+        try:
+            result = self._run_with_restarts(record)
+        except ExecutionCancelled as exc:
+            record.context.mark_unfinished("cancelled", str(exc))
+            record.set_state("cancelled", error=str(exc))
+        except Exception as exc:
+            error = f"{type(exc).__name__}: {exc}"
+            record.context.mark_unfinished("failed", error)
+            record.set_state("failed", error=error)
+        else:
+            record.set_state(result.status, result=result)
+
+    def _run_with_restarts(self, record: ExecutionRecord) -> AlgorithmResult:
+        assert record.context is not None
+        for attempt in range(1, self.max_attempts + 1):
+            with record.condition:
+                record.attempt = attempt
+                record.condition.notify_all()
+
+            runner, generation = self._acquire_runner()
+            failure: Exception | None = None
+            try:
+                return runner.run(
+                    record.request,
+                    context=record.context,
+                    mark_unfinished_on_error=False,
+                )
+            except ExecutionCancelled:
+                raise
+            except Exception as exc:
+                record.context.raise_if_cancelled()
+                failure = exc
+                if attempt < self.max_attempts:
+                    delay = random.uniform(
+                        self.retry_seconds * (1 - self.retry_jitter_ratio),
+                        self.retry_seconds * (1 + self.retry_jitter_ratio),
+                    )
+                    logger.warning(
+                        "algorithm execution %s failed on attempt %d/%d; "
+                        "releasing and restarting algorithm in %.3f seconds: %s: %s",
+                        record.execution_id,
+                        attempt,
+                        self.max_attempts,
+                        delay,
+                        type(exc).__name__,
+                        exc,
+                    )
+                else:
+                    delay = 0.0
+                    logger.error(
+                        "algorithm execution %s exhausted %d attempts; "
+                        "releasing and restarting algorithm for future requests: %s: %s",
+                        record.execution_id,
+                        self.max_attempts,
+                        type(exc).__name__,
+                        exc,
+                    )
+            finally:
+                self._release_runner()
+
+            self._restart_runner(runner, generation, delay)
+            if attempt >= self.max_attempts:
+                assert failure is not None
+                raise failure
+            record.context.raise_if_cancelled()
+
+        raise AssertionError("retry loop exited unexpectedly")
+
+    def _acquire_runner(self) -> tuple[AlgorithmRunner, int]:
+        with self._runner_condition:
+            while self._restart_pending:
+                self._runner_condition.wait()
+            if self._closed:
+                raise ExecutionCancelled("service shutdown")
+            self._runner_users += 1
+            return self.runner, self._runner_generation
+
+    def _release_runner(self) -> None:
+        with self._runner_condition:
+            self._runner_users -= 1
+            self._runner_condition.notify_all()
+
+    def _restart_runner(
+        self,
+        failed_runner: AlgorithmRunner,
+        failed_generation: int,
+        delay: float,
+    ) -> None:
+        with self._runner_condition:
+            while (
+                self._restart_pending
+                and failed_generation == self._runner_generation
+            ):
+                self._runner_condition.wait()
+            if (
+                failed_generation != self._runner_generation
+                or failed_runner is not self.runner
+            ):
+                return
+            if self._closed:
+                raise ExecutionCancelled("service shutdown")
+            self._restart_pending = True
+            while self._runner_users:
+                self._runner_condition.wait()
+
+        try:
+            try:
+                failed_runner.close()
+            except Exception:
+                logger.exception("failed to release algorithm before restart")
+            time.sleep(delay)
+            replacement = self._runner_factory()
+            replacement.start()
+        except Exception:
+            with self._runner_condition:
+                self._restart_pending = False
+                self._runner_condition.notify_all()
+            raise
+
+        with self._runner_condition:
+            self.runner = replacement
+            self._runner_generation += 1
+            self._restart_pending = False
+            self._runner_condition.notify_all()
+
+
+def create_app(
+    manager: ExecutionManager,
+    *,
+    webui: bool = False,
+    webui_url: str | None = None,
+    token: str | None = None,
+    lifecycle: Any | None = None,
+) -> Any:
+    """Create the optional generic HTTP adapter."""
+    try:
+        from fastapi import FastAPI, Query, Request
+        from fastapi.responses import JSONResponse
+    except ImportError as exc:
+        raise RuntimeError(
+            'HTTP service dependencies are missing; install "ls-algorithm-plugin-sdk[service]"'
+        ) from exc
+
+    # FastAPI resolves postponed annotations through module globals. Request is
+    # imported lazily so the base SDK keeps FastAPI optional; expose it here
+    # before route decorators inspect the endpoint signatures.
+    globals()["Request"] = Request
+
+    @asynccontextmanager
+    async def lifespan(_: Any):
+        manager.start()
+        if lifecycle is not None:
+            lifecycle.start()
+        announce_webui(application)
+        try:
+            yield
+        finally:
+            if lifecycle is not None:
+                lifecycle.stop()
+            manager.close()
+
+    application = FastAPI(
+        title="Algorithm Plugin Service",
+        version="0.1.0",
+        lifespan=lifespan,
+    )
+
+    @application.middleware("http")
+    async def authenticate(request: Request, call_next: Any) -> Any:
+        public_paths = {"/v1/health", "/v1/metadata"}
+        if (
+            token
+            and request.url.path.startswith("/v1/")
+            and request.url.path not in public_paths
+            and not hmac.compare_digest(
+                request.headers.get("X-LDP-Plugin-Token", ""), token
+            )
+        ):
+            return JSONResponse(status_code=401, content={"error": "unauthorized"})
+        return await call_next(request)
+
+    @application.exception_handler(ExecutionNotFound)
+    async def handle_not_found(_: Request, exc: ExecutionNotFound) -> JSONResponse:
+        return JSONResponse(
+            status_code=404,
+            content={"error": f"execution not found: {exc.args[0]}"},
+        )
+
+    @application.exception_handler(ExecutionNotFinished)
+    async def handle_not_finished(_: Request, exc: ExecutionNotFinished) -> JSONResponse:
+        return JSONResponse(
+            status_code=409,
+            content={"error": f"execution is not finished: {exc.args[0]}"},
+        )
+
+    @application.exception_handler(IdempotencyConflict)
+    async def handle_idempotency_conflict(
+        _: Request, exc: IdempotencyConflict
+    ) -> JSONResponse:
+        return JSONResponse(
+            status_code=409,
+            content={"error": f"idempotency key reused with another request: {exc.args[0]}"},
+        )
+
+    @application.get("/v1/metadata")
+    def metadata() -> dict[str, Any]:
+        return manager.runner.algorithm.metadata().to_dict()
+
+    @application.get("/v1/health")
+    def health() -> dict[str, str]:
+        return {"status": "ready"}
+
+    @application.get("/v1/manifest")
+    def manifest() -> dict[str, Any]:
+        return manager.manifest()
+
+    @application.get("/v1/heartbeat")
+    def heartbeat() -> dict[str, Any]:
+        return manager.heartbeat()
+
+    @application.post("/v1/executions", status_code=202)
+    async def create_execution(request: Request) -> Any:
+        try:
+            body = await request.json()
+            execution_request = AlgorithmRequest.from_dict(body)
+            if execution_request.workspace is None:
+                raise ValueError("workspace is required")
+        except Exception as exc:
+            return JSONResponse(
+                status_code=422,
+                content={"error": f"{type(exc).__name__}: {exc}"},
+            )
+        return manager.submit(
+            execution_request,
+            idempotency_key=request.headers.get("Idempotency-Key"),
+        )
+
+    @application.get("/v1/executions/{execution_id}")
+    def execution_status(execution_id: str) -> dict[str, Any]:
+        return manager.status(execution_id)
+
+    @application.get("/v1/executions/{execution_id}/events")
+    def execution_events(
+        execution_id: str,
+        after_sequence: int = Query(0, alias="afterSequence", ge=0),
+        wait_seconds: float = Query(0, alias="waitSeconds", ge=0, le=30),
+    ) -> dict[str, Any]:
+        return manager.events(
+            execution_id,
+            after_sequence=after_sequence,
+            wait_seconds=wait_seconds,
+        )
+
+    @application.get("/v1/executions/{execution_id}/result")
+    def execution_result(execution_id: str) -> dict[str, Any]:
+        return manager.result(execution_id)
+
+    @application.post("/v1/executions/{execution_id}:cancel", status_code=202)
+    async def cancel_execution(execution_id: str, request: Request) -> dict[str, Any]:
+        try:
+            body = await request.json()
+        except Exception:
+            body = {}
+        return manager.cancel(execution_id, str(body.get("reason", "cancelled by user")))
+
+    if webui:
+        mount_webui(application, webui_url=webui_url)
+
+    application.state.execution_manager = manager
+    return application
+
+
+def run_server(
+    manager: ExecutionManager,
+    *,
+    host: str = "0.0.0.0",
+    port: int = 8000,
+    webui: bool = False,
+    token: str | None = None,
+    lifecycle: Any | None = None,
+) -> None:
+    try:
+        import uvicorn
+    except ImportError as exc:
+        raise RuntimeError(
+            'HTTP service dependencies are missing; install "ls-algorithm-plugin-sdk[service]"'
+        ) from exc
+    display_host = "127.0.0.1" if host in {"0.0.0.0", "::"} else host
+    webui_url = f"http://{display_host}:{port}/ui/" if webui else None
+    uvicorn.run(
+        create_app(
+            manager,
+            webui=webui,
+            webui_url=webui_url,
+            token=token,
+            lifecycle=lifecycle,
+        ),
+        host=host,
+        port=port,
+    )

algorithm_plugin_sdk/webui/__init__.py

@@ -0,0 +1 @@
+"""Static assets for the optional internal processing UI."""