Skip to content

Worker API Documentation

This document describes how workers interact with the MicroHub server to process jobs.

Authentication

Workers authenticate using Bearer tokens obtained during registration:

Authorization: Bearer <worker_token>

API Base URL

All worker API endpoints use the /api/v1/workers prefix:

  • Production: https://api.microdc.ai/api/v1/workers
  • Development: http://localhost:8000/api/v1/workers

Job Assignment System

The MicroHub platform uses a claim-based job assignment system to prevent multiple workers from competing for the same job. This ensures:

  • No race conditions with hundreds of workers
  • Jobs are only shown to workers that can handle them
  • Fair distribution of work
  • Atomic job claiming (only one worker gets each job)

How It Works

  1. Check Available Jobs - Worker requests list of jobs matching its capabilities
  2. Claim a Job - Worker atomically claims a specific job (first claim wins)
  3. Process - Only the successful claimer processes the job
  4. Complete/Fail - Submit results or report failure

API Endpoints

1. Get Available Jobs

Endpoint: GET /api/v1/workers/jobs/available

Query Parameters:

  • limit (optional, default=5): Maximum number of jobs to return

Description: Returns jobs that match the worker's capabilities and are available to claim.

response = requests.get(
    f"{server_url}/api/v1/workers/jobs/available",
    headers={"Authorization": f"Bearer {worker_token}"},
    params={"limit": 5}
)

data = response.json()
print(f"Found {data['count']} available jobs")

Response:

{
    "jobs": [
        {
            "job_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
            "model": "llama2-7b",
            "platform": "ollama",
            "type": "LLM",
            "priority": "high",
            "estimated_reward": 0.0125,
            "created_at": "2024-01-01T12:00:00Z",
            "retry_count": 0,
            "expires_in": 30
        }
    ],
    "count": 1,
    "message": "Found 1 available jobs"
}

Fields:

  • platform: The inference platform to use (e.g., "ollama", "transformers", "vllm"). Derived from the worker's registered model support.

Important Notes:

  • Jobs returned are filtered by worker's supported models
  • Jobs with retry_count >= 3 are not shown
  • Jobs are ordered by priority (descending) then creation time (ascending)
  • These jobs are NOT reserved - you must claim them atomically

2. Claim a Job

Endpoint: POST /api/v1/workers/jobs/{job_id}/claim

Description: Atomically claim a job. Uses database row-level locking to ensure only one worker can claim each job.

response = requests.post(
    f"{server_url}/api/v1/workers/jobs/{job_id}/claim",
    headers={"Authorization": f"Bearer {worker_token}"}
)

result = response.json()
if result["status"] == "claimed":
    job_data = result["job_data"]
    print(f"Successfully claimed job {job_data['job_id']}")
    # Start processing the job...
else:
    print(f"Could not claim job: {result['status']}")

Success Response (status=claimed) — Unencrypted Job:

{
    "status": "claimed",
    "message": "Job successfully claimed",
    "job_data": {
        "job_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
        "assignment_id": "7b3d5f64-1234-4562-b3fc-2c963f66afa6",
        "type": "LLM",
        "model": "llama2-7b",
        "platform": "ollama",
        "is_encrypted": false,
        "payload": {
            "prompt": "Hello, how are you?",
            "max_tokens": 512,
            "temperature": 0.7
        },
        "llm_interaction_type": "generation",
        "input_modalities": ["text"],
        "output_modalities": ["text"],
        "priority": "high",
        "timeout": 300,
        "started_at": "2024-01-01T12:00:00Z",
        "must_complete_by": "2024-01-01T12:05:00Z"
    }
}

Success Response (status=claimed) — Encrypted Job:

{
    "status": "claimed",
    "message": "Job successfully claimed",
    "job_data": {
        "job_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
        "assignment_id": "7b3d5f64-1234-4562-b3fc-2c963f66afa6",
        "type": "LLM",
        "model": "llama2-7b",
        "platform": "ollama",
        "is_encrypted": true,
        "payload": {
            "_encrypted": true,
            "ciphertext": "<base64-encoded AES-256-GCM ciphertext>",
            "iv": "<base64-encoded IV>",
            "algorithm": "AES-256-GCM"
        },
        "encryption_key": "<base64-encoded 32-byte AES symmetric key>",
        "payload_iv": "<base64-encoded 12-byte IV>",
        "customer_public_key": "-----BEGIN PUBLIC KEY-----\n...\n-----END PUBLIC KEY-----",
        "priority": "high",
        "timeout": 300,
        "started_at": "2024-01-01T12:00:00Z",
        "must_complete_by": "2024-01-01T12:05:00Z"
    }
}

