submmit/modelhub_submmit_api/common.py

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