data_server.py

from __future__ import annotations

import argparse
import array
import bisect
import concurrent.futures
import gzip
import hashlib
import json
import os
import random
import threading
import time
import uuid
from dataclasses import dataclass
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any
from urllib.parse import parse_qs, urlparse

from omegaconf import OmegaConf


def load_api_keys_from_file(file_path: str | None) -> list[str]:
    if not file_path:
        return []
    expanded = os.path.expanduser(file_path)
    if not os.path.exists(expanded):
        return []
    keys: list[str] = []
    with open(expanded, "r", encoding="utf-8") as key_file:
        for line in key_file:
            value = line.strip()
            if value:
                keys.append(value)
    return keys


@dataclass(frozen=True)
class IndexedFile:
    path: str
    size: int
    mtime_ns: int
    index_path: str
    count: int
    cumulative_end: int


def _iter_jsonl_files(data_root: str) -> list[str]:
    result: list[str] = []
    for root, _, names in os.walk(data_root):
        for name in names:
            if name.endswith(".jsonl"):
                result.append(os.path.join(root, name))
    result.sort()
    return result


def _index_path_for(index_root: str, file_path: str) -> str:
    digest = hashlib.sha1(file_path.encode("utf-8")).hexdigest()
    return os.path.join(index_root, f"{digest}.idx")


def _build_offsets_file(file_path: str, index_path: str) -> tuple[int, int, int]:
    offsets = array.array("Q")
    with open(file_path, "rb") as source:
        while True:
            offset = source.tell()
            line = source.readline()
            if not line:
                break
            offsets.append(offset)

    os.makedirs(os.path.dirname(index_path), exist_ok=True)
    with open(index_path, "wb") as index_file:
        offsets.tofile(index_file)

    stat = os.stat(file_path)
    return len(offsets), stat.st_size, stat.st_mtime_ns