Fields:

  • type: Job type - LLM, EMBED, or DOCUMENT
  • platform: The inference platform to use (e.g., "ollama", "transformers", "vllm")
  • is_encrypted: Whether the job payload is encrypted. When true, the worker must decrypt the payload before processing and encrypt the result before submitting.
  • encryption_key: (Encrypted jobs only) Base64-encoded AES-256 symmetric key for decrypting the payload
  • payload_iv: (Encrypted jobs only) Base64-encoded initialization vector for AES-GCM decryption
  • customer_public_key: (Encrypted jobs only) PEM-encoded RSA public key for encrypting the result
  • llm_interaction_type: Only for LLM jobs - "generation" or "chat". EMBED and DOCUMENT jobs don't use this field. Not present in the encrypted claim response — available inside the decrypted payload.

Failure Response Statuses:

  • already_claimed - Another worker claimed the job
  • incompatible - Worker doesn't support the required model
  • worker_busy - Worker is already processing another job
  • 404 - Job not found

3. Complete Job

Endpoint: POST /api/v1/workers/jobs/{job_id}/complete

Description: Submit successful job completion with results. Supports both plaintext and encrypted result modes.

import time

start_time = time.time()
result = process_job(job_data)  # Your processing logic
execution_time = time.time() - start_time

response = requests.post(
    f"{server_url}/api/v1/workers/jobs/{job_id}/complete",
    headers={"Authorization": f"Bearer {worker_token}"},
    json={
        "output": result,
        "execution_time": execution_time,
        "tokens_used": 150  # Optional
    }
)

Request Body — Unencrypted (existing):

{
    "output": {
        "text": "Generated response...",
        "finish_reason": "stop",
        "reasoning_content": "optional — chain-of-thought from a thinking model"
    },
    "execution_time": 2.5,
    "tokens_used": 150
}

output.reasoning_content is included only for reasoning ("thinking") models that return a chain-of-thought (DeepSeek-R1 / o1 convention); it is omitted for ordinary responses. The answer text stays in output.text.

Request Body — Encrypted:

{
    "encrypted_output": "<base64-encoded encrypted result ciphertext>",
    "result_iv": "<base64-encoded IV for result encryption>",
    "execution_time": 2.5,
    "input_tokens": 150,
    "output_tokens": 200,
    "total_tokens": 350
}

Important: Provide either output (plaintext) OR encrypted_output + result_iv (encrypted) — never both. For encrypted jobs, include input_tokens/output_tokens since the server cannot count tokens from encrypted payloads. If token counts are omitted for encrypted jobs, the server records zeros — billing and usage stats will be inaccurate.

Response:

{
    "status": "completed",
    "earnings": 0.0125,
    "next_poll": 60
}

4. Report Job Failure

Endpoint: POST /api/v1/workers/jobs/{job_id}/fail

Description: Report that a job could not be completed.

response = requests.post(
    f"{server_url}/api/v1/workers/jobs/{job_id}/fail",
    headers={"Authorization": f"Bearer {worker_token}"},
    json={
        "reason": "Model loading failed: Out of memory",
        "error_details": traceback.format_exc()  # Optional
    }
)

Request Body:

{
    "reason": "Out of memory",
    "error_details": "Detailed error traceback (optional)"
}

5. Worker Status Heartbeat

Endpoint: POST /api/v1/workers/heartbeat

Description: Send periodic heartbeat to update worker status, resource usage, and system metrics. The worker sends this on the monitoring.heartbeat_interval cadence (default 30 seconds, see config/default.yaml) to maintain active status.

IMPORTANT FOR LONG-RUNNING JOBS: If a worker is processing a job that may exceed the timeout (default 300 seconds / 5 minutes), it MUST include job_progress in the heartbeat to reset the timeout and prevent the job from being reassigned. This allows jobs to run indefinitely as long as the worker reports progress.

Request Body:

{
    "status": "idle",  // or "busy", "maintenance"
    "cpu_usage": 25.5,
    "memory_usage": 45.2,
    "gpu_usage": 0.0,  // Optional, null if no GPU
    "active_jobs": 0,
    "completed_jobs_session": 5,
    "worker_version": "1.0.0",  // Optional: Worker software version
    "current_models": [  // Optional: Currently available models with platform info
        {"model_name": "llama2-7b", "platform": "ollama"},
        {"model_name": "sentence-transformers/all-minilm-l6-v2", "platform": "transformers"}
    ],
    "job_progress": {  // IMPORTANT: Include this when processing a job to reset timeout
        "job_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
        "assignment_id": "7b3d5f64-1234-4562-b3fc-2c963f66afa6",
        "progress_percentage": 45.0,  // Optional: 0-100
        "message": "Loading model weights",  // Optional: Human-readable status
        "estimated_remaining_seconds": 120  // Optional: Estimated time to completion
    },
    "system_metrics": {  // Optional: Detailed system metrics
        "load_average": [1.5, 2.0, 1.8],  // 1, 5, and 15 minute load averages
        "cpu_count": 16,
        "memory_total_gb": 64.0,
        "memory_available_gb": 32.5,
        "memory_percent": 49.2,
        "disk_total_gb": 1000.0,
        "disk_available_gb": 450.0,
        "disk_percent": 55.0,
        "gpu_memory_total_gb": 24.0,  // Optional
        "gpu_memory_used_gb": 8.5,     // Optional
        "gpu_temperature_c": 65.0,      // Optional: GPU temperature in Celsius
        "cpu_temperature_c": 58.0,      // Optional: CPU temperature in Celsius
        "network_sent_mb": 1024.5,      // Optional: Network bytes sent in MB
        "network_recv_mb": 2048.3,      // Optional: Network bytes received in MB
        "uptime_seconds": 86400
    },
    "capabilities": ["encryption"],  // Optional: Worker capabilities. Include "encryption" to receive encrypted jobs.
    "supported_models": ["llama2-7b", "mistral-7b"]  // Legacy: Use current_models instead
}

