43 lines
1.6 KiB
Python
43 lines
1.6 KiB
Python
"""Utilities for managing scratch space used by frame sampling."""
|
|
|
|
import logging
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ScratchManager:
|
|
"""Create and clean per-video frame scratch directories."""
|
|
|
|
def __init__(self, base_path: str = "/scratch", video_id: Optional[str] = None, auto_cleanup: bool = True):
|
|
self.base_path = Path(base_path)
|
|
self.video_id = video_id or "unknown"
|
|
self.auto_cleanup = auto_cleanup
|
|
self.frame_dir = self.base_path / self.video_id / "frames"
|
|
|
|
def ensure_frame_dir(self) -> Path:
|
|
"""Create the per-video scratch directory if it does not exist."""
|
|
self.frame_dir.mkdir(parents=True, exist_ok=True)
|
|
return self.frame_dir
|
|
|
|
def cleanup(self) -> None:
|
|
"""Remove the frame scratch directory only."""
|
|
if self.auto_cleanup and self.frame_dir.exists():
|
|
shutil.rmtree(self.frame_dir)
|
|
logger.info("Cleaned scratch frames for %s", self.video_id)
|
|
|
|
def cleanup_all(self) -> None:
|
|
"""Remove the entire per-video scratch directory (frames + crops)."""
|
|
video_root = self.base_path / self.video_id
|
|
if self.auto_cleanup and video_root.exists():
|
|
shutil.rmtree(video_root)
|
|
logger.info("Cleaned all scratch data for %s", self.video_id)
|
|
|
|
def usage_bytes(self) -> int:
|
|
"""Return the total size of the scratch directory in bytes."""
|
|
if not self.frame_dir.exists():
|
|
return 0
|
|
return sum(path.stat().st_size for path in self.frame_dir.rglob("*") if path.is_file())
|