Tonic/g-android-control
0
1# Originally copied from https://github.com/huggingface/huggingface_hub/blob/d0a948fc2a32ed6e557042a95ef3e4af97ec4a7c/src/huggingface_hub/_commit_scheduler.py2 3import atexit4import logging5import os6import time7from concurrent.futures import Future8from dataclasses import dataclass9from io import SEEK_END, SEEK_SET, BytesIO10from pathlib import Path11from threading import Lock, Thread12from typing import Callable, Dict, List, Optional, Union13 14from huggingface_hub.hf_api import (15 DEFAULT_IGNORE_PATTERNS,16 CommitInfo,17 CommitOperationAdd,18 HfApi,19)20from huggingface_hub.utils import filter_repo_objects21 22logger = logging.getLogger(__name__)23 24 25@dataclass(frozen=True)26class _FileToUpload:27 """Temporary dataclass to store info about files to upload. Not meant to be used directly."""28 29 local_path: Path30 path_in_repo: str31 size_limit: int32 last_modified: float33 34 35class CommitScheduler:36 """37 Scheduler to upload a local folder to the Hub at regular intervals (e.g. push to hub every 5 minutes).38 39 The recommended way to use the scheduler is to use it as a context manager. This ensures that the scheduler is40 properly stopped and the last commit is triggered when the script ends. The scheduler can also be stopped manually41 with the `stop` method. Checkout the [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#scheduled-uploads)42 to learn more about how to use it.43 44 Args:45 repo_id (`str`):46 The id of the repo to commit to.47 folder_path (`str` or `Path`):48 Path to the local folder to upload regularly.49 every (`int` or `float`, *optional*):50 The number of minutes between each commit. Defaults to 5 minutes.51 path_in_repo (`str`, *optional*):52 Relative path of the directory in the repo, for example: `"checkpoints/"`. Defaults to the root folder53 of the repository.54 repo_type (`str`, *optional*):55 The type of the repo to commit to. Defaults to `model`.56 revision (`str`, *optional*):57 The revision of the repo to commit to. Defaults to `main`.58 private (`bool`, *optional*):59 Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.60 token (`str`, *optional*):61 The token to use to commit to the repo. Defaults to the token saved on the machine.62 allow_patterns (`List[str]` or `str`, *optional*):63 If provided, only files matching at least one pattern are uploaded.64 ignore_patterns (`List[str]` or `str`, *optional*):65 If provided, files matching any of the patterns are not uploaded.66 squash_history (`bool`, *optional*):67 Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is68 useful to avoid degraded performances on the repo when it grows too large.69 hf_api (`HfApi`, *optional*):70 The [`HfApi`] client to use to commit to the Hub. Can be set with custom settings (user agent, token,...).71 on_before_commit (`Callable[[], None]`, *optional*):72 If specified, a function that will be called before the CommitScheduler lists files to create a commit.73 74 Example:75 ```py76 >>> from pathlib import Path77 >>> from huggingface_hub import CommitScheduler78 79 # Scheduler uploads every 10 minutes80 >>> csv_path = Path("watched_folder/data.csv")81 >>> CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path=csv_path.parent, every=10)82 83 >>> with csv_path.open("a") as f:84 ... f.write("first line")85 86 # Some time later (...)87 >>> with csv_path.open("a") as f:88 ... f.write("second line")89 ```90 91 Example using a context manager:92 ```py93 >>> from pathlib import Path94 >>> from huggingface_hub import CommitScheduler95 96 >>> with CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path="watched_folder", every=10) as scheduler:97 ... csv_path = Path("watched_folder/data.csv")98 ... with csv_path.open("a") as f:99 ... f.write("first line")100 ... (...)101 ... with csv_path.open("a") as f:102 ... f.write("second line")103 104 # Scheduler is now stopped and last commit have been triggered105 ```106 """107 108 def __init__(109 self,110 *,111 repo_id: str,112 folder_path: Union[str, Path],113 every: Union[int, float] = 5,114 path_in_repo: Optional[str] = None,115 repo_type: Optional[str] = None,116 revision: Optional[str] = None,117 private: Optional[bool] = None,118 token: Optional[str] = None,119 allow_patterns: Optional[Union[List[str], str]] = None,120 ignore_patterns: Optional[Union[List[str], str]] = None,121 squash_history: bool = False,122 hf_api: Optional["HfApi"] = None,123 on_before_commit: Optional[Callable[[], None]] = None,124 ) -> None:125 self.api = hf_api or HfApi(token=token)126 self.on_before_commit = on_before_commit127 128 # Folder129 self.folder_path = Path(folder_path).expanduser().resolve()130 self.path_in_repo = path_in_repo or ""131 self.allow_patterns = allow_patterns132 133 if ignore_patterns is None:134 ignore_patterns = []135 elif isinstance(ignore_patterns, str):136 ignore_patterns = [ignore_patterns]137 self.ignore_patterns = ignore_patterns + DEFAULT_IGNORE_PATTERNS138 139 if self.folder_path.is_file():140 raise ValueError(141 f"'folder_path' must be a directory, not a file: '{self.folder_path}'."142 )143 self.folder_path.mkdir(parents=True, exist_ok=True)144 145 # Repository146 repo_url = self.api.create_repo(147 repo_id=repo_id, private=private, repo_type=repo_type, exist_ok=True148 )149 self.repo_id = repo_url.repo_id150 self.repo_type = repo_type151 self.revision = revision152 self.token = token153 154 # Keep track of already uploaded files155 self.last_uploaded: Dict[156 Path, float157 ] = {} # key is local path, value is timestamp158 159 # Scheduler160 if not every > 0:161 raise ValueError(f"'every' must be a positive integer, not '{every}'.")162 self.lock = Lock()163 self.every = every164 self.squash_history = squash_history165 166 logger.info(167 f"Scheduled job to push '{self.folder_path}' to '{self.repo_id}' every {self.every} minutes."168 )169 self._scheduler_thread = Thread(target=self._run_scheduler, daemon=True)170 self._scheduler_thread.start()171 atexit.register(self._push_to_hub)172 173 self.__stopped = False174 175 def stop(self) -> None:176 """Stop the scheduler.177 178 A stopped scheduler cannot be restarted. Mostly for tests purposes.179 """180 self.__stopped = True181 182 def __enter__(self) -> "CommitScheduler":183 return self184 185 def __exit__(self, exc_type, exc_value, traceback) -> None:186 # Upload last changes before exiting187 self.trigger().result()188 self.stop()189 return190 191 def _run_scheduler(self) -> None:192 """Dumb thread waiting between each scheduled push to Hub."""193 while True:194 self.last_future = self.trigger()195 time.sleep(self.every * 60)196 if self.__stopped:197 break198 199 def trigger(self) -> Future:200 """Trigger a `push_to_hub` and return a future.201 202 This method is automatically called every `every` minutes. You can also call it manually to trigger a commit203 immediately, without waiting for the next scheduled commit.204 """205 return self.api.run_as_future(self._push_to_hub)206 207 def _push_to_hub(self) -> Optional[CommitInfo]:208 if self.__stopped: # If stopped, already scheduled commits are ignored209 return None210 211 logger.info("(Background) scheduled commit triggered.")212 try:213 value = self.push_to_hub()214 if self.squash_history:215 logger.info("(Background) squashing repo history.")216 self.api.super_squash_history(217 repo_id=self.repo_id, repo_type=self.repo_type, branch=self.revision218 )219 return value220 except Exception as e:221 logger.error(222 f"Error while pushing to Hub: {e}"223 ) # Depending on the setup, error might be silenced224 raise225 226 def push_to_hub(self) -> Optional[CommitInfo]:227 """228 Push folder to the Hub and return the commit info.229 230 <Tip warning={true}>231 232 This method is not meant to be called directly. It is run in the background by the scheduler, respecting a233 queue mechanism to avoid concurrent commits. Making a direct call to the method might lead to concurrency234 issues.235 236 </Tip>237 238 The default behavior of `push_to_hub` is to assume an append-only folder. It lists all files in the folder and239 uploads only changed files. If no changes are found, the method returns without committing anything. If you want240 to change this behavior, you can inherit from [`CommitScheduler`] and override this method. This can be useful241 for example to compress data together in a single file before committing. For more details and examples, check242 out our [integration guide](https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#scheduled-uploads).243 """244 # Check files to upload (with lock)245 with self.lock:246 if self.on_before_commit is not None:247 self.on_before_commit()248 249 logger.debug("Listing files to upload for scheduled commit.")250 251 # List files from folder (taken from `_prepare_upload_folder_additions`)252 relpath_to_abspath = {253 path.relative_to(self.folder_path).as_posix(): path254 for path in sorted(255 self.folder_path.glob("**/*")256 ) # sorted to be deterministic257 if path.is_file()258 }259 prefix = f"{self.path_in_repo.strip('/')}/" if self.path_in_repo else ""260 261 # Filter with pattern + filter out unchanged files + retrieve current file size262 files_to_upload: List[_FileToUpload] = []263 for relpath in filter_repo_objects(264 relpath_to_abspath.keys(),265 allow_patterns=self.allow_patterns,266 ignore_patterns=self.ignore_patterns,267 ):268 local_path = relpath_to_abspath[relpath]269 stat = local_path.stat()270 if (271 self.last_uploaded.get(local_path) is None272 or self.last_uploaded[local_path] != stat.st_mtime273 ):274 files_to_upload.append(275 _FileToUpload(276 local_path=local_path,277 path_in_repo=prefix + relpath,278 size_limit=stat.st_size,279 last_modified=stat.st_mtime,280 )281 )282 283 # Return if nothing to upload284 if len(files_to_upload) == 0:285 logger.debug("Dropping schedule commit: no changed file to upload.")286 return None287 288 # Convert `_FileToUpload` as `CommitOperationAdd` (=> compute file shas + limit to file size)289 logger.debug("Removing unchanged files since previous scheduled commit.")290 add_operations = [291 CommitOperationAdd(292 # TODO: Cap the file to its current size, even if the user append data to it while a scheduled commit is happening293 # (requires an upstream fix for XET-535: `hf_xet` should support `BinaryIO` for upload)294 path_or_fileobj=file_to_upload.local_path,295 path_in_repo=file_to_upload.path_in_repo,296 )297 for file_to_upload in files_to_upload298 ]299 300 # Upload files (append mode expected - no need for lock)301 logger.debug("Uploading files for scheduled commit.")302 commit_info = self.api.create_commit(303 repo_id=self.repo_id,304 repo_type=self.repo_type,305 operations=add_operations,306 commit_message="Scheduled Commit",307 revision=self.revision,308 )309 310 # Successful commit: keep track of the latest "last_modified" for each file311 for file in files_to_upload:312 self.last_uploaded[file.local_path] = file.last_modified313 return commit_info314 315 316class PartialFileIO(BytesIO):317 """A file-like object that reads only the first part of a file.318 319 Useful to upload a file to the Hub when the user might still be appending data to it. Only the first part of the320 file is uploaded (i.e. the part that was available when the filesystem was first scanned).321 322 In practice, only used internally by the CommitScheduler to regularly push a folder to the Hub with minimal323 disturbance for the user. The object is passed to `CommitOperationAdd`.324 325 Only supports `read`, `tell` and `seek` methods.326 327 Args:328 file_path (`str` or `Path`):329 Path to the file to read.330 size_limit (`int`):331 The maximum number of bytes to read from the file. If the file is larger than this, only the first part332 will be read (and uploaded).333 """334 335 def __init__(self, file_path: Union[str, Path], size_limit: int) -> None:336 self._file_path = Path(file_path)337 self._file = self._file_path.open("rb")338 self._size_limit = min(size_limit, os.fstat(self._file.fileno()).st_size)339 340 def __del__(self) -> None:341 self._file.close()342 return super().__del__()343 344 def __repr__(self) -> str:345 return (346 f"<PartialFileIO file_path={self._file_path} size_limit={self._size_limit}>"347 )348 349 def __len__(self) -> int:350 return self._size_limit351 352 def __getattribute__(self, name: str):353 if name.startswith("_") or name in (354 "read",355 "tell",356 "seek",357 ): # only 3 public methods supported358 return super().__getattribute__(name)359 raise NotImplementedError(f"PartialFileIO does not support '{name}'.")360 361 def tell(self) -> int:362 """Return the current file position."""363 return self._file.tell()364 365 def seek(self, __offset: int, __whence: int = SEEK_SET) -> int:366 """Change the stream position to the given offset.367 368 Behavior is the same as a regular file, except that the position is capped to the size limit.369 """370 if __whence == SEEK_END:371 # SEEK_END => set from the truncated end372 __offset = len(self) + __offset373 __whence = SEEK_SET374 375 pos = self._file.seek(__offset, __whence)376 if pos > self._size_limit:377 return self._file.seek(self._size_limit)378 return pos379 380 def read(self, __size: Optional[int] = -1) -> bytes:381 """Read at most `__size` bytes from the file.382 383 Behavior is the same as a regular file, except that it is capped to the size limit.384 """385 current = self._file.tell()386 if __size is None or __size < 0:387 # Read until file limit388 truncated_size = self._size_limit - current389 else:390 # Read until file limit or __size391 truncated_size = min(__size, self._size_limit - current)392 return self._file.read(truncated_size)393