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

algorithm_plugin_sdk/registration.py

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+import hashlib
+import ipaddress
+import json
+import logging
+import os
+import socket
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from urllib.error import HTTPError
+from urllib.parse import urlparse
+from urllib.request import Request, urlopen
+
+logger = logging.getLogger(__name__)
+ManifestProvider = Callable[[], dict[str, Any]]
+HeartbeatProvider = Callable[[], dict[str, Any]]
+
+
+def _env_disabled() -> bool:
+    return (os.getenv("LDP_PLUGIN_REGISTRATION_DISABLED") or "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
+
+
+def _absolute_directory(name: str, value: str, *, writable: bool) -> Path:
+    path = Path(value).expanduser()
+    if not path.is_absolute():
+        raise ValueError(f"{name} must be an absolute path")
+    resolved = path.resolve(strict=True)
+    if not resolved.is_dir():
+        raise ValueError(f"{name} must be a directory: {resolved}")
+    if writable and not os.access(resolved, os.W_OK):
+        raise ValueError(f"{name} must be writable: {resolved}")
+    return resolved
+
+
+def _validate_public_url(value: str) -> str:
+    endpoint = value.strip().rstrip("/")
+    parsed = urlparse(endpoint)
+    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
+        raise ValueError("LDP_PLUGIN_PUBLIC_URL must be an http or https URL")
+    hostname = parsed.hostname.lower()
+    if hostname in {"0.0.0.0", "localhost"}:
+        raise ValueError("LDP_PLUGIN_PUBLIC_URL must be reachable by the scheduler")
+    try:
+        address = ipaddress.ip_address(hostname)
+    except ValueError:
+        address = None
+    if address is not None and (address.is_unspecified or address.is_loopback):
+        raise ValueError("LDP_PLUGIN_PUBLIC_URL must be reachable by the scheduler")
+    return endpoint
+
+
+def _machine_identity() -> str:
+    try:
+        value = Path("/etc/machine-id").read_text(encoding="utf-8").strip()
+    except OSError:
+        value = ""
+    return value or socket.gethostname()
+
+
+def stable_instance_key(algorithm_type: str) -> str:
+    identity = f"{_machine_identity()}\0{algorithm_type}".encode()
+    return f"{algorithm_type}-{hashlib.sha256(identity).hexdigest()[:24]}"
+
+
+def _private_address() -> str:
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
+            sock.connect(("192.0.2.1", 9))
+            return str(sock.getsockname()[0])
+    except OSError:
+        return ""
+
+
+@dataclass(frozen=True, slots=True)
+class RegistrationSettings:
+    compute_url: str
+    api_key: str
+    public_url: str
+    input_root: Path
+    workspace_root: Path
+    cluster: str
+    namespace: str
+    node_name: str
+    instance_key: str | None = None
+    heartbeat_seconds: float = 5.0
+
+    @classmethod
+    def from_env(cls) -> "RegistrationSettings":
+        required = {
+            "LDP_COMPUTE_URL": os.getenv("LDP_COMPUTE_URL"),
+            "LDP_INTERNAL_API_KEY": os.getenv("LDP_INTERNAL_API_KEY"),
+            "LDP_PLUGIN_PUBLIC_URL": os.getenv("LDP_PLUGIN_PUBLIC_URL"),
+            "LDP_JOB_INPUT_ROOT": os.getenv("LDP_JOB_INPUT_ROOT"),
+            "LDP_JOB_WORKSPACE_ROOT": os.getenv("LDP_JOB_WORKSPACE_ROOT"),
+            "LDP_CLUSTER": os.getenv("LDP_CLUSTER"),
+            "POD_NAMESPACE": os.getenv("POD_NAMESPACE"),
+            "NODE_NAME": os.getenv("NODE_NAME"),
+        }
+        missing = [name for name, value in required.items() if not str(value or "").strip()]
+        if missing:
+            raise ValueError(
+                "missing plugin registration settings: " + ", ".join(missing)
+            )
+        compute_url = str(required["LDP_COMPUTE_URL"]).strip().rstrip("/")
+        parsed_compute = urlparse(compute_url)
+        if parsed_compute.scheme not in {"http", "https"} or not parsed_compute.hostname:
+            raise ValueError("LDP_COMPUTE_URL must be an http or https URL")
+        return cls(
+            compute_url=compute_url,
+            api_key=str(required["LDP_INTERNAL_API_KEY"]).strip(),
+            public_url=_validate_public_url(str(required["LDP_PLUGIN_PUBLIC_URL"])),
+            input_root=_absolute_directory(
+                "LDP_JOB_INPUT_ROOT", str(required["LDP_JOB_INPUT_ROOT"]), writable=False
+            ),
+            workspace_root=_absolute_directory(
+                "LDP_JOB_WORKSPACE_ROOT",
+                str(required["LDP_JOB_WORKSPACE_ROOT"]),
+                writable=True,
+            ),
+            cluster=str(required["LDP_CLUSTER"]).strip(),
+            namespace=str(required["POD_NAMESPACE"]).strip(),
+            node_name=str(required["NODE_NAME"]).strip(),
+            instance_key=(os.getenv("LDP_PLUGIN_INSTANCE_KEY") or "").strip() or None,
+        )
+
+
+class JsonClient:
+    def __init__(self, base_url: str, api_key: str, *, timeout: float = 10.0):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.timeout = timeout
+
+    def post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        request = Request(
+            self.base_url + path,
+            data=json.dumps(payload, separators=(",", ":")).encode(),
+            method="POST",
+            headers={
+                "Content-Type": "application/json",
+                "X-API-Key": self.api_key,
+            },
+        )
+        with urlopen(request, timeout=self.timeout) as response:
+            body = response.read()
+        if not body:
+            return {}
+        decoded = json.loads(body)
+        if not isinstance(decoded, dict):
+            raise RuntimeError("compute API returned a non-object response")
+        return decoded
+
+
+class PluginRegistrationAgent:
+    """Register one bare-metal plugin instance and maintain its heartbeat."""
+
+    def __init__(
+        self,
+        settings: RegistrationSettings,
+        manifest_provider: ManifestProvider,
+        heartbeat_provider: HeartbeatProvider,
+        *,
+        client: JsonClient | None = None,
+    ):
+        self.settings = settings
+        self.manifest_provider = manifest_provider
+        self.heartbeat_provider = heartbeat_provider
+        self.client = client or JsonClient(settings.compute_url, settings.api_key)
+        self.instance_id = ""
+        self.stop_event = threading.Event()
+        self.thread: threading.Thread | None = None
+
+    @classmethod
+    def from_env(
+        cls,
+        manifest_provider: ManifestProvider,
+        heartbeat_provider: HeartbeatProvider,
+    ) -> "PluginRegistrationAgent | None":
+        if _env_disabled():
+            return None
+        names = (
+            "LDP_COMPUTE_URL",
+            "LDP_INTERNAL_API_KEY",
+            "LDP_PLUGIN_PUBLIC_URL",
+            "LDP_JOB_INPUT_ROOT",
+            "LDP_JOB_WORKSPACE_ROOT",
+        )
+        if not any(str(os.getenv(name) or "").strip() for name in names):
+            return None
+        return cls(RegistrationSettings.from_env(), manifest_provider, heartbeat_provider)
+
+    def start(self) -> None:
+        if self.thread is not None:
+            return
+        self.stop_event.clear()
+        self.thread = threading.Thread(
+            target=self._run,
+            name="plugin-registration-agent",
+            daemon=True,
+        )
+        self.thread.start()
+
+    def stop(self) -> None:
+        self.stop_event.set()
+        if self.thread is not None:
+            self.thread.join(timeout=self.settings.heartbeat_seconds + 2)
+        if self.instance_id:
+            try:
+                self._heartbeat("offline")
+            except Exception:
+                logger.warning("failed to mark plugin instance offline", exc_info=True)
+
+    def _run(self) -> None:
+        while not self.stop_event.is_set():
+            try:
+                if not self.instance_id:
+                    self._register()
+                self._heartbeat()
+            except HTTPError as exc:
+                if exc.code == 404:
+                    self.instance_id = ""
+                logger.warning("plugin registration request failed: %s", exc)
+            except Exception:
+                logger.exception("plugin registration iteration failed")
+            self.stop_event.wait(self.settings.heartbeat_seconds)
+
+    def _register(self) -> None:
+        manifest = self.manifest_provider()
+        algorithm_type = str(manifest.get("algorithmType") or "")
+        implementation_digest = str(manifest.get("implementationDigest") or "")
+        if not algorithm_type or not implementation_digest:
+            raise RuntimeError(
+                "plugin manifest requires algorithmType and implementationDigest"
+            )
+        registration_manifest = dict(manifest)
+        registration_manifest["imageDigest"] = implementation_digest
+        registration_manifest.pop("implementationDigest", None)
+        registration_manifest.pop("metadata", None)
+        response = self.client.post(
+            "/api/internal/compute-schedule/plugin-instances/register",
+            {
+                "instanceKey": self.settings.instance_key
+                or stable_instance_key(algorithm_type),
+                "endpoint": self.settings.public_url,
+                "inputRoot": str(self.settings.input_root),
+                "workspaceRoot": str(self.settings.workspace_root),
+                "cluster": self.settings.cluster,
+                "namespace": self.settings.namespace,
+                "nodeName": self.settings.node_name,
+                "manifest": registration_manifest,
+            },
+        )
+        self.instance_id = str(response.get("instanceId") or "")
+        if not self.instance_id:
+            raise RuntimeError("compute API registration response has no instanceId")
+        logger.info(
+            "plugin instance registered instance_id=%s implementation_id=%s",
+            self.instance_id,
+            response.get("implementationId", ""),
+        )
+
+    def _heartbeat(self, status: str | None = None) -> None:
+        if not self.instance_id:
+            return
+        heartbeat = self.heartbeat_provider()
+        self.client.post(
+            f"/api/internal/compute-schedule/plugin-instances/{self.instance_id}/heartbeat",
+            {
+                "status": status or heartbeat["status"],
+                "activeExecutions": int(heartbeat.get("activeExecutions", 0)),
+                "maxConcurrency": int(heartbeat["maxConcurrency"]),
+                "resources": dict(heartbeat.get("resources", {})),
+            },
+        )

algorithm_plugin_sdk/release.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import importlib
+import json
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+_SHA256 = re.compile(r"^sha256:[0-9a-f]{64}$")
+
+
+@dataclass(frozen=True, slots=True)
+class ReleaseManifest:
+    """Immutable identity and protocol metadata for a plugin release."""
+
+    algorithm_type: str
+    name: str
+    version: str
+    implementation_digest: str
+    build_revision: str
+    protocol_version: str
+    input_schema: str
+    output_schema: str
+    model_components: tuple[dict[str, Any], ...] = field(default_factory=tuple)
+
+    @classmethod
+    def discover(
+        cls,
+        algorithm: Any,
+        *,
+        start: str | Path | None = None,
+        required: bool = True,
+    ) -> "ReleaseManifest | None":
+        """Locate release-manifest.json from the checkout or Algorithm module."""
+        roots: list[Path] = []
+        if start is not None:
+            roots.append(Path(start).expanduser().resolve())
+        try:
+            module = importlib.import_module(type(algorithm).__module__)
+            module_file = getattr(module, "__file__", None)
+            if module_file:
+                roots.append(Path(module_file).resolve().parent)
+        except (ImportError, OSError):
+            pass
+        roots.append(Path.cwd().resolve())
+
+        seen: set[Path] = set()
+        for root in roots:
+            if root.is_file():
+                root = root.parent
+            for candidate in (root, *root.parents):
+                if candidate in seen:
+                    continue
+                seen.add(candidate)
+                path = candidate / "release-manifest.json"
+                if path.is_file():
+                    return cls.load(
+                        path,
+                        expected_algorithm_type=algorithm.metadata().name.replace("-", "_"),
+                    )
+        if required:
+            raise RuntimeError("release-manifest.json was not found from the current repository")
+        return None
+
+    @classmethod
+    def load(
+        cls,
+        path: str | Path,
+        *,
+        expected_algorithm_type: str | None = None,
+    ) -> "ReleaseManifest":
+        manifest_path = Path(path)
+        try:
+            payload = json.loads(manifest_path.read_text(encoding="utf-8"))
+        except FileNotFoundError as exc:
+            raise RuntimeError(f"release manifest is missing: {manifest_path}") from exc
+        except json.JSONDecodeError as exc:
+            raise RuntimeError(
+                f"release manifest is not valid JSON: {manifest_path}: {exc}"
+            ) from exc
+        if not isinstance(payload, dict):
+            raise RuntimeError("release manifest must contain a JSON object")
+
+        required = (
+            "algorithmType",
+            "name",
+            "version",
+            "implementationDigest",
+            "buildRevision",
+            "protocolVersion",
+            "inputSchema",
+            "outputSchema",
+        )
+        missing = [name for name in required if not str(payload.get(name) or "").strip()]
+        if missing:
+            raise RuntimeError(
+                "release manifest is missing required fields: " + ", ".join(missing)
+            )
+
+        algorithm_type = str(payload["algorithmType"]).strip()
+        if expected_algorithm_type and algorithm_type != expected_algorithm_type:
+            raise RuntimeError(
+                "release manifest algorithmType must be "
+                f"{expected_algorithm_type}, got {algorithm_type}"
+            )
+        digest = str(payload["implementationDigest"]).strip().lower()
+        if not _SHA256.fullmatch(digest):
+            raise RuntimeError(
+                "release manifest implementationDigest must be sha256 plus 64 hex digits"
+            )
+        components = payload.get("modelComponents", [])
+        if not isinstance(components, list) or not all(
+            isinstance(component, dict) for component in components
+        ):
+            raise RuntimeError("release manifest modelComponents must be a list of objects")
+        for component in components:
+            component_digest = str(component.get("digest") or "").lower()
+            if not component.get("name") or not _SHA256.fullmatch(component_digest):
+                raise RuntimeError(
+                    "each release manifest model component requires name and sha256 digest"
+                )
+
+        return cls(
+            algorithm_type=algorithm_type,
+            name=str(payload["name"]).strip(),
+            version=str(payload["version"]).strip(),
+            implementation_digest=digest,
+            build_revision=str(payload["buildRevision"]).strip(),
+            protocol_version=str(payload["protocolVersion"]).strip(),
+            input_schema=str(payload["inputSchema"]).strip(),
+            output_schema=str(payload["outputSchema"]).strip(),
+            model_components=tuple(dict(component) for component in components),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "algorithmType": self.algorithm_type,
+            "name": self.name,
+            "version": self.version,
+            "implementationDigest": self.implementation_digest,
+            "buildRevision": self.build_revision,
+            "protocolVersion": self.protocol_version,
+            "inputSchema": self.input_schema,
+            "outputSchema": self.output_schema,
+            "modelComponents": [dict(component) for component in self.model_components],
+        }
+
+    def plugin_manifest(
+        self,
+        *,
+        max_concurrency: int,
+        capabilities: dict[str, Any],
+        resources: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        if max_concurrency < 1:
+            raise ValueError("max_concurrency must be positive")
+        return {
+            "apiVersion": self.protocol_version,
+            "algorithmType": self.algorithm_type,
+            "name": self.name,
+            "version": self.version,
+            "implementationDigest": self.implementation_digest,
+            "buildRevision": self.build_revision,
+            "inputSchema": self.input_schema,
+            "outputSchema": self.output_schema,
+            "maxConcurrency": max_concurrency,
+            "capabilities": dict(capabilities),
+            "resources": dict(resources or {}),
+            "modelComponents": [dict(component) for component in self.model_components],
+        }

algorithm_plugin_sdk/runner.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import threading
+import uuid
+
+from .algorithm import Algorithm
+from .context import ExecutionContext, ProgressCallback
+from .errors import ExecutionCancelled, InvalidAlgorithmResult
+from .models import AlgorithmRequest, AlgorithmResult
+
+
+class AlgorithmRunner:
+    """Runs one algorithm implementation without owning its scheduling policy."""
+
+    def __init__(self, algorithm: Algorithm) -> None:
+        if not isinstance(algorithm, Algorithm):
+            raise TypeError("algorithm must be a Algorithm instance")
+        self.algorithm = algorithm
+        self._started = False
+        self._closed = False
+        self._lifecycle_lock = threading.Lock()
+
+    def start(self) -> None:
+        with self._lifecycle_lock:
+            if self._closed:
+                raise RuntimeError("runner is closed")
+            if self._started:
+                return
+            self.algorithm.startup()
+            self._started = True
+
+    def close(self) -> None:
+        with self._lifecycle_lock:
+            if self._closed:
+                return
+            if self._started:
+                self.algorithm.shutdown()
+            self._closed = True
+
+    def run(
+        self,
+        request: AlgorithmRequest,
+        *,
+        context: ExecutionContext | None = None,
+        execution_id: str | None = None,
+        scratch_dir: str | None = None,
+        progress_callback: ProgressCallback | None = None,
+        mark_unfinished_on_error: bool = True,
+    ) -> AlgorithmResult:
+        if not isinstance(request, AlgorithmRequest):
+            raise TypeError("request must be an AlgorithmRequest")
+        self.start()
+        datasets = [item.dataset for item in request.inputs]
+        selected_context = context or ExecutionContext(
+            execution_id or f"exec-{uuid.uuid4().hex}",
+            datasets,
+            scratch_dir=(
+                scratch_dir
+                or (
+                    request.workspace["scratchRoot"]
+                    if request.workspace is not None
+                    else None
+                )
+            ),
+            progress_callback=progress_callback,
+        )
+        selected_context.raise_if_cancelled()
+
+        try:
+            result = self.algorithm.execute(request, selected_context)
+            selected_context.raise_if_cancelled()
+            self._validate_result(request, result)
+            for dataset_result in result.datasets:
+                selected_context.mark_dataset_result(dataset_result)
+        except ExecutionCancelled as exc:
+            if mark_unfinished_on_error:
+                selected_context.mark_unfinished("cancelled", str(exc))
+            raise
+        except Exception as exc:
+            if mark_unfinished_on_error:
+                selected_context.mark_unfinished(
+                    "failed", f"{type(exc).__name__}: {exc}"
+                )
+            raise
+
+        return result
+
+    @staticmethod
+    def _validate_result(
+        request: AlgorithmRequest,
+        result: AlgorithmResult,
+    ) -> None:
+        if not isinstance(result, AlgorithmResult):
+            raise InvalidAlgorithmResult("execute() must return AlgorithmResult")
+        result.validate()
+        actual = [dataset_result.dataset for dataset_result in result.datasets]
+        if len(actual) != len(set(actual)):
+            raise InvalidAlgorithmResult(
+                "algorithm returned duplicate dataset results"
+            )
+        expected = [item.dataset for item in request.inputs]
+        if set(actual) != set(expected):
+            missing = [dataset for dataset in expected if dataset not in actual]
+            unknown = [dataset for dataset in actual if dataset not in expected]
+            raise InvalidAlgorithmResult(
+                f"algorithm result does not match request; missing={missing}, unknown={unknown}"
+            )
+
+    def __enter__(self) -> "AlgorithmRunner":
+        self.start()
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        self.close()