Spaces:
Running
Running
File size: 7,102 Bytes
fba6023 3493993 fba6023 3493993 fba6023 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | from __future__ import annotations
import asyncio
import os
import shutil
import time
from dataclasses import dataclass
from pathlib import Path
from uuid import UUID
from app.core.config import Settings
from app.core.exceptions import NotFoundError, ProcessingError
from app.core.logger import get_logger
logger = get_logger(__name__)
@dataclass(frozen=True, slots=True)
class RequestWorkspace:
request_id: str
root: Path
uploads: Path
outputs: Path
logs: Path
class CleanupService:
"""Owns per-request directories and expires completed or abandoned work."""
def __init__(self, settings: Settings) -> None:
self.settings = settings
self._active: set[str] = set()
self._lock = asyncio.Lock()
async def create_workspace(self, request_id: str) -> RequestWorkspace:
self._validate_request_id(request_id)
root = self.settings.temp_dir / request_id
workspace = RequestWorkspace(
request_id=request_id,
root=root,
uploads=root / "uploads",
outputs=root / "outputs",
logs=root / "logs",
)
for directory in (workspace.uploads, workspace.outputs, workspace.logs):
directory.mkdir(parents=True, exist_ok=True)
async with self._lock:
self._active.add(request_id)
return workspace
async def complete(self, request_id: str) -> None:
async with self._lock:
self._active.discard(request_id)
for base in (self.settings.temp_dir, self.settings.output_dir):
path = base / request_id
if path.exists():
await asyncio.to_thread(os.utime, path, None)
async def publish(self, request_id: str, source: Path, filename: str) -> Path:
self._validate_request_id(request_id)
safe_name = Path(filename).name
if not safe_name or safe_name in {".", ".."}:
raise ProcessingError("The generated output filename is invalid")
destination_dir = self.settings.output_dir / request_id
destination_dir.mkdir(parents=True, exist_ok=True)
destination = destination_dir / safe_name
try:
await asyncio.to_thread(os.replace, source, destination)
except OSError:
await asyncio.to_thread(shutil.move, str(source), str(destination))
return destination
async def publish_new(self, request_id: str, source: Path, filename: str) -> Path | None:
"""Publish a generated output without replacing an existing file.
This is used by durable generation-job reconciliation, where a retry
must never overwrite a canonical output created by an earlier worker
attempt. Existing media operations continue to use ``publish`` and
retain their established replacement semantics.
"""
self._validate_request_id(request_id)
safe_name = Path(filename).name
if not safe_name or safe_name in {".", ".."}:
raise ProcessingError("The generated output filename is invalid")
destination_dir = self.settings.output_dir / request_id
destination_dir.mkdir(parents=True, exist_ok=True)
destination = destination_dir / safe_name
created = await asyncio.to_thread(self._publish_new_sync, source, destination)
return destination if created else None
@staticmethod
def _publish_new_sync(source: Path, destination: Path) -> bool:
"""Atomically claim a destination, with a cross-device fallback."""
try:
os.link(source, destination)
except FileExistsError:
return False
except OSError:
# ``temp_dir`` and ``output_dir`` can be different mounts. An
# exclusive create still prevents overwrite in that arrangement.
try:
with source.open("rb") as input_stream, destination.open("xb") as output_stream:
shutil.copyfileobj(input_stream, output_stream, length=1024 * 1024)
except FileExistsError:
return False
except OSError:
# Never leave a partial file reachable from the output
# directory when cross-device publication fails.
try:
destination.unlink(missing_ok=True)
except OSError:
pass
raise
try:
source.unlink()
except FileNotFoundError:
pass
return True
def resolve_download(self, request_id: str, filename: str) -> Path:
self._validate_request_id(request_id)
if filename != Path(filename).name:
raise NotFoundError("Output file not found")
root = (self.settings.output_dir / request_id).resolve()
candidate = (root / filename).resolve()
if candidate.parent != root or not candidate.is_file():
raise NotFoundError("Output file not found")
return candidate
async def cleanup_expired(self) -> int:
cutoff = time.time() - self.settings.cleanup_minutes * 60
async with self._lock:
active = self._active.copy()
removed = 0
for base in (self.settings.temp_dir, self.settings.output_dir):
if not base.exists():
continue
for path in list(base.iterdir()):
if not path.is_dir() or path.name in active:
continue
try:
if path.stat().st_mtime < cutoff:
await asyncio.to_thread(shutil.rmtree, path)
removed += 1
logger.info("expired workspace removed", extra={"path": str(path)})
except FileNotFoundError:
continue
except OSError as exc:
logger.warning(
"workspace cleanup failed",
extra={"path": str(path), "error": str(exc)},
)
return removed
async def remove_request(self, request_id: str) -> None:
self._validate_request_id(request_id)
async with self._lock:
self._active.discard(request_id)
for base in (self.settings.temp_dir, self.settings.output_dir):
path = base / request_id
if path.is_dir():
await asyncio.to_thread(shutil.rmtree, path)
async def remove_temporary_request(self, request_id: str) -> None:
"""Remove only bounded staging data while retaining published output."""
self._validate_request_id(request_id)
async with self._lock:
self._active.discard(request_id)
path = self.settings.temp_dir / request_id
if path.is_dir():
await asyncio.to_thread(shutil.rmtree, path)
@staticmethod
def _validate_request_id(request_id: str) -> None:
try:
UUID(request_id)
except (ValueError, AttributeError) as exc:
raise NotFoundError("Output file not found") from exc
|