Response:

{
    "status": "acknowledged",
    "server_time": "2024-01-01T12:00:00Z",
    "next_heartbeat": 60,
    "model_sync": {
        "added": ["new-model"],
        "removed": ["old-model"],
        "updated": ["existing-model"]
    },
    "metrics_received": true,
    "progress_recorded": true  // true if job_progress was successfully recorded
}

Notes:

  • The current_models field should list all models currently available on the worker with their platform
  • Each model entry should include model_name and platform (e.g., "ollama", "transformers", "vllm")
  • The platform is used to route jobs correctly - workers receive jobs with matching platform
  • System metrics help the server monitor worker health and resource availability
  • All temperature values should be in Celsius
  • Load average follows Unix convention: [1-minute, 5-minute, 15-minute]
  • Network metrics are cumulative since worker start
  • Job Progress: When processing a job, include job_progress to extend the timeout. Each progress update resets the timeout window, allowing jobs to run longer than the initial timeout limit.
  • Timeout Handling: If no progress is reported, jobs will timeout based on timeout_seconds (default 300s). Including progress updates prevents timeout and job reassignment.

6. Optional Job Endpoints

Send Job Heartbeat

Endpoint: POST /api/v1/workers/jobs/{job_id}/heartbeat

Keep a long-running job alive and report progress. The worker calls this during streaming/long generations (server_client.send_job_heartbeat) and reads the response for a cancel_requested flag, which lets the server cancel an in-flight job.

Release Job

Endpoint: POST /api/v1/workers/jobs/{job_id}/release

Release a claimed job back to the queue when the worker can't process it (e.g. insufficient VRAM, over cost budget, or a parse failure). Known issue: the server currently returns 404 for this endpoint, so releases fail silently; the worker mitigates this by gating jobs before claiming (VRAM/RAM admission) so a release is rarely needed. Tracked for a server-side fix.

Job Processing

Execute the job based on its type, platform, interaction type, and modalities. Check is_encrypted to determine whether the payload needs decryption first.

def process_job(job):
    job_type = job["type"]
    model = job["model"]
    platform = job["platform"]  # The inference platform to use (ollama, transformers, vllm, etc.)
    llm_interaction_type = job.get("llm_interaction_type")  # Only for LLM jobs
    input_modalities = job.get("input_modalities", ["text"])
    output_modalities = job.get("output_modalities", ["text"])

    # Handle encrypted payloads
    if job.get("is_encrypted", False):
        payload = decrypt_payload(job)
    else:
        payload = job["payload"]

    if job_type == "LLM":
        # Route based on LLM interaction type
        if llm_interaction_type == "chat":
            # Chat-based interaction (messages format)
            result = run_chat_inference(
                model=model,
                messages=payload["messages"],
                max_tokens=payload.get("max_tokens", 512),
                temperature=payload.get("temperature", 0.7),
                input_modalities=input_modalities,
                output_modalities=output_modalities
            )
        elif llm_interaction_type == "generation":
            # Text generation (completion format)
            result = run_llm_inference(
                model=model,
                prompt=payload["prompt"],
                max_tokens=payload.get("max_tokens", 512),
                temperature=payload.get("temperature", 0.7),
                input_modalities=input_modalities,
                output_modalities=output_modalities
            )
        return result

    elif job_type == "EMBED":
        # Generate embeddings using the specified platform
        result = generate_embeddings(
            model=model,
            platform=platform,  # Use platform to route to correct inference backend
            texts=payload.get("texts") or [payload.get("text")]
        )
        return result

    elif job_type == "DOCUMENT":
        # Document processing (docling, surya, etc.)
        result = process_document(
            model=model,
            platform=platform,
            files=payload.get("files", [])
        )
        return result

    # Add more job types as needed

Multimodal Job Examples

Text-to-Image Generation

if job_type == "LLM" and "image" in output_modalities:
    # Input: text, Output: image
    result = generate_image(
        model=model,
        prompt=payload["prompt"],
        size=payload.get("size", "1024x1024")
    )
    return {
        "image_url": result.url,
        "image_data": result.base64  # or URL
    }

