166 lines
4.7 KiB
Python
166 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import socket
|
|
import tempfile
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Iterator
|
|
|
|
|
|
def ensure_utc(value: datetime) -> datetime:
|
|
if value.tzinfo is None:
|
|
return value.replace(tzinfo=timezone.utc)
|
|
return value.astimezone(timezone.utc)
|
|
|
|
|
|
def utc_now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def parse_datetime(value: Any) -> datetime | None:
|
|
if value in (None, ""):
|
|
return None
|
|
if isinstance(value, datetime):
|
|
return ensure_utc(value)
|
|
text = str(value).strip()
|
|
candidates = [
|
|
text,
|
|
text.replace("Z", "+00:00"),
|
|
text.replace(" ", "T"),
|
|
text.replace(" ", "T").replace("Z", "+00:00"),
|
|
]
|
|
for candidate in candidates:
|
|
try:
|
|
return ensure_utc(datetime.fromisoformat(candidate))
|
|
except ValueError:
|
|
continue
|
|
for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
|
|
try:
|
|
parsed = datetime.strptime(text, fmt)
|
|
return parsed.replace(tzinfo=timezone.utc)
|
|
except ValueError:
|
|
continue
|
|
return None
|
|
|
|
|
|
def isoformat_z(value: datetime) -> str:
|
|
return ensure_utc(value).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def format_modelhub_datetime(value: datetime) -> str:
|
|
return ensure_utc(value).strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
|
def json_dumps(value: Any) -> str:
|
|
return json.dumps(value, ensure_ascii=False, sort_keys=True)
|
|
|
|
|
|
def runtime_instance_id() -> str:
|
|
"""Return a non-secret identity that differs between concurrent workers."""
|
|
base = (
|
|
os.getenv("MODELHUB_AGENT_INSTANCE_ID")
|
|
or os.getenv("STRATEGY_ID")
|
|
or os.getenv("HOSTNAME")
|
|
or socket.gethostname()
|
|
)
|
|
raw = f"{base}:{os.getpid()}"
|
|
return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:20]
|
|
|
|
|
|
def _lock_path(path: Path) -> Path:
|
|
return path.with_name(f".{path.name}.lock")
|
|
|
|
|
|
@contextmanager
|
|
def file_lock(path: Path, *, shared: bool = False) -> Iterator[None]:
|
|
"""Coordinate access to a path between processes on the same filesystem."""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
lock_path = _lock_path(path)
|
|
with lock_path.open("a+b") as handle:
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_SH if shared else fcntl.LOCK_EX)
|
|
try:
|
|
yield
|
|
finally:
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
|
|
|
|
|
def _atomic_write_text(path: Path, content: str) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
descriptor, temporary_name = tempfile.mkstemp(
|
|
dir=str(path.parent),
|
|
prefix=f".{path.name}.",
|
|
suffix=".tmp",
|
|
)
|
|
try:
|
|
with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
|
|
handle.write(content)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary_name, path)
|
|
except Exception:
|
|
try:
|
|
os.unlink(temporary_name)
|
|
except FileNotFoundError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def _read_jsonl_unlocked(path: Path) -> list[dict[str, Any]]:
|
|
if not path.exists():
|
|
return []
|
|
rows: list[dict[str, Any]] = []
|
|
for line in path.read_text(encoding="utf-8").splitlines():
|
|
if not line.strip():
|
|
continue
|
|
rows.append(json.loads(line))
|
|
return rows
|
|
|
|
|
|
def _jsonl_content(rows: list[dict[str, Any]]) -> str:
|
|
return "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows)
|
|
|
|
|
|
def read_json(path: Path) -> Any:
|
|
with file_lock(path, shared=True):
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def write_json(path: Path, value: Any) -> None:
|
|
content = json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
|
with file_lock(path):
|
|
_atomic_write_text(path, content)
|
|
|
|
|
|
def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
|
with file_lock(path, shared=True):
|
|
return _read_jsonl_unlocked(path)
|
|
|
|
|
|
def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
|
|
with file_lock(path):
|
|
_atomic_write_text(path, _jsonl_content(rows))
|
|
|
|
|
|
def append_jsonl(path: Path, row: dict[str, Any]) -> None:
|
|
line = json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n"
|
|
with file_lock(path):
|
|
with path.open("a", encoding="utf-8") as handle:
|
|
handle.write(line)
|
|
handle.flush()
|
|
|
|
|
|
def update_jsonl(
|
|
path: Path,
|
|
updater: Callable[[list[dict[str, Any]]], list[dict[str, Any]]],
|
|
) -> list[dict[str, Any]]:
|
|
"""Atomically read, update, and replace a JSONL file under one lock."""
|
|
with file_lock(path):
|
|
updated = updater(_read_jsonl_unlocked(path))
|
|
_atomic_write_text(path, _jsonl_content(updated))
|
|
return updated
|