task.py

"""Per-request task isolation envelope.

`GenerationTask` is a thin, pickle-safe dataclass that carries just enough
metadata to log/diagnose a single `generate()` call. The `_TASK_CTX` ContextVar
threads the active task through any code path that needs it — including the
HTTP-yielding `encode_remote` call inside `_generate_impl`, where a module
global would race.

`task_id` is a `uuid4().hex` — opaque, unguessable, and unique per request.
The work_dir is task-scoped under `<_TASK_TEMP_ROOT>/h3-<task_id>/` so concurrent
requests never collide on fixed filenames.
"""
from __future__ import annotations

import contextvars
import os
import tempfile
import time
import uuid
from dataclasses import dataclass, field
from typing import Optional

from state import _log


# Per-request task context — ContextVar is isolated across concurrent Gradio handlers
# in the same process (unlike a module global, which races on encode_remote HTTP yield).
_TASK_CTX: contextvars.ContextVar["GenerationTask | None"] = contextvars.ContextVar(
    "h3_generation_task", default=None
)


# Task-scoped temp root. All intermediates and outputs for a request live under
# <temp>/h3-<task_id>/ so concurrent requests never collide on fixed filenames.
_TASK_TEMP_ROOT = os.environ.get(
    "H3_TASK_TEMP_ROOT", os.path.join(tempfile.gettempdir(), "h3-tasks")
)


@dataclass
class GenerationTask:
    """Lightweight, pickle-safe task envelope for isolation and diagnostics.

    Holds only primitives / small serializable data. Never stores tensors, models,
    CUDA objects, locks, open files, or Gradio objects.
    """
    task_id: str
    created_at: float
    owner_id: Optional[str] = None
    # Request parameters (lightweight)
    prompt_len: int = 0
    prompt_hash: str = ""
    canvas: str = ""
    duration: float = 0.0
    steps: int = 0
    seed: int = -1
    turbo: bool = False
    nsfw_keys: tuple = ()
    custom_lora_url: str = ""
    # Lifecycle
    status: str = "created"  # created|validating|loading_loras|conditioning|gpu_running|finalizing|completed|failed|cancelled
    # Paths (task-local)
    work_dir: str = ""
    output_path: str = ""
    # Applied adapters (authoritative for this task; do not use global *_merged flags)
    applied_loras: list = field(default_factory=list)

    @staticmethod
    def create(owner_id: Optional[str] = None, **kwargs) -> "GenerationTask":
        tid = uuid.uuid4().hex
        work = os.path.join(_TASK_TEMP_ROOT, f"h3-{tid}")
        os.makedirs(work, exist_ok=True)
        os.makedirs(os.path.join(work, "intermediates"), exist_ok=True)
        return GenerationTask(
            task_id=tid,
            created_at=time.time(),
            owner_id=owner_id,
            work_dir=work,
            **kwargs,
        )


def _task_log(task: GenerationTask | None, msg: str, level: str = "INFO") -> None:
    prefix = f"task={task.task_id[:12]}" if task else "task=?"
    _log(f"{prefix} {msg}", level)