wehub-resource-sync
917eedffcf
Main / Python 3.11 - Docs (push) Waiting to run
Main / Python 3.11 - Build (push) Waiting to run
Main / Python 3.11 - Lint (push) Waiting to run
Main / Python 3.11 - Style (push) Waiting to run
Main / Python 3.11 - Test (push) Waiting to run
Main / GPU CI (push) Blocked by required conditions
Main / Release (push) Blocked by required conditions
Main / Build and Push Docker Images (push) Blocked by required conditions
474 lines
18 KiB
Python
474 lines
18 KiB
Python
import abc
|
|
import asyncio
|
|
import csv
|
|
import datetime
|
|
import hashlib
|
|
import io
|
|
import logging
|
|
import os
|
|
import random
|
|
from asyncio import Queue, QueueEmpty
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List, Optional, Set
|
|
|
|
import zstandard
|
|
|
|
from olmocr.s3_utils import (
|
|
download_zstd_csv,
|
|
expand_s3_glob,
|
|
parse_s3_path,
|
|
upload_zstd_csv,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Shared directory names for both local and S3 backends
|
|
WORKER_LOCKS_DIR = "worker_locks"
|
|
DONE_FLAGS_DIR = "done_flags"
|
|
|
|
|
|
@dataclass
|
|
class WorkItem:
|
|
"""Represents a single work item in the queue."""
|
|
|
|
hash: str
|
|
work_paths: List[str]
|
|
|
|
|
|
class Backend(abc.ABC):
|
|
"""Abstract backend for storage operations."""
|
|
|
|
@abc.abstractmethod
|
|
async def load_index_lines(self) -> List[str]:
|
|
"""Load raw index lines from storage."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def save_index_lines(self, lines: List[str]) -> None:
|
|
"""Save raw index lines to storage."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def get_completed_hashes(self) -> Set[str]:
|
|
"""Get set of completed work hashes."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def is_completed(self, work_hash: str) -> bool:
|
|
"""Check if a work item has been completed."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def is_worker_lock_taken(self, work_hash: str, worker_lock_timeout_secs: int = 1800) -> bool:
|
|
"""Check if a worker lock is taken and not stale."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def create_worker_lock(self, work_hash: str) -> None:
|
|
"""Create a worker lock for a work hash."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def delete_worker_lock(self, work_hash: str) -> None:
|
|
"""Delete the worker lock for a work hash if it exists."""
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
async def create_done_flag(self, work_hash: str) -> None:
|
|
"""Create a done flag for a work hash."""
|
|
pass
|
|
|
|
|
|
class WorkQueue:
|
|
"""
|
|
Manages a work queue with pluggable storage backends (e.g., local or S3).
|
|
"""
|
|
|
|
def __init__(self, backend: Backend):
|
|
self.backend = backend
|
|
self._queue: Queue[WorkItem] = Queue()
|
|
self._completed_hash_cache = set()
|
|
|
|
@staticmethod
|
|
def _encode_csv_row(row: List[str]) -> str:
|
|
"""Encodes a row of data for CSV storage with proper escaping."""
|
|
output = io.StringIO()
|
|
writer = csv.writer(output, quoting=csv.QUOTE_MINIMAL)
|
|
writer.writerow(row)
|
|
return output.getvalue().strip()
|
|
|
|
@staticmethod
|
|
def _decode_csv_row(line: str) -> List[str]:
|
|
"""Decodes a CSV row with proper unescaping."""
|
|
return next(csv.reader([line]))
|
|
|
|
@staticmethod
|
|
def _compute_workgroup_hash(work_paths: List[str]) -> str:
|
|
"""Compute a deterministic hash for a group of paths."""
|
|
sha1 = hashlib.sha1()
|
|
for path in sorted(work_paths):
|
|
sha1.update(path.encode("utf-8"))
|
|
return sha1.hexdigest()
|
|
|
|
def _parse_index_lines(self, lines: List[str]) -> Dict[str, List[str]]:
|
|
"""Parse index lines into a dict of hash to paths."""
|
|
result = {}
|
|
for line in lines:
|
|
if line.strip():
|
|
parts = self._decode_csv_row(line)
|
|
if parts:
|
|
result[parts[0]] = parts[1:]
|
|
return result
|
|
|
|
def _make_index_lines(self, groups: Dict[str, List[str]]) -> List[str]:
|
|
"""Create encoded lines from groups dict."""
|
|
return [self._encode_csv_row([group_hash] + group_paths) for group_hash, group_paths in groups.items()]
|
|
|
|
async def populate_queue(self, work_paths: List[str], items_per_group: int) -> None:
|
|
"""
|
|
Add new items to the work queue.
|
|
"""
|
|
all_paths = set(work_paths)
|
|
|
|
lines = await self.backend.load_index_lines()
|
|
existing_groups = self._parse_index_lines(lines)
|
|
|
|
existing_path_set = {p for paths in existing_groups.values() for p in paths}
|
|
new_paths = sorted(all_paths - existing_path_set)
|
|
|
|
if not new_paths:
|
|
return
|
|
|
|
new_groups = []
|
|
for i in range(0, len(new_paths), items_per_group):
|
|
group = new_paths[i : i + items_per_group]
|
|
group_hash = self._compute_workgroup_hash(group)
|
|
new_groups.append((group_hash, group))
|
|
|
|
combined_groups = {**existing_groups, **dict(new_groups)}
|
|
combined_lines = self._make_index_lines(combined_groups)
|
|
|
|
await self.backend.save_index_lines(combined_lines)
|
|
|
|
async def initialize_queue(self) -> int:
|
|
"""
|
|
Load the work queue and initialize it for processing.
|
|
Removes already completed work items and randomizes the order.
|
|
"""
|
|
lines = await self.backend.load_index_lines()
|
|
work_queue = self._parse_index_lines(lines)
|
|
|
|
done_hashes = await self.backend.get_completed_hashes()
|
|
|
|
remaining_hashes = set(work_queue) - done_hashes
|
|
remaining_items = [WorkItem(hash=h, work_paths=work_queue[h]) for h in remaining_hashes]
|
|
random.shuffle(remaining_items)
|
|
|
|
self._queue = Queue()
|
|
for item in remaining_items:
|
|
await self._queue.put(item)
|
|
|
|
logger.info(f"Initialized queue with {self.size:,} work items")
|
|
return self.size
|
|
|
|
async def get_work(self, worker_lock_timeout_secs: int = 1800) -> Optional[WorkItem]:
|
|
"""
|
|
Get the next available work item that isn't completed or locked.
|
|
"""
|
|
REFRESH_COMPLETED_HASH_CACHE_MAX_ATTEMPTS = 3
|
|
refresh_completed_hash_attempt = 0
|
|
|
|
while True:
|
|
try:
|
|
work_item = self._queue.get_nowait()
|
|
except QueueEmpty:
|
|
return None
|
|
|
|
if work_item.hash in self._completed_hash_cache:
|
|
logger.debug(f"Work item {work_item.hash} already completed (cache hit), skipping")
|
|
self._queue.task_done()
|
|
continue
|
|
|
|
if await self.backend.is_completed(work_item.hash):
|
|
logger.debug(f"Work item {work_item.hash} already completed, skipping")
|
|
self._queue.task_done()
|
|
|
|
refresh_completed_hash_attempt += 1
|
|
|
|
if refresh_completed_hash_attempt >= REFRESH_COMPLETED_HASH_CACHE_MAX_ATTEMPTS:
|
|
logger.info(f"More than {REFRESH_COMPLETED_HASH_CACHE_MAX_ATTEMPTS} queue items already done, refreshing local completed cache fully")
|
|
self._completed_hash_cache = await self.backend.get_completed_hashes()
|
|
refresh_completed_hash_attempt = 0
|
|
|
|
continue
|
|
|
|
if await self.backend.is_worker_lock_taken(work_item.hash, worker_lock_timeout_secs):
|
|
logger.debug(f"Work item {work_item.hash} is locked by another worker, skipping")
|
|
self._queue.task_done()
|
|
continue
|
|
|
|
# Create lock (overwrites if stale)
|
|
try:
|
|
await self.backend.create_worker_lock(work_item.hash)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to create lock for {work_item.hash}: {e}")
|
|
self._queue.task_done()
|
|
continue
|
|
|
|
refresh_completed_hash_attempt = 0
|
|
return work_item
|
|
|
|
async def mark_done(self, work_item: WorkItem) -> None:
|
|
"""
|
|
Mark a work item as done by removing its lock file and creating a done flag.
|
|
"""
|
|
# Create done flag in done_flags_dir
|
|
await self.backend.create_done_flag(work_item.hash)
|
|
|
|
# Remove the worker lock
|
|
await self.backend.delete_worker_lock(work_item.hash)
|
|
self._queue.task_done()
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
"""Get current size of work queue."""
|
|
return self._queue.qsize()
|
|
|
|
|
|
class LocalBackend(Backend):
|
|
"""Local file system backend."""
|
|
|
|
# Default timeout for filesystem ops, not to block queue, in seconds
|
|
DEFAULT_FS_TIMEOUT = 120
|
|
|
|
def __init__(self, workspace_path: str):
|
|
self.workspace_path = os.path.abspath(workspace_path)
|
|
self._index_path = os.path.join(self.workspace_path, "work_index_list.csv.zstd")
|
|
self._done_flags_dir = os.path.join(self.workspace_path, DONE_FLAGS_DIR)
|
|
self._locks_dir = os.path.join(self.workspace_path, WORKER_LOCKS_DIR)
|
|
|
|
os.makedirs(self.workspace_path, exist_ok=True)
|
|
os.makedirs(self._done_flags_dir, exist_ok=True)
|
|
os.makedirs(self._locks_dir, exist_ok=True)
|
|
|
|
async def _run_with_timeout(self, func, *args, operation_name: str = "filesystem operation"):
|
|
"""Run a blocking function in a thread with timeout and logging."""
|
|
logger.debug(f"LocalBackend: Starting {operation_name}")
|
|
try:
|
|
result = await asyncio.wait_for(asyncio.to_thread(func, *args), timeout=LocalBackend.DEFAULT_FS_TIMEOUT)
|
|
logger.debug(f"LocalBackend: Completed {operation_name}")
|
|
return result
|
|
except asyncio.TimeoutError:
|
|
logger.error(f"LocalBackend: TIMEOUT after {LocalBackend.DEFAULT_FS_TIMEOUT}s on {operation_name}")
|
|
raise TimeoutError(f"Filesystem operation '{operation_name}' timed out after {LocalBackend.DEFAULT_FS_TIMEOUT} seconds")
|
|
|
|
def _download_zstd_csv_local(self, local_path: str) -> List[str]:
|
|
"""
|
|
Read a zstd-compressed CSV from a local path.
|
|
If the file doesn't exist, returns an empty list.
|
|
"""
|
|
if not os.path.exists(local_path):
|
|
return []
|
|
|
|
with open(local_path, "rb") as f:
|
|
dctx = zstandard.ZstdDecompressor()
|
|
data = dctx.decompress(f.read())
|
|
lines = data.decode("utf-8").splitlines()
|
|
return lines
|
|
|
|
def _upload_zstd_csv_local(self, local_path: str, lines: List[str]) -> None:
|
|
"""
|
|
Write a zstd-compressed CSV to a local path.
|
|
"""
|
|
data = "\n".join(lines).encode("utf-8")
|
|
cctx = zstandard.ZstdCompressor()
|
|
compressed_data = cctx.compress(data)
|
|
|
|
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
|
with open(local_path, "wb") as f:
|
|
f.write(compressed_data)
|
|
|
|
async def load_index_lines(self) -> List[str]:
|
|
return await asyncio.to_thread(self._download_zstd_csv_local, self._index_path)
|
|
|
|
async def save_index_lines(self, lines: List[str]) -> None:
|
|
await asyncio.to_thread(self._upload_zstd_csv_local, self._index_path, lines)
|
|
|
|
async def get_completed_hashes(self) -> Set[str]:
|
|
def _list_completed() -> Set[str]:
|
|
if not os.path.isdir(self._done_flags_dir):
|
|
return set()
|
|
return {f[len("done_") : -len(".flag")] for f in os.listdir(self._done_flags_dir) if f.startswith("done_") and f.endswith(".flag")}
|
|
|
|
return await asyncio.to_thread(_list_completed)
|
|
|
|
def _get_worker_lock_path(self, work_hash: str) -> str:
|
|
"""Internal method to get worker lock path."""
|
|
return os.path.join(self._locks_dir, f"worker_{work_hash}.lock")
|
|
|
|
def _get_done_flag_path(self, work_hash: str) -> str:
|
|
"""Internal method to get done flag path."""
|
|
return os.path.join(self._done_flags_dir, f"done_{work_hash}.flag")
|
|
|
|
async def _get_object_mtime(self, path: str) -> Optional[datetime.datetime]:
|
|
"""Internal method to get object mtime."""
|
|
|
|
def _get_mtime() -> Optional[datetime.datetime]:
|
|
if not os.path.exists(path):
|
|
return None
|
|
return datetime.datetime.fromtimestamp(os.path.getmtime(path), datetime.timezone.utc)
|
|
|
|
return await self._run_with_timeout(_get_mtime, operation_name=f"get_mtime({path})")
|
|
|
|
async def is_worker_lock_taken(self, work_hash: str, worker_lock_timeout_secs: int = 1800) -> bool:
|
|
"""Check if a worker lock is taken and not stale."""
|
|
lock_path = self._get_worker_lock_path(work_hash)
|
|
try:
|
|
lock_mtime = await self._get_object_mtime(lock_path)
|
|
except (TimeoutError, Exception) as e:
|
|
# On timeout or error, return False to let this worker attempt the work
|
|
logger.warning(f"is_worker_lock_taken({work_hash}) failed: {e}, returning False")
|
|
return False
|
|
|
|
if not lock_mtime:
|
|
return False
|
|
|
|
now = datetime.datetime.now(datetime.timezone.utc)
|
|
return (now - lock_mtime).total_seconds() <= worker_lock_timeout_secs
|
|
|
|
async def create_worker_lock(self, work_hash: str) -> None:
|
|
"""Create a worker lock for a work hash. Best-effort, does not fail on timeout."""
|
|
lock_path = self._get_worker_lock_path(work_hash)
|
|
|
|
def _create() -> None:
|
|
with open(lock_path, "wb"):
|
|
pass
|
|
|
|
try:
|
|
await self._run_with_timeout(_create, operation_name=f"create_worker_lock({work_hash})")
|
|
except (TimeoutError, Exception) as e:
|
|
# Best effort - if we can't create lock, still proceed (duplicate work is okay)
|
|
logger.warning(f"create_worker_lock({work_hash}) failed: {e}, proceeding anyway")
|
|
|
|
async def delete_worker_lock(self, work_hash: str) -> None:
|
|
"""Delete the worker lock for a work hash if it exists. Best-effort, does not fail on timeout."""
|
|
lock_path = self._get_worker_lock_path(work_hash)
|
|
|
|
def _delete() -> None:
|
|
if os.path.exists(lock_path):
|
|
os.remove(lock_path)
|
|
|
|
try:
|
|
await self._run_with_timeout(_delete, operation_name=f"delete_worker_lock({work_hash})")
|
|
except (TimeoutError, Exception) as e:
|
|
# Best effort - if we can't delete lock, just log and continue
|
|
logger.warning(f"delete_worker_lock({work_hash}) failed: {e}, proceeding anyway")
|
|
|
|
async def is_completed(self, work_hash: str) -> bool:
|
|
"""Check if a work item has been completed."""
|
|
done_flag_path = self._get_done_flag_path(work_hash)
|
|
try:
|
|
return await self._get_object_mtime(done_flag_path) is not None
|
|
except (TimeoutError, Exception) as e:
|
|
# On timeout or error, return False to let worker attempt the work
|
|
logger.warning(f"is_completed({work_hash}) failed: {e}, returning False")
|
|
return False
|
|
|
|
async def create_done_flag(self, work_hash: str) -> None:
|
|
"""Create a done flag for a work hash. Best-effort, does not fail on timeout."""
|
|
done_flag_path = self._get_done_flag_path(work_hash)
|
|
|
|
def _create() -> None:
|
|
with open(done_flag_path, "wb"):
|
|
pass
|
|
|
|
try:
|
|
await self._run_with_timeout(_create, operation_name=f"create_done_flag({work_hash})")
|
|
except (TimeoutError, Exception) as e:
|
|
# Best effort - work was done, just couldn't mark it. Log and continue.
|
|
logger.warning(f"create_done_flag({work_hash}) failed: {e}, proceeding anyway")
|
|
|
|
|
|
class S3Backend(Backend):
|
|
"""S3 backend."""
|
|
|
|
def __init__(self, s3_client: Any, workspace_path: str):
|
|
self.s3_client = s3_client
|
|
self.workspace_path = workspace_path.rstrip("/")
|
|
self._index_path = os.path.join(self.workspace_path, "work_index_list.csv.zstd")
|
|
self._output_glob = os.path.join(self.workspace_path, DONE_FLAGS_DIR, "*.flag")
|
|
|
|
async def load_index_lines(self) -> List[str]:
|
|
return await asyncio.to_thread(download_zstd_csv, self.s3_client, self._index_path)
|
|
|
|
async def save_index_lines(self, lines: List[str]) -> None:
|
|
await asyncio.to_thread(upload_zstd_csv, self.s3_client, self._index_path, lines)
|
|
|
|
async def get_completed_hashes(self) -> Set[str]:
|
|
def _list_completed() -> Set[str]:
|
|
done_work_items = expand_s3_glob(self.s3_client, self._output_glob)
|
|
return {
|
|
os.path.basename(item)[len("done_") : -len(".flag")]
|
|
for item in done_work_items
|
|
if os.path.basename(item).startswith("done_") and os.path.basename(item).endswith(".flag")
|
|
}
|
|
|
|
return await asyncio.to_thread(_list_completed)
|
|
|
|
def _get_worker_lock_path(self, work_hash: str) -> str:
|
|
"""Internal method to get worker lock path."""
|
|
return os.path.join(self.workspace_path, WORKER_LOCKS_DIR, f"worker_{work_hash}.lock")
|
|
|
|
def _get_done_flag_path(self, work_hash: str) -> str:
|
|
"""Internal method to get done flag path."""
|
|
return os.path.join(self.workspace_path, DONE_FLAGS_DIR, f"done_{work_hash}.flag")
|
|
|
|
async def _get_object_mtime(self, path: str) -> Optional[datetime.datetime]:
|
|
"""Internal method to get object mtime."""
|
|
bucket, key = parse_s3_path(path)
|
|
|
|
def _head_object() -> Optional[datetime.datetime]:
|
|
try:
|
|
response = self.s3_client.head_object(Bucket=bucket, Key=key)
|
|
return response["LastModified"]
|
|
except self.s3_client.exceptions.ClientError as e:
|
|
if e.response["Error"]["Code"] == "404":
|
|
return None
|
|
raise
|
|
|
|
return await asyncio.to_thread(_head_object)
|
|
|
|
async def is_worker_lock_taken(self, work_hash: str, worker_lock_timeout_secs: int = 1800) -> bool:
|
|
"""Check if a worker lock is taken and not stale."""
|
|
lock_path = self._get_worker_lock_path(work_hash)
|
|
lock_mtime = await self._get_object_mtime(lock_path)
|
|
|
|
if not lock_mtime:
|
|
return False
|
|
|
|
now = datetime.datetime.now(datetime.timezone.utc)
|
|
return (now - lock_mtime).total_seconds() <= worker_lock_timeout_secs
|
|
|
|
async def create_worker_lock(self, work_hash: str) -> None:
|
|
"""Create a worker lock for a work hash."""
|
|
lock_path = self._get_worker_lock_path(work_hash)
|
|
bucket, key = parse_s3_path(lock_path)
|
|
await asyncio.to_thread(self.s3_client.put_object, Bucket=bucket, Key=key, Body=b"")
|
|
|
|
async def delete_worker_lock(self, work_hash: str) -> None:
|
|
"""Delete the worker lock for a work hash if it exists."""
|
|
lock_path = self._get_worker_lock_path(work_hash)
|
|
bucket, key = parse_s3_path(lock_path)
|
|
await asyncio.to_thread(self.s3_client.delete_object, Bucket=bucket, Key=key)
|
|
|
|
async def is_completed(self, work_hash: str) -> bool:
|
|
"""Check if a work item has been completed."""
|
|
done_flag_path = self._get_done_flag_path(work_hash)
|
|
return await self._get_object_mtime(done_flag_path) is not None
|
|
|
|
async def create_done_flag(self, work_hash: str) -> None:
|
|
"""Create a done flag for a work hash."""
|
|
done_flag_path = self._get_done_flag_path(work_hash)
|
|
bucket, key = parse_s3_path(done_flag_path)
|
|
await asyncio.to_thread(self.s3_client.put_object, Bucket=bucket, Key=key, Body=b"")
|