File size: 6,834 Bytes
3825ff2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f000a32
 
3825ff2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
# -*- coding: utf-8 -*-
"""The submission ledger, kept in a private Hugging Face dataset.



A Space's container filesystem is ephemeral. Every rebuild, restart or sleep wipes it,

and with it every submission and every score. That is survivable for a demo and not

survivable for a contest with a prize attached, so the ledger lives outside the container

in a dataset repo that the entrant's work outlives the service.



Layout - one file per record, never appended to:



    submissions/<id>.json   written by the web service when an entry is accepted

    results/<id>.json       written by the GPU worker when it finishes scoring

    leaderboard.json        rolled up by the worker so the page reads one file



One file per record is what makes concurrent writers safe. Two entrants submitting at the

same moment touch different paths, so neither commit can clobber the other - which an

append to a shared JSONL absolutely would.



The rollup exists because a page load must not fan out into one request per entry. The

worker already holds every score at the moment it writes one, so it is the natural place

to rebuild the table.

"""
import base64
import json
import os
import threading
import time
import urllib.error
import urllib.request

REPO = os.environ.get("OMC_DATASET", "FINAL-Bench/omc-submissions")
TOKEN = os.environ.get("HF_TOKEN", "")
API = "https://huggingface.co/api/datasets/%s" % REPO
RESOLVE = "https://huggingface.co/datasets/%s/resolve/main" % REPO


def _hdr(extra=None):
    h = {"User-Agent": "VIDRAFT-OMC/1.0"}
    if TOKEN:
        h["Authorization"] = "Bearer " + TOKEN
    if extra:
        h.update(extra)
    return h


def _get_json(url, timeout=30):
    req = urllib.request.Request(url, headers=_hdr())
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode())


def read(path, default=None):
    """Fetch one record. Missing files are a normal state, not an error."""
    try:
        return _get_json("%s/%s" % (RESOLVE, path))
    except urllib.error.HTTPError as e:
        if e.code in (404, 401, 403):
            return default
        raise
    except Exception:
        return default


def write(path, obj, summary=None):
    """Commit one record. NDJSON is the format this endpoint takes - a plain JSON body

    is accepted and then quietly does nothing, which is a long way to debug."""
    blob = base64.b64encode(json.dumps(obj, ensure_ascii=False).encode()).decode()
    lines = [
        json.dumps({"key": "header", "value": {"summary": summary or ("write " + path)}}),
        json.dumps({"key": "file", "value": {"path": path, "content": blob,
                                             "encoding": "base64"}}),
    ]
    body = ("\n".join(lines) + "\n").encode()
    req = urllib.request.Request(API + "/commit/main", data=body,
                                 headers=_hdr({"Content-Type": "application/x-ndjson"}))
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.loads(r.read().decode())


MAX_PAGES = int(os.environ.get("OMC_MAX_PAGES", "200"))


def _tree_pages(prefix):
    """Every page of a tree listing, following the Link cursor.



    The Hub caps a page at 1,000 entries. Reading only the first page was silent while

    the dataset was small and started dropping records the moment it was not."""
    url = "%s/tree/main/%s" % (API, prefix)
    seen = 0
    for _ in range(MAX_PAGES):
        req = urllib.request.Request(url, headers=_hdr())
        with urllib.request.urlopen(req, timeout=60) as r:
            page = json.loads(r.read().decode())
            link = r.headers.get("Link") or ""
        yield page
        seen += len(page)
        nxt = ""
        for part in link.split(","):
            if 'rel="next"' in part and "<" in part:
                nxt = part[part.index("<") + 1:part.index(">")]
        if not nxt:
            return
        url = nxt
    raise RuntimeError("tree listing exceeded %d pages at %s (%d entries)"
                       % (MAX_PAGES, prefix, seen))


def listdir(prefix):
    """Record ids under a prefix. Absent directory means nothing has been written yet."""
    out = []
    try:
        for page in _tree_pages(prefix):
            for e in page:
                q = e.get("path", "")
                if q.endswith(".json") and not os.path.basename(q).startswith("_"):
                    out.append(os.path.basename(q)[:-5])
    except urllib.error.HTTPError as e:
        if e.code == 404:
            return []
        raise
    except Exception:
        # a partial listing is worse than none: the caller would treat the missing ids as
        # unscored and the rollup would drop them, which is exactly the failure this fixes
        raise
    return out


class Cached:
    """A small TTL cache so a page refresh does not become a round trip to the Hub.



    Staleness is bounded and harmless here: the worst case is a leaderboard a few seconds

    behind, and the page polls anyway.



    **Single flight.** One refresher at a time; everyone else is served the value already

    held, so an expiry does not send every concurrent request to the Hub at once.



    **Stale beats blocking, and stale beats an error.** A leaderboard a minute old is a

    working page. A timeout is not.

    """

    def __init__(self, ttl=20):
        self.ttl = ttl
        self._v = {}
        self._locks = {}
        self._guard = threading.Lock()

    def _lock_for(self, key):
        with self._guard:
            lk = self._locks.get(key)
            if lk is None:
                lk = self._locks[key] = threading.Lock()
            return lk

    def get(self, key, produce):
        hit = self._v.get(key)
        if hit and time.time() - hit[0] < self.ttl:
            return hit[1]

        lk = self._lock_for(key)
        # Block only when there is nothing at all to serve. If a refresh is already in
        # flight and we hold a stale value, hand that back instead of joining the queue.
        if not lk.acquire(blocking=(hit is None)):
            return hit[1]
        try:
            hit = self._v.get(key)          # the winner may have filled it while we waited
            if hit and time.time() - hit[0] < self.ttl:
                return hit[1]
            try:
                val = produce()
            except Exception:
                if hit:
                    return hit[1]
                raise
            self._v[key] = (time.time(), val)
            return val
        finally:
            lk.release()

    def drop(self, key=None):
        if key is None:
            self._v.clear()
        else:
            self._v.pop(key, None)