src/evaluation/storage.py

"""Private submission persistence; local roots should be the ignored submissions/.

Use SubmissionStore(Path('submissions')) locally, or pass the private dataset
repo_id='XLearning-SCU/ARK-Bench-Evaluation' and an explicit token. Remote
operations never report local success or write a local fallback.
"""
from __future__ import annotations

from datetime import datetime, timezone
from functools import wraps
import fcntl
import json
import os
from pathlib import Path
import re
import shutil
import tempfile


class StorageError(Exception):
    """A private storage operation failed (details deliberately withheld)."""


def _private_errors(function):
    @wraps(function)
    def wrapped(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except Exception:
            raise StorageError('Submission operation failed.') from None
    return wrapped


def _submission_id(value):
    if not isinstance(value, str) or re.fullmatch(r'[0-9a-f]{32}', value) is None:
        raise StorageError('Submission operation failed.')
    return value


def _json_bytes(value):
    return json.dumps(value, allow_nan=False, sort_keys=True, ensure_ascii=False).encode('utf-8')


def _markdown(value):
    """Render untrusted model metadata as plain Markdown text."""
    text = str(value or '').replace('\r', ' ').replace('\n', ' ').strip()
    for character in r'\\`*_{}[]<>|':
        text = text.replace(character, '\\' + character)
    return text.replace('@', '＠') or '—'


def _issue_content(record):
    submission_id = record['submission_id']
    score = record['score']
    model = score['model']
    metadata = score['meta_data']
    model_name = str(model['name']).replace('@', '＠')
    title = f"ARK-Bench submission {submission_id[:8]} — {model_name}"[:200]
    rows = []
    for family, metrics in (('Recall', ('R@1', 'R@5', 'R@10')),
                            ('NDCG', ('nDCG@5', 'nDCG@10', 'nDCG@20'))):
        results = score['results'][family]
        for metric in metrics:
            knowledge = results['knowledge_domain']['macro_avg'][metric]
            reasoning = results['reasoning_skill']['macro_avg'][metric]
            rows.append(f"| {metric} | {knowledge:.2f} | {reasoning:.2f} | {(knowledge + reasoning) / 2:.2f} |")
    description = '\n'.join([
        'A participant submitted an evaluated ARK-Bench report for maintainer review.', '',
        f"- **Submission ID:** `{submission_id}`",
        f"- **Model:** {_markdown(model.get('name'))}",
        f"- **Model URL:** {_markdown(model.get('repo_link'))}",
        f"- **Parameters (B):** {_markdown(model.get('params_b'))}",
        f"- **Evaluation date:** {_markdown(metadata.get('date'))}",
        f"- **Benchmark version:** {_markdown(metadata.get('benchmark_version'))}",
        f"- **Dataset revision:** `{_markdown(metadata.get('dataset_revision'))}`",
        f"- **Protocol:** `{_markdown(metadata.get('evaluation_protocol'))}`", '',
        '| Metric | Knowledge macro | Reasoning macro | Overview |',
        '| --- | ---: | ---: | ---: |', *rows, '',
        f"Private artifacts: `submissions/{submission_id}/`", '',
        'Review and recompute this submission with the maintainer CLI before publication.',
    ])
    return title, description


class SubmissionStore:
    @_private_errors
    def __init__(self, root: Path, repo_id: str | None = None, token: str | None = None):
        self.root = Path(root)
        self.repo_id = repo_id
        self.token = token
        self.api = None
        if repo_id is not None:
            if not repo_id or not token:
                raise ValueError()
            from huggingface_hub import HfApi
            self.api = HfApi(token=token)

    def _private_info(self):
        info = self.api.dataset_info(repo_id=self.repo_id, token=self.token)
        if info.private is not True:
            raise ValueError()
        return info

    def _local_path(self, submission_id, filename=None):
        path = self.root / _submission_id(submission_id)
        if path.is_symlink():
            raise ValueError()
        if filename:
            path = path / filename
            if path.is_symlink():
                raise ValueError()
        return path

    @_private_errors
    def save(self, record: dict, predictions: bytes) -> None:
        submission_id = _submission_id(record['submission_id'])
        if (record.get('status') not in ('evaluated', 'pending_review') or not isinstance(record.get('score'), dict)
                or not isinstance(record.get('created_at'), str) or not record['created_at']
                or not isinstance(predictions, bytes)):
            raise ValueError()
        payload = _json_bytes(record)
        if self.api is not None:
            from huggingface_hub import CommitOperationAdd
            info = self._private_info()
            prefix = f'submissions/{submission_id}/'
            files = self.api.list_repo_files(repo_id=self.repo_id, repo_type='dataset',
                                             token=self.token, revision=info.sha)
            if any(name.startswith(prefix) for name in files):
                raise ValueError()
            self.api.create_commit(
                repo_id=self.repo_id, repo_type='dataset', token=self.token,
                parent_commit=info.sha, commit_message='Save private submission',
                operations=[CommitOperationAdd(path_in_repo=prefix + 'predictions.jsonl',
                                               path_or_fileobj=predictions),
                            CommitOperationAdd(path_in_repo=prefix + 'record.json', path_or_fileobj=payload)])
            return
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        destination = self._local_path(submission_id)
        staging = Path(tempfile.mkdtemp(prefix='.staging-', dir=self.root))
        try:
            for name, content in [('predictions.jsonl', predictions), ('record.json', payload)]:
                with (staging / name).open('xb') as handle:
                    os.chmod(staging / name, 0o600)
                    handle.write(content)
                    handle.flush()
                    os.fsync(handle.fileno())
            if destination.exists():
                raise ValueError()
            staging.rename(destination)
        finally:
            if staging.exists():
                shutil.rmtree(staging)

    def _read(self, submission_id, filename):
        submission_id = _submission_id(submission_id)
        if self.api is None:
            return self._local_path(submission_id, filename).read_bytes()
        from huggingface_hub import hf_hub_download
        info = self._private_info()
        path = hf_hub_download(repo_id=self.repo_id, repo_type='dataset', token=self.token,
                               revision=info.sha, filename=f'submissions/{submission_id}/{filename}')
        return Path(path).read_bytes()

    @_private_errors
    def load(self, submission_id: str) -> dict:
        record = json.loads(self._read(submission_id, 'record.json'))
        if record['submission_id'] != submission_id:
            raise ValueError()
        return record

    @_private_errors
    def read_predictions(self, submission_id: str) -> bytes:
        return self._read(submission_id, 'predictions.jsonl')

    @_private_errors
    def list_records(self, *, include_closed: bool = False) -> list[dict]:
        if self.api is not None:
            info = self._private_info()
            files = self.api.list_repo_files(repo_id=self.repo_id, repo_type='dataset',
                                             token=self.token, revision=info.sha)
            ids = [name.split('/')[1] for name in files
                   if re.fullmatch(r'submissions/[0-9a-f]{32}/record.json', name)]
        else:
            ids = [path.name for path in self.root.iterdir()
                   if re.fullmatch(r'[0-9a-f]{32}', path.name)] if self.root.exists() else []
        records = [self.load(value) for value in sorted(ids)]
        if self.api is not None and not include_closed and any(
                record.get('review_request', {}).get('issue_number') is not None for record in records):
            closed = {discussion.num for discussion in self.api.get_repo_discussions(
                repo_id=self.repo_id, repo_type='dataset', token=self.token)
                if not discussion.is_pull_request and discussion.status == 'closed'}
            records = [record for record in records
                       if record.get('review_request', {}).get('issue_number') not in closed]
        return records

    @_private_errors
    def submit_report(self, submission_id: str) -> dict:
        """Move an evaluated draft to review and create one private Hub Issue."""
        submission_id = _submission_id(submission_id)
        if self.api is None:
            directory_fd = os.open(self._local_path(submission_id), os.O_RDONLY)
            try:
                fcntl.flock(directory_fd, fcntl.LOCK_EX)
                record = self.load(submission_id)
                if record['status'] == 'pending_review':
                    return record.get('review_request', {})
                if record['status'] != 'evaluated':
                    raise ValueError()
                request = {'submitted_at': datetime.now(timezone.utc).isoformat(),
                           'issue_url': None, 'issue_number': None}
                record.update(status='pending_review', review_request=request)
                self._write_local_record(submission_id, record)
                return request
            finally:
                os.close(directory_fd)

        from huggingface_hub import CommitOperationAdd, hf_hub_download
        info = self._private_info()
        path = hf_hub_download(repo_id=self.repo_id, repo_type='dataset', token=self.token,
                               revision=info.sha, filename=f'submissions/{submission_id}/record.json')
        record = json.loads(Path(path).read_bytes())
        if record['submission_id'] != submission_id:
            raise ValueError()
        if record['status'] == 'pending_review':
            return record.get('review_request', {})
        if record['status'] != 'evaluated':
            raise ValueError()
        title, description = _issue_content(record)
        issue = next((item for item in self.api.get_repo_discussions(
            repo_id=self.repo_id, repo_type='dataset', token=self.token)
                      if not item.is_pull_request and item.title == title), None)
        if issue is None:
            issue = self.api.create_discussion(
                repo_id=self.repo_id, repo_type='dataset', token=self.token,
                title=title, description=description, pull_request=False)
        request = {'submitted_at': datetime.now(timezone.utc).isoformat(),
                   'issue_url': issue.url, 'issue_number': issue.num}
        record.update(status='pending_review', review_request=request)
        self.api.create_commit(
            repo_id=self.repo_id, repo_type='dataset', token=self.token,
            parent_commit=info.sha, commit_message='Submit evaluation report for review',
            operations=[CommitOperationAdd(
                path_in_repo=f'submissions/{submission_id}/record.json',
                path_or_fileobj=_json_bytes(record))])
        return request

    def _write_local_record(self, submission_id, record):
        directory = self._local_path(submission_id)
        fd, temporary = tempfile.mkstemp(prefix='.record-', dir=directory)
        try:
            with os.fdopen(fd, 'wb') as handle:
                handle.write(_json_bytes(record))
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temporary, self._local_path(submission_id, 'record.json'))
        finally:
            Path(temporary).unlink(missing_ok=True)

    @_private_errors
    def decide(self, submission_id: str, status: str, reviewer: str,
               parent_commit: str | None = None) -> None:
        if self.api is not None:
            return self._decide(submission_id, status, reviewer, parent_commit)
        directory_fd = os.open(self._local_path(submission_id), os.O_RDONLY)
        try:
            fcntl.flock(directory_fd, fcntl.LOCK_EX)
            return self._decide(submission_id, status, reviewer, parent_commit)
        finally:
            os.close(directory_fd)

    def _decide(self, submission_id, status, reviewer, parent_commit):
        if status not in ('rejected', 'approved_for_publication') or not reviewer.strip():
            raise ValueError()
        # Pin remote read and write to the same revision to avoid lost decisions.
        info = self._private_info() if self.api is not None else None
        if info is not None:
            from huggingface_hub import CommitOperationAdd, hf_hub_download
            revision = parent_commit or info.sha
            path = hf_hub_download(repo_id=self.repo_id, repo_type='dataset', token=self.token,
                                   revision=revision, filename=f'submissions/{_submission_id(submission_id)}/record.json')
            record = json.loads(Path(path).read_bytes())
        else:
            record = self.load(submission_id)
        if record['submission_id'] != submission_id or record['status'] != 'pending_review':
            raise ValueError()
        record.update(status=status, review={'reviewer': reviewer,
                      'reviewed_at': datetime.now(timezone.utc).isoformat()})
        payload = _json_bytes(record)
        if info is not None:
            self._private_info()
            self.api.create_commit(repo_id=self.repo_id, repo_type='dataset', token=self.token,
                                   parent_commit=revision, commit_message='Record private review decision',
                                   operations=[CommitOperationAdd(
                                       path_in_repo=f'submissions/{submission_id}/record.json',
                                       path_or_fileobj=payload)])
        else:
            self._write_local_record(submission_id, record)