working for larger batches now

notes.md

@@ -0,0 +1,4 @@
+
+
+**Manually Clear the GPU Lock from REDIS:
+ redis-cli -u "$REDIS_URL" DEL gpu:lock

stream3r/worker/__init__.py

@@ -1,6 +1,14 @@
 """Worker utilities for running STream3R jobs via RQ."""
 
-from .tasks import model_build_job, pose_pointmap_job
+from rq.queue import Queue
+
+from .config import WorkerSettings
+
+_settings = WorkerSettings.from_env()
+if _settings.default_job_timeout and _settings.default_job_timeout > 0:
+    Queue.DEFAULT_TIMEOUT = _settings.default_job_timeout
+
+from .tasks import model_build_job, pose_pointmap_job  # noqa: E402
 
 __all__ = [
     "pose_pointmap_job",

stream3r/worker/config.py

@@ -114,8 +114,9 @@ class WorkerSettings:
     scene_media_api_base_url: str | None = None
     scene_media_api_token: str | None = None
     scene_media_page_size: int = 200
-    stream_window_size: int = 10
-    max_frames_per_job: int = 10
+    stream_window_size: int = 20
+    max_frames_per_job: int = 0
+    default_job_timeout: int = 15 * 60
 
     @classmethod
     def from_env(cls) -> "WorkerSettings":
@@ -208,6 +209,9 @@ class WorkerSettings:
             "max_frames_per_job": _env_int(
                 "STREAM3R_MAX_FRAMES", base.max_frames_per_job
             ),
+            "default_job_timeout": _env_int(
+                "STREAM3R_JOB_TIMEOUT", base.default_job_timeout
+            ),
         }
 
         return cls(**kwargs)

stream3r/worker/main.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import argparse
 import logging
+import os
 from typing import Sequence
 
 from rq import Queue, Worker
@@ -40,8 +41,42 @@ def _parse_args(default_queues: Sequence[str]) -> argparse.Namespace:
     return args
 
 
+class Stream3RWorker(Worker):
+    """RQ worker that enforces configured default timeouts before execution."""
+
+    def __init__(self, *args, default_timeout: int | None = None, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._default_timeout = default_timeout if default_timeout and default_timeout > 0 else None
+
+    def _normalize_timeout(self, timeout_value: object) -> int | None:
+        if timeout_value is None:
+            return None
+        try:
+            return int(timeout_value)
+        except (TypeError, ValueError):
+            return None
+
+    def execute_job(self, job, queue):  # type: ignore[override]
+        if self._default_timeout is not None:
+            current = self._normalize_timeout(getattr(job, "timeout", None))
+            if current is None or current < self._default_timeout:
+                job.timeout = self._default_timeout
+        runtime = get_runtime()
+        lock_ctx = runtime.gpu_lock()
+        self.log.debug("Worker %s acquiring GPU lock for job %s", self.name, job.id)
+        with lock_ctx:
+            os.environ["STREAM3R_GPU_LOCK_HELD"] = "1"
+            try:
+                return super().execute_job(job, queue)
+            finally:
+                os.environ.pop("STREAM3R_GPU_LOCK_HELD", None)
+
+
 def main() -> None:
     settings = WorkerSettings.from_env()
+    if settings.default_job_timeout and settings.default_job_timeout > 0:
+        Queue.DEFAULT_TIMEOUT = settings.default_job_timeout
+
     args = _parse_args([settings.pose_queue, settings.model_queue])
     logging.basicConfig(level=getattr(logging, str(args.log_level).upper(), logging.INFO))
 
@@ -51,7 +86,11 @@ def main() -> None:
     for queue in queues:
         logger.info("Listening on queue '%s'", queue.name)
 
-    worker = Worker(queues, name=settings.worker_name)
+    worker = Stream3RWorker(
+        queues,
+        name=settings.worker_name,
+        default_timeout=settings.default_job_timeout,
+    )
     worker.work(burst=args.burst)
 
 

stream3r/worker/pipeline.py

@@ -86,6 +86,7 @@ def run_stream3r_inference(
                 frame = images[idx : idx + 1].to(device)
                 with autocast_ctx:
                     session.forward_stream(frame)
+                print(f"Processed frame {idx + 1}/{total_frames}")
                 if progress_cb is not None:
                     progress_cb(idx + 1, total_frames)
 

stream3r/worker/tasks.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import base64
 import json
 import logging
+import os
 import re
 import shutil
 import tempfile
@@ -13,6 +14,7 @@ import uuid
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from pathlib import Path
+from contextlib import nullcontext
 from typing import Any, Callable, Mapping
 
 import numpy as np
@@ -797,16 +799,35 @@ def _execute_job(job_type: str, payload: Mapping[str, Any], handler: JobHandler)
             window_size = None
 
     payload["mode"] = mode
+    desired_timeout = runtime.settings.default_job_timeout
     timeout_override = payload.get("timeout")
+    applied_timeout: int | None = None
+
     if timeout_override is not None:
         try:
-            job.timeout = int(timeout_override)
+            applied_timeout = int(timeout_override)
+            if job is not None:
+                job.timeout = applied_timeout
         except (TypeError, ValueError):
-            pass
+            applied_timeout = None
+
+    if applied_timeout is None and desired_timeout and desired_timeout > 0:
+        if job is not None:
+            current_timeout = getattr(job, "timeout", None)
+            try:
+                current_timeout_value = int(current_timeout) if current_timeout is not None else None
+            except (TypeError, ValueError):
+                current_timeout_value = None
+            if current_timeout_value is None or current_timeout_value < desired_timeout:
+                job.timeout = desired_timeout
+                applied_timeout = desired_timeout
+            else:
+                applied_timeout = current_timeout_value
+        else:
+            applied_timeout = desired_timeout
 
-    # Default to 15 minutes if no timeout already applied
-    if job.timeout is None:
-        job.timeout = 15 * 60
+    if applied_timeout is not None:
+        payload["timeout"] = applied_timeout
 
     sanitized_payload = _sanitize_payload(payload)
 
@@ -834,8 +855,10 @@ def _execute_job(job_type: str, payload: Mapping[str, Any], handler: JobHandler)
         },
     )
 
+    lock_ctx = nullcontext() if os.getenv("STREAM3R_GPU_LOCK_HELD") == "1" else runtime.gpu_lock()
+
     try:
-        with runtime.gpu_lock():
+        with lock_ctx:
             with tempfile.TemporaryDirectory(prefix=f"stream3r_{job_id}_") as tmp_dir:
                 temp_path = Path(tmp_dir)
                 frame_records = _collect_frames(runtime, scene_id, payload, temp_path)