116 lines
3.7 KiB
Python
116 lines
3.7 KiB
Python
|
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||
|
|
"""Upload utilities for Ultralytics, mirroring downloads.py patterns."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
from time import sleep
|
||
|
|
|
||
|
|
from ultralytics.utils import LOGGER, TQDM
|
||
|
|
|
||
|
|
|
||
|
|
class _ProgressReader:
|
||
|
|
"""File wrapper that reports read progress for upload monitoring."""
|
||
|
|
|
||
|
|
def __init__(self, file_path, pbar):
|
||
|
|
self.file = open(file_path, "rb")
|
||
|
|
self.pbar = pbar
|
||
|
|
self._size = os.path.getsize(file_path)
|
||
|
|
|
||
|
|
def read(self, size=-1):
|
||
|
|
"""Read data and update progress bar."""
|
||
|
|
data = self.file.read(size)
|
||
|
|
if data and self.pbar:
|
||
|
|
self.pbar.update(len(data))
|
||
|
|
return data
|
||
|
|
|
||
|
|
def __len__(self):
|
||
|
|
"""Return file size for Content-Length header."""
|
||
|
|
return self._size
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
"""Close the file."""
|
||
|
|
self.file.close()
|
||
|
|
|
||
|
|
|
||
|
|
def safe_upload(
|
||
|
|
file: str | Path,
|
||
|
|
url: str,
|
||
|
|
headers: dict | None = None,
|
||
|
|
retry: int = 2,
|
||
|
|
timeout: int = 600,
|
||
|
|
progress: bool = False,
|
||
|
|
) -> bool:
|
||
|
|
"""Upload a file to a URL with retry logic and optional progress bar.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
file (str | Path): Path to the file to upload.
|
||
|
|
url (str): The URL endpoint to upload the file to (e.g., signed GCS URL).
|
||
|
|
headers (dict, optional): Additional headers to include in the request.
|
||
|
|
retry (int, optional): Number of retry attempts on failure (default: 2 for 3 total attempts).
|
||
|
|
timeout (int, optional): Request timeout in seconds.
|
||
|
|
progress (bool, optional): Whether to display a progress bar during upload.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
(bool): True if upload succeeded, False otherwise.
|
||
|
|
|
||
|
|
Examples:
|
||
|
|
>>> from ultralytics.utils.uploads import safe_upload
|
||
|
|
>>> success = safe_upload("model.pt", "https://storage.googleapis.com/...", progress=True)
|
||
|
|
"""
|
||
|
|
import requests
|
||
|
|
|
||
|
|
file = Path(file)
|
||
|
|
if not file.exists():
|
||
|
|
raise FileNotFoundError(f"File not found: {file}")
|
||
|
|
|
||
|
|
file_size = file.stat().st_size
|
||
|
|
desc = f"Uploading {file.name}"
|
||
|
|
|
||
|
|
# Prepare headers (Content-Length set automatically from file size)
|
||
|
|
upload_headers = {"Content-Type": "application/octet-stream"}
|
||
|
|
if headers:
|
||
|
|
upload_headers.update(headers)
|
||
|
|
|
||
|
|
last_error = None
|
||
|
|
for attempt in range(retry + 1):
|
||
|
|
pbar = None
|
||
|
|
reader = None
|
||
|
|
try:
|
||
|
|
if progress:
|
||
|
|
pbar = TQDM(total=file_size, desc=desc, unit="B", unit_scale=True, unit_divisor=1024)
|
||
|
|
reader = _ProgressReader(file, pbar)
|
||
|
|
|
||
|
|
r = requests.put(url, data=reader, headers=upload_headers, timeout=timeout)
|
||
|
|
r.raise_for_status()
|
||
|
|
reader.close()
|
||
|
|
reader = None # Prevent double-close in finally
|
||
|
|
if pbar:
|
||
|
|
pbar.close()
|
||
|
|
pbar = None
|
||
|
|
LOGGER.info(f"Uploaded {file.name} ✅")
|
||
|
|
return True
|
||
|
|
|
||
|
|
except requests.exceptions.HTTPError as e:
|
||
|
|
status = e.response.status_code if e.response is not None else 0
|
||
|
|
if 400 <= status < 500 and status not in {408, 429}:
|
||
|
|
LOGGER.warning(f"{desc} failed: {status} {getattr(e.response, 'reason', '')}")
|
||
|
|
return False
|
||
|
|
last_error = f"HTTP {status}"
|
||
|
|
except Exception as e:
|
||
|
|
last_error = str(e)
|
||
|
|
finally:
|
||
|
|
if reader:
|
||
|
|
reader.close()
|
||
|
|
if pbar:
|
||
|
|
pbar.close()
|
||
|
|
|
||
|
|
if attempt < retry:
|
||
|
|
wait_time = 2 ** (attempt + 1)
|
||
|
|
LOGGER.warning(f"{desc} failed ({last_error}), retrying {attempt + 1}/{retry} in {wait_time}s...")
|
||
|
|
sleep(wait_time)
|
||
|
|
|
||
|
|
LOGGER.warning(f"{desc} failed after {retry + 1} attempts: {last_error}")
|
||
|
|
return False
|