heimdex-worker-sdk 0.1.0__tar.gz

heimdex_worker_sdk-0.1.0/.claude/CLAUDE.md

@@ -0,0 +1,62 @@
+# Heimdex Worker SDK
+
+Shared worker utilities for the Heimdex ecosystem.
+
+## Quick Reference
+
+```bash
+pip install -e ".[dev]"
+pytest -v
+```
+
+## Design Philosophy
+
+This package provides shared infrastructure for Heimdex workers:
+- S3/MinIO dual-mode client
+- SQS consumer with heartbeat and visibility extension
+- Aircloud GPU worker lifecycle management (start/stop/scale)
+- Internal API client for livecommerce endpoints
+- Deterministic S3 key generation helpers
+- MIME type classification
+
+Dependencies are intentionally minimal: `pydantic-settings`, `boto3`, `requests`.
+
+## Package Structure
+
+```
+src/heimdex_worker_sdk/
+├── aircloud_client.py    # Aircloud External API (start/stop/scale/status)
+├── gpu_orchestrator.py   # Auto start/stop GPU workers based on SQS queue depth
+├── sqs_client.py         # Thin boto3 SQS wrapper
+├── sqs_consumer.py       # Background SQS polling loop with heartbeat
+├── s3.py                 # S3/MinIO dual-mode client
+├── settings.py           # WorkerSettings (pydantic-settings, 60+ env vars)
+├── internal_api.py       # HTTP client for livecommerce /internal/drive/* endpoints
+├── drive_keys.py         # Deterministic S3 key generation (Drive files)
+├── youtube_keys.py       # S3 key generation (YouTube files)
+├── youtube_api.py        # HTTP client for YouTube internal endpoints
+├── content_type.py       # MIME type classification
+└── message_adapters.py   # SQS message → dataclass converters
+```
+
+## Consumers
+
+Changes here affect ALL downstream repos. Always verify impact before merging.
+
+| Consumer | What it uses |
+|---------|-------------|
+| dev-heimdex-for-livecommerce (8 workers + API) | Full SDK |
+| dev-heimdex-playground | aircloud_client, gpu_orchestrator |
+
+## Release
+
+- Trigger: Git tag `v*`
+- Pipeline: Test -> Build -> Publish to PyPI -> GitHub Release
+- Current version: 0.1.0
+
+## Rules
+
+- Keep dependencies minimal (no torch, cv2, or heavy ML libs)
+- All Aircloud calls must be fire-and-forget (never block ingest)
+- GPU orchestrator settings are injected via `configure_settings_provider()`, not hardcoded imports
+- Version bumps must consider downstream consumer compatibility

heimdex_worker_sdk-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,24 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install dependencies
+        run: pip install -e ".[dev]"
+
+      - name: Run tests
+        run: pytest -v

heimdex_worker_sdk-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,62 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - 'v*'
+  workflow_dispatch:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install dependencies
+        run: pip install -e ".[dev]"
+
+      - name: Run tests
+        run: pytest
+
+  publish:
+    needs: test
+    runs-on: ubuntu-latest
+    if: startsWith(github.ref, 'refs/tags/v')
+    permissions:
+      contents: read
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install build dependencies
+        run: pip install build
+
+      - name: Build distribution
+        run: python -m build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    needs: publish
+    runs-on: ubuntu-latest
+    if: startsWith(github.ref, 'refs/tags/v')
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Create GitHub Release
+        uses: softprops/action-gh-release@v2
+        with:
+          generate_release_notes: true

heimdex_worker_sdk-0.1.0/.gitignore

@@ -0,0 +1,41 @@
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+*.egg-info/
+*.egg
+.eggs/
+dist/
+build/
+*.whl
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+.mypy_cache/
+
+# Environment / secrets
+.env
+.env.*
+!.env.example
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# oh-my-claudecode (local agent state)
+.omc/