Vision Model (Image + Text → Text)

if "image" in input_modalities and "text" in input_modalities:
    # Multi-modal input processing
    result = process_vision_model(
        model=model,
        text=payload["prompt"],
        images=payload["images"],  # Base64 or URLs
        max_tokens=payload.get("max_tokens", 512)
    )
    return {
        "text": result.text,
        "finish_reason": "stop"
    }

Audio Transcription

if "audio" in input_modalities:
    # Audio to text
    result = transcribe_audio(
        model=model,
        audio_data=payload["audio"],  # Base64 or URL
        language=payload.get("language", "en")
    )
    return {
        "text": result.transcription,
        "language": result.detected_language
    }

Video Analysis

if "video" in input_modalities:
    # Video to text/image analysis
    result = analyze_video(
        model=model,
        video_url=payload["video_url"],
        analysis_type=payload.get("analysis_type", "description")
    )
    return {
        "analysis": result.text,
        "keyframes": result.keyframes if "image" in output_modalities else None
    }

Encrypted Job Processing

Two gates control access to encrypted jobs:

  1. Admin approval — A platform admin must enable the encryption capability for the worker via the Workers UI (lock icon on the worker card). This creates a WorkerCapability record.
  2. Worker confirmation — The worker must report "encryption" in its heartbeat capabilities list to confirm it can actually handle encrypted jobs at runtime.

Both must be true. If an admin approves a worker but the worker hasn't reported the capability yet, the worker card shows a yellow "E2E?" badge. Once the worker confirms, it turns blue "E2E".

// Heartbeat with encryption capability
{
    "status": "idle",
    "cpu_usage": 10.0,
    "memory_usage": 30.0,
    "active_jobs": 0,
    "completed_jobs_session": 5,
    "capabilities": ["encryption"]
}

Workers without both approvals will not see encrypted jobs in /jobs/available and will be rejected if they attempt to claim one.

Workers must support encrypted jobs when is_encrypted is true in the claimed job data. The flow:

  1. Decrypt the payload using the encryption_key (AES-256-GCM symmetric key) and payload_iv
  2. Process the job as normal (the decrypted payload has the same structure as a plaintext payload)
  3. Encrypt the result using the customer_public_key (RSA-OAEP + AES-256-GCM hybrid encryption)
  4. Submit the encrypted result with encrypted_output and result_iv instead of output

How it works, end to end

Three keys are involved, and they travel different paths:

Key Generated by Seen by the platform? Seen by the worker?
Payload AES-256 key (encryption_key) client yes — held in platform key custody, delivered to the approved worker in the claim response over TLS yes (used to decrypt)
Result AES-256 key worker — fresh, one-time, per result only RSA-wrapped (opaque) yes (generated, used, discarded)
Customer RSA keypair client public half only (customer_public_key) public half only

Step by step (worker implementation: src/utils/encryption.py):

  1. Submit — the client encrypts the job payload with AES-256-GCM under a symmetric key and sends ciphertext, key, and IV to the platform.
  2. Claim — the platform delivers all three to an approved worker in the claim response (payload.ciphertext, encryption_key, payload_iv), plus the customer's RSA public key for the return path.
  3. Decrypt — the worker decrypts the payload in memory (decrypt_payload). The plaintext prompt then flows through the normal engine path; the inference engines are encryption-unaware.
  4. Encrypt the resultencrypt_result generates a fresh one-time 32-byte AES-256-GCM key and a random 96-bit IV, encrypts the result JSON with them, then wraps the AES key with the customer's RSA public key (RSA-OAEP, MGF1 + SHA-256). The hybrid construction is required, not optional: RSA cannot encrypt bulk data (a 4096-bit key caps out at 446 bytes per OAEP-SHA-256 operation), so the RSA key encrypts only the 32-byte AES key, and the AES key encrypts the arbitrarily-sized result.
  5. Submit the result — the wrapped key is embedded in the result blob itself (format below), base64-encoded into encrypted_output; the GCM IV goes alongside as result_iv. Estimated input_tokens/output_tokens are sent as plaintext metadata because the server cannot count tokens in ciphertext.
  6. Client decrypt — mirror the packing: read the 4-byte length prefix, split off the wrapped key, unwrap it with the RSA private key (OAEP-SHA-256), then AES-256-GCM-decrypt the remainder using the recovered key and result_iv.

Result blob format

encrypted_output = base64(
    <4-byte big-endian unsigned int: length of wrapped_key>
    <wrapped_key:  RSA-OAEP(MGF1+SHA-256) encryption of the 32-byte AES key>
    <ciphertext:   AES-256-GCM(result JSON) — 16-byte GCM tag appended>
)
result_iv = base64(<12-byte GCM IV>)

Security model

