SignerX/SignVerse-2M

1from __future__ import annotations2 3import contextlib4import fcntl5from pathlib import Path6from typing import Iterable, Iterator, Sequence7 8VIDEO_EXTENSIONS = {".mp4", ".mkv", ".webm", ".mov"}9 10 11def existing_raw_dirs(*dirs: Path | None) -> list[Path]:12    result: list[Path] = []13    seen: set[str] = set()14    for directory in dirs:15        if directory is None:16            continue17        key = str(directory)18        if key in seen:19            continue20        seen.add(key)21        result.append(directory)22    return result23 24 25def collect_raw_videos(*dirs: Path | None) -> dict[str, Path]:26    videos: dict[str, Path] = {}27    for directory in existing_raw_dirs(*dirs):28        if not directory.exists():29            continue30        for path in sorted(directory.iterdir()):31            if not path.is_file() or path.suffix.lower() not in VIDEO_EXTENSIONS:32                continue33            videos.setdefault(path.stem, path)34    return videos35 36 37def count_raw_videos(*dirs: Path | None) -> int:38    return len(collect_raw_videos(*dirs))39 40 41def sum_raw_video_sizes(*dirs: Path | None) -> int:42    return sum(path.stat().st_size for path in collect_raw_videos(*dirs).values() if path.exists())43 44 45def iter_raw_video_files(*dirs: Path | None) -> Iterator[Path]:46    for path in collect_raw_videos(*dirs).values():47        yield path48 49 50def find_video_file(video_id: str, *dirs: Path | None) -> Path | None:51    for directory in existing_raw_dirs(*dirs):52        if not directory.exists():53            continue54        candidates = []55        for path in directory.glob(f"{video_id}.*"):56            if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS:57                candidates.append(path)58        if candidates:59            return sorted(candidates)[0]60    return None61 62 63def iter_partial_download_files(video_id: str, *dirs: Path | None) -> Iterator[Path]:64    seen: set[Path] = set()65    for directory in existing_raw_dirs(*dirs):66        if not directory.exists():67            continue68        for path in directory.glob(f"{video_id}*"):69            if not path.is_file():70                continue71            suffixes = set(path.suffixes)72            if '.part' in suffixes or '.ytdl' in suffixes or path.suffix in {'.part', '.ytdl'}:73                resolved = path.resolve()74                if resolved in seen:75                    continue76                seen.add(resolved)77                yield path78 79 80def cleanup_partial_downloads(video_id: str, *dirs: Path | None) -> None:81    for partial_path in iter_partial_download_files(video_id, *dirs):82        partial_path.unlink(missing_ok=True)83 84 85def _count_reservations(reservation_dir: Path | None, pool_name: str) -> int:86    if reservation_dir is None:87        return 088    pool_dir = reservation_dir / pool_name89    if not pool_dir.exists():90        return 091    return sum(1 for path in pool_dir.iterdir() if path.is_file() and path.suffix == ".reserve")92 93 94def _create_reservation(reservation_dir: Path | None, pool_name: str, reservation_key: str | None) -> Path | None:95    if reservation_dir is None or not reservation_key:96        return None97    pool_dir = reservation_dir / pool_name98    pool_dir.mkdir(parents=True, exist_ok=True)99    reservation_path = pool_dir / f"{reservation_key}.reserve"100    reservation_path.write_text(f"pool={pool_name}\nkey={reservation_key}\n", encoding="utf-8")101    return reservation_path102 103 104def release_download_reservation(reservation_path: Path | None) -> None:105    if reservation_path is not None:106        reservation_path.unlink(missing_ok=True)107 108 109def choose_download_target(110    primary_dir: Path,111    scratch_dir: Path | None,112    primary_limit: int,113    scratch_limit: int,114    reservation_dir: Path | None = None,115    reservation_key: str | None = None,116) -> tuple[Path, Path | None]:117    primary_dir.mkdir(parents=True, exist_ok=True)118    if reservation_dir is None:119        primary_count = count_raw_videos(primary_dir)120        if primary_count < primary_limit:121            return primary_dir, None122        if scratch_dir is None:123            raise RuntimeError(124                f"raw backlog full in primary pool ({primary_count}/{primary_limit}) and no scratch raw pool configured"125            )126        scratch_dir.mkdir(parents=True, exist_ok=True)127        scratch_count = count_raw_videos(scratch_dir)128        if scratch_count < scratch_limit:129            return scratch_dir, None130        raise RuntimeError(131            f"raw backlog full in both pools: primary {primary_count}/{primary_limit}, scratch {scratch_count}/{scratch_limit}"132        )133 134    reservation_dir.mkdir(parents=True, exist_ok=True)135    lock_path = reservation_dir / ".target_selection.lock"136    with lock_path.open("a+", encoding="utf-8") as lock_handle:137        fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX)138        try:139            primary_count = count_raw_videos(primary_dir) + _count_reservations(reservation_dir, "home")140            if primary_count < primary_limit:141                return primary_dir, _create_reservation(reservation_dir, "home", reservation_key)142            if scratch_dir is None:143                raise RuntimeError(144                    f"raw backlog full in primary pool ({primary_count}/{primary_limit}) and no scratch raw pool configured"145                )146            scratch_dir.mkdir(parents=True, exist_ok=True)147            scratch_count = count_raw_videos(scratch_dir) + _count_reservations(reservation_dir, "scratch")148            if scratch_count < scratch_limit:149                return scratch_dir, _create_reservation(reservation_dir, "scratch", reservation_key)150            raise RuntimeError(151                f"raw backlog full in both pools: primary {primary_count}/{primary_limit}, scratch {scratch_count}/{scratch_limit}"152            )153        finally:154            with contextlib.suppress(OSError):155                fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN)156 157 158def remove_video_files(video_id: str, *dirs: Path | None) -> None:159    for directory in existing_raw_dirs(*dirs):160        if not directory.exists():161            continue162        for path in directory.glob(f"{video_id}.*"):163            if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS:164                path.unlink(missing_ok=True)165