WalisonCruz/function-gemma
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, 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: str | None = None,115 repo_type: str | None = None,116 revision: str | None = None,117 private: bool | None = None,118 token: str | None = None,119 allow_patterns: list[str] | str | None = None,120 ignore_patterns: list[str] | str | None = None,121 squash_history: bool = False,122 hf_api: HfApi | None = None,123 on_before_commit: Callable[[], None] | 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 self.last_uploaded: Dict[Path, float] = {}155 self.last_push_time: float | None = None156 157 if not every > 0:158 raise ValueError(f"'every' must be a positive integer, not '{every}'.")159 self.lock = Lock()160 self.every = every161 self.squash_history = squash_history162 163 logger.info(164 f"Scheduled job to push '{self.folder_path}' to '{self.repo_id}' every {self.every} minutes."165 )166 self._scheduler_thread = Thread(target=self._run_scheduler, daemon=True)167 self._scheduler_thread.start()168 atexit.register(self._push_to_hub)169 170 self.__stopped = False171 172 def stop(self) -> None:173 """Stop the scheduler.174 175 A stopped scheduler cannot be restarted. Mostly for tests purposes.176 """177 self.__stopped = True178 179 def __enter__(self) -> "CommitScheduler":180 return self181 182 def __exit__(self, exc_type, exc_value, traceback) -> None:183 # Upload last changes before exiting184 self.trigger().result()185 self.stop()186 return187 188 def _run_scheduler(self) -> None:189 """Dumb thread waiting between each scheduled push to Hub."""190 while True:191 self.last_future = self.trigger()192 time.sleep(self.every * 60)193 if self.__stopped:194 break195 196 def trigger(self) -> Future:197 """Trigger a `push_to_hub` and return a future.198 199 This method is automatically called every `every` minutes. You can also call it manually to trigger a commit200 immediately, without waiting for the next scheduled commit.201 """202 return self.api.run_as_future(self._push_to_hub)203 204 def _push_to_hub(self) -> CommitInfo | None:205 if self.__stopped: # If stopped, already scheduled commits are ignored206 return None207 208 logger.info("(Background) scheduled commit triggered.")209 try:210 value = self.push_to_hub()211 if self.squash_history:212 logger.info("(Background) squashing repo history.")213 self.api.super_squash_history(214 repo_id=self.repo_id, repo_type=self.repo_type, branch=self.revision215 )216 return value217 except Exception as e:218 logger.error(219 f"Error while pushing to Hub: {e}"220 ) # Depending on the setup, error might be silenced221 raise222 223 def push_to_hub(self) -> CommitInfo | None:224 """225 Push folder to the Hub and return the commit info.226 227 <Tip warning={true}>228 229 This method is not meant to be called directly. It is run in the background by the scheduler, respecting a230 queue mechanism to avoid concurrent commits. Making a direct call to the method might lead to concurrency231 issues.232 233 </Tip>234 235 The default behavior of `push_to_hub` is to assume an append-only folder. It lists all files in the folder and236 uploads only changed files. If no changes are found, the method returns without committing anything. If you want237 to change this behavior, you can inherit from [`CommitScheduler`] and override this method. This can be useful238 for example to compress data together in a single file before committing. For more details and examples, check239 out our [integration guide](https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#scheduled-uploads).240 """241 # Check files to upload (with lock)242 with self.lock:243 if self.on_before_commit is not None:244 self.on_before_commit()245 246 logger.debug("Listing files to upload for scheduled commit.")247 248 # List files from folder (taken from `_prepare_upload_folder_additions`)249 relpath_to_abspath = {250 path.relative_to(self.folder_path).as_posix(): path251 for path in sorted(252 self.folder_path.glob("**/*")253 ) # sorted to be deterministic254 if path.is_file()255 }256 prefix = f"{self.path_in_repo.strip('/')}/" if self.path_in_repo else ""257 258 # Filter with pattern + filter out unchanged files + retrieve current file size259 files_to_upload: List[_FileToUpload] = []260 for relpath in filter_repo_objects(261 relpath_to_abspath.keys(),262 allow_patterns=self.allow_patterns,263 ignore_patterns=self.ignore_patterns,264 ):265 local_path = relpath_to_abspath[relpath]266 stat = local_path.stat()267 if (268 self.last_uploaded.get(local_path) is None269 or self.last_uploaded[local_path] != stat.st_mtime270 ):271 files_to_upload.append(272 _FileToUpload(273 local_path=local_path,274 path_in_repo=prefix + relpath,275 size_limit=stat.st_size,276 last_modified=stat.st_mtime,277 )278 )279 280 # Return if nothing to upload281 if len(files_to_upload) == 0:282 logger.debug("Dropping schedule commit: no changed file to upload.")283 return None284 285 # Convert `_FileToUpload` as `CommitOperationAdd` (=> compute file shas + limit to file size)286 logger.debug("Removing unchanged files since previous scheduled commit.")287 add_operations = [288 CommitOperationAdd(289 # TODO: Cap the file to its current size, even if the user append data to it while a scheduled commit is happening290 # (requires an upstream fix for XET-535: `hf_xet` should support `BinaryIO` for upload)291 path_or_fileobj=file_to_upload.local_path,292 path_in_repo=file_to_upload.path_in_repo,293 )294 for file_to_upload in files_to_upload295 ]296 297 # Upload files (append mode expected - no need for lock)298 logger.debug("Uploading files for scheduled commit.")299 commit_info = self.api.create_commit(300 repo_id=self.repo_id,301 repo_type=self.repo_type,302 operations=add_operations,303 commit_message="Scheduled Commit",304 revision=self.revision,305 )306 307 for file in files_to_upload:308 self.last_uploaded[file.local_path] = file.last_modified309 310 self.last_push_time = time.time()311 312 return commit_info313 314 315class PartialFileIO(BytesIO):316 """A file-like object that reads only the first part of a file.317 318 Useful to upload a file to the Hub when the user might still be appending data to it. Only the first part of the319 file is uploaded (i.e. the part that was available when the filesystem was first scanned).320 321 In practice, only used internally by the CommitScheduler to regularly push a folder to the Hub with minimal322 disturbance for the user. The object is passed to `CommitOperationAdd`.323 324 Only supports `read`, `tell` and `seek` methods.325 326 Args:327 file_path (`str` or `Path`):328 Path to the file to read.329 size_limit (`int`):330 The maximum number of bytes to read from the file. If the file is larger than this, only the first part331 will be read (and uploaded).332 """333 334 def __init__(self, file_path: Union[str, Path], size_limit: int) -> None:335 self._file_path = Path(file_path)336 self._file = self._file_path.open("rb")337 self._size_limit = min(size_limit, os.fstat(self._file.fileno()).st_size)338 339 def __del__(self) -> None:340 self._file.close()341 return super().__del__()342 343 def __repr__(self) -> str:344 return (345 f"<PartialFileIO file_path={self._file_path} size_limit={self._size_limit}>"346 )347 348 def __len__(self) -> int:349 return self._size_limit350 351 def __getattribute__(self, name: str):352 if name.startswith("_") or name in (353 "read",354 "tell",355 "seek",356 ): # only 3 public methods supported357 return super().__getattribute__(name)358 raise NotImplementedError(f"PartialFileIO does not support '{name}'.")359 360 def tell(self) -> int:361 """Return the current file position."""362 return self._file.tell()363 364 def seek(self, __offset: int, __whence: int = SEEK_SET) -> int:365 """Change the stream position to the given offset.366 367 Behavior is the same as a regular file, except that the position is capped to the size limit.368 """369 if __whence == SEEK_END:370 # SEEK_END => set from the truncated end371 __offset = len(self) + __offset372 __whence = SEEK_SET373 374 pos = self._file.seek(__offset, __whence)375 if pos > self._size_limit:376 return self._file.seek(self._size_limit)377 return pos378 379 def read(self, __size: int | None = -1) -> bytes:380 """Read at most `__size` bytes from the file.381 382 Behavior is the same as a regular file, except that it is capped to the size limit.383 """384 current = self._file.tell()385 if __size is None or __size < 0:386 # Read until file limit387 truncated_size = self._size_limit - current388 else:389 # Read until file limit or __size390 truncated_size = min(__size, self._size_limit - current)391 return self._file.read(truncated_size)392