420 lines
13 KiB
Python
420 lines
13 KiB
Python
"""
|
|
S3 client for RunPod network volume operations.
|
|
|
|
Provides:
|
|
- File uploads with progress tracking
|
|
- Log file tailing via byte-range requests
|
|
- Model download
|
|
- Volume inventory listing
|
|
"""
|
|
|
|
import hashlib
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
import boto3
|
|
from botocore.config import Config
|
|
from botocore.exceptions import ClientError
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
BarColumn,
|
|
DownloadColumn,
|
|
Progress,
|
|
TaskID,
|
|
TextColumn,
|
|
TimeRemainingColumn,
|
|
TransferSpeedColumn,
|
|
)
|
|
|
|
from .config import RunPodConfig, get_config
|
|
from .errors import S3Error, S3ObjectNotFoundError
|
|
|
|
console = Console()
|
|
|
|
|
|
class S3Client:
|
|
"""
|
|
S3 client for RunPod network volume operations.
|
|
|
|
Handles all S3 interactions with proper retry logic and error handling.
|
|
"""
|
|
|
|
def __init__(self, config: RunPodConfig | None = None):
|
|
"""
|
|
Initialize S3 client.
|
|
|
|
Args:
|
|
config: RunPodConfig instance (uses global config if None)
|
|
"""
|
|
self.config = config or get_config()
|
|
|
|
# Initialize boto3 S3 client
|
|
self.s3_client = boto3.client(
|
|
"s3",
|
|
aws_access_key_id=self.config.runpod_s3_access_key,
|
|
aws_secret_access_key=self.config.runpod_s3_secret,
|
|
region_name=self.config.runpod_s3_region,
|
|
endpoint_url=self.config.runpod_s3_endpoint,
|
|
config=Config(
|
|
signature_version="s3v4",
|
|
s3={"addressing_style": "path"},
|
|
retries={
|
|
"max_attempts": self.config.max_retries,
|
|
"mode": "adaptive",
|
|
},
|
|
connect_timeout=30,
|
|
read_timeout=self.config.api_timeout,
|
|
),
|
|
)
|
|
|
|
self.bucket = self.config.get_s3_bucket()
|
|
|
|
def compute_md5(self, file_path: Path) -> str:
|
|
"""
|
|
Compute MD5 checksum of local file.
|
|
|
|
Args:
|
|
file_path: Path to local file
|
|
|
|
Returns:
|
|
MD5 hex digest
|
|
|
|
Raises:
|
|
S3Error: If file cannot be read
|
|
"""
|
|
try:
|
|
md5 = hashlib.md5()
|
|
with open(file_path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(8192), b""):
|
|
md5.update(chunk)
|
|
return md5.hexdigest()
|
|
except OSError as e:
|
|
raise S3Error(f"Failed to read file {file_path}: {e}") from e
|
|
|
|
def object_exists(self, s3_key: str) -> bool:
|
|
"""
|
|
Check if S3 object exists.
|
|
|
|
Args:
|
|
s3_key: S3 object key
|
|
|
|
Returns:
|
|
True if object exists, False otherwise
|
|
"""
|
|
try:
|
|
self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
|
|
return True
|
|
except ClientError as e:
|
|
if e.response["Error"]["Code"] == "404":
|
|
return False
|
|
raise S3Error(f"Failed to check object existence: {e}", self.bucket, s3_key) from e
|
|
|
|
def get_object_size(self, s3_key: str) -> int:
|
|
"""
|
|
Get size of S3 object in bytes.
|
|
|
|
Args:
|
|
s3_key: S3 object key
|
|
|
|
Returns:
|
|
Object size in bytes
|
|
|
|
Raises:
|
|
S3ObjectNotFoundError: If object doesn't exist
|
|
S3Error: On other S3 errors
|
|
"""
|
|
try:
|
|
response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
|
|
return response["ContentLength"]
|
|
except ClientError as e:
|
|
if e.response["Error"]["Code"] == "404":
|
|
raise S3ObjectNotFoundError(self.bucket, s3_key)
|
|
raise S3Error(f"Failed to get object size: {e}", self.bucket, s3_key) from e
|
|
|
|
def needs_upload(self, local_file: Path, s3_key: str) -> tuple[bool, str]:
|
|
"""
|
|
Check if file needs upload based on MD5 checksum.
|
|
|
|
Args:
|
|
local_file: Path to local file
|
|
s3_key: Destination S3 key
|
|
|
|
Returns:
|
|
Tuple of (needs_upload, reason)
|
|
"""
|
|
local_md5 = self.compute_md5(local_file)
|
|
|
|
try:
|
|
response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
|
|
remote_etag = response["ETag"].strip('"')
|
|
|
|
if local_md5 == remote_etag:
|
|
return False, "up-to-date (checksum match)"
|
|
else:
|
|
return True, "changed (checksum mismatch)"
|
|
except ClientError as e:
|
|
if e.response["Error"]["Code"] == "404":
|
|
return True, "new file (doesn't exist on S3)"
|
|
else:
|
|
return True, f"error checking: {e}"
|
|
|
|
def upload_binary(
|
|
self,
|
|
local_file: Path,
|
|
s3_key: str,
|
|
force: bool = False,
|
|
progress_callback: Callable[[int], None] | None = None,
|
|
) -> bool:
|
|
"""
|
|
Upload binary file with progress tracking.
|
|
|
|
Args:
|
|
local_file: Path to local file
|
|
s3_key: Destination S3 key
|
|
force: Skip checksum check and force upload
|
|
progress_callback: Optional callback(bytes_transferred)
|
|
|
|
Returns:
|
|
True if upload succeeded, False otherwise
|
|
|
|
Raises:
|
|
S3Error: On upload failure
|
|
"""
|
|
if not local_file.exists():
|
|
raise S3Error(f"Local file not found: {local_file}")
|
|
|
|
file_size = local_file.stat().st_size
|
|
|
|
# Check if upload needed (unless forced)
|
|
if not force:
|
|
needs_upload, reason = self.needs_upload(local_file, s3_key)
|
|
if not needs_upload:
|
|
console.print(f"[green]✓ {s3_key}[/green]: {reason}")
|
|
return True
|
|
console.print(f"[cyan]↻ {s3_key}[/cyan]: {reason}")
|
|
|
|
try:
|
|
# Upload with optional progress callback
|
|
self.s3_client.upload_file(
|
|
str(local_file),
|
|
self.bucket,
|
|
s3_key,
|
|
Callback=progress_callback,
|
|
ExtraArgs={"ContentType": "application/octet-stream"},
|
|
)
|
|
|
|
# Verify upload
|
|
response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
|
|
remote_size = response["ContentLength"]
|
|
|
|
if remote_size == file_size:
|
|
console.print(f"[green]✓ Uploaded: {s3_key} ({file_size:,} bytes)[/green]")
|
|
return True
|
|
else:
|
|
# Size mismatch - delete and raise error
|
|
console.print(
|
|
f"[red]✗ Size mismatch: local={file_size}, remote={remote_size}[/red]"
|
|
)
|
|
self.s3_client.delete_object(Bucket=self.bucket, Key=s3_key)
|
|
raise S3Error(
|
|
f"Upload verification failed: size mismatch", self.bucket, s3_key
|
|
)
|
|
|
|
except ClientError as e:
|
|
raise S3Error(f"Upload failed: {e}", self.bucket, s3_key) from e
|
|
|
|
def tail_log_file(
|
|
self, s3_key: str, start_byte: int = 0, max_bytes: int = 1024 * 1024
|
|
) -> tuple[bytes, int]:
|
|
"""
|
|
Tail log file using byte-range requests.
|
|
|
|
Args:
|
|
s3_key: S3 key for log file
|
|
start_byte: Starting byte position
|
|
max_bytes: Maximum bytes to read (default 1MB)
|
|
|
|
Returns:
|
|
Tuple of (log_content, new_position)
|
|
|
|
Raises:
|
|
S3ObjectNotFoundError: If log file doesn't exist
|
|
S3Error: On other S3 errors
|
|
"""
|
|
try:
|
|
# Get current file size
|
|
file_size = self.get_object_size(s3_key)
|
|
|
|
# If file hasn't grown, return empty
|
|
if file_size <= start_byte:
|
|
return b"", start_byte
|
|
|
|
# Calculate byte range to fetch
|
|
end_byte = min(start_byte + max_bytes - 1, file_size - 1)
|
|
|
|
# Fetch byte range
|
|
response = self.s3_client.get_object(
|
|
Bucket=self.bucket, Key=s3_key, Range=f"bytes={start_byte}-{end_byte}"
|
|
)
|
|
|
|
content = response["Body"].read()
|
|
new_position = start_byte + len(content)
|
|
|
|
return content, new_position
|
|
|
|
except ClientError as e:
|
|
if e.response["Error"]["Code"] == "NoSuchKey":
|
|
raise S3ObjectNotFoundError(self.bucket, s3_key)
|
|
raise S3Error(f"Failed to tail log file: {e}", self.bucket, s3_key) from e
|
|
|
|
def download_results(self, s3_prefix: str, local_dir: Path) -> list[Path]:
|
|
"""
|
|
Download all files with given S3 prefix to local directory.
|
|
|
|
Args:
|
|
s3_prefix: S3 prefix (directory) to download
|
|
local_dir: Local destination directory
|
|
|
|
Returns:
|
|
List of downloaded file paths
|
|
|
|
Raises:
|
|
S3Error: On download failure
|
|
"""
|
|
local_dir.mkdir(parents=True, exist_ok=True)
|
|
downloaded_files = []
|
|
|
|
try:
|
|
# List all objects with prefix
|
|
paginator = self.s3_client.get_paginator("list_objects_v2")
|
|
pages = paginator.paginate(Bucket=self.bucket, Prefix=s3_prefix)
|
|
|
|
objects = []
|
|
for page in pages:
|
|
objects.extend(page.get("Contents", []))
|
|
|
|
if not objects:
|
|
console.print(f"[yellow]⚠ No files found with prefix: {s3_prefix}[/yellow]")
|
|
return downloaded_files
|
|
|
|
console.print(f"Found {len(objects)} file(s) to download")
|
|
|
|
# Download each object with progress
|
|
with Progress(
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
DownloadColumn(),
|
|
TransferSpeedColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as progress:
|
|
for obj in objects:
|
|
s3_key = obj["Key"]
|
|
file_size = obj["Size"]
|
|
|
|
# Skip directory markers
|
|
if s3_key.endswith("/"):
|
|
continue
|
|
|
|
# Determine local path (preserve directory structure)
|
|
rel_path = Path(s3_key).relative_to(s3_prefix)
|
|
local_file = local_dir / rel_path
|
|
local_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Download with progress
|
|
task = progress.add_task(f"Downloading {rel_path}", total=file_size)
|
|
|
|
def callback(bytes_transferred):
|
|
progress.update(task, completed=bytes_transferred)
|
|
|
|
self.s3_client.download_file(
|
|
self.bucket, s3_key, str(local_file), Callback=callback
|
|
)
|
|
|
|
downloaded_files.append(local_file)
|
|
console.print(f"[green]✓ Downloaded: {rel_path}[/green]")
|
|
|
|
return downloaded_files
|
|
|
|
except ClientError as e:
|
|
raise S3Error(f"Download failed: {e}", self.bucket, s3_prefix) from e
|
|
|
|
def list_binaries(self) -> list[dict[str, Any]]:
|
|
"""
|
|
List all binaries in binaries/ directory.
|
|
|
|
Returns:
|
|
List of dicts with keys: name, size, last_modified
|
|
"""
|
|
return self._list_prefix("binaries/")
|
|
|
|
def list_models(self) -> list[dict[str, Any]]:
|
|
"""
|
|
List all models in models/ directory.
|
|
|
|
Returns:
|
|
List of dicts with keys: name, size, last_modified
|
|
"""
|
|
models_prefix = self.config.get_models_s3_prefix()
|
|
return self._list_prefix(models_prefix)
|
|
|
|
def _list_prefix(self, prefix: str) -> list[dict[str, Any]]:
|
|
"""
|
|
List all objects with given prefix.
|
|
|
|
Args:
|
|
prefix: S3 prefix to list
|
|
|
|
Returns:
|
|
List of object metadata dicts
|
|
"""
|
|
try:
|
|
response = self.s3_client.list_objects_v2(Bucket=self.bucket, Prefix=prefix)
|
|
|
|
objects = []
|
|
for obj in response.get("Contents", []):
|
|
# Skip directory markers
|
|
if obj["Key"].endswith("/"):
|
|
continue
|
|
|
|
objects.append(
|
|
{
|
|
"name": obj["Key"],
|
|
"size": obj["Size"],
|
|
"last_modified": obj["LastModified"],
|
|
}
|
|
)
|
|
|
|
return objects
|
|
|
|
except ClientError as e:
|
|
raise S3Error(f"Failed to list objects: {e}", self.bucket, prefix) from e
|
|
|
|
def create_directory(self, s3_key: str) -> bool:
|
|
"""
|
|
Create directory marker in S3.
|
|
|
|
Args:
|
|
s3_key: Directory path (should end with /)
|
|
|
|
Returns:
|
|
True if created or already exists
|
|
"""
|
|
if not s3_key.endswith("/"):
|
|
s3_key += "/"
|
|
|
|
try:
|
|
# Check if already exists
|
|
if self.object_exists(s3_key):
|
|
return True
|
|
|
|
# Create empty directory marker
|
|
self.s3_client.put_object(Bucket=self.bucket, Key=s3_key)
|
|
console.print(f"[green]✓ Created directory: {s3_key}[/green]")
|
|
return True
|
|
|
|
except ClientError as e:
|
|
console.print(f"[yellow]⚠ Failed to create directory {s3_key}: {e}[/yellow]")
|
|
return False
|