diff --git a/.github/workflows/remote-real-model-e2e.yml b/.github/workflows/remote-real-model-e2e.yml new file mode 100644 index 0000000..85e75a8 --- /dev/null +++ b/.github/workflows/remote-real-model-e2e.yml @@ -0,0 +1,108 @@ +name: Remote real model E2E + +on: + pull_request: + branches: [develop] + paths: + - "anylabeling/services/auto_labeling/remote_client.py" + - "anylabeling/services/auto_labeling/remote_model.py" + - "anylabeling/services/auto_labeling/model_manager.py" + - "scripts/validate_remote_inference.py" + - ".github/workflows/remote-real-model-e2e.yml" + push: + branches: [develop] + paths: + - "anylabeling/services/auto_labeling/remote_client.py" + - "anylabeling/services/auto_labeling/remote_model.py" + - "anylabeling/services/auto_labeling/model_manager.py" + - "scripts/validate_remote_inference.py" + - ".github/workflows/remote-real-model-e2e.yml" + schedule: + - cron: "41 3 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: remote-real-model-e2e-${{ github.ref }} + cancel-in-progress: true + +jobs: + yolox-s: + name: Authenticated YOLOX-S / ${{ matrix.os }} + runs-on: ${{ matrix.os }} + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + env: + ASSET_DIR: _anylearning/tests/fixtures/inference/real_models + MODEL_SHA256: c5c2d13e59ae883e6af3b45daea64af4833a4951c92d116ec270d9ddbe998063 + IMAGE_SHA256: 5a9522051c3cec2bbd2f6323fccba32e8fbf3ddcc2b3e2fd46b04c720bc6f866 + OPENBLAS_NUM_THREADS: "1" + OMP_NUM_THREADS: "1" + QT_QPA_PLATFORM: offscreen + steps: + - uses: actions/checkout@v7 + + - name: Check out the pinned AnyLearning server contract + uses: actions/checkout@v7 + with: + repository: nrl-ai/anylearning-oss + ref: 7cabc1e8caaec070410fbc47e8ee250ae50454ce + path: _anylearning + + - uses: actions/setup-python@v7 + with: + python-version: "3.13" + cache: pip + + - name: Install Qt system libraries (Linux) + if: runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libegl1 libxkbcommon-x11-0 libdbus-1-3 libxcb-cursor0 \ + libxcb-icccm4 libxcb-image0 libxcb-keysyms1 libxcb-randr0 \ + libxcb-render-util0 libxcb-shape0 libxcb-xinerama0 libxcb-xkb1 + + - name: Install PyQt6 (macOS) + if: runner.os == 'macOS' + run: python -m pip install "PyQt6>=6.7.0" + + - name: Install client and isolated server dependencies + run: | + python -m pip install --upgrade pip + python -m pip install . + python -m pip install "fastapi>=0.141,<1" "pydantic>=2.13,<3" "argon2-cffi>=25.1,<26" "psutil>=7.2,<8" "uvicorn>=0.41,<1" + + - name: Download immutable Apache-2.0 test assets + shell: bash + run: | + set -euo pipefail + python _anylearning/scripts/download_verified_file.py \ + "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.onnx" \ + "$ASSET_DIR/yolox_s.onnx" --sha256 "$MODEL_SHA256" --max-bytes 104857600 + python _anylearning/scripts/download_verified_file.py \ + "https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/6ddff4824372906469a7fae2dc3206c7aa4bbaee/assets/dog.jpg" \ + "$ASSET_DIR/dog.jpg" --sha256 "$IMAGE_SHA256" --max-bytes 10485760 + + - name: Run AnyLabeling through the real authenticated server + run: >- + python scripts/validate_remote_inference.py + --anylearning-root _anylearning + --manifest _anylearning/tests/fixtures/inference/real_models/yolox_s_official.json + --model _anylearning/tests/fixtures/inference/real_models/yolox_s.onnx + --image _anylearning/tests/fixtures/inference/real_models/dog.jpg + --output-root validation-results/${{ runner.os }} + + - name: Retain visual, result, and timing evidence + if: always() + uses: actions/upload-artifact@v6 + with: + name: anylabeling-remote-yolox-s-${{ runner.os }} + path: validation-results/ + if-no-files-found: error + retention-days: 30 diff --git a/README.md b/README.md index f57a0e2..51b1565 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,7 @@ - **SAM 3** (ViT-H) — open-vocabulary segmentation with text prompts - [x] Text detection, recognition and KIE (Key Information Extraction) labeling. - [x] Hardware acceleration with CUDA, CoreML, DirectML, OpenVINO, and vendor NPU providers. +- [x] Authenticated shared ONNX inference through an AnyLearning server. - [x] Multiple languages available: English, Vietnamese, Chinese. ### Supported Models @@ -57,6 +58,8 @@ Required model weights are downloaded automatically on first use. +For centrally hosted models, see the [authenticated remote inference guide](docs/remote_inference.md). + ## Latest Release [AnyLabeling v0.4.43](https://github.com/vietanhdev/anylabeling/releases/tag/v0.4.43) is the current stable release. It fixes recovery after model download/load failures, skips invalid images during SAM preload, persists grouped shapes with undo support, and closes label files reliably after saving and loading. diff --git a/anylabeling/services/auto_labeling/__init__.py b/anylabeling/services/auto_labeling/__init__.py index b334713..d296bbe 100644 --- a/anylabeling/services/auto_labeling/__init__.py +++ b/anylabeling/services/auto_labeling/__init__.py @@ -1,4 +1,5 @@ # Import models to ensure they register themselves via @ModelRegistry.register +from . import remote_model as remote_model # noqa: F401 from . import segment_anything as segment_anything # noqa: F401 from . import yolov5 as yolov5 # noqa: F401 from . import yolov8 as yolov8 # noqa: F401 diff --git a/anylabeling/services/auto_labeling/model.py b/anylabeling/services/auto_labeling/model.py index 0e11c28..6731063 100644 --- a/anylabeling/services/auto_labeling/model.py +++ b/anylabeling/services/auto_labeling/model.py @@ -1,7 +1,6 @@ import logging import os import socket -import ssl from abc import abstractmethod import yaml @@ -18,11 +17,6 @@ socket.setdefaulttimeout(240) # Prevent timeout when downloading models -ssl._create_default_https_context = ( - ssl._create_unverified_context -) # Prevent issue when downloading models behind a proxy - - class Model(QObject): BASE_DOWNLOAD_URL = "https://github.com/vietanhdev/anylabeling-assets/raw/main/" diff --git a/anylabeling/services/auto_labeling/model_manager.py b/anylabeling/services/auto_labeling/model_manager.py index 769fad7..f5824d3 100644 --- a/anylabeling/services/auto_labeling/model_manager.py +++ b/anylabeling/services/auto_labeling/model_manager.py @@ -4,7 +4,6 @@ import os import pathlib import shutil -import ssl import tempfile import time import urllib.request @@ -22,10 +21,6 @@ from .registry import ModelRegistry -ssl._create_default_https_context = ( - ssl._create_unverified_context -) # Prevent issue when downloading models behind a proxy - class ModelManager(QObject): """Model manager""" @@ -162,7 +157,11 @@ def on_model_download_finished(self): self.model_loaded.emit(self.loaded_model_config) self.output_modes_changed.emit( self.loaded_model_config["model"].Meta.output_modes, - self.loaded_model_config["model"].Meta.default_output_mode, + getattr( + self.loaded_model_config["model"], + "output_mode", + self.loaded_model_config["model"].Meta.default_output_mode, + ), ) else: self.model_loaded.emit({}) @@ -208,7 +207,8 @@ def load_custom_model(self, config_file): "type" not in model_config or "display_name" not in model_config or "name" not in model_config - or model_config["type"] not in ["segment_anything", "yolov5", "yolov8"] + or model_config["type"] + not in ["remote", "segment_anything", "yolov5", "yolov8"] ): self._report_model_load_error( self.tr("Error in loading custom model: Invalid config file format.") @@ -439,7 +439,9 @@ def _load_model(self, model_id): # Specific logic for interactive models (like SAM) vs detection models # Ideally this should be a property of the model class (capabilities) - if model_type == "segment_anything": + if model_type == "segment_anything" or getattr( + model_config["model"], "supports_interactive_prompts", False + ): self.auto_segmentation_model_selected.emit() # Request next files for prediction self.request_next_files_requested.emit() @@ -460,9 +462,13 @@ def set_auto_labeling_marks(self, marks): """Set auto labeling marks (For example, for segment_anything model, it is the marks for) """ - if ( - self.loaded_model_config is None - or self.loaded_model_config["type"] != "segment_anything" + if self.loaded_model_config is None or not ( + self.loaded_model_config["type"] == "segment_anything" + or getattr( + self.loaded_model_config["model"], + "supports_interactive_prompts", + False, + ) ): return self.loaded_model_config["model"].set_auto_labeling_marks(marks) @@ -534,8 +540,11 @@ def predict_shapes_threading(self, image, filename=None): self.model_execution_thread is not None and self.model_execution_thread.isRunning() ): - if hasattr(self.loaded_model_config["model"], "unload"): - self.loaded_model_config["model"].unload() + model = self.loaded_model_config["model"] + if hasattr(model, "cancel_prediction"): + model.cancel_prediction() + elif hasattr(model, "unload"): + model.unload() # Wait for the thread to finish self.model_execution_thread.quit() @@ -573,7 +582,14 @@ def on_next_files_changed(self, next_files): return # Currently only segment_anything model supports this feature - if self.loaded_model_config["type"] != "segment_anything": + if not ( + self.loaded_model_config["type"] == "segment_anything" + or getattr( + self.loaded_model_config["model"], + "supports_interactive_prompts", + False, + ) + ): return self.loaded_model_config["model"].on_next_files_changed(next_files) diff --git a/anylabeling/services/auto_labeling/remote_client.py b/anylabeling/services/auto_labeling/remote_client.py new file mode 100644 index 0000000..172c545 --- /dev/null +++ b/anylabeling/services/auto_labeling/remote_client.py @@ -0,0 +1,577 @@ +"""Bounded client for the public AnyLearning inference protocol.""" + +from __future__ import annotations + +import base64 +import hashlib +import ipaddress +import json +import math +import secrets +import ssl +import threading +import time +from dataclasses import dataclass +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.parse import quote, urlsplit, urlunsplit +from urllib.request import ( + HTTPHandler, + HTTPRedirectHandler, + HTTPSHandler, + ProxyHandler, + Request, + build_opener, +) + +PROTOCOL_VERSION = "1.0" +_REQUEST_HEADER = "X-AnyLearning-Request" +_MAX_METADATA_BYTES = 8 * 1024 +_MAX_JSON_RESPONSE_BYTES = 9 * 1024**2 +_MAX_IMAGE_BYTES = 32 * 1024**2 +_MAX_SHAPES = 10_000 +_MAX_POINTS = 100_000 +_STATES = frozenset( + {"queued", "running", "succeeded", "failed", "cancelled", "timed_out"} +) + + +class RemoteInferenceError(RuntimeError): + """A public, credential-free remote inference failure.""" + + +class _NoRedirects(HTTPRedirectHandler): + def redirect_request(self, request, file_pointer, code, message, headers, url): + del request, file_pointer, code, message, headers, url + return None + + +@dataclass(frozen=True) +class RemoteModelCapabilities: + model_id: str + model_revision: str + tasks: tuple[str, ...] + metadata: dict[str, Any] + + @property + def promptable(self) -> bool: + return "promptable_segmentation" in self.tasks + + +class RemoteInferenceClient: + """Authenticate, discover one model, and run token-owned prediction jobs.""" + + def __init__( + self, + server_url: str, + model_id: str, + password: str, + *, + prediction_timeout_seconds: float = 120, + poll_interval_seconds: float = 0.1, + max_image_bytes: int = _MAX_IMAGE_BYTES, + max_response_bytes: int = _MAX_JSON_RESPONSE_BYTES, + ) -> None: + self._server_url = _validate_server_url(server_url) + if not isinstance(model_id, str) or not 1 <= len(model_id) <= 512: + raise ValueError("remote model_id must contain 1 to 512 characters") + if not isinstance(password, str) or not 12 <= len(password.encode()) <= 1_024: + raise ValueError("remote password must contain 12 to 1024 UTF-8 bytes") + if not 1 <= prediction_timeout_seconds <= 3_600: + raise ValueError("prediction timeout must be between 1 and 3600 seconds") + if not 0.02 <= poll_interval_seconds <= 2: + raise ValueError("poll interval must be between 0.02 and 2 seconds") + if not 1_024 <= max_image_bytes <= 512 * 1024**2: + raise ValueError("remote image byte limit is invalid") + if not 1_024 <= max_response_bytes <= 256 * 1024**2: + raise ValueError("remote response byte limit is invalid") + + context = ssl.create_default_context() + self._opener = build_opener( + ProxyHandler({}), + HTTPHandler(), + HTTPSHandler(context=context), + _NoRedirects(), + ) + self._model_id = model_id + self._password: str | None = password + self._prediction_timeout = float(prediction_timeout_seconds) + self._poll_interval = float(poll_interval_seconds) + self._max_image_bytes = max_image_bytes + self._max_response_bytes = max_response_bytes + self._token: str | None = None + self._token_expires_at = 0.0 + self._cancelled = threading.Event() + self._prediction_lock = threading.Lock() + self._authenticate() + self.capabilities = self._discover_model() + + def close(self) -> None: + self._cancelled.set() + self._token = None + self._password = None + self._token_expires_at = 0.0 + + def cancel(self) -> None: + self._cancelled.set() + + def predict( + self, + encoded_image: bytes, + media_type: str, + *, + prompts: list[dict[str, Any]] | None = None, + output_shape: str | None = None, + parameters: dict[str, Any] | None = None, + ) -> dict[str, Any]: + if not isinstance(encoded_image, bytes) or not encoded_image: + raise ValueError("remote image must be non-empty bytes") + if len(encoded_image) > self._max_image_bytes: + raise RemoteInferenceError("Encoded image exceeds the client limit") + if media_type not in {"image/jpeg", "image/png", "image/webp"}: + raise ValueError("remote image media type is unsupported") + if not self._prediction_lock.acquire(blocking=False): + raise RemoteInferenceError("A remote prediction is already running") + + job_id: str | None = None + try: + self._cancelled.clear() + self._ensure_token_lifetime() + request_id = secrets.token_hex(16) + source_id = f"content-sha256:{hashlib.sha256(encoded_image).hexdigest()}" + request_payload: dict[str, Any] = { + "protocol_version": PROTOCOL_VERSION, + "request_id": request_id, + "source_id": source_id, + "model_id": self.capabilities.model_id, + "model_revision": self.capabilities.model_revision, + "prompts": prompts or [], + "parameters": parameters or {}, + } + if output_shape is not None: + request_payload["output_shape"] = output_shape + metadata = json.dumps( + request_payload, + ensure_ascii=True, + allow_nan=False, + separators=(",", ":"), + ).encode("ascii") + if len(metadata) > _MAX_METADATA_BYTES: + raise RemoteInferenceError("Inference request metadata is too large") + encoded_metadata = base64.urlsafe_b64encode(metadata).rstrip(b"=") + submitted = self._json_request( + "POST", + "/v1/predictions", + body=encoded_image, + headers={ + "Authorization": f"Bearer {self._required_token()}", + "Content-Type": media_type, + _REQUEST_HEADER: encoded_metadata.decode("ascii"), + }, + expected_status=202, + ) + job_id = _bounded_text(submitted.get("job_id"), "job_id", 512) + state = self._validate_job(submitted, job_id, request_id) + if state == "succeeded": + result = submitted.get("result") + elif state in {"queued", "running"}: + result = self._poll(job_id, request_id) + else: + raise RemoteInferenceError( + _bounded_text(submitted.get("error"), "prediction error", 2_048) + if isinstance(submitted.get("error"), str) + else "Remote prediction failed" + ) + return _validate_result( + result, + request_id=request_id, + source_id=source_id, + capabilities=self.capabilities, + ) + finally: + if job_id is not None: + self._delete_job(job_id) + self._prediction_lock.release() + + def _poll(self, job_id: str, request_id: str) -> Any: + deadline = time.monotonic() + self._prediction_timeout + while True: + if self._cancelled.is_set(): + raise RemoteInferenceError("Remote prediction was cancelled") + if time.monotonic() >= deadline: + raise RemoteInferenceError("Remote prediction exceeded its deadline") + snapshot = self._json_request( + "GET", + f"/v1/predictions/{quote(job_id, safe='')}", + headers={"Authorization": f"Bearer {self._required_token()}"}, + expected_status=200, + ) + state = self._validate_job(snapshot, job_id, request_id) + if state == "succeeded": + return snapshot.get("result") + if state not in {"queued", "running"}: + raise RemoteInferenceError( + _bounded_text(snapshot.get("error"), "prediction error", 2_048) + if isinstance(snapshot.get("error"), str) + else "Remote prediction failed" + ) + self._cancelled.wait(self._poll_interval) + + def _validate_job(self, value: Any, job_id: str, request_id: str) -> str: + if not isinstance(value, dict): + raise RemoteInferenceError("Server returned an invalid prediction job") + if value.get("job_id") != job_id or value.get("request_id") != request_id: + raise RemoteInferenceError("Server prediction identity did not match") + state = value.get("state") + if state not in _STATES: + raise RemoteInferenceError("Server returned an invalid prediction state") + return state + + def _authenticate(self) -> None: + if self._password is None: + raise RemoteInferenceError("Remote inference client is closed") + payload = json.dumps( + {"password": self._password}, ensure_ascii=True, separators=(",", ":") + ).encode("utf-8") + response = self._json_request( + "POST", + "/v1/auth/token", + body=payload, + headers={"Content-Type": "application/json"}, + expected_status=200, + maximum=16 * 1024, + ) + token = _bounded_text(response.get("access_token"), "access token", 4_096) + if len(token.encode("ascii", errors="ignore")) != len(token) or not token: + raise RemoteInferenceError("Server returned an invalid access token") + if response.get("token_type") != "bearer": + raise RemoteInferenceError("Server returned an invalid token type") + expires_in = response.get("expires_in") + if type(expires_in) is not int or not 30 <= expires_in <= 3_600: + raise RemoteInferenceError("Server returned an invalid token lifetime") + self._token = token + self._token_expires_at = time.monotonic() + expires_in + + def _ensure_token_lifetime(self) -> None: + required = self._prediction_timeout + 5 + if self._token is None or self._token_expires_at - time.monotonic() < required: + self._authenticate() + if self._token_expires_at - time.monotonic() < required: + raise RemoteInferenceError( + "Server token lifetime is shorter than the prediction timeout" + ) + + def _discover_model(self) -> RemoteModelCapabilities: + response = self._json_request( + "GET", + f"/v1/models/{quote(self._model_id, safe='')}", + headers={"Authorization": f"Bearer {self._required_token()}"}, + expected_status=200, + maximum=256 * 1024, + ) + model_id = _bounded_text(response.get("model_id"), "model_id", 512) + revision = _bounded_text(response.get("model_revision"), "model_revision", 512) + if ( + model_id != self._model_id + or response.get("protocol_version") != PROTOCOL_VERSION + ): + raise RemoteInferenceError("Server model identity is incompatible") + raw_tasks = response.get("tasks") + if ( + not isinstance(raw_tasks, list) + or not 1 <= len(raw_tasks) <= 16 + or any( + not isinstance(item, str) or not 1 <= len(item) <= 128 + for item in raw_tasks + ) + or len(raw_tasks) != len(set(raw_tasks)) + ): + raise RemoteInferenceError("Server returned invalid model tasks") + metadata = response.get("metadata", {}) + if not isinstance(metadata, dict) or len(metadata) > 128: + raise RemoteInferenceError("Server returned invalid model metadata") + _validate_metadata(metadata, "model metadata") + return RemoteModelCapabilities(model_id, revision, tuple(raw_tasks), metadata) + + def _delete_job(self, job_id: str) -> None: + if self._token is None: + return + try: + self._raw_request( + "DELETE", + f"/v1/predictions/{quote(job_id, safe='')}", + headers={"Authorization": f"Bearer {self._token}"}, + expected_status=204, + maximum=1_024, + ) + except RemoteInferenceError: + pass + + def _required_token(self) -> str: + if self._token is None: + raise RemoteInferenceError("Remote inference client is not authenticated") + return self._token + + def _json_request(self, method: str, path: str, **kwargs: Any) -> dict[str, Any]: + payload = self._raw_request(method, path, **kwargs) + try: + value = json.loads( + payload, + parse_constant=_reject_json_constant, + object_pairs_hook=_unique_json_object, + ) + except (UnicodeError, json.JSONDecodeError, ValueError) as error: + raise RemoteInferenceError("Server returned invalid JSON") from error + if not isinstance(value, dict): + raise RemoteInferenceError("Server returned an invalid JSON object") + return value + + def _raw_request( + self, + method: str, + path: str, + *, + body: bytes | None = None, + headers: dict[str, str] | None = None, + expected_status: int, + maximum: int | None = None, + ) -> bytes: + if not path.startswith("/") or "?" in path or "#" in path: + raise ValueError("remote API path is invalid") + request = Request( + self._server_url + path, + data=body, + headers={"Accept": "application/json", **(headers or {})}, + method=method, + ) + try: + response = self._opener.open( + request, timeout=min(30.0, max(1.0, self._prediction_timeout)) + ) + except HTTPError as error: + status = error.code + error.close() + if status == 401: + raise RemoteInferenceError("Remote authentication failed") from error + if status == 404: + raise RemoteInferenceError( + "Remote model or prediction was not found" + ) from error + if status == 429: + raise RemoteInferenceError( + "Remote inference capacity was reached" + ) from error + raise RemoteInferenceError( + f"Remote server rejected the request (HTTP {status})" + ) from error + except (OSError, TimeoutError, URLError) as error: + raise RemoteInferenceError( + "Could not reach the remote inference server" + ) from error + + with response: + if response.status != expected_status: + raise RemoteInferenceError( + f"Remote server returned HTTP {response.status}" + ) + limit = self._max_response_bytes if maximum is None else maximum + lengths = response.headers.get_all("Content-Length", failobj=[]) + if len(lengths) > 1: + raise RemoteInferenceError("Server returned invalid response framing") + if lengths: + declared = lengths[0] + if not declared.isascii() or not declared.isdecimal(): + raise RemoteInferenceError( + "Server returned invalid response framing" + ) + if int(declared) > limit: + raise RemoteInferenceError( + "Server response exceeds the client limit" + ) + payload = response.read(limit + 1) + if len(payload) > limit: + raise RemoteInferenceError("Server response exceeds the client limit") + if ( + expected_status != 204 + and response.headers.get_content_type() != "application/json" + ): + raise RemoteInferenceError("Server response is not JSON") + return payload + + +def _validate_server_url(value: str) -> str: + if not isinstance(value, str) or not 1 <= len(value) <= 2_048: + raise ValueError("remote server URL is invalid") + parsed = urlsplit(value) + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.path not in {"", "/"} + or parsed.query + or parsed.fragment + ): + raise ValueError("remote server URL must be an exact HTTP(S) origin") + try: + port = parsed.port + except ValueError as error: + raise ValueError("remote server URL port is invalid") from error + if parsed.scheme == "http" and not _is_loopback(parsed.hostname): + raise ValueError("non-loopback remote inference requires HTTPS") + host = f"[{parsed.hostname}]" if ":" in parsed.hostname else parsed.hostname + authority = f"{host}:{port}" if port is not None else host + return urlunsplit((parsed.scheme, authority, "", "", "")) + + +def _is_loopback(host: str) -> bool: + if host.lower() == "localhost": + return True + try: + return ipaddress.ip_address(host).is_loopback + except ValueError: + return False + + +def _bounded_text(value: Any, name: str, maximum: int) -> str: + if not isinstance(value, str) or not 1 <= len(value) <= maximum: + raise RemoteInferenceError(f"Server returned an invalid {name}") + return value + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"Non-finite JSON constant is not accepted: {value}") + + +def _unique_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + value: dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise ValueError("Duplicate JSON object keys are not accepted") + value[key] = item + return value + + +def _validate_result( + value: Any, + *, + request_id: str, + source_id: str, + capabilities: RemoteModelCapabilities, +) -> dict[str, Any]: + if not isinstance(value, dict) or value.get("protocol_version") != PROTOCOL_VERSION: + raise RemoteInferenceError("Server returned an invalid inference result") + expected = { + "request_id": request_id, + "source_id": source_id, + "model_id": capabilities.model_id, + "model_revision": capabilities.model_revision, + } + if any( + value.get(field) != expected_value for field, expected_value in expected.items() + ): + raise RemoteInferenceError("Server inference result identity did not match") + shapes = value.get("shapes") + if not isinstance(shapes, list) or len(shapes) > _MAX_SHAPES: + raise RemoteInferenceError("Server returned an invalid shape list") + for shape in shapes: + _validate_shape(shape) + warnings = value.get("warnings", []) + if ( + not isinstance(warnings, list) + or len(warnings) > 128 + or any(not isinstance(item, str) or len(item) > 2_048 for item in warnings) + ): + raise RemoteInferenceError("Server returned invalid inference warnings") + timings = value.get("timings_ms", {}) + if not isinstance(timings, dict) or len(timings) > 64: + raise RemoteInferenceError("Server returned invalid inference timings") + for key, timing in timings.items(): + if ( + not isinstance(key, str) + or not 1 <= len(key) <= 128 + or not isinstance(timing, (int, float)) + or isinstance(timing, bool) + or not math.isfinite(timing) + or timing < 0 + ): + raise RemoteInferenceError("Server returned invalid inference timings") + return value + + +def _validate_shape(value: Any) -> None: + if not isinstance(value, dict): + raise RemoteInferenceError("Server returned an invalid shape") + if not {"type", "points"} <= set(value) or not set(value) <= { + "type", + "points", + "label", + "score", + "group_id", + "attributes", + }: + raise RemoteInferenceError("Server returned an invalid shape") + expected_points = { + "point": (1, 1), + "rectangle": (2, 2), + "polygon": (3, _MAX_POINTS), + "rotated_rectangle": (4, 4), + }.get(value.get("type")) + points = value.get("points") + if expected_points is None or not isinstance(points, list): + raise RemoteInferenceError("Server returned an unsupported shape") + if not expected_points[0] <= len(points) <= expected_points[1]: + raise RemoteInferenceError("Server returned invalid shape points") + for point in points: + if not isinstance(point, dict) or set(point) != {"x", "y"}: + raise RemoteInferenceError("Server returned invalid point coordinates") + if any( + not isinstance(point[axis], (int, float)) + or isinstance(point[axis], bool) + or not math.isfinite(point[axis]) + or abs(point[axis]) > 1_000_000_000 + for axis in ("x", "y") + ): + raise RemoteInferenceError("Server returned invalid point coordinates") + label = value.get("label") + if label is not None and (not isinstance(label, str) or len(label) > 1_024): + raise RemoteInferenceError("Server returned an invalid shape label") + score = value.get("score") + if score is not None and ( + not isinstance(score, (int, float)) + or isinstance(score, bool) + or not math.isfinite(score) + or not 0 <= score <= 1 + ): + raise RemoteInferenceError("Server returned an invalid shape score") + group_id = value.get("group_id") + if group_id is not None and not ( + (isinstance(group_id, str) and len(group_id) <= 2_048) + or ( + isinstance(group_id, int) + and not isinstance(group_id, bool) + and -(2**63) <= group_id <= 2**63 - 1 + ) + ): + raise RemoteInferenceError("Server returned an invalid shape group_id") + attributes = value.get("attributes", {}) + if not isinstance(attributes, dict) or len(attributes) > 128: + raise RemoteInferenceError("Server returned invalid shape attributes") + _validate_metadata(attributes, "shape attributes") + + +def _validate_metadata(value: dict[Any, Any], name: str) -> None: + for key, item in value.items(): + if not isinstance(key, str) or not 1 <= len(key) <= 128: + raise RemoteInferenceError(f"Server returned invalid {name}") + if item is None or isinstance(item, (str, bool)): + if isinstance(item, str) and len(item) > 2_048: + raise RemoteInferenceError(f"Server returned invalid {name}") + continue + if isinstance(item, int): + if not -(2**63) <= item <= 2**63 - 1: + raise RemoteInferenceError(f"Server returned invalid {name}") + continue + if isinstance(item, float) and math.isfinite(item): + continue + raise RemoteInferenceError(f"Server returned invalid {name}") diff --git a/anylabeling/services/auto_labeling/remote_model.py b/anylabeling/services/auto_labeling/remote_model.py new file mode 100644 index 0000000..5f64f8f --- /dev/null +++ b/anylabeling/services/auto_labeling/remote_model.py @@ -0,0 +1,251 @@ +"""AnyLabeling adapter for authenticated AnyLearning inference servers.""" + +from __future__ import annotations + +import math +import os +import re +from typing import Any + +from PyQt6 import QtCore +from PyQt6.QtCore import QBuffer, QCoreApplication, QIODevice + +from anylabeling.views.labeling.shape import Shape + +from .model import Model +from .registry import ModelRegistry +from .remote_client import RemoteInferenceClient, RemoteInferenceError +from .types import AutoLabelingResult + +_ENVIRONMENT_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,127}$") + + +@ModelRegistry.register("remote") +class RemoteModel(Model): + """Run auto-labeling through a preconfigured remote ONNX model.""" + + class Meta: + required_config_names = [ + "type", + "name", + "display_name", + "server_url", + "model_id", + "password_env", + ] + widgets = ["button_run"] + output_modes = { + "polygon": QCoreApplication.translate("Model", "Polygon"), + "rectangle": QCoreApplication.translate("Model", "Rectangle"), + } + default_output_mode = "rectangle" + + _PROMPT_WIDGETS = [ + "output_label", + "output_select_combobox", + "button_run", + "button_add_point", + "button_remove_point", + "button_add_rect", + "button_clear", + "button_finish_object", + ] + + def __init__(self, model_config, on_message) -> None: + super().__init__(model_config, on_message) + environment_name = self.config["password_env"] + if not isinstance(environment_name, str) or not _ENVIRONMENT_NAME.fullmatch( + environment_name + ): + raise ValueError("Remote password_env is not a valid environment name") + password = os.environ.get(environment_name) + if password is None: + raise ValueError( + "Remote inference password environment variable is not set: " + f"{environment_name}" + ) + parameters = _wire_parameters(self.config.get("parameters", {})) + self.client: RemoteInferenceClient | None = RemoteInferenceClient( + self.config["server_url"], + self.config["model_id"], + password, + prediction_timeout_seconds=self.config.get( + "prediction_timeout_seconds", 120 + ), + poll_interval_seconds=self.config.get("poll_interval_seconds", 0.1), + max_image_bytes=self.config.get("max_image_bytes", 32 * 1024**2), + ) + self.supports_interactive_prompts = self.client.capabilities.promptable + self.marks: list[dict[str, Any]] = [] + self.parameters = parameters + self.output_mode = ( + "polygon" if self.supports_interactive_prompts else "rectangle" + ) + + def get_required_widgets(self): + return list( + self._PROMPT_WIDGETS + if self.supports_interactive_prompts + else self.Meta.widgets + ) + + def set_auto_labeling_marks(self, marks): + if not isinstance(marks, list) or len(marks) > 10_000: + raise ValueError("Remote prompt list is invalid") + self.marks = marks + + def predict_shapes(self, image, filename=None) -> AutoLabelingResult: + del filename + if image is None or image.isNull(): + return AutoLabelingResult([], replace=True) + client = self.client + if client is None: + raise RemoteInferenceError("Remote inference model is unloaded") + result = client.predict( + _encode_png(image), + "image/png", + prompts=_wire_prompts(self.marks) + if self.supports_interactive_prompts + else [], + output_shape=self.output_mode + if self.supports_interactive_prompts + else None, + parameters=self.parameters, + ) + return AutoLabelingResult( + [_labeling_shape(value) for value in result["shapes"]], replace=True + ) + + def cancel_prediction(self) -> None: + if self.client is not None: + self.client.cancel() + + def unload(self) -> None: + if self.client is not None: + self.client.close() + self.client = None + + +def _encode_png(image) -> bytes: + buffer = QBuffer() + if not buffer.open(QIODevice.OpenModeFlag.WriteOnly): + raise RemoteInferenceError("Could not allocate the remote image buffer") + try: + if not image.save(buffer, "PNG"): + raise RemoteInferenceError("Could not encode the image for inference") + encoded = bytes(buffer.data()) + finally: + buffer.close() + if not encoded: + raise RemoteInferenceError("Could not encode the image for inference") + return encoded + + +def _wire_prompts(marks: list[dict[str, Any]]) -> list[dict[str, Any]]: + prompts = [] + for index, mark in enumerate(marks): + if not isinstance(mark, dict): + raise ValueError(f"Remote prompt {index} must be an object") + kind, data = mark.get("type"), mark.get("data") + if kind == "point": + if not _finite_coordinates(data, 2): + raise ValueError(f"Remote point prompt {index} is invalid") + label = mark.get("label") + if label not in (0, 1): + raise ValueError(f"Remote point prompt {index} label must be 0 or 1") + prompts.append( + { + "type": "point", + "point": {"x": float(data[0]), "y": float(data[1])}, + "foreground": label == 1, + } + ) + elif kind == "rectangle": + if not _finite_coordinates(data, 4): + raise ValueError(f"Remote rectangle prompt {index} is invalid") + x1, y1, x2, y2 = (float(item) for item in data) + if x2 <= x1 or y2 <= y1: + raise ValueError( + f"Remote rectangle prompt {index} must have positive area" + ) + prompts.append( + { + "type": "box", + "top_left": {"x": x1, "y": y1}, + "bottom_right": {"x": x2, "y": y2}, + } + ) + else: + raise ValueError(f"Remote prompt {index} has an unsupported type") + return prompts + + +def _finite_coordinates(value: Any, length: int) -> bool: + return ( + isinstance(value, (list, tuple)) + and len(value) == length + and all( + isinstance(item, (int, float)) + and not isinstance(item, bool) + and math.isfinite(item) + for item in value + ) + ) + + +def _wire_parameters(value: Any) -> dict[str, Any]: + if not isinstance(value, dict) or len(value) > 128: + raise ValueError("Remote inference parameters must be a bounded mapping") + parameters = {} + for key, item in value.items(): + if not isinstance(key, str) or not 1 <= len(key) <= 128: + raise ValueError("Remote inference parameter names are invalid") + if item is None or isinstance(item, (str, bool)): + if isinstance(item, str) and len(item) > 2_048: + raise ValueError(f"Remote inference parameter {key!r} is too long") + parameters[key] = item + elif isinstance(item, int): + if not -(2**63) <= item <= 2**63 - 1: + raise ValueError(f"Remote inference parameter {key!r} is invalid") + parameters[key] = item + elif isinstance(item, float) and math.isfinite(item): + parameters[key] = item + elif ( + isinstance(item, (list, tuple)) + and len(item) <= 256 + and ( + all(isinstance(member, str) and len(member) <= 2_048 for member in item) + or all( + isinstance(member, int) + and not isinstance(member, bool) + and -(2**63) <= member <= 2**63 - 1 + for member in item + ) + ) + ): + parameters[key] = list(item) + else: + raise ValueError(f"Remote inference parameter {key!r} is invalid") + return parameters + + +def _labeling_shape(value: dict[str, Any]) -> Shape: + shape_type = "polygon" if value["type"] == "rotated_rectangle" else value["type"] + shape = Shape( + label=value.get("label") or "AUTOLABEL_OBJECT", + shape_type=shape_type, + flags={}, + group_id=value.get("group_id"), + ) + for point in value["points"]: + shape.add_point(QtCore.QPointF(point["x"], point["y"])) + if shape_type == "polygon": + shape.close() + if value.get("score") is not None: + shape.other_data["score"] = value["score"] + if value.get("attributes"): + shape.other_data["attributes"] = value["attributes"] + return shape + + +__all__ = ["RemoteModel"] diff --git a/docs/remote_inference.md b/docs/remote_inference.md new file mode 100644 index 0000000..e1eccbd --- /dev/null +++ b/docs/remote_inference.md @@ -0,0 +1,72 @@ +# Authenticated remote inference + +AnyLabeling can use a centrally hosted AnyLearning ONNX model without copying +the model to every labeling computer. The desktop client sends the encoded +image and bounded inference metadata; the server returns editable shapes. + +## Configure the server + +Follow the [AnyLearning server guide](https://github.com/nrl-ai/anylearning-oss/blob/develop/docs/server.md) +to configure the inference service and its startup model manifest. The server +chooses every model path and backend setting; clients cannot upload models. + +Use direct TLS or a trusted TLS-terminating reverse proxy for every network +deployment. Plain HTTP is accepted only for localhost and numeric loopback +addresses. + +## Configure AnyLabeling + +Keep the plaintext password out of YAML. Put it in the environment before +launching AnyLabeling: + +```shell +export ANYLABELING_REMOTE_PASSWORD='use-a-long-random-password' +anylabeling +``` + +On PowerShell: + +```powershell +$env:ANYLABELING_REMOTE_PASSWORD = 'use-a-long-random-password' +anylabeling +``` + +Create a custom model YAML file and select it in the auto-labeling model picker: + +```yaml +type: remote +name: shared-yolox +display_name: Shared YOLOX detector +server_url: https://inference.example.com +model_id: shared-detector +password_env: ANYLABELING_REMOTE_PASSWORD +prediction_timeout_seconds: 120 +poll_interval_seconds: 0.1 +parameters: + confidence: 0.5 + iou: 0.45 +``` + +`model_id` must match the server's immutable startup manifest. `parameters` are +optional bounded values; the selected AnyLearning backend decides which names +it supports. + +The client exchanges the password for a short-lived token and keeps both in +memory only. It ignores proxy environment variables, rejects redirects, +verifies HTTPS certificates, hashes the exact encoded image into the request +identity, bounds responses, and deletes completed or cancelled jobs. Never put +passwords or tokens in YAML, URLs, command lines, screenshots, or issue reports. + +Interactive models expose point and rectangle controls when the server +advertises `promptable_segmentation`. Detection models expose the Run button. + +## Troubleshooting + +- `non-loopback remote inference requires HTTPS`: use HTTPS, or test locally + through `http://127.0.0.1:`. +- `Remote authentication failed`: ensure the configured environment variable + exists in the process that launched AnyLabeling and matches the server hash. +- `Server token lifetime is shorter than the prediction timeout`: raise the + server token TTL or lower `prediction_timeout_seconds`. +- `Remote inference capacity was reached`: the bounded server queue is full; + wait for current jobs or deliberately raise measured server limits. diff --git a/scripts/validate_remote_inference.py b/scripts/validate_remote_inference.py new file mode 100644 index 0000000..ffb3b5e --- /dev/null +++ b/scripts/validate_remote_inference.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +"""Validate AnyLabeling against an authenticated AnyLearning TCP server.""" + +from __future__ import annotations + +import argparse +import hashlib +import html +import json +import os +import secrets +import socket +import sys +import threading +import time +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import cv2 +import psutil +import uvicorn +from PyQt6.QtGui import QImage + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from anylabeling.services.auto_labeling.remote_model import RemoteModel + +_PASSWORD = "real-remote-validation-password" +_PASSWORD_ENV = "ANYLABELING_REAL_REMOTE_PASSWORD" + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _write_json(path: Path, value: object) -> None: + path.write_text( + json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + +def _start_server(app: Any) -> tuple[uvicorn.Server, threading.Thread, int]: + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("127.0.0.1", 0)) + listener.listen(128) + port = listener.getsockname()[1] + server = uvicorn.Server(uvicorn.Config(app, log_level="warning", access_log=False)) + thread = threading.Thread( + target=server.run, + kwargs={"sockets": [listener]}, + name="anylearning-real-remote-server", + daemon=True, + ) + thread.start() + deadline = time.monotonic() + 30 + while not server.started: + if not thread.is_alive(): + raise RuntimeError("AnyLearning validation server stopped during startup") + if time.monotonic() >= deadline: + server.should_exit = True + thread.join(timeout=10) + raise TimeoutError("AnyLearning validation server did not start in time") + time.sleep(0.02) + return server, thread, port + + +def _stop_server(server: uvicorn.Server, thread: threading.Thread) -> None: + server.should_exit = True + thread.join(timeout=20) + if thread.is_alive(): + server.force_exit = True + thread.join(timeout=10) + if thread.is_alive(): + raise RuntimeError("AnyLearning validation server did not stop cleanly") + + +def _adapt_result(model, result, *, request_id: str): + from anylearning.inference import InferenceResult, InferenceShape, Point + + capabilities = model.client.capabilities + return InferenceResult( + request_id=request_id, + source_id="anylabeling-qimage-png", + model_id=capabilities.model_id, + model_revision=capabilities.model_revision, + shapes=tuple( + InferenceShape( + type=shape.shape_type, + points=tuple(Point(x=point.x(), y=point.y()) for point in shape.points), + label=shape.label, + score=shape.other_data.get("score"), + group_id=shape.group_id, + attributes=shape.other_data.get("attributes", {}), + ) + for shape in result.shapes + ), + ) + + +def run_validation( + anylearning_root: Path, + manifest_path: Path, + model_path: Path, + image_path: Path, + output_root: Path, +) -> Path: + sys.path.insert(0, str(anylearning_root.resolve(strict=True))) + from anylearning.inference.validation import ( + _annotate, + _check_expectations, + _load_rgb, + _result_digest, + load_validation_manifest, + ) + from anylearning.server import ( + ServerModelDefinition, + ServerSettings, + create_server_app, + hash_password, + ) + + manifest_path = manifest_path.resolve(strict=True) + model_path = model_path.resolve(strict=True) + image_path = image_path.resolve(strict=True) + manifest = load_validation_manifest(manifest_path) + if manifest.backend != "yolo_onnx" or len(manifest.images) != 1: + raise ValueError("remote validation requires one YOLO ONNX image case") + config = dict(manifest.config) + config.update(config_file=str(manifest_path), model_path=str(model_path)) + definition = ServerModelDefinition(backend=manifest.backend, config=config) + settings = ServerSettings( + password_hash=hash_password(_PASSWORD), + token_secret=secrets.token_bytes(32), + token_ttl_seconds=300, + prediction_timeout_seconds=120, + prediction_result_ttl_seconds=300, + ) + server, thread, port = _start_server( + create_server_app(settings, model_definitions=(definition,)) + ) + remote_model = None + previous_password = os.environ.get(_PASSWORD_ENV) + stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + output_dir = ( + output_root.resolve() / f"{stamp}-anylabeling-remote-{secrets.token_hex(4)}" + ) + output_dir.mkdir(parents=True, mode=0o700) + process = psutil.Process() + peak_rss = process.memory_info().rss + started = time.perf_counter() + try: + os.environ[_PASSWORD_ENV] = _PASSWORD + image_case = manifest.images[0] + remote_model = RemoteModel( + { + "type": "remote", + "name": "real-remote-yolox", + "display_name": "Real remote YOLOX-S", + "server_url": f"http://127.0.0.1:{port}", + "model_id": config["name"], + "password_env": _PASSWORD_ENV, + "prediction_timeout_seconds": 120, + "poll_interval_seconds": 0.02, + "parameters": dict(image_case.request_parameters), + }, + lambda _message: None, + ) + image = QImage(str(image_path)) + if image.isNull(): + raise ValueError("validation image could not be decoded by Qt") + results, round_trip_ms = [], [] + for run in range(manifest.runs): + run_started = time.perf_counter() + labeling_result = remote_model.predict_shapes(image, str(image_path)) + round_trip_ms.append((time.perf_counter() - run_started) * 1000) + results.append( + _adapt_result(remote_model, labeling_result, request_id=f"run-{run}") + ) + peak_rss = max(peak_rss, process.memory_info().rss) + digests = [ + _result_digest(result.model_copy(update={"request_id": "canonical"})) + for result in results + ] + failures = _check_expectations(results[0], image_case.expected) + if len(set(digests)) != 1: + failures.append("AnyLabeling remote results changed across identical runs") + annotated_name = "000-dog-anylabeling-remote.png" + if not cv2.imwrite( + str(output_dir / annotated_name), + _annotate(_load_rgb(image_path), results[0]), + ): + raise OSError("could not write annotated remote validation image") + summary = { + "schema_version": 1, + "passed": not failures, + "created_at": datetime.now(UTC).isoformat(), + "transport": "authenticated TCP HTTP loopback", + "client": "AnyLabeling RemoteModel", + "server": "AnyLearning inference server", + "manifest": manifest_path.name, + "model_sha256": _sha256(model_path), + "image_sha256": _sha256(image_path), + "provenance": manifest.provenance.model_dump(mode="json"), + "runs": manifest.runs, + "shape_count": len(results[0].shapes), + "labels": [shape.label for shape in results[0].shapes], + "consistent_runs": len(set(digests)) == 1, + "consistency_digest": digests[0], + "round_trip_ms": round_trip_ms, + "total_elapsed_ms": (time.perf_counter() - started) * 1000, + "peak_observed_rss_bytes": peak_rss, + "annotated_image": annotated_name, + "failures": failures, + } + _write_json(output_dir / "summary.json", summary) + _write_json( + output_dir / "results.json", + [result.model_dump(mode="json") for result in results], + ) + (output_dir / "index.html").write_text( + "AnyLabeling remote validation" + f"