heimdex_worker_sdk-0.1.0/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: heimdex-worker-sdk
+Version: 0.1.0
+Summary: Shared S3 client, SQS consumer, Aircloud orchestrator, and worker utilities for Heimdex
+Requires-Python: >=3.11
+Requires-Dist: boto3>=1.28
+Requires-Dist: pydantic-settings>=2.0
+Requires-Dist: requests>=2.31
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == 'dev'

heimdex_worker_sdk-0.1.0/README.md

@@ -0,0 +1,46 @@
+# heimdex-worker-sdk
+
+Shared worker utilities for the Heimdex ecosystem: S3 client, SQS consumer, Aircloud GPU orchestrator, and internal API helpers.
+
+## Install
+
+```bash
+pip install heimdex-worker-sdk
+```
+
+## Modules
+
+| Module | Purpose |
+|--------|---------|
+| `aircloud_client` | HTTP client for Aircloud External API (start/stop/scale/status) |
+| `gpu_orchestrator` | Auto start/stop GPU workers based on SQS queue depth |
+| `sqs_client` | Thin boto3 SQS wrapper |
+| `sqs_consumer` | Background SQS polling loop with heartbeat |
+| `s3` | S3/MinIO dual-mode client |
+| `settings` | WorkerSettings via pydantic-settings |
+| `internal_api` | HTTP client for livecommerce internal endpoints |
+| `drive_keys` | Deterministic S3 key generation for Drive files |
+| `youtube_keys` | S3 key generation for YouTube files |
+| `youtube_api` | HTTP client for YouTube internal endpoints |
+| `content_type` | MIME type classification |
+| `message_adapters` | SQS message to dataclass converters |
+
+## Consumers
+
+| Product | What it uses |
+|---------|-------------|
+| dev-heimdex-for-livecommerce | Full SDK (all modules) |
+| dev-heimdex-playground | `aircloud_client`, `gpu_orchestrator` |
+
+## Release
+
+1. Update `version` in `pyproject.toml`
+2. `git tag vX.Y.Z && git push origin vX.Y.Z`
+3. CI runs tests, builds wheel, publishes to PyPI, creates GitHub Release
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+pytest -v
+```

heimdex_worker_sdk-0.1.0/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "heimdex-worker-sdk"
+version = "0.1.0"
+description = "Shared S3 client, SQS consumer, Aircloud orchestrator, and worker utilities for Heimdex"
+requires-python = ">=3.11"
+dependencies = [
+    "pydantic-settings>=2.0",
+    "boto3>=1.28",
+    "requests>=2.31",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/heimdex_worker_sdk"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = "test_*.py"

heimdex_worker_sdk-0.1.0/src/heimdex_worker_sdk/__init__.py

@@ -0,0 +1,76 @@
+"""Heimdex Worker SDK — shared settings, S3 client, drive key helpers, and internal API client."""
+
+from heimdex_worker_sdk.content_type import (
+    classify_mime,
+    is_image,
+    is_supported_mime,
+    is_video,
+)
+from heimdex_worker_sdk.drive_keys import (
+    audio_s3_key,
+    drive_video_id,
+    enrichment_keyframe_s3_key,
+    enrichment_keyframe_s3_prefix,
+    proxy_s3_key,
+    scene_manifest_s3_key,
+    thumbnail_s3_key,
+    thumbnail_s3_prefix,
+)
+from heimdex_worker_sdk.s3 import S3Client
+from heimdex_worker_sdk.settings import WorkerSettings, get_worker_settings
+from heimdex_worker_sdk.internal_api import (
+    AccessToken,
+    ClaimedConnection,
+    ClaimedFile,
+    ClaimedProcessingFile,
+    InternalAPIClient,
+    UpsertResult,
+)
+
+from heimdex_worker_sdk.sqs_client import SQSJobClient, SQSMessage
+from heimdex_worker_sdk.sqs_consumer import (
+    InvalidMessageError,
+    SQSConsumerLoop,
+    VisibilityHeartbeat,
+)
+from heimdex_worker_sdk.message_adapters import (
+    sqs_to_claimed_file,
+    sqs_to_claimed_processing_file,
+)
+from heimdex_worker_sdk.gpu_orchestrator import (
+    configure_settings_provider,
+    ensure_worker_running,
+)
+
+__all__ = [
+    "WorkerSettings",
+    "get_worker_settings",
+    "S3Client",
+    "InternalAPIClient",
+    "ClaimedConnection",
+    "ClaimedFile",
+    "UpsertResult",
+    "AccessToken",
+    "ClaimedProcessingFile",
+    "SQSJobClient",
+    "SQSMessage",
+    "InvalidMessageError",
+    "SQSConsumerLoop",
+    "VisibilityHeartbeat",
+    "sqs_to_claimed_file",
+    "sqs_to_claimed_processing_file",
+    "configure_settings_provider",
+    "ensure_worker_running",
+    "audio_s3_key",
+    "classify_mime",
+    "drive_video_id",
+    "enrichment_keyframe_s3_key",
+    "enrichment_keyframe_s3_prefix",
+    "is_image",
+    "is_supported_mime",
+    "is_video",
+    "proxy_s3_key",
+    "scene_manifest_s3_key",
+    "thumbnail_s3_key",
+    "thumbnail_s3_prefix",
+]

heimdex_worker_sdk-0.1.0/src/heimdex_worker_sdk/aircloud_client.py

@@ -0,0 +1,198 @@
+"""
+Aircloud External API client for GPU worker lifecycle management.
+
+Thin HTTP wrapper over the Aircloud External API (start/stop/scale/status).
+Used by:
+  - sqs_producer.py (API): wake up workers when publishing SQS jobs
+  - gpu_orchestrator.py (drive-worker): shut down idle workers periodically
+
+All methods are synchronous and fire-and-forget safe.  Errors are logged,
+never raised to callers (unless explicitly requested via raise_on_error).
+
+API docs: https://external.aieev.cloud:5007/external/api/v1
+Auth: Bearer token via API key generated in Aircloud web console.
+"""
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_BASE_URL = "https://external.aieev.cloud:5007/external/api/v1"
+_DEFAULT_TIMEOUT = 10  # seconds
+
+
+@dataclass(frozen=True)
+class EndpointStatus:
+    """Parsed response from GET /endpoints/{endpoint_id}."""
+
+    endpoint_id: str
+    name: str
+    is_active: bool
+    num_replicas: int
+    replica_status_summary: dict[str, int]
+    enable_autoscaling: bool
+    instance_type_name: str
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "EndpointStatus":
+        return cls(
+            endpoint_id=data.get("endpoint_id", ""),
+            name=data.get("name", ""),
+            is_active=data.get("is_active", False),
+            num_replicas=data.get("num_replicas", 0),
+            replica_status_summary=data.get("replica_status_summary", {}),
+            enable_autoscaling=data.get("enable_autoscaling", False),
+            instance_type_name=data.get("instance_type_name", ""),
+        )
+
+
+class AircloudClient:
+    """Synchronous HTTP client for the Aircloud External API.
+
+    Args:
+        api_key: Bearer token for authentication.
+        base_url: API base URL (default: production Aircloud endpoint).
+        timeout: Request timeout in seconds.
+        max_retries: Number of retries on transient failures (5xx, connection errors).
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = _DEFAULT_BASE_URL,
+        timeout: int = _DEFAULT_TIMEOUT,
+        max_retries: int = 2,
+    ) -> None:
+        self._api_key = api_key
+        self._base_url = base_url.rstrip("/")
+        self._timeout = timeout
+
+        self._session = requests.Session()
+        self._session.headers.update(
+            {
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+            }
+        )
+
+        retry = Retry(
+            total=max_retries,
+            backoff_factor=0.5,
+            status_forcelist=[500, 502, 503, 504],
+            allowed_methods=["GET", "POST"],
+        )
+        adapter = HTTPAdapter(max_retries=retry)
+        self._session.mount("https://", adapter)
+        self._session.mount("http://", adapter)
+
+    # ── Public API ─────────────────────────────────────────────────
+
+    def get_status(self, endpoint_id: str) -> Optional[EndpointStatus]:
+        """Get current endpoint status.  Returns None on failure."""
+        try:
+            resp = self._session.get(
+                f"{self._base_url}/endpoints/{endpoint_id}",
+                timeout=self._timeout,
+            )
+            resp.raise_for_status()
+            return EndpointStatus.from_dict(resp.json())
+        except Exception:
+            logger.exception(
+                "aircloud_get_status_failed",
+                extra={"endpoint_id": endpoint_id},
+            )
+            return None
+
+    def start(self, endpoint_id: str) -> bool:
+        """Start an inactive endpoint.  Returns True on success.
+
+        Idempotent: starting an already-active endpoint is a no-op on
+        the Aircloud side (returns success).
+        """
+        try:
+            resp = self._session.post(
+                f"{self._base_url}/endpoints/{endpoint_id}/start",
+                timeout=self._timeout,
+            )
+            resp.raise_for_status()
+            data = resp.json()
+            logger.info(
+                "aircloud_endpoint_started",
+                extra={
+                    "endpoint_id": endpoint_id,
+                    "is_active": data.get("is_active"),
+                    "response_message": data.get("message", ""),
+                },
+            )
+            return True
+        except Exception:
+            logger.exception(
+                "aircloud_start_failed",
+                extra={"endpoint_id": endpoint_id},
+            )
+            return False
+
+    def stop(self, endpoint_id: str) -> bool:
+        """Stop an active endpoint.  Returns True on success."""
+        try:
+            resp = self._session.post(
+                f"{self._base_url}/endpoints/{endpoint_id}/stop",
+                timeout=self._timeout,
+            )
+            resp.raise_for_status()
+            data = resp.json()
+            logger.info(
+                "aircloud_endpoint_stopped",
+                extra={
+                    "endpoint_id": endpoint_id,
+                    "is_active": data.get("is_active"),
+                    "response_message": data.get("message", ""),
+                },
+            )
+            return True
+        except Exception:
+            logger.exception(
+                "aircloud_stop_failed",
+                extra={"endpoint_id": endpoint_id},
+            )
+            return False
+
+    def scale(self, endpoint_id: str, num_replicas: int) -> bool:
+        """Scale replicas for an active endpoint.  Returns True on success.
+
+        Requires autoscaling to be DISABLED on the Aircloud web console.
+        """
+        try:
+            resp = self._session.post(
+                f"{self._base_url}/endpoints/{endpoint_id}/scale",
+                json={"num_replicas": num_replicas},
+                timeout=self._timeout,
+            )
+            resp.raise_for_status()
+            data = resp.json()
+            logger.info(
+                "aircloud_endpoint_scaled",
+                extra={
+                    "endpoint_id": endpoint_id,
+                    "previous_replicas": data.get("previous_replicas"),
+                    "current_replicas": data.get("current_replicas"),
+                    "response_message": data.get("message", ""),
+                },
+            )
+            return True
+        except Exception:
+            logger.exception(
+                "aircloud_scale_failed",
+                extra={
+                    "endpoint_id": endpoint_id,
+                    "num_replicas": num_replicas,
+                },
+            )
+            return False

heimdex_worker_sdk-0.1.0/src/heimdex_worker_sdk/content_type.py

@@ -0,0 +1,52 @@
+"""Centralized MIME-type classification for the content pipeline.
+
+Shared between API (via re-export at ``app.modules.content_type``) and
+drive-worker (direct import from ``heimdex_worker_sdk.content_type``).
+
+All MIME-type decisions must flow through these helpers — no scattered
+``startswith("video/")`` checks elsewhere in the codebase.
+"""
+
+from __future__ import annotations
+
+IMAGE_MIME_TYPES: frozenset[str] = frozenset({
+    "image/jpeg",
+    "image/png",
+    "image/webp",
+    "image/heic",
+    "image/heif",
+})
+
+VIDEO_MIME_PREFIX: str = "video/"
+
+
+def classify_mime(mime_type: str) -> str:
+    """Classify a MIME type string as ``"video"``, ``"image"``, or ``"unknown"``.
+
+    >>> classify_mime("video/mp4")
+    'video'
+    >>> classify_mime("image/jpeg")
+    'image'
+    >>> classify_mime("application/pdf")
+    'unknown'
+    """
+    if mime_type in IMAGE_MIME_TYPES:
+        return "image"
+    if mime_type.startswith(VIDEO_MIME_PREFIX):
+        return "video"
+    return "unknown"
+
+
+def is_supported_mime(mime_type: str) -> bool:
+    """Return ``True`` if *mime_type* is a supported video or image type."""
+    return classify_mime(mime_type) != "unknown"
+
+
+def is_image(mime_type: str) -> bool:
+    """Return ``True`` if *mime_type* is a supported image type."""
+    return mime_type in IMAGE_MIME_TYPES
+
+
+def is_video(mime_type: str) -> bool:
+    """Return ``True`` if *mime_type* is any video type."""
+    return mime_type.startswith(VIDEO_MIME_PREFIX)

heimdex_worker_sdk-0.1.0/src/heimdex_worker_sdk/drive_keys.py

@@ -0,0 +1,53 @@
+"""Drive S3 key helpers — standalone copy of ``app.modules.drive.keys``.
+
+Pure functions, only ``hashlib`` dependency.  Kept 1:1 with the API copy so
+that workers and the API always generate identical keys.
+"""
+
+import hashlib
+
+
+def drive_video_id(org_id: str, google_file_id: str) -> str:
+    """Deterministic video_id for Drive files. Collision-resistant, idempotent."""
+    digest = hashlib.sha256(f"{org_id}:{google_file_id}".encode()).hexdigest()[:16]
+    return f"gd_{digest}"
+
+
+def proxy_s3_key(org_id: str, drive_id: str, google_file_id: str) -> str:
+    return f"{org_id}/drive/{drive_id}/{google_file_id}/proxy.mp4"
+
+
+def thumbnail_s3_key(org_id: str, video_id: str, scene_id: str) -> str:
+    return f"{org_id}/drive/thumbs/{video_id}/{scene_id}.jpg"
+
+
+def thumbnail_s3_prefix(org_id: str, video_id: str) -> str:
+    return f"{org_id}/drive/thumbs/{video_id}/"
+
+
+def audio_s3_key(org_id: str, video_id: str) -> str:
+    return f"{org_id}/drive/audio/{video_id}/audio.wav"
+
+
+def enrichment_keyframe_s3_prefix(org_id: str, video_id: str) -> str:
+    return f"{org_id}/drive/keyframes/{video_id}/"
+
+
+def enrichment_keyframe_s3_key(
+    org_id: str, video_id: str, scene_id: str,
+) -> str:
+    return f"{org_id}/drive/keyframes/{video_id}/{scene_id}.jpg"
+
+
+def scene_manifest_s3_key(org_id: str, video_id: str) -> str:
+    return f"{org_id}/drive/manifests/{video_id}/scenes.json"
+
+
+def original_s3_key(org_id: str, drive_id: str, google_file_id: str) -> str:
+    """S3 key for the original (pre-transcode) file uploaded by drive-worker."""
+    return f"{org_id}/drive/{drive_id}/{google_file_id}/original"
+
+
+def stt_result_s3_key(org_id: str, video_id: str) -> str:
+    """S3 key for STT result JSON used by speech-aware scene splitting."""
+    return f"{org_id}/drive/stt/{video_id}/stt_result.json"