Spaces:
Paused
Paused
File size: 4,827 Bytes
3be8894 | 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 | import copy
import io
import json
import os
import tempfile
import threading
from datetime import datetime, timezone
def utc_now_iso():
return datetime.now(timezone.utc).isoformat(timespec="seconds")
class HubStateStore:
"""Small JSON state store backed by a private Hugging Face dataset."""
def __init__(self, data_dir):
self.repo_id = os.getenv("STATE_HUB_REPO_ID")
self.repo_type = os.getenv("STATE_HUB_REPO_TYPE", "dataset")
self.filename = os.getenv("STATE_HUB_FILENAME", "monitor_state.json")
token = os.getenv("HF_TOKEN")
self.path = os.path.join(data_dir, self.filename)
self._lock = threading.RLock()
self._api = None
self.last_error = None
self.remote_synced = False
if self.repo_id and token:
try:
from huggingface_hub import HfApi
self._api = HfApi(token=token)
except Exception as exc:
self.last_error = str(exc)
self._state = self._load()
@property
def enabled(self):
return self._api is not None and bool(self.repo_id)
def _empty(self):
return {"updated_at": "1970-01-01T00:00:00+00:00", "paused_monitor_ids": [], "values": {}}
def _decode(self, path):
with open(path, "r", encoding="utf-8") as handle:
state = json.load(handle)
if not isinstance(state, dict):
raise ValueError("state root must be an object")
state.setdefault("paused_monitor_ids", [])
state.setdefault("values", {})
if not isinstance(state["paused_monitor_ids"], list):
state["paused_monitor_ids"] = []
if not isinstance(state["values"], dict):
state["values"] = {}
return state
def _write_local(self, state):
os.makedirs(os.path.dirname(self.path), exist_ok=True)
fd, temp_path = tempfile.mkstemp(prefix=".monitor-state-", dir=os.path.dirname(self.path))
try:
with os.fdopen(fd, "w", encoding="utf-8") as handle:
json.dump(state, handle, ensure_ascii=False, indent=2, sort_keys=True)
os.replace(temp_path, self.path)
except Exception:
try:
os.unlink(temp_path)
except OSError:
pass
raise
def _load(self):
local = self._empty()
if os.path.exists(self.path):
try:
local = self._decode(self.path)
except Exception as exc:
self.last_error = str(exc)
if not self.enabled:
return local
try:
from huggingface_hub import hf_hub_download
remote_path = hf_hub_download(
repo_id=self.repo_id,
repo_type=self.repo_type,
filename=self.filename,
force_download=True,
)
remote = self._decode(remote_path)
if str(remote.get("updated_at") or "") >= str(local.get("updated_at") or ""):
self.remote_synced = True
self._write_local(remote)
return remote
return local
except Exception as exc:
self.last_error = str(exc)
return local
def _upload(self, state):
if not self.enabled:
return False
try:
payload = json.dumps(state, ensure_ascii=False, indent=2, sort_keys=True).encode("utf-8")
self._api.upload_file(
path_or_fileobj=io.BytesIO(payload),
path_in_repo=self.filename,
repo_id=self.repo_id,
repo_type=self.repo_type,
commit_message="Update monitor state",
)
self.last_error = None
self.remote_synced = True
return True
except Exception as exc:
self.last_error = str(exc)
self.remote_synced = False
return False
def get(self, key, default=None):
with self._lock:
if key == "paused_monitor_ids":
return copy.deepcopy(self._state.get(key, []))
return copy.deepcopy(self._state.get("values", {}).get(key, default))
def update(self, **updates):
with self._lock:
self._state.update(updates)
self._state["updated_at"] = utc_now_iso()
state = copy.deepcopy(self._state)
self._write_local(state)
self._upload(state)
def set_kv(self, key, value):
with self._lock:
self._state.setdefault("values", {})[key] = copy.deepcopy(value)
self._state["updated_at"] = utc_now_iso()
state = copy.deepcopy(self._state)
self._write_local(state)
self._upload(state)
|