{'PASS' if not failures else 'FAIL'}: authenticated remote YOLOX-S

" + f"

{'
'.join(html.escape(item) for item in failures)}

" + f'', + encoding="utf-8", + ) + if failures: + raise AssertionError("; ".join(failures)) + return output_dir + except Exception as error: + _write_json( + output_dir / "failure.json", + { + "schema_version": 1, + "passed": False, + "created_at": datetime.now(UTC).isoformat(), + "error_type": type(error).__name__, + "error": str(error)[:2_048], + }, + ) + raise + finally: + if remote_model is not None: + remote_model.unload() + if previous_password is None: + os.environ.pop(_PASSWORD_ENV, None) + else: + os.environ[_PASSWORD_ENV] = previous_password + _stop_server(server, thread) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--anylearning-root", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--image", type=Path, required=True) + parser.add_argument("--output-root", type=Path, default=Path("validation-results")) + args = parser.parse_args() + print( + run_validation( + args.anylearning_root, + args.manifest, + args.model, + args.image, + args.output_root, + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_model_manager.py b/tests/test_model_manager.py index 73e6c7a..45cf061 100644 --- a/tests/test_model_manager.py +++ b/tests/test_model_manager.py @@ -1,4 +1,5 @@ import os +import subprocess import sys import tempfile import unittest @@ -11,6 +12,27 @@ class TestModelManager(unittest.TestCase): + def test_import_does_not_disable_process_tls_verification(self): + source = """ +import ssl +original = ssl._create_default_https_context +import anylabeling.services.auto_labeling.model_manager +assert ssl._create_default_https_context is original +assert ssl._create_default_https_context is not ssl._create_unverified_context +""" + environment = dict(os.environ) + environment["QT_QPA_PLATFORM"] = "offscreen" + completed = subprocess.run( + [sys.executable, "-c", source], + cwd=os.path.dirname(os.path.dirname(__file__)), + env=environment, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + self.assertEqual(completed.returncode, 0, completed.stderr) + @patch( "anylabeling.services.auto_labeling.model_manager.ModelManager.load_model_configs" ) @@ -35,6 +57,30 @@ def test_unload_model(self, mock_load): mock_model.unload.assert_called_once() self.assertIsNone(manager.loaded_model_config) + @patch( + "anylabeling.services.auto_labeling.model_manager.ModelManager.load_model_configs" + ) + @patch("anylabeling.services.auto_labeling.model_manager.QThread") + @patch("anylabeling.services.auto_labeling.model_manager.GenericWorker") + def test_new_prediction_cancels_without_unloading_remote_model( + self, worker_cls, thread_cls, mock_load + ): + del mock_load + manager = ModelManager() + model = MagicMock() + manager.loaded_model_config = {"model": model, "type": "remote"} + old_thread = MagicMock() + old_thread.isRunning.return_value = True + old_thread.wait.return_value = True + manager.model_execution_thread = old_thread + manager.predict_shapes_threading(MagicMock(), "image.png") + model.cancel_prediction.assert_called_once_with() + model.unload.assert_not_called() + old_thread.quit.assert_called_once_with() + old_thread.wait.assert_called_once_with(1000) + worker_cls.assert_called_once() + thread_cls.return_value.start.assert_called_once_with() + @patch( "anylabeling.services.auto_labeling.model_manager.ModelManager.load_model_configs" ) diff --git a/tests/test_registry.py b/tests/test_registry.py index b8679ae..41af2dc 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -72,6 +72,11 @@ def test_segment_anything_registered(self): import anylabeling.services.auto_labeling.segment_anything # noqa: F401 self.assertIn("segment_anything", ModelRegistry.list_models()) + def test_remote_model_registered(self): + import anylabeling.services.auto_labeling.remote_model # noqa: F401 + + self.assertIn("remote", ModelRegistry.list_models()) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_remote_client.py b/tests/test_remote_client.py new file mode 100644 index 0000000..02acdc9 --- /dev/null +++ b/tests/test_remote_client.py @@ -0,0 +1,263 @@ +import base64 +import hashlib +import json +import threading +import unittest +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +from anylabeling.services.auto_labeling.remote_client import ( + RemoteInferenceClient, + RemoteInferenceError, +) + + +class _State: + def __init__( + self, *, redirect_auth=False, invalid_attributes=False, hold_running=False + ): + self.redirect_auth = redirect_auth + self.invalid_attributes = invalid_attributes + self.hold_running = hold_running + self.deleted = threading.Event() + self.polled = threading.Event() + self.request = None + self.image = None + self.authorization = [] + + +class _Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, _format, *args): + del args + + @property + def state(self): + return self.server.protocol_state + + def _json(self, status, value): + body = json.dumps(value, separators=(",", ":")).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _authorized(self): + value = self.headers.get("Authorization") + self.state.authorization.append(value) + return value == "Bearer test-token" + + def do_POST(self): + body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if self.path == "/v1/auth/token": + if self.state.redirect_auth: + self.send_response(307) + self.send_header("Location", "/redirected-auth") + self.send_header("Content-Length", "0") + self.end_headers() + return + if json.loads(body) != {"password": "correct horse battery staple"}: + self._json(401, {"detail": "denied"}) + return + self._json( + 200, + { + "access_token": "test-token", + "token_type": "bearer", + "expires_in": 300, + }, + ) + return + if self.path == "/v1/predictions" and self._authorized(): + encoded = self.headers["X-AnyLearning-Request"] + encoded += "=" * (-len(encoded) % 4) + self.state.request = json.loads(base64.urlsafe_b64decode(encoded)) + self.state.image = body + self._json( + 202, + { + "job_id": "job-1", + "request_id": self.state.request["request_id"], + "state": "queued", + }, + ) + return + self._json(404, {"detail": "not found"}) + + def do_GET(self): + if self.path == "/v1/models/model-1" and self._authorized(): + self._json( + 200, + { + "protocol_version": "1.0", + "model_id": "model-1", + "model_revision": "sha256:revision", + "tasks": ["detection"], + "supports_batch": False, + "supports_cancellation": True, + "max_batch_size": 1, + "metadata": {"backend": "yolo_onnx"}, + }, + ) + return + if self.path == "/v1/predictions/job-1" and self._authorized(): + request = self.state.request + if self.state.hold_running: + self.state.polled.set() + self._json( + 200, + { + "job_id": "job-1", + "request_id": request["request_id"], + "state": "running", + }, + ) + return + attributes = ( + {"nested": ["rejected"]} + if self.state.invalid_attributes + else {"class_id": 16} + ) + self._json( + 200, + { + "job_id": "job-1", + "request_id": request["request_id"], + "state": "succeeded", + "result": { + "protocol_version": "1.0", + "request_id": request["request_id"], + "source_id": request["source_id"], + "model_id": "model-1", + "model_revision": "sha256:revision", + "shapes": [ + { + "type": "rectangle", + "points": [ + {"x": 10.5, "y": 20.25}, + {"x": 30.75, "y": 40.5}, + ], + "label": "dog", + "score": 0.9, + "group_id": 0, + "attributes": attributes, + } + ], + "warnings": [], + "timings_ms": {"inference": 1.25}, + }, + }, + ) + return + self._json(404, {"detail": "not found"}) + + def do_DELETE(self): + if self.path == "/v1/predictions/job-1" and self._authorized(): + self.state.deleted.set() + self.send_response(204) + self.send_header("Content-Length", "0") + self.end_headers() + return + self._json(404, {"detail": "not found"}) + + +@contextmanager +def _server(**options): + server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + server.protocol_state = _State(**options) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server, server.protocol_state + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class TestRemoteInferenceClient(unittest.TestCase): + def test_authenticated_round_trip_validates_identity_and_deletes_job(self): + image = b"\x89PNG\r\n\x1a\nreal-image-payload" + with _server() as (server, state): + client = RemoteInferenceClient( + f"http://127.0.0.1:{server.server_port}", + "model-1", + "correct horse battery staple", + prediction_timeout_seconds=2, + poll_interval_seconds=0.02, + ) + result = client.predict(image, "image/png") + self.assertEqual(result["shapes"][0]["label"], "dog") + self.assertEqual(state.image, image) + self.assertEqual( + state.request["source_id"], + f"content-sha256:{hashlib.sha256(image).hexdigest()}", + ) + self.assertTrue(state.deleted.is_set()) + self.assertTrue( + all(value == "Bearer test-token" for value in state.authorization) + ) + + def test_non_loopback_plain_http_is_rejected(self): + with self.assertRaisesRegex(ValueError, "requires HTTPS"): + RemoteInferenceClient( + "http://example.com", "model-1", "correct horse battery staple" + ) + + def test_redirected_authentication_is_not_followed(self): + with _server(redirect_auth=True) as (server, _state): + with self.assertRaisesRegex( + RemoteInferenceError, r"rejected the request \(HTTP 307\)" + ): + RemoteInferenceClient( + f"http://127.0.0.1:{server.server_port}", + "model-1", + "correct horse battery staple", + ) + + def test_nested_shape_attributes_are_rejected_and_job_is_deleted(self): + with _server(invalid_attributes=True) as (server, state): + client = RemoteInferenceClient( + f"http://127.0.0.1:{server.server_port}", + "model-1", + "correct horse battery staple", + prediction_timeout_seconds=2, + poll_interval_seconds=0.02, + ) + with self.assertRaisesRegex(RemoteInferenceError, "shape attributes"): + client.predict(b"valid bytes", "image/png") + self.assertTrue(state.deleted.is_set()) + + def test_cancellation_interrupts_polling_and_deletes_job(self): + with _server(hold_running=True) as (server, state): + client = RemoteInferenceClient( + f"http://127.0.0.1:{server.server_port}", + "model-1", + "correct horse battery staple", + prediction_timeout_seconds=5, + poll_interval_seconds=0.02, + ) + captured = [] + + def run_prediction(): + try: + client.predict(b"valid bytes", "image/png") + except Exception as error: # noqa: BLE001 + captured.append(error) + + thread = threading.Thread(target=run_prediction) + thread.start() + self.assertTrue(state.polled.wait(timeout=2)) + client.cancel() + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + self.assertEqual(len(captured), 1) + self.assertIsInstance(captured[0], RemoteInferenceError) + self.assertIn("cancelled", str(captured[0])) + self.assertTrue(state.deleted.is_set()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_remote_model.py b/tests/test_remote_model.py new file mode 100644 index 0000000..0d964e1 --- /dev/null +++ b/tests/test_remote_model.py @@ -0,0 +1,123 @@ +import os +import unittest +from unittest.mock import MagicMock, patch + +from PyQt6.QtGui import QColor, QImage + +from anylabeling.services.auto_labeling.remote_model import RemoteModel, _wire_prompts + + +class TestRemoteModel(unittest.TestCase): + @patch("anylabeling.services.auto_labeling.remote_model.RemoteInferenceClient") + def test_real_qimage_is_losslessly_encoded_and_shapes_are_editable( + self, client_cls + ): + client = client_cls.return_value + client.capabilities.promptable = False + client.predict.return_value = { + "shapes": [ + { + "type": "rectangle", + "points": [{"x": 1.5, "y": 2.5}, {"x": 8.5, "y": 7.5}], + "label": "cat", + "score": 0.75, + "group_id": 0, + "attributes": {"class_id": 15}, + } + ] + } + config = { + "type": "remote", + "name": "shared-model", + "display_name": "Shared model", + "server_url": "https://inference.example.com", + "model_id": "model-1", + "password_env": "ANYLABELING_TEST_REMOTE_PASSWORD", + "parameters": {"confidence": 0.5}, + } + image = QImage(10, 10, QImage.Format.Format_RGB32) + image.fill(QColor("red")) + with patch.dict( + os.environ, + {"ANYLABELING_TEST_REMOTE_PASSWORD": "correct horse battery staple"}, + ): + model = RemoteModel(config, MagicMock()) + result = model.predict_shapes(image) + encoded, media_type = client.predict.call_args.args[:2] + self.assertTrue(encoded.startswith(b"\x89PNG\r\n\x1a\n")) + self.assertEqual(media_type, "image/png") + self.assertEqual( + client.predict.call_args.kwargs["parameters"], {"confidence": 0.5} + ) + shape = result.shapes[0] + self.assertEqual( + (shape.shape_type, shape.label, shape.group_id), ("rectangle", "cat", 0) + ) + self.assertEqual( + shape.other_data, {"score": 0.75, "attributes": {"class_id": 15}} + ) + + def test_prompt_conversion_rejects_invalid_geometry(self): + self.assertEqual( + _wire_prompts( + [ + {"type": "point", "data": [1, 2], "label": 1}, + {"type": "rectangle", "data": [3, 4, 8, 9]}, + ] + ), + [ + {"type": "point", "point": {"x": 1.0, "y": 2.0}, "foreground": True}, + { + "type": "box", + "top_left": {"x": 3.0, "y": 4.0}, + "bottom_right": {"x": 8.0, "y": 9.0}, + }, + ], + ) + with self.assertRaisesRegex(ValueError, "positive area"): + _wire_prompts([{"type": "rectangle", "data": [3, 4, 3, 9]}]) + + @patch("anylabeling.services.auto_labeling.remote_model.RemoteInferenceClient") + def test_cancel_does_not_unload_client(self, client_cls): + client = client_cls.return_value + client.capabilities.promptable = False + config = { + "type": "remote", + "name": "shared-model", + "display_name": "Shared model", + "server_url": "https://inference.example.com", + "model_id": "model-1", + "password_env": "ANYLABELING_TEST_REMOTE_PASSWORD", + } + with patch.dict( + os.environ, + {"ANYLABELING_TEST_REMOTE_PASSWORD": "correct horse battery staple"}, + ): + model = RemoteModel(config, MagicMock()) + model.cancel_prediction() + client.cancel.assert_called_once_with() + self.assertIs(model.client, client) + + @patch("anylabeling.services.auto_labeling.remote_model.RemoteInferenceClient") + def test_promptable_remote_model_defaults_to_polygon(self, client_cls): + client_cls.return_value.capabilities.promptable = True + config = { + "type": "remote", + "name": "shared-promptable", + "display_name": "Shared promptable model", + "server_url": "https://inference.example.com", + "model_id": "model-1", + "password_env": "ANYLABELING_TEST_REMOTE_PASSWORD", + } + with patch.dict( + os.environ, + {"ANYLABELING_TEST_REMOTE_PASSWORD": "correct horse battery staple"}, + ): + model = RemoteModel(config, MagicMock()) + self.assertTrue(model.supports_interactive_prompts) + self.assertEqual(model.output_mode, "polygon") + self.assertIn("button_add_point", model.get_required_widgets()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_remote_workflow.py b/tests/test_remote_workflow.py new file mode 100644 index 0000000..7161dd5 --- /dev/null +++ b/tests/test_remote_workflow.py @@ -0,0 +1,31 @@ +import unittest +from pathlib import Path + + +class TestRemoteRealModelWorkflow(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.source = ( + Path(__file__).parents[1] / ".github/workflows/remote-real-model-e2e.yml" + ).read_text(encoding="utf-8") + + def test_runs_real_authenticated_model_on_all_desktop_operating_systems(self): + self.assertIn("os: [ubuntu-latest, windows-latest, macos-latest]", self.source) + self.assertIn("scripts/validate_remote_inference.py", self.source) + self.assertIn("if: always()", self.source) + self.assertIn("retention-days: 30", self.source) + + def test_pins_server_contract_and_model_checksums(self): + self.assertIn("ref: 7cabc1e8caaec070410fbc47e8ee250ae50454ce", self.source) + self.assertIn( + "c5c2d13e59ae883e6af3b45daea64af4833a4951c92d116ec270d9ddbe998063", + self.source, + ) + self.assertIn( + "5a9522051c3cec2bbd2f6323fccba32e8fbf3ddcc2b3e2fd46b04c720bc6f866", + self.source, + ) + + +if __name__ == "__main__": + unittest.main()