stream3r/worker/config.py

"""Configuration helpers for the STream3R RQ worker."""

from __future__ import annotations

from dataclasses import dataclass, field
import os
from pathlib import Path


def _env_bool(name: str, default: bool) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    value = value.strip().lower()
    if value in {"1", "true", "yes", "y", "on"}:
        return True
    if value in {"0", "false", "no", "n", "off"}:
        return False
    return default


def _env_bool_any(default: bool, *names: str) -> bool:
    for name in names:
        if not name:
            continue
        value = os.getenv(name)
        if value is None:
            continue
        value = value.strip().lower()
        if value in {"1", "true", "yes", "y", "on"}:
            return True
        if value in {"0", "false", "no", "n", "off"}:
            return False
    return default


def _env_int(name: str, default: int) -> int:
    value = os.getenv(name)
    if value is None:
        return default
    try:
        return int(value)
    except ValueError:
        return default


def _env_value(primary: str, *aliases: str, default: str | None = None) -> str | None:
    """Best-effort environment lookup with fallbacks for legacy names."""

    names = (primary, *aliases)
    for name in names:
        if not name:
            continue
        value = os.getenv(name)
        if value:
            return value
    return default


@dataclass(slots=True)
class WorkerSettings:
    """Runtime configuration derived from environment variables."""

    redis_url: str = "redis://localhost:6379/0"
    redis_events_stream: str = "stream3r:events"
    redis_stream_maxlen: int = 50_000
    redis_healthcheck_interval: int = 30

    pose_queue: str = "pose_pointmap"
    model_queue: str = "model_build"
    keyframe_queue: str = "keyframe_selection"

    gpu_lock_key: str = "gpu:lock"
    gpu_lock_timeout: int = 3600
    gpu_lock_blocking_timeout: int = 600

    storage_prefix: str = "scenes"
    s3_bucket: str | None = None
    s3_endpoint_url: str | None = None
    s3_region: str | None = None
    s3_profile: str | None = None
    s3_force_path_style: bool = False
    aws_access_key_id: str | None = None
    aws_secret_access_key: str | None = None
    aws_session_token: str | None = None
    local_storage_root: Path = field(default_factory=lambda: Path("storage"))

    db_dsn: str | None = None
    job_table: str = "stream3r_jobs"

    model_id: str = "yslan/STream3R"
    model_revision: str | None = None
    model_device_preference: str | None = None
    model_dtype: str | None = None

    default_mode: str = "window"
    default_streaming: bool = True
    download_workers: int = 4

    worker_name: str | None = None

    session_cache_filename: str = "kv_cache.pt"
    predictions_filename: str = "predictions.npz"
    poses_filename: str = "poses.jsonl"
    result_filename: str = "summary.json"
    selected_frames_filename: str = "selected_frames.json"
    scene_glb_filename: str = "scene.glb"
    session_settings_filename: str = "session_settings.json"
    summary_results_filename: str = "summary.json"

    pointmap_dir: str = "pointmaps"
    models_dir: str = "models"
    results_dir: str = "results"

    scene_media_api_base_url: str | None = None
    scene_media_api_token: str | None = None
    scene_media_api_secret: str | None = None
    scene_media_page_size: int = 200
    stream_window_size: int = 14
    max_frames_per_job: int = 0
    default_job_timeout: int = 45 * 60
    upload_session_cache: bool = True
    keyframe_prepass_enabled: bool = True
    keyframe_motion_threshold: float = 0.4
    keyframe_rotation_weight: float = 0.5
    keyframe_min_gap_frames: int = 2
    keyframe_max_gap_frames: int = 45
    keyframe_default_top_k: int = 16
    keyframe_coverage_confidence: float = 0.3
    keyframe_coverage_voxel_size: float = 0.05
    keyframe_coverage_max_points: int = 5000
    keyframe_min_gain_ratio: float = 0.01
    keyframe_full_mode_max_frames: int = 24
    keyframe_extract_fps: float = 6.0
    keyframe_extract_max_frames: int = 1200
    keyframe_upload_dir: str = "images/keyframes"

    @classmethod
    def from_env(cls) -> "WorkerSettings":
        base = cls()

        kwargs: dict[str, object] = {
            "redis_url": _env_value("STREAM3R_REDIS_URL", "REDIS_URL", default=base.redis_url),
            "redis_events_stream": os.getenv("STREAM3R_EVENTS_STREAM", base.redis_events_stream),
            "redis_stream_maxlen": _env_int("STREAM3R_EVENTS_MAXLEN", base.redis_stream_maxlen),
            "redis_healthcheck_interval": _env_int(
                "STREAM3R_REDIS_HEALTHCHECK", base.redis_healthcheck_interval
            ),
            "pose_queue": os.getenv("STREAM3R_QUEUE_POSE", base.pose_queue),
            "model_queue": os.getenv("STREAM3R_QUEUE_MODEL", base.model_queue),
            "keyframe_queue": os.getenv("STREAM3R_QUEUE_KEYFRAME", base.keyframe_queue),
            "gpu_lock_key": os.getenv("STREAM3R_GPU_LOCK_KEY", base.gpu_lock_key),
            "gpu_lock_timeout": _env_int("STREAM3R_GPU_LOCK_TIMEOUT", base.gpu_lock_timeout),
            "gpu_lock_blocking_timeout": _env_int(
                "STREAM3R_GPU_LOCK_BLOCK", base.gpu_lock_blocking_timeout
            ),
            "storage_prefix": os.getenv("STREAM3R_STORAGE_PREFIX", base.storage_prefix),
            "s3_bucket": _env_value("STREAM3R_STORAGE_BUCKET", "S3_BUCKET", default=base.s3_bucket) or None,
            "s3_endpoint_url": _env_value("STREAM3R_S3_ENDPOINT", "S3_ENDPOINT", default=base.s3_endpoint_url) or None,
            "s3_region": _env_value("STREAM3R_S3_REGION", "AWS_REGION", default=base.s3_region) or None,
            "s3_profile": os.getenv("STREAM3R_S3_PROFILE", base.s3_profile or "") or None,
            "s3_force_path_style": _env_bool_any(
                base.s3_force_path_style,
                "STREAM3R_S3_FORCE_PATH",
                "S3_FORCE_PATH_STYLE",
            ),
            "aws_access_key_id": _env_value("AWS_ACCESS_KEY_ID", default=base.aws_access_key_id) or None,
            "aws_secret_access_key": _env_value("AWS_SECRET_ACCESS_KEY", default=base.aws_secret_access_key) or None,
            "aws_session_token": _env_value("AWS_SESSION_TOKEN", default=base.aws_session_token) or None,
            "local_storage_root": Path(
                os.getenv("STREAM3R_LOCAL_STORAGE", str(base.local_storage_root))
            ).resolve(),
            "db_dsn": _env_value("STREAM3R_DB_DSN", "DATABASE_URL", default=base.db_dsn) or None,
            "job_table": os.getenv("STREAM3R_JOBS_TABLE", base.job_table),
            "model_id": os.getenv("STREAM3R_MODEL_ID", base.model_id),
            "model_revision": os.getenv("STREAM3R_MODEL_REVISION", base.model_revision or "") or None,
            "model_device_preference": os.getenv(
                "STREAM3R_MODEL_DEVICE", base.model_device_preference or ""
            )
            or None,
            "model_dtype": os.getenv("STREAM3R_MODEL_DTYPE", base.model_dtype or "") or None,
            "default_mode": os.getenv("STREAM3R_DEFAULT_MODE", base.default_mode),
            "default_streaming": _env_bool(
                "STREAM3R_DEFAULT_STREAMING", base.default_streaming
            ),
            "download_workers": _env_int("STREAM3R_DOWNLOAD_WORKERS", base.download_workers),
            "worker_name": os.getenv("STREAM3R_WORKER_NAME", base.worker_name or "") or None,
            "session_cache_filename": os.getenv(
                "STREAM3R_SESSION_CACHE", base.session_cache_filename
            ),
            "predictions_filename": os.getenv(
                "STREAM3R_PREDICTIONS_FILE", base.predictions_filename
            ),
            "poses_filename": os.getenv("STREAM3R_POSES_FILE", base.poses_filename),
            "result_filename": os.getenv("STREAM3R_RESULT_FILE", base.result_filename),
            "selected_frames_filename": os.getenv(
                "STREAM3R_SELECTED_FRAMES_FILE", base.selected_frames_filename
            ),
            "scene_glb_filename": os.getenv("STREAM3R_SCENE_GLB_FILE", base.scene_glb_filename),
            "session_settings_filename": os.getenv(
                "STREAM3R_SESSION_SETTINGS_FILE", base.session_settings_filename
            ),
            "summary_results_filename": os.getenv(
                "STREAM3R_SUMMARY_FILE", base.summary_results_filename
            ),
            "pointmap_dir": os.getenv("STREAM3R_POINTMAP_DIR", base.pointmap_dir),
            "models_dir": os.getenv("STREAM3R_MODELS_DIR", base.models_dir),
            "results_dir": os.getenv("STREAM3R_RESULTS_DIR", base.results_dir),
            "scene_media_api_base_url": _env_value(
                "STREAM3R_MEDIA_API_BASE_URL",
                "API_BASE_URL",
                default=base.scene_media_api_base_url,
            )
            or None,
            "scene_media_api_token": _env_value(
                "STREAM3R_MEDIA_API_TOKEN",
                "MEDIA_API_TOKEN",
                default=base.scene_media_api_token,
            )
            or None,
            "scene_media_api_secret": _env_value(
                "STREAM3R_MEDIA_API_SECRET",
                "MEDIA_API_SECRET",
                "INTERNAL_NOTIFY_SECRET",
                default=base.scene_media_api_secret,
            )
            or None,
            "scene_media_page_size": _env_int(
                "STREAM3R_MEDIA_PAGE_SIZE", base.scene_media_page_size
            ),
            "stream_window_size": _env_int(
                "STREAM3R_WINDOW_SIZE", base.stream_window_size
            ),
            "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
            ),
            "upload_session_cache": _env_bool(
                "STREAM3R_UPLOAD_CACHE", base.upload_session_cache
            ),
            "keyframe_prepass_enabled": _env_bool(
                "STREAM3R_KEYFRAME_PREPASS", base.keyframe_prepass_enabled
            ),
            "keyframe_motion_threshold": float(
                os.getenv("STREAM3R_KEYFRAME_MOTION_THRESH", base.keyframe_motion_threshold)
            ),
            "keyframe_rotation_weight": float(
                os.getenv("STREAM3R_KEYFRAME_ROT_WEIGHT", base.keyframe_rotation_weight)
            ),
            "keyframe_min_gap_frames": _env_int(
                "STREAM3R_KEYFRAME_MIN_GAP_FRAMES", base.keyframe_min_gap_frames
            ),
            "keyframe_max_gap_frames": _env_int(
                "STREAM3R_KEYFRAME_MAX_GAP_FRAMES", base.keyframe_max_gap_frames
            ),
            "keyframe_default_top_k": _env_int(
                "STREAM3R_KEYFRAME_TOP_K", base.keyframe_default_top_k
            ),
            "keyframe_coverage_confidence": float(
                os.getenv("STREAM3R_KEYFRAME_CONF_THRESH", base.keyframe_coverage_confidence)
            ),
            "keyframe_coverage_voxel_size": float(
                os.getenv("STREAM3R_KEYFRAME_VOXEL_SIZE", base.keyframe_coverage_voxel_size)
            ),
            "keyframe_coverage_max_points": _env_int(
                "STREAM3R_KEYFRAME_MAX_POINTS", base.keyframe_coverage_max_points
            ),
            "keyframe_min_gain_ratio": float(
                os.getenv("STREAM3R_KEYFRAME_MIN_GAIN", base.keyframe_min_gain_ratio)
            ),
            "keyframe_full_mode_max_frames": _env_int(
                "STREAM3R_KEYFRAME_FULL_MAX_FRAMES", base.keyframe_full_mode_max_frames
            ),
            "keyframe_extract_fps": float(
                os.getenv("STREAM3R_KEYFRAME_EXTRACT_FPS", base.keyframe_extract_fps)
            ),
            "keyframe_extract_max_frames": _env_int(
                "STREAM3R_KEYFRAME_EXTRACT_MAX_FRAMES", base.keyframe_extract_max_frames
            ),
            "keyframe_upload_dir": os.getenv(
                "STREAM3R_KEYFRAME_UPLOAD_DIR", base.keyframe_upload_dir
            ),
        }

        return cls(**kwargs)