Spaces:
Running on Zero
Running on Zero
File size: 2,960 Bytes
64db20c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """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)
|