Tonic/g-android-control

1import threading2import time3 4import huggingface_hub5from gradio_client import Client6 7from trackio.sqlite_storage import SQLiteStorage8from trackio.typehints import LogEntry9from trackio.utils import RESERVED_KEYS, fibo, generate_readable_name10 11 12class Run:13    def __init__(14        self,15        url: str,16        project: str,17        client: Client | None,18        name: str | None = None,19        config: dict | None = None,20    ):21        self.url = url22        self.project = project23        self._client_lock = threading.Lock()24        self._client_thread = None25        self._client = client26        self.name = name or generate_readable_name(SQLiteStorage.get_runs(project))27        self.config = config or {}28        self._queued_logs: list[LogEntry] = []29        self._stop_flag = threading.Event()30 31        self._client_thread = threading.Thread(target=self._init_client_background)32        self._client_thread.daemon = True33        self._client_thread.start()34 35    def _batch_sender(self):36        """Send batched logs every 500ms."""37        while not self._stop_flag.is_set():38            time.sleep(0.5)39 40            with self._client_lock:41                if self._queued_logs and self._client is not None:42                    logs_to_send = self._queued_logs.copy()43                    self._queued_logs.clear()44 45                    self._client.predict(46                        api_name="/bulk_log",47                        logs=logs_to_send,48                        hf_token=huggingface_hub.utils.get_token(),49                    )50 51    def _init_client_background(self):52        if self._client is None:53            fib = fibo()54            for sleep_coefficient in fib:55                try:56                    client = Client(self.url, verbose=False)57                    with self._client_lock:58                        self._client = client59                    break60                except Exception:61                    pass62                if sleep_coefficient is not None:63                    time.sleep(0.1 * sleep_coefficient)64 65        self._batch_sender()66 67    def log(self, metrics: dict, step: int | None = None):68        for k in metrics.keys():69            if k in RESERVED_KEYS or k.startswith("__"):70                raise ValueError(71                    f"Please do not use this reserved key as a metric: {k}"72                )73 74        log_entry: LogEntry = {75            "project": self.project,76            "run": self.name,77            "metrics": metrics,78            "step": step,79        }80 81        with self._client_lock:82            self._queued_logs.append(log_entry)83 84    def finish(self):85        """Cleanup when run is finished."""86        self._stop_flag.set()87 88        with self._client_lock:89            if self._queued_logs and self._client is not None:90                logs_to_send = self._queued_logs.copy()91                self._queued_logs.clear()92                self._client.predict(93                    api_name="/bulk_log",94                    logs=logs_to_send,95                    hf_token=huggingface_hub.utils.get_token(),96                )97 98        if self._client_thread is not None:99            print(f"* Uploading logs to Trackio Space: {self.url} (please wait...)")100            self._client_thread.join(timeout=30)101