Skip to content

GPU Worker (VLM2Vec + Qwen2-VL)

Metadata

  • System type: service

System Intent

  • What this is: An SQS pull-mode worker that runs on an EC2 GPU instance and performs two model tasks — a visual embedding model for semantic video search and a captioning/generation model for scene description, NER, and triple extraction. It has zero inbound network exposure: it long-polls the encache-gpu-requests SQS queue, fetches job payloads from S3 (claim-check pattern — base64 frames exceed SQS's 256KB message limit), runs inference with the already-loaded models, and writes results back to S3. It is the GPU-side counterpart to VLM2VecClient and GPULLMClient, which submit jobs and poll for results instead of making HTTP calls.

Architecture

  • Entry point: main/server/worldmm/gpu_worker/sqs_worker.py (main()), run under systemd as the vlm2vec unit.
  • main/server/worldmm/gpu_worker/server.py no longer serves traffic in production. sqs_worker.py imports it only for load_models(), _start_watchdog(), _touch_activity(), and the four handler functions (encode_video, encode_text, caption, generate), calling them directly as Python functions rather than over HTTP. The FastAPI app in server.py is kept only as a local-dev harness (see its module docstring).
  • Job protocol: main/server/worldmm/gpu_worker/protocol.py — the shared claim-check contract used by both the Lambda callers (VLM2VecClient in main/server/worldmm/memory/visual/encoder.py) and the worker.
  • No inbound ports: the worker's security group (gpu-worker-sg, main/devops/main.tf) has egress only — no ingress rules at all. All coordination happens through the encache-gpu-requests SQS queue and gpu-jobs/ S3 prefixes.

Mermaid Diagram

flowchart TD
  Caller["API Lambda / Ingest Lambda\n(VLM2VecClient)"] -->|"submit_job(): put gpu-jobs/requests/{id}.json\n+ SQS send_message"| Queue["SQS: encache-gpu-requests"]
  Queue -->|"long-poll receive_message\nWaitTimeSeconds=20"| Worker["sqs_worker.process_message()"]
  Worker -->|"fetch_job(): get gpu-jobs/requests/{id}.json"| S3Req["S3: gpu-jobs/requests/"]
  Worker -->|"dispatch by op"| EV["encode_video()"]
  Worker -->|"dispatch by op"| ET["encode_text()"]
  Worker -->|"dispatch by op"| CAP["caption()"]
  Worker -->|"dispatch by op"| GEN["generate()"]

  EV --> EM["_embed_model\nQwen2-VL-2B-Instruct\n+ TIGER-Lab/VLM2Vec-Qwen2VL-2B LoRA"]
  ET --> EM
  CAP --> CM["_caption_model\nQwen/Qwen2-VL-7B-Instruct"]
  GEN --> CM

  EM -->|result| WriteResult["write_result(): put gpu-jobs/results/{id}.json"]
  CM -->|result| WriteResult
  Worker -->|"delete_message() on success\n(deterministic errors also delete — no redrive)"| Queue
  Worker -.->|"transient failure before delete →\nredelivered; 3rd receive → DLQ"| DLQ["SQS DLQ: encache-gpu-requests-dlq"]

  Caller -->|"poll_result(): get gpu-jobs/results/{id}.json"| S3Res["S3: gpu-jobs/results/"]

  Worker -->|"write_heartbeat() every 30s\n(only after models loaded)"| Heartbeat["S3: gpu-jobs/heartbeat.json"]
  Caller -->|"heartbeat_age_s() < 90s?"| Heartbeat

Request Flow

  1. Caller (VLM2VecClient._submit) writes the job payload to gpu-jobs/requests/{request_id}.json and sends an SQS message {request_id, op} to encache-gpu-requests.
  2. The worker long-polls the queue (receive_message, WaitTimeSeconds=20). On a message it fetches the request object from S3, dispatches to the matching handler function (the same encode_video/encode_text/caption/generate functions server.py used to expose over HTTP), and writes the result to gpu-jobs/results/{request_id}.json.
  3. The caller polls for gpu-jobs/results/{request_id}.json until it appears or its per-op timeout elapses: encode_video 300s, encode_text 240s, caption 420s, generate 300s (JOB_TIMEOUTS_S in encoder.py).
  4. Failure semantics — deterministic vs. transient:
  5. Deterministic failures (bad input, model/inference error raised by a handler) are written as {"error": "<ExceptionType>: <message>"} to the result object, and the SQS message is deleted — redriving a deterministic failure would only waste queue slots and repeat the same error. The caller raises GpuJobError on this shape.
  6. Transient failures (an S3/SQS call itself fails, or a poison message with no parseable request_id) are not followed by delete_message — the in-flight message is left for SQS to redeliver after the 300s visibility timeout. After 3 receives it is routed to the encache-gpu-requests-dlq DLQ (14-day retention).
  7. gpu-jobs/ request and result objects expire after 1 day via an S3 lifecycle rule on encache-raw-memory (aws_s3_bucket_lifecycle_configuration.raw_data_gpu_jobs).

Health (heartbeat, not HTTP)

  • The worker writes an empty JSON object to gpu-jobs/heartbeat.json every 30 seconds (HEARTBEAT_INTERVAL_S in sqs_worker.py), starting only once load_models() has completed — so a fresh heartbeat means "ready for work," the same semantics the old GET /health had.
  • Callers treat the worker as healthy when protocol.heartbeat_age_s() < 90 seconds (HEARTBEAT_FRESH_S in encoder.py — allows 3 missed heartbeats before declaring the worker dead). Both VLM2VecClient._is_healthy and the chat Lambda's _gpu_is_healthy() (main/server/api/memories/chat/app.py) check this instead of making an HTTP request.
  • The FastAPI app's GET /health route still exists for local-dev use only — it is not reachable in production (no inbound port).

No Inbound Ports

  • aws_security_group.gpu_worker in main/devops/main.tf has an egress-only rule set — no ingress block at all.
  • There is no GPU_WORKER_TOKEN/shared-secret auth layer; there is nothing to authenticate against because the worker never accepts a connection.

Flows

The flows below describe the handler functions themselves (unchanged from the HTTP era). In production sqs_worker.py invokes them directly as Python function calls dispatched from an SQS message, not as HTTP requests; a raised HTTPException (e.g. "not-loaded", "no-frames") is caught by sqs_worker.process_message and reported as a deterministic {"error": ...} result object rather than an actual HTTP response. The FastAPI routes and status codes below only apply when running server.py directly as a local-dev harness.

Flow: load_models

  • Core files: main/server/worldmm/gpu_worker/server.py

Types

(no HTTP types — called internally at startup)

Paths

path input output path-type notes
load_models.success models in globals happy path logs GPU memory used
load_models.cuda-oom OOM error error requires g5.xlarge (24 GB VRAM)

Pseudocode

load Qwen/Qwen2-VL-7B-Instruct → _caption_model (fp16, device_map=auto)
load Qwen/Qwen2-VL-2B-Instruct base → base_2b (fp16, device_map=auto)
wrap base_2b with PeftModel.from_pretrained(base_2b, "TIGER-Lab/VLM2Vec-Qwen2VL-2B") → _embed_model
  NOTE: adapter must be VLM2Vec-Qwen2VL-2B (Qwen2-VL-compatible).
        VLM2Vec-LoRA targets Phi-3.5 and must NOT be used here (see bug #403).
load AutoProcessor for 7B → _processor
load AutoProcessor for 2B → _embed_processor
start keepalive thread during load to prevent idle-shutdown

Flow: encode_video

  • Core files: main/server/worldmm/gpu_worker/server.py

Types

EncodeVideoRequest {
  frames: list[string]  (base64-encoded JPEG)
}

EmbeddingResponse {
  embedding: list[float]  (1536-dim, last hidden state of final token)
}

Paths

path input output path-type notes
encode_video.success EncodeVideoRequest EmbeddingResponse happy path prompt: "Represent the given video clip."
encode_video.not-loaded EncodeVideoRequest HTTP 503 error model not yet loaded
encode_video.no-frames EncodeVideoRequest HTTP 400 error empty frames list

Flow: encode_text

  • Core files: main/server/worldmm/gpu_worker/server.py

Types

EncodeTextRequest {
  text: string
}

EmbeddingResponse {
  embedding: list[float]  (1536-dim)
}

Paths

path input output path-type notes
encode_text.success EncodeTextRequest EmbeddingResponse happy path prefix: "Represent the given query for retrieving relevant video clips: {text}"
encode_text.not-loaded EncodeTextRequest HTTP 503 error
encode_text.no-text EncodeTextRequest HTTP 400 error

Flow: caption

  • Core files: main/server/worldmm/gpu_worker/server.py

Types

CaptionRequest {
  frames: list[string]  (base64-encoded JPEG)
  transcript: string    (optional, audio transcript for context)
}

CaptionResponse {
  caption: string
}

Paths

path input output path-type notes
caption.success CaptionRequest CaptionResponse happy path max 512 new tokens
caption.not-loaded CaptionRequest HTTP 503 error
caption.no-frames CaptionRequest HTTP 400 error

Flow: generate

  • Core files: main/server/worldmm/gpu_worker/server.py

Types

GenerateRequest {
  messages: list[dict]    (OpenAI-style chat messages)
  max_new_tokens: int     (default 512)
}

GenerateResponse {
  text: string
}

Paths

path input output path-type notes
generate.success GenerateRequest GenerateResponse happy path used for NER, triple extraction
generate.not-loaded GenerateRequest HTTP 503 error

Models

Role Base model Adapter Processor
Embedding (_embed_model) Qwen/Qwen2-VL-2B-Instruct TIGER-Lab/VLM2Vec-Qwen2VL-2B (PEFT LoRA) Qwen/Qwen2-VL-2B-Instruct
Captioning/generation (_caption_model) Qwen/Qwen2-VL-7B-Instruct none Qwen/Qwen2-VL-7B-Instruct

The embedding LoRA must be VLM2Vec-Qwen2VL-2B. The adapter VLM2Vec-LoRA targets Phi-3.5 and is incompatible with Qwen2-VL-2B-Instruct (issue #403).

Idle Watchdog

The server tracks the last-activity timestamp in /tmp/vlm2vec_last_activity. A background thread checks every 30 seconds; if idle for more than 480 seconds the instance calls sudo shutdown -h now. Activity is touched on every request and during model loading.

Known SSM write gap: The watchdog shuts the instance down without updating the SSM parameter /encache/gpu/instance_id. Callers (chat Lambda, ingest Lambda) that read the stale ID will address traffic to the terminated instance. Both Lambdas include a tag-based fallback — they scan EC2 for a running instance tagged Name=encache-gpu-worker and update SSM when they find one — so the stale ID is corrected on the first invocation that hits this path. See memories-chat.md and ingest-window.md for the full recovery logic.

Logs

Source Location
GPU worker stdout systemd journal on the EC2 instance (journalctl -u vlm2vec -f)

Deployment

  • Mechanism: EC2 (g5.xlarge, Deep Learning AMI — PyTorch 2.8, Amazon Linux 2023)
  • AMI: ami-0e72acaa1863957cd
  • Deploy command:
    # Upload server code to S3, then launch from launch template or run user data manually:
    aws s3 cp main/server/worldmm/gpu_worker/server.py s3://encache-raw-memory/gpu-worker/server.py
    aws s3 cp main/server/worldmm/gpu_worker/protocol.py s3://encache-raw-memory/gpu-worker/protocol.py
    aws s3 cp main/server/worldmm/gpu_worker/sqs_worker.py s3://encache-raw-memory/gpu-worker/sqs_worker.py
    # EC2 user data at main/server/worldmm/gpu_worker/ec2_user_data.sh installs deps and starts the vlm2vec systemd service automatically.
    
  • Notes: No inbound ports — the worker has an egress-only security group (gpu-worker-sg). Callers resolve the running instance via EC2 describe_instances using GPU_INSTANCE_ID (or the tag-based fallback if SSM is stale) to decide whether to start/launch it, but coordination with the worker itself is entirely through the encache-gpu-requests SQS queue and gpu-jobs/ S3 prefixes (see Architecture, above) — not a direct HTTP connection to the instance. The instance is started on demand by the API Lambda if stopped; a new instance is launched from the launch template if terminated.