class JsonlIndexStore:
    def __init__(self, data_root: str, index_root: str, workers: int | None = None) -> None:
        self.data_root = os.path.abspath(data_root)
        self.index_root = os.path.abspath(index_root)
        os.makedirs(self.index_root, exist_ok=True)
        self.files: list[IndexedFile] = []
        self._offset_cache: dict[str, array.array] = {}
        self.workers = max(1, workers or (os.cpu_count() or 4))
        self.manifest_path = os.path.join(self.index_root, "manifest.json")
        self._build_manifest()

    def _load_manifest(self) -> dict[str, dict[str, int | str]] | None:
        if not os.path.exists(self.manifest_path):
            return None
        with open(self.manifest_path, "r", encoding="utf-8") as manifest_file:
            payload = json.load(manifest_file)
        files = payload.get("files")
        if not isinstance(files, dict):
            return None
        return files

    def _load_or_build_offsets(self, file_path: str) -> array.array:
        cached = self._offset_cache.get(file_path)
        if cached is not None:
            return cached

        index_path = _index_path_for(self.index_root, file_path)
        if os.path.exists(index_path):
            offsets = array.array("Q")
            with open(index_path, "rb") as index_file:
                offsets.frombytes(index_file.read())
            self._offset_cache[file_path] = offsets
            return offsets

        _build_offsets_file(file_path, index_path)
        offsets = array.array("Q")
        with open(index_path, "rb") as index_file:
            offsets.frombytes(index_file.read())
        self._offset_cache[file_path] = offsets
        return offsets

    def _save_manifest(self, manifest_files: dict[str, dict[str, int | str]]) -> None:
        payload = {
            "data_root": self.data_root,
            "index_root": self.index_root,
            "files": manifest_files,
        }
        tmp_path = f"{self.manifest_path}.tmp"
        with open(tmp_path, "w", encoding="utf-8") as manifest_file:
            json.dump(payload, manifest_file, ensure_ascii=False)
        os.replace(tmp_path, self.manifest_path)

    def _build_manifest(self) -> None:
        data_files = _iter_jsonl_files(self.data_root)
        if not data_files:
            raise FileNotFoundError(f"no .jsonl files found under {self.data_root}")

        cached_manifest = self._load_manifest() or {}
        manifest_files: dict[str, dict[str, int | str]] = {}
        rebuild_jobs: list[tuple[str, str]] = []

        for file_path in data_files:
            stat = os.stat(file_path)
            index_path = _index_path_for(self.index_root, file_path)
            cached = cached_manifest.get(file_path)
            if (
                cached is not None
                and cached.get("size") == stat.st_size
                and cached.get("mtime_ns") == stat.st_mtime_ns
                and cached.get("index_path") == index_path
                and os.path.exists(index_path)
            ):
                manifest_files[file_path] = cached
            else:
                rebuild_jobs.append((file_path, index_path))

        if rebuild_jobs:
            with concurrent.futures.ProcessPoolExecutor(max_workers=self.workers) as executor:
                future_map = {
                    executor.submit(_build_offsets_file, file_path, index_path): (file_path, index_path)
                    for file_path, index_path in rebuild_jobs
                }
                for future in concurrent.futures.as_completed(future_map):
                    file_path, index_path = future_map[future]
                    count, size, mtime_ns = future.result()
                    manifest_files[file_path] = {
                        "index_path": index_path,
                        "count": count,
                        "size": size,
                        "mtime_ns": mtime_ns,
                    }

        self._save_manifest(manifest_files)

        cumulative = 0
        for file_path in data_files:
            info = manifest_files[file_path]
            count = int(info["count"])
            cumulative += count
            stat = os.stat(file_path)
            self.files.append(
                IndexedFile(
                    path=file_path,
                    size=stat.st_size,
                    mtime_ns=stat.st_mtime_ns,
                    index_path=str(info["index_path"]),
                    count=count,
                    cumulative_end=cumulative,
                )
            )
        if cumulative == 0:
            raise ValueError(f"no records found under {self.data_root}")

    @property
    def total_records(self) -> int:
        return self.files[-1].cumulative_end

    def resolve_index(self, global_index: int) -> tuple[IndexedFile, int]:
        if global_index < 0 or global_index >= self.total_records:
            raise IndexError(f"global_index out of range: {global_index}")
        cumulative_ends = [item.cumulative_end for item in self.files]
        file_idx = bisect.bisect_right(cumulative_ends, global_index)
        indexed_file = self.files[file_idx]
        previous_end = 0 if file_idx == 0 else self.files[file_idx - 1].cumulative_end
        local_index = global_index - previous_end
        return indexed_file, local_index

    def read_record(self, global_index: int) -> dict[str, Any]:
        indexed_file, local_index = self.resolve_index(global_index)
        offsets = self._load_or_build_offsets(indexed_file.path)
        offset = offsets[local_index]
        with open(indexed_file.path, "rb") as source:
            source.seek(offset)
            line = source.readline()
        record = json.loads(line.decode("utf-8"))
        record["_global_index"] = global_index
        return record


class StreamSession:
    def __init__(self, session_id: str, store: JsonlIndexStore, seed: int, shuffle_buffer_size: int) -> None:
        self.session_id = session_id
        self.store = store
        self.seed = seed
        self.shuffle_buffer_size = max(1, shuffle_buffer_size)
        self.file_order = list(range(len(store.files)))
        self.rng = random.Random(seed)
        self.rng.shuffle(self.file_order)
        self.file_order_cursor = 0
        self.current_local_index = 0
        self.buffer: list[int] = []
        self.exhausted = False
        self.lock = threading.Lock()

    def _next_global_index(self) -> int | None:
        while self.file_order_cursor < len(self.file_order):
            file_idx = self.file_order[self.file_order_cursor]
            indexed_file = self.store.files[file_idx]
            if self.current_local_index < indexed_file.count:
                file_start = indexed_file.cumulative_end - indexed_file.count
                global_index = file_start + self.current_local_index
                self.current_local_index += 1
                return global_index

            self.file_order_cursor += 1
            self.current_local_index = 0
        return None

    def _fill_buffer(self) -> None:
        while len(self.buffer) < self.shuffle_buffer_size:
            next_index = self._next_global_index()
            if next_index is None:
                break
            self.buffer.append(next_index)
        if not self.buffer and self._next_global_index() is None:
            self.exhausted = True

    def next_batch(self, batch_size: int) -> dict[str, Any]:
        texts: list[str] = []
        indices: list[int] = []
        with self.lock:
            self._fill_buffer()
            while len(texts) < batch_size and self.buffer:
                picked_idx = self.rng.randrange(len(self.buffer))
                global_index = self.buffer[picked_idx]
                replacement = self._next_global_index()
                if replacement is None:
                    self.buffer.pop(picked_idx)
                else:
                    self.buffer[picked_idx] = replacement

                record = self.store.read_record(global_index)
                text = record.get("text")
                if not text:
                    continue
                texts.append(text)
                indices.append(global_index)

            if not self.buffer and self.file_order_cursor >= len(self.file_order):
                self.exhausted = True

        return {
            "session_id": self.session_id,
            "texts": texts,
            "indices": indices,
            "batch_size": len(texts),
            "exhausted": self.exhausted and not texts,
        }


class DataServerState:
    def __init__(
        self,
        data_root: str,
        index_root: str,
        api_key_file: str | None = None,
        workers: int | None = None,
    ) -> None:
        self.store = JsonlIndexStore(data_root=data_root, index_root=index_root, workers=workers)
        self.sessions: dict[str, StreamSession] = {}
        self.sessions_lock = threading.Lock()
        self.started_at = time.time()
        self.api_keys = self._resolve_api_keys(api_key_file=api_key_file)
        if not self.api_keys:
            raise ValueError(
                "no API keys loaded; please configure data.data_server_api_key_file "
                "with at least one non-empty key"
            )

    @staticmethod
    def _resolve_api_keys(api_key_file: str | None) -> set[str]:
        keys: set[str] = set()
        for key in load_api_keys_from_file(api_key_file):
            keys.add(key)
        return keys

    def open_session(self, seed: int, shuffle_buffer_size: int) -> dict[str, Any]:
        session_id = uuid.uuid4().hex
        session = StreamSession(
            session_id=session_id,
            store=self.store,
            seed=seed,
            shuffle_buffer_size=shuffle_buffer_size,
        )
        with self.sessions_lock:
            self.sessions[session_id] = session
        return {
            "session_id": session_id,
            "seed": seed,
            "shuffle_buffer_size": shuffle_buffer_size,
            "total_records": self.store.total_records,
        }

    def get_session(self, session_id: str) -> StreamSession:
        with self.sessions_lock:
            session = self.sessions.get(session_id)
        if session is None:
            raise KeyError(f"unknown session_id: {session_id}")
        return session


def build_handler(state: DataServerState):
    class DataServerHandler(BaseHTTPRequestHandler):
        server_version = "NestedLearningDataServer/1.0"

        def log_message(self, fmt: str, *args: object) -> None:
            return

        def _supports_gzip(self) -> bool:
            accept_encoding = self.headers.get("Accept-Encoding", "")
            return "gzip" in accept_encoding.lower()

        def _write_json(self, payload: dict[str, Any], status: int = HTTPStatus.OK) -> None:
            body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
            use_gzip = self._supports_gzip() and len(body) >= 1024
            if use_gzip:
                body = gzip.compress(body, compresslevel=5)
            self.send_response(status)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            if use_gzip:
                self.send_header("Content-Encoding", "gzip")
                self.send_header("Vary", "Accept-Encoding")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def _check_api_key(self) -> bool:
            if not state.api_keys:
                return True
            provided = self.headers.get("X-Infimory-Api-Key", "")
            return provided in state.api_keys

        def do_GET(self) -> None:
            parsed = urlparse(self.path)
            query = parse_qs(parsed.query)
            try:
                if not self._check_api_key():
                    self._write_json({"error": "unauthorized"}, status=HTTPStatus.UNAUTHORIZED)
                    return

                if parsed.path == "/health":
                    self._write_json(
                        {
                            "status": "ok",
                            "total_records": state.store.total_records,
                            "uptime_sec": round(time.time() - state.started_at, 3),
                        }
                    )
                    return

                if parsed.path == "/v1/meta":
                    self._write_json(
                        {
                            "total_records": state.store.total_records,
                            "files": len(state.store.files),
                            "data_root": state.store.data_root,
                            "index_root": state.store.index_root,
                            "workers": state.store.workers,
                        }
                    )
                    return

                if parsed.path == "/v1/sample":
                    global_index = int(query["idx"][0])
                    record = state.store.read_record(global_index)
                    self._write_json(record)
                    return

                if parsed.path == "/v1/session/open":
                    seed = int(query.get("seed", ["1234"])[0])
                    shuffle_buffer_size = int(query.get("shuffle_buffer_size", ["8192"])[0])
                    self._write_json(state.open_session(seed=seed, shuffle_buffer_size=shuffle_buffer_size))
                    return

                if parsed.path == "/v1/session/next":
                    session_id = query["session_id"][0]
                    batch_size = int(query.get("batch_size", ["1"])[0])
                    session = state.get_session(session_id)
                    payload = session.next_batch(batch_size=batch_size)
                    print(
                        json.dumps(
                            {
                                "event": "batch_sent",
                                "time": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                                "session_id": session_id,
                                "requested_batch_size": batch_size,
                                "sent_batch_size": payload.get("batch_size", 0),
                            },
                            ensure_ascii=False,
                        ),
                        flush=True,
                    )
                    self._write_json(payload)
                    return

                self._write_json({"error": f"unknown path: {parsed.path}"}, status=HTTPStatus.NOT_FOUND)
            except KeyError as exc:
                self._write_json({"error": str(exc)}, status=HTTPStatus.NOT_FOUND)
            except IndexError as exc:
                self._write_json({"error": str(exc)}, status=HTTPStatus.BAD_REQUEST)
            except Exception as exc:
                self._write_json({"error": str(exc)}, status=HTTPStatus.INTERNAL_SERVER_ERROR)

    return DataServerHandler


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Serve cleaned JSONL training data over HTTP.")
    parser.add_argument(
        "--config-file",
        default="configs/hope_tiny.yaml",
        help="YAML config file that contains data.data_server_* settings.",
    )
    parser.add_argument(
        "--data-root",
        default=None,
        help="Root directory containing cleaned JSONL files.",
    )
    parser.add_argument(
        "--index-root",
        default=None,
        help="Directory for persistent line offset indexes.",
    )
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=18080)
    parser.add_argument("--workers", type=int, default=max(1, os.cpu_count() or 4))
    return parser.parse_args()


def load_server_config(config_file: str) -> dict[str, Any]:
    expanded = os.path.expanduser(config_file)
    if not os.path.exists(expanded):
        return {}
    cfg = OmegaConf.load(expanded)
    data_cfg = cfg.get("data")
    if data_cfg is None:
        return {}
    return OmegaConf.to_container(data_cfg, resolve=True)  # type: ignore[return-value]


def main() -> None:
    args = parse_args()
    data_cfg = load_server_config(args.config_file)
    data_root = args.data_root or data_cfg.get(
        "data_server_data_root",
        "/data/datasets/nested-learning/extract_data/cleaned_data",
    )
    index_root = args.index_root or data_cfg.get(
        "data_server_index_root",
        "/data/datasets/nested-learning/extract_data/data_index",
    )
    api_key_file = data_cfg.get(
        "data_server_api_key_file",
        "~/.config/infimory/api_key.txt",
    )

    state = DataServerState(
        data_root=data_root,
        index_root=index_root,
        api_key_file=api_key_file,
        workers=args.workers,
    )
    server = ThreadingHTTPServer((args.host, args.port), build_handler(state))
    print(
        json.dumps(
            {
                "event": "data_server_started",
                "host": args.host,
                "port": args.port,
                "config_file": os.path.abspath(os.path.expanduser(args.config_file)),
                "data_root": os.path.abspath(os.path.expanduser(data_root)),
                "index_root": os.path.abspath(os.path.expanduser(index_root)),
                "workers": args.workers,
                "api_key_enabled": bool(state.api_keys),
                "api_key_count": len(state.api_keys),
                "total_records": state.store.total_records,
            },
            ensure_ascii=False,
        ),
        flush=True,
    )
    server.serve_forever()


if __name__ == "__main__":
    main()