Encrypted jobs are protected by multiple layers, each defending a different boundary:

  1. TLS in transit. All worker ↔ platform API traffic runs over HTTPS in production, terminated and enforced at the platform edge — there is no plaintext-HTTP path to the API. (The worker itself follows whatever scheme MICRODC_SERVER_URL specifies, so plain HTTP stays available for local development against localhost.)
  2. Encrypted payloads. The prompt is encrypted client-side with AES-256-GCM before submission. The platform stores and routes ciphertext — plaintext prompts never touch the platform's database or queues. The payload key is held in the platform's key-custody layer and released only to the approved worker at claim time, over TLS.
  3. Results sealed to the customer. The worker encrypts every result under a fresh one-time AES-256-GCM key, wrapped with the customer's RSA public key. Only the customer's private key — which never leaves their machine — can open a result; the platform stores and routes a blob it cannot decrypt. Result keys are never reused, so each result is independently protected.

Trust boundaries (implementers should know exactly what each layer defends):

  • Payload confidentiality rests on TLS plus platform key custody: it protects against network attackers and at-rest data exposure, with the platform acting as the trusted key custodian between client and worker.
  • Result confidentiality is the strongest guarantee — it holds independently of the platform, because only the customer holds the unwrapping key.
  • customer_public_key is delivered by the platform in the claim response; the worker uses it as supplied.
  • encrypt_result accepts any RSA public key; 4096-bit is the platform convention.
  • Job metadata (type, model, timings, estimated token counts) is intentionally cleartext so scheduling, routing, and billing work without ever decrypting content.

Worker obligations: decrypt only in memory; never log or persist the payload key, decrypted payload, or plaintext result. Worker logs record only sizes and non-sensitive metadata (job id/type/model/encrypted flag, input/output lengths).

Dependencies

pip install cryptography

Decrypting the Payload

import json
from base64 import b64decode
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

def decrypt_payload(job_data: dict) -> dict:
    """Decrypt an encrypted job payload using the provided symmetric key."""
    symmetric_key = b64decode(job_data["encryption_key"])
    iv = b64decode(job_data["payload_iv"])
    ciphertext = b64decode(job_data["payload"]["ciphertext"])

    aesgcm = AESGCM(symmetric_key)
    plaintext_bytes = aesgcm.decrypt(iv, ciphertext, None)
    return json.loads(plaintext_bytes.decode("utf-8"))

Encrypting the Result

The result is encrypted using hybrid encryption: a new AES-256-GCM key encrypts the data, and the customer's RSA public key encrypts that AES key.

import json
import struct
import secrets
from base64 import b64encode
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes, serialization

def encrypt_result(result: dict, customer_public_key_pem: str) -> tuple[str, str]:
    """
    Encrypt job result for the customer.

    Returns:
        (encrypted_output_b64, result_iv_b64) — both base64-encoded strings
    """
    result_bytes = json.dumps(result).encode("utf-8")

    # Generate a one-time AES key for this result
    result_key = AESGCM.generate_key(bit_length=256)
    result_iv = secrets.token_bytes(12)  # 96-bit IV for AES-GCM
    aesgcm = AESGCM(result_key)
    encrypted_data = aesgcm.encrypt(result_iv, result_bytes, None)

    # Encrypt the AES key with the customer's RSA public key
    customer_pubkey = serialization.load_pem_public_key(
        customer_public_key_pem.encode("utf-8")
    )
    encrypted_key = customer_pubkey.encrypt(
        result_key,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )

    # Pack: <4-byte key length><encrypted_key><encrypted_data>
    packed = struct.pack(">I", len(encrypted_key)) + encrypted_key + encrypted_data

    return b64encode(packed).decode("utf-8"), b64encode(result_iv).decode("utf-8")

Submitting an Encrypted Result

def complete_encrypted_job(self, job_id: str, result: dict, job_data: dict,
                           execution_time: float, input_tokens: int, output_tokens: int) -> bool:
    """Submit an encrypted job result."""
    encrypted_output, result_iv = encrypt_result(result, job_data["customer_public_key"])

    response = requests.post(
        f"{self.server_url}/api/v1/workers/jobs/{job_id}/complete",
        headers=self.headers,
        json={
            "encrypted_output": encrypted_output,
            "result_iv": result_iv,
            "execution_time": execution_time,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
        }
    )

    return response.status_code == 200

Updated process_job with Encryption Support

def process_job(self, job_data: dict):
    """Process a job — handles both encrypted and plaintext."""
    is_encrypted = job_data.get("is_encrypted", False)

    # Decrypt payload if needed
    if is_encrypted:
        payload = decrypt_payload(job_data)
    else:
        payload = job_data["payload"]

    # Process as normal (same logic for both modes)
    start_time = time.time()
    result = self.execute(payload, job_data)
    execution_time = time.time() - start_time

    # Submit result — encrypted or plaintext
    if is_encrypted:
        return self.complete_encrypted_job(
            job_id=job_data["job_id"],
            result=result,
            job_data=job_data,
            execution_time=execution_time,
            input_tokens=result.get("input_tokens", 0),
            output_tokens=result.get("output_tokens", 0),
        )
    else:
        return self.complete_job(job_data["job_id"], result, execution_time)

Complete Python Worker Example

import json
import secrets
import struct
import requests
import time
import logging
import traceback
from base64 import b64encode, b64decode
from typing import Dict, Any, Optional, List

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.asymmetric import padding as asym_padding
from cryptography.hazmat.primitives import hashes, serialization


def decrypt_payload(job_data: dict) -> dict:
    """Decrypt an encrypted job payload using the provided symmetric key."""
    symmetric_key = b64decode(job_data["encryption_key"])
    iv = b64decode(job_data["payload_iv"])
    ciphertext = b64decode(job_data["payload"]["ciphertext"])

    aesgcm = AESGCM(symmetric_key)
    plaintext_bytes = aesgcm.decrypt(iv, ciphertext, None)
    return json.loads(plaintext_bytes.decode("utf-8"))


def encrypt_result(result: dict, customer_public_key_pem: str) -> tuple[str, str]:
    """
    Encrypt job result for the customer using hybrid encryption.

    Returns:
        (encrypted_output_b64, result_iv_b64) — both base64-encoded strings
    """
    result_bytes = json.dumps(result).encode("utf-8")

    # Generate a one-time AES key for this result
    result_key = AESGCM.generate_key(bit_length=256)
    result_iv = secrets.token_bytes(12)  # 96-bit IV for AES-GCM
    aesgcm = AESGCM(result_key)
    encrypted_data = aesgcm.encrypt(result_iv, result_bytes, None)

    # Encrypt the AES key with the customer's RSA public key
    customer_pubkey = serialization.load_pem_public_key(
        customer_public_key_pem.encode("utf-8")
    )
    encrypted_key = customer_pubkey.encrypt(
        result_key,
        asym_padding.OAEP(
            mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )

    # Pack: <4-byte key length><encrypted_key><encrypted_data>
    packed = struct.pack(">I", len(encrypted_key)) + encrypted_key + encrypted_data

    return b64encode(packed).decode("utf-8"), b64encode(result_iv).decode("utf-8")


class MicroHubWorker:
    def __init__(self, server_url: str, worker_token: str):
        self.server_url = server_url.rstrip('/')
        self.worker_token = worker_token
        self.headers = {
            "Authorization": f"Bearer {worker_token}",
            "Content-Type": "application/json"
        }
        self.logger = logging.getLogger(__name__)

    def get_available_jobs(self, limit: int = 5) -> List[Dict[str, Any]]:
        """Get list of jobs this worker can handle"""
        try:
            response = requests.get(
                f"{self.server_url}/api/v1/workers/jobs/available",
                headers=self.headers,
                params={"limit": limit}
            )

            if response.status_code == 200:
                data = response.json()
                return data.get("jobs", [])

            return []

        except Exception as e:
            self.logger.error(f"Error getting available jobs: {e}")
            return []

    def claim_job(self, job_id: str) -> Optional[Dict[str, Any]]:
        """Attempt to claim a specific job"""
        try:
            response = requests.post(
                f"{self.server_url}/api/v1/workers/jobs/{job_id}/claim",
                headers=self.headers
            )

            if response.status_code == 200:
                data = response.json()
                if data["status"] == "claimed":
                    self.logger.info(f"Successfully claimed job {job_id}")
                    return data["job_data"]
                else:
                    self.logger.info(f"Could not claim job {job_id}: {data['status']}")

            return None

        except Exception as e:
            self.logger.error(f"Error claiming job: {e}")
            return None

    def send_worker_heartbeat(self, status: str = "idle", active_jobs: int = 0,
                             cpu_usage: float = 0, memory_usage: float = 0,
                             gpu_usage: Optional[float] = None,
                             completed_jobs_session: int = 0,
                             supported_models: Optional[List[str]] = None,
                             job_progress: Optional[Dict[str, Any]] = None) -> bool:
        """Send worker status heartbeat (every 60 seconds)

        Args:
            job_progress: Optional dict with job_id, assignment_id, and optional
                         progress_percentage, message, estimated_remaining_seconds.
                         Including this resets the job timeout.
        """
        try:
            payload = {
                "status": status,  # "idle", "busy", or "maintenance"
                "cpu_usage": cpu_usage,
                "memory_usage": memory_usage,
                "gpu_usage": gpu_usage,
                "active_jobs": active_jobs,
                "completed_jobs_session": completed_jobs_session
            }

            if supported_models:
                payload["supported_models"] = supported_models

            if job_progress:
                payload["job_progress"] = job_progress

            response = requests.post(
                f"{self.server_url}/api/v1/workers/heartbeat",
                headers=self.headers,
                json=payload
            )

            if response.status_code == 200:
                self.logger.debug("Worker heartbeat sent successfully")
                return True
            else:
                self.logger.error(f"Worker heartbeat failed: {response.text}")
                return False

        except Exception as e:
            self.logger.error(f"Error sending worker heartbeat: {e}")
            return False

    def send_job_heartbeat(self, job_id: str, progress: float = None, message: str = None) -> bool:
        """Send heartbeat to keep job alive (not yet implemented in server)"""
        try:
            payload = {}
            if progress is not None:
                payload["progress"] = progress
            if message:
                payload["message"] = message

            response = requests.post(
                f"{self.server_url}/api/v1/workers/jobs/{job_id}/heartbeat",
                headers=self.headers,
                json=payload
            )
            return response.status_code == 200

        except Exception as e:
            self.logger.error(f"Error sending job heartbeat: {e}")
            return False

    def complete_job(self, job_id: str, result: Dict[str, Any],
                    execution_time: float, tokens_used: Optional[int] = None,
                    encrypted: bool = False, job_data: Optional[Dict] = None) -> bool:
        """Submit successful job completion (plaintext or encrypted)"""
        try:
            if encrypted and job_data:
                encrypted_output, result_iv = encrypt_result(result, job_data["customer_public_key"])
                payload = {
                    "encrypted_output": encrypted_output,
                    "result_iv": result_iv,
                    "execution_time": execution_time,
                    "input_tokens": result.get("input_tokens", 0),
                    "output_tokens": result.get("output_tokens", 0),
                }
            else:
                payload = {
                    "output": result,
                    "execution_time": execution_time
                }
                if tokens_used:
                    payload["tokens_used"] = tokens_used

            response = requests.post(
                f"{self.server_url}/api/v1/workers/jobs/{job_id}/complete",
                headers=self.headers,
                json=payload
            )

            if response.status_code == 200:
                data = response.json()
                self.logger.info(f"Job completed! Earned: ${data.get('earnings', 0)}")
                return True
            else:
                self.logger.error(f"Failed to complete job: {response.text}")
                return False

        except Exception as e:
            self.logger.error(f"Error completing job: {e}")
            return False

    def fail_job(self, job_id: str, reason: str, details: Optional[str] = None) -> bool:
        """Report job failure"""
        try:
            payload = {"reason": reason}
            if details:
                payload["error_details"] = details

            response = requests.post(
                f"{self.server_url}/api/v1/workers/jobs/{job_id}/fail",
                headers=self.headers,
                json=payload
            )

            return response.status_code == 200

        except Exception as e:
            self.logger.error(f"Error reporting job failure: {e}")
            return False

    def process_job(self, job: Dict[str, Any]) -> Dict[str, Any]:
        """Process a job based on its type, platform, interaction type, and modalities - Override this method"""
        job_type = job["type"]
        model = job["model"]
        platform = job["platform"]  # Use this to route to correct inference backend
        llm_interaction_type = job.get("llm_interaction_type")  # Only for LLM jobs
        input_modalities = job.get("input_modalities", ["text"])
        output_modalities = job.get("output_modalities", ["text"])

        # Decrypt payload if encrypted
        if job.get("is_encrypted", False):
            payload = decrypt_payload(job)
        else:
            payload = job["payload"]

        # This is where you implement your actual processing logic
        if job_type == "LLM":
            # Route based on LLM interaction type and platform
            if llm_interaction_type == "chat":
                # Chat-based interaction
                return {
                    "text": f"Chat response to: {payload.get('messages', [])}",
                    "finish_reason": "stop"
                }
            elif llm_interaction_type == "generation":
                # Text generation
                return {
                    "text": f"Generated from: {payload.get('prompt', '')}",
                    "finish_reason": "stop"
                }

            # Handle multimodal cases
            if "image" in output_modalities:
                # Text-to-image generation
                return {"image_url": "https://example.com/generated.png"}

        elif job_type == "EMBED":
            # Embedding processing - use platform to route to correct backend
            # platform could be "transformers", "ollama", "sentence-transformers", etc.
            return {
                "embeddings": [[0.1] * 768]  # Replace with actual embeddings
            }

        elif job_type == "DOCUMENT":
            # Document processing (docling, surya, etc.)
            return {
                "processed": True,
                "content": "Document content..."
            }
        else:
            raise ValueError(f"Unsupported job type: {job_type}")

    def run(self, poll_interval: int = 10, heartbeat_interval: int = 60):
        """Main worker loop with claim-based assignment"""
        self.logger.info("Worker started (claim-based system)")
        last_heartbeat = 0
        current_job = None

        while True:
            try:
                # Send periodic heartbeat (every 60s)
                if time.time() - last_heartbeat > heartbeat_interval:
                    job_progress = None
                    if current_job:
                        # Include job progress to reset timeout for long-running jobs
                        job_progress = {
                            "job_id": current_job["job_id"],
                            "assignment_id": current_job["assignment_id"],
                            "progress_percentage": current_job.get("progress", 0),
                            "message": current_job.get("status_message", "Processing"),
                        }

                    self.send_worker_heartbeat(
                        status="busy" if current_job else "idle",
                        active_jobs=1 if current_job else 0,
                        cpu_usage=0,  # TODO: Get actual metrics
                        memory_usage=0,
                        job_progress=job_progress
                    )
                    last_heartbeat = time.time()

                # Get available jobs (not assigned)
                available_jobs = self.get_available_jobs(limit=5)

                if available_jobs:
                    self.logger.info(f"Found {len(available_jobs)} available jobs")

                    # Try to claim jobs in priority order
                    for job_info in available_jobs:
                        job_id = job_info["job_id"]

                        # Attempt to claim this job
                        job_data = self.claim_job(job_id)

                        if job_data:
                            # We got the job!
                            self.logger.info(f"Processing job {job_id}")
                            current_job = job_data
                            current_job["progress"] = 0

                            # Process the job
                            start_time = time.time()
                            try:
                                result = self.process_job(job_data)
                                execution_time = time.time() - start_time

                                # Submit result (encrypted or plaintext)
                                is_encrypted = job_data.get("is_encrypted", False)
                                if self.complete_job(
                                    job_id, result, execution_time,
                                    encrypted=is_encrypted, job_data=job_data if is_encrypted else None
                                ):
                                    self.logger.info(f"Successfully completed job: {job_id}")
                                    current_job = None
                                    # Short delay before looking for next job
                                    time.sleep(2)
                                else:
                                    self.logger.error(f"Failed to submit job result: {job_id}")

                            except Exception as e:
                                # Report failure
                                self.logger.error(f"Job processing failed: {e}")
                                self.fail_job(job_id, str(e), traceback.format_exc())
                                current_job = None

                            # After processing one job, check for more
                            break
                        else:
                            # Someone else claimed this job, try next
                            continue
                else:
                    self.logger.debug("No available jobs")

                # Wait before next poll
                time.sleep(poll_interval)

            except KeyboardInterrupt:
                self.logger.info("Worker stopped by user")
                break
            except Exception as e:
                self.logger.error(f"Unexpected error: {e}")
                time.sleep(poll_interval)

# Usage
if __name__ == "__main__":
    import os

    logging.basicConfig(level=logging.INFO)

    worker = MicroHubWorker(
        server_url=os.getenv("MICROHUB_SERVER_URL", "http://localhost:8000"),
        worker_token=os.getenv("WORKER_TOKEN")
    )

    # Use claim-based system with fast polling
    worker.run(poll_interval=10)

Important Notes

  1. Execution Time: Always track and report execution time accurately for billing purposes

  2. Error Handling: Always use try-catch blocks and report failures properly to avoid job getting stuck

  3. Polling Interval: Respect the next_poll value in responses to avoid overloading the server

  4. Token Management: Keep your worker token secure and never share it

  5. Result Format: Ensure your output matches the expected format for the job type

  6. Timeouts: Respect the timeout value provided with each job

  7. Idempotency: Jobs may be retried, ensure your processing is idempotent where possible

  8. Encrypted Jobs: Check is_encrypted on every claimed job. Encrypted jobs require the cryptography Python package. Always report input_tokens and output_tokens for encrypted jobs since the server cannot count tokens from encrypted payloads. Never log or persist decrypted payloads or encryption keys.

HTTP Status Codes

  • 200: Success
  • 401: Authentication failed (invalid or expired token)
  • 404: Job not found or not assigned to worker
  • 400: Bad request (invalid data format)
  • 429: Rate limited (too many requests)
  • 500: Server error

Rate Limits

  • Poll requests: Max 60 per minute
  • Job completions: No limit
  • Other endpoints: 100 requests per minute

Model Management

Model Registration and Approval Process

Workers can register support for models, which go through an approval process:

  1. Worker Registration: When registering, workers can specify models they support
  2. Model Request: Models start with status REQUESTED
  3. Admin Review: Admins review model requests in the /model-requests UI
  4. Approval Process:
  5. Admin clicks "Approve" on a model request
  6. Fills in configuration (size, requirements, pricing)
  7. Can use quick presets (Small/Medium/Large/XL)
  8. Model becomes AVAILABLE and appears on /models page
  9. Worker Support: Once approved, workers can process jobs for that model

Model Status Values

  • REQUESTED: Initial request from user/worker
  • PENDING: Awaiting admin review
  • IN_REVIEW: Admin is actively reviewing
  • APPROVED: Approved with configuration (transitions to AVAILABLE)
  • AVAILABLE: Active and available for use
  • REJECTED: Request was rejected with reason
  • UNAVAILABLE: Model is offline/not currently hosted
  • DEPRECATED: No longer supported
  • BETA: In beta testing

Support

For issues or questions, contact support@microdc.ai or visit the documentation at https://docs.microdc.ai