import fnmatch import logging import os import pathlib import time import zipfile from collections.abc import Iterable, Sequence from concurrent.futures import ThreadPoolExecutor from datetime import datetime from tempfile import TemporaryDirectory import requests from kagglesdk.blobs.types.blob_api_service import ApiBlobType, ApiStartBlobUploadRequest from kagglesdk.datasets.types.dataset_api_service import ApiDatasetNewFile, ApiUploadDirectoryInfo from requests.exceptions import Timeout from tqdm import tqdm from tqdm.utils import CallbackIOWrapper from kagglehub.clients import build_kaggle_client from kagglehub.exceptions import BackendError, handle_call logger = logging.getLogger(__name__) MAX_FILES_TO_UPLOAD = 50 TEMP_ARCHIVE_FILE = "archive.zip" MAX_RETRIES = 5 REQUEST_TIMEOUT = 600 MAX_PARALLEL_UPLOADS = 8 class UploadDirectoryInfo: def __init__( self, name: str, files: list[str] | None = None, directories: list["UploadDirectoryInfo"] | None = None, ): self.name = name self.files = files if files is not None else [] self.directories = directories if directories is not None else [] def to_proto(self) -> ApiUploadDirectoryInfo: d = ApiUploadDirectoryInfo() d.name = self.name d.files = [] for file in self.files: f = ApiDatasetNewFile() f.token = file d.files.append(f) d.directories = [d.to_proto() for d in self.directories] return d def _get_or_create_upload_directory(root: UploadDirectoryInfo, path: str) -> UploadDirectoryInfo: current = root if path == ".": return current for part in path.split(os.sep): for subdir in current.directories: if subdir.name == part: current = subdir break else: new_dir = UploadDirectoryInfo(name=part) current.directories.append(new_dir) current = new_dir return current def parse_datetime_string(string: str) -> datetime | str: time_formats = ["%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S.%fZ"] for t in time_formats: try: return datetime.strptime(string[:26], t).replace(microsecond=0) # noqa: DTZ007 except: # noqa: E722, S110 pass return string class File(object): # noqa: UP004 def __init__(self, init_dict: dict) -> None: parsed_dict = {k: parse_datetime_string(v) for k, v in init_dict.items()} self.__dict__.update(parsed_dict) @staticmethod def get_size(size: float, precision: int = 0) -> str: suffixes = ["B", "KB", "MB", "GB", "TB"] suffix_index = 0 while size >= 1024 and suffix_index < 4: # noqa: PLR2004 suffix_index += 1 size /= 1024.0 return f"{size:.{precision}f}{suffixes[suffix_index]}" def filtered_walk(*, base_dir: str, ignore_patterns: Sequence[str]) -> Iterable[tuple[str, list[str], list[str]]]: """An `os.walk` like directory tree generator with filtering. This method filters out files matching any ignore pattern. Args: base_dir (str): The base dir to walk in. ignore_patterns (Sequence[str]): The patterns for ignored files. These are standard wildcards relative to base_dir. Yields: Iterable[tuple[str, list[str], list[str]]]: (base_dir_path, list[dir_names], list[filtered_file_names]) """ for dir_path, dir_names, file_names in os.walk(base_dir): dir_p = pathlib.Path(dir_path) filtered_files = [] for file_name in file_names: rel_file_p = (dir_p / file_name).relative_to(base_dir) if not any(fnmatch.fnmatch(name=str(rel_file_p), pat=pat) for pat in ignore_patterns): filtered_files.append(file_name) if filtered_files: yield (dir_path, dir_names, filtered_files) def _check_uploaded_size(session_uri: str, file_size: int, backoff_factor: int = 1) -> int: """Check the status of the resumable upload.""" headers = {"Content-Length": "0", "Content-Range": f"bytes */{file_size}"} retry_count = 0 while retry_count < MAX_RETRIES: try: response = requests.put(session_uri, headers=headers, timeout=REQUEST_TIMEOUT) if response.status_code == 308: # Resume Incomplete # noqa: PLR2004 range_header = response.headers.get("Range") if range_header: bytes_uploaded = int(range_header.split("-")[1]) + 1 return bytes_uploaded return 0 # If no Range header, assume no bytes were uploaded else: return file_size except (requests.ConnectionError, Timeout): logger.info(f"Network issue while checking uploaded size, retrying in {backoff_factor} seconds...") time.sleep(backoff_factor) backoff_factor = min(backoff_factor * 2, 60) retry_count += 1 return 0 # Return 0 if all retries fail def _upload_blob(file_path: str, item_type: ApiBlobType) -> str: """Uploads a file to a remote server as a blob and returns an upload token. Args: file_path: The path to the file to be uploaded. item_type : The type of the item associated with the file. Returns: A str token of uploaded blob. """ file_size = os.path.getsize(file_path) with build_kaggle_client() as api_client: r = ApiStartBlobUploadRequest() r.type = item_type r.name = os.path.basename(file_path) r.content_length = file_size r.last_modified_epoch_seconds = int(os.path.getmtime(file_path)) response = handle_call(lambda: api_client.blobs.blob_api_client.start_blob_upload(r)) session_uri = response.create_url headers = {"Content-Type": "application/octet-stream"} retry_count = 0 uploaded_bytes = 0 backoff_factor = 1 # Initial backoff duration in seconds with open(file_path, "rb") as f, tqdm(total=file_size, desc="Uploading", unit="B", unit_scale=True) as pbar: while retry_count < MAX_RETRIES and (file_size == 0 or uploaded_bytes < file_size): try: # Special case for empty files. if file_size == 0: headers["Content-Length"] = "0" upload_data = None # Resumable upload for non-empty files. else: f.seek(uploaded_bytes) headers["Content-Range"] = f"bytes {uploaded_bytes}-{file_size - 1}/{file_size}" upload_data = CallbackIOWrapper(pbar.update, f, "read") upload_response = requests.put(session_uri, headers=headers, data=upload_data, timeout=REQUEST_TIMEOUT) if upload_response.status_code in [200, 201]: return response.token elif upload_response.status_code == 308: # Resume Incomplete # noqa: PLR2004 uploaded_bytes = _check_uploaded_size(session_uri, file_size) else: upload_failed_exception = ( f"Upload failed with status code {upload_response.status_code}: {upload_response.text}" ) raise BackendError(upload_failed_exception) except (requests.ConnectionError, requests.Timeout) as e: logger.info(f"Network issue: {e}, retrying in {backoff_factor} seconds...") time.sleep(backoff_factor) backoff_factor = min(backoff_factor * 2, 60) retry_count += 1 uploaded_bytes = _check_uploaded_size(session_uri, file_size) pbar.n = uploaded_bytes # Update progress bar to reflect actual uploaded bytes return response.token def upload_files_and_directories( folder: str, *, ignore_patterns: Sequence[str], item_type: ApiBlobType, quiet: bool = False, ) -> UploadDirectoryInfo: # Count the total number of files file_count = 0 for _, _, files in filtered_walk(base_dir=folder, ignore_patterns=ignore_patterns): file_count += len(files) if file_count > MAX_FILES_TO_UPLOAD: if not quiet: logger.info(f"More than {MAX_FILES_TO_UPLOAD} files detected, creating a zip archive...") with TemporaryDirectory() as temp_dir: zip_path = os.path.join(temp_dir, TEMP_ARCHIVE_FILE) with zipfile.ZipFile(zip_path, "w") as zipf: for root, _, files in filtered_walk(base_dir=folder, ignore_patterns=ignore_patterns): for file in files: file_path = os.path.join(root, file) zipf.write(file_path, os.path.relpath(file_path, folder)) archive_tokens = [ token for token in [_upload_file(file_path=zip_path, item_type=item_type, quiet=quiet)] if token is not None ] return UploadDirectoryInfo(name="archive", files=archive_tokens) if os.path.isfile(folder): root_dict = UploadDirectoryInfo(name="root") # Directly upload the file if the path is a file token = _upload_file(file_path=folder, item_type=item_type, quiet=quiet) if token: root_dict.files.append(token) else: root_dict = UploadDirectoryInfo(name="root") files_to_upload = [] for root, _, files in filtered_walk(base_dir=folder, ignore_patterns=ignore_patterns): # Path of the current folder relative to the base folder path = os.path.relpath(root, folder) # Build the directory structure before parallel uploads so only the main thread mutates it. current_dict = _get_or_create_upload_directory(root_dict, path) for file in files: files_to_upload.append((current_dict, os.path.join(root, file))) if files_to_upload: max_workers = min(MAX_PARALLEL_UPLOADS, len(files_to_upload)) with ThreadPoolExecutor(max_workers=max_workers) as executor: upload_tokens = executor.map( lambda item: _upload_file(file_path=item[1], item_type=item_type, quiet=quiet), files_to_upload, ) for (current_dict, _), token in zip(files_to_upload, upload_tokens, strict=False): if token: current_dict.files.append(token) return root_dict def _upload_file(file_path: str, *, quiet: bool, item_type: ApiBlobType) -> str | None: """Helper function to upload a single file. Args: full_path: path to the file to upload quiet: suppress verbose output item_type: Type of the item that is being uploaded. Returns: A str token of uploaded file if successful, otherwise None. """ if not quiet: logger.info("Starting upload for file " + file_path) if not os.path.isfile(file_path): logger.warn("Skip uploading %s because it is not a file.", file_path) return None content_length = os.path.getsize(file_path) token = _upload_blob(file_path, item_type) if not quiet: logger.info("Upload successful: " + file_path + " (" + File.get_size(content_length) + ")") return token def normalize_patterns(*, default: list[str], additional: list[str] | str | None) -> list[str]: """Merges additional patterns with the default, and normalize the dir pattern with wildcard.""" def add_wildcard_to_dir(pattern: str) -> str: return pattern + "*" if pattern.endswith("/") else pattern if additional is None: additional = [] elif isinstance(additional, str): additional = [additional] return [add_wildcard_to_dir(pattern) for pattern in default + additional]