Object Storage in Machine Learning

import boto3
import os
import math
from concurrent.futures import ThreadPoolExecutor

def upload_model_checkpoint(
    file_path: str,
    bucket: str,
    key: str,
    part_size_mb: int = 100,
    max_workers: int = 8,
) -> str:
    """Upload a large model checkpoint using S3 multipart upload.
    
    Args:
        file_path: Local path to the model checkpoint file.
        bucket: S3 bucket name.
        key: S3 object key (e.g., 'checkpoints/v2/model.safetensors').
        part_size_mb: Size of each upload part in MB (min 5 MB).
        max_workers: Number of parallel upload threads.
    
    Returns:
        The ETag of the completed upload.
    """
    s3 = boto3.client("s3", region_name="ap-south-1")  # Mumbai region
    file_size = os.path.getsize(file_path)
    part_size = part_size_mb * 1024 * 1024
    num_parts = math.ceil(file_size / part_size)

    # Initiate multipart upload with metadata
    mpu = s3.create_multipart_upload(
        Bucket=bucket,
        Key=key,
        Metadata={
            "model-framework": "pytorch",
            "checkpoint-epoch": "42",
            "training-run-id": "run-2026-02-07-abc123",
        },
        StorageClass="STANDARD",  # Use INTELLIGENT_TIERING for uncertain access
    )
    upload_id = mpu["UploadId"]

    def upload_part(part_number: int) -> dict:
        offset = (part_number - 1) * part_size
        length = min(part_size, file_size - offset)
        with open(file_path, "rb") as f:
            f.seek(offset)
            data = f.read(length)
        response = s3.upload_part(
            Bucket=bucket,
            Key=key,
            PartNumber=part_number,
            UploadId=upload_id,
            Body=data,
        )
        return {"PartNumber": part_number, "ETag": response["ETag"]}

    try:
        # Upload parts in parallel
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            parts = list(executor.map(upload_part, range(1, num_parts + 1)))

        # Complete the multipart upload
        result = s3.complete_multipart_upload(
            Bucket=bucket,
            Key=key,
            UploadId=upload_id,
            MultipartUpload={"Parts": sorted(parts, key=lambda p: p["PartNumber"])},
        )
        print(f"Upload complete: {result['ETag']}")
        return result["ETag"]

    except Exception as e:
        # Abort on failure to avoid orphaned parts (and storage charges)
        s3.abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)
        raise RuntimeError(f"Multipart upload failed: {e}") from e


# Usage: upload a 15 GB SafeTensors model file
upload_model_checkpoint(
    file_path="/data/checkpoints/llama-3-8b-finetuned.safetensors",
    bucket="ml-models-prod",
    key="models/llama-3-8b-ft/v2/model.safetensors",
    part_size_mb=100,
    max_workers=8,
)

import boto3
from datetime import datetime

def generate_presigned_upload_url(
    bucket: str,
    key: str,
    expiry_seconds: int = 3600,
    max_size_bytes: int = 5 * 1024 * 1024 * 1024,  # 5 GB
) -> dict:
    """Generate a presigned URL for direct-to-S3 upload.
    
    Use this when clients (browsers, edge devices, labeling tools)
    need to upload data directly to S3 without routing through your backend.
    
    Args:
        bucket: Target S3 bucket.
        key: Object key for the upload.
        expiry_seconds: URL validity period.
        max_size_bytes: Maximum allowed upload size.
    
    Returns:
        Dict with presigned URL and required fields.
    """
    s3 = boto3.client("s3", region_name="ap-south-1")
    
    # Generate presigned POST (more secure than presigned PUT)
    conditions = [
        ["content-length-range", 1, max_size_bytes],
        {"x-amz-meta-uploaded-by": "labeling-tool"},
        {"x-amz-meta-upload-timestamp": datetime.utcnow().isoformat()},
    ]
    
    presigned = s3.generate_presigned_post(
        Bucket=bucket,
        Key=key,
        Fields={
            "x-amz-meta-uploaded-by": "labeling-tool",
            "x-amz-meta-upload-timestamp": datetime.utcnow().isoformat(),
        },
        Conditions=conditions,
        ExpiresIn=expiry_seconds,
    )
    
    return presigned


def generate_presigned_download_url(
    bucket: str,
    key: str,
    expiry_seconds: int = 3600,
    version_id: str = None,
) -> str:
    """Generate a presigned URL for downloading an object.
    
    Useful for serving model artifacts to inference services
    without exposing bucket credentials.
    """
    s3 = boto3.client("s3", region_name="ap-south-1")
    params = {"Bucket": bucket, "Key": key}
    if version_id:
        params["VersionId"] = version_id
    
    url = s3.generate_presigned_url(
        "get_object",
        Params=params,
        ExpiresIn=expiry_seconds,
    )
    return url


# Generate upload URL for a data labeling tool
upload_info = generate_presigned_upload_url(
    bucket="ml-datasets-raw",
    key=f"uploads/labeling/batch-2026-02-07/image-{42:06d}.jpg",
)
print(f"Upload URL: {upload_info['url']}")
print(f"Fields: {upload_info['fields']}")

# Generate download URL for a specific model version
download_url = generate_presigned_download_url(
    bucket="ml-models-prod",
    key="models/fraud-detector/v3/model.onnx",
    version_id="abc123def456",
)
print(f"Download URL: {download_url}")

# Terraform configuration for ML artifact buckets with versioning
# and cost-optimized lifecycle policies

resource "aws_s3_bucket" "ml_artifacts" {
  bucket = "mycompany-ml-artifacts-prod"

  tags = {
    Environment = "production"
    Team        = "ml-platform"
    CostCenter  = "ml-infra"
  }
}

resource "aws_s3_bucket_versioning" "ml_artifacts" {
  bucket = aws_s3_bucket.ml_artifacts.id
  versioning_configuration {
    status = "Enabled"
  }
}

resource "aws_s3_bucket_lifecycle_configuration" "ml_artifacts" {
  bucket = aws_s3_bucket.ml_artifacts.id

  # Rule 1: Training datasets -- move to IA after 30 days, Glacier after 90
  rule {
    id     = "training-data-tiering"
    status = "Enabled"

    filter {
      prefix = "datasets/training/"
    }

    transition {
      days          = 30
      storage_class = "STANDARD_IA"
    }

    transition {
      days          = 90
      storage_class = "GLACIER"
    }

    expiration {
      days = 365  # Delete training data older than 1 year
    }

    noncurrent_version_transition {
      noncurrent_days = 7
      storage_class   = "STANDARD_IA"
    }

    noncurrent_version_expiration {
      noncurrent_days = 30  # Keep old versions for 30 days only
    }
  }

  # Rule 2: Model checkpoints -- keep recent, archive old
  rule {
    id     = "checkpoint-lifecycle"
    status = "Enabled"

    filter {
      prefix = "checkpoints/"
    }

    transition {
      days          = 14
      storage_class = "STANDARD_IA"
    }

    transition {
      days          = 60
      storage_class = "GLACIER"
    }

    noncurrent_version_expiration {
      noncurrent_days = 14
    }
  }

  # Rule 3: Production model artifacts -- keep hot for serving
  rule {
    id     = "production-models"
    status = "Enabled"

    filter {
      prefix = "models/production/"
    }

    # No transition -- production models stay in Standard
    # Only expire noncurrent versions after 90 days
    noncurrent_version_expiration {
      noncurrent_days = 90
    }
  }

  # Rule 4: Abort incomplete multipart uploads after 7 days
  rule {
    id     = "abort-incomplete-uploads"
    status = "Enabled"

    filter {
      prefix = ""
    }

    abort_incomplete_multipart_upload {
      days_after_initiation = 7
    }
  }
}

# Block public access (ML data should never be public)
resource "aws_s3_bucket_public_access_block" "ml_artifacts" {
  bucket = aws_s3_bucket.ml_artifacts.id

  block_public_acls       = true
  block_public_policy     = true
  ignore_public_acls      = true
  restrict_public_buckets = true
}

# Enable server-side encryption
resource "aws_s3_bucket_server_side_encryption_configuration" "ml_artifacts" {
  bucket = aws_s3_bucket.ml_artifacts.id

  rule {
    apply_server_side_encryption_by_default {
      sse_algorithm = "aws:kms"
    }
    bucket_key_enabled = true  # Reduces KMS costs by up to 99%
  }
}

import fsspec
import pyarrow.parquet as pq
import torch
from torch.utils.data import Dataset, DataLoader


class ObjectStorageMLDataset(Dataset):
    """PyTorch Dataset that reads data shards from any object store.
    
    Supports S3, GCS, Azure Blob, and local filesystem via fsspec.
    """

    def __init__(
        self,
        data_uri: str,  # e.g., "s3://bucket/datasets/train/"
        storage_options: dict = None,
        shard_pattern: str = "*.parquet",
    ):
        self.fs, self.base_path = fsspec.core.url_to_fs(
            data_uri, **(storage_options or {})
        )
        # List all shard files matching the pattern
        self.shard_paths = sorted(
            self.fs.glob(f"{self.base_path}/{shard_pattern}")
        )
        if not self.shard_paths:
            raise FileNotFoundError(f"No shards found at {data_uri}")

        # Read schema from first shard to get column names
        with self.fs.open(self.shard_paths[0], "rb") as f:
            schema = pq.read_schema(f)
        self.columns = schema.names
        
        # Build an index mapping global row index -> (shard, local_row)
        self._index_map = []
        for shard_path in self.shard_paths:
            with self.fs.open(shard_path, "rb") as f:
                metadata = pq.read_metadata(f)
                num_rows = metadata.num_rows
            self._index_map.extend(
                [(shard_path, i) for i in range(num_rows)]
            )

    def __len__(self) -> int:
        return len(self._index_map)

    def __getitem__(self, idx: int) -> dict:
        shard_path, local_row = self._index_map[idx]
        with self.fs.open(shard_path, "rb") as f:
            table = pq.read_table(f)
        row = {col: table.column(col)[local_row].as_py() for col in self.columns}
        return row


# Usage with S3 (Mumbai region)
dataset = ObjectStorageMLDataset(
    data_uri="s3://ml-datasets-prod/training/image-classification/v3/",
    storage_options={
        "anon": False,
        "client_kwargs": {"region_name": "ap-south-1"},
    },
)

loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True)

for batch in loader:
    # batch is a dict of tensors
    features = batch["features"]
    labels = batch["labels"]
    # ... training step ...
    break  # Demo: process one batch

# Same code works with GCS:
# dataset = ObjectStorageMLDataset("gs://my-bucket/training/v3/")

# Same code works with Azure Blob:
# dataset = ObjectStorageMLDataset("az://container/training/v3/")

# Same code works locally (for development):
# dataset = ObjectStorageMLDataset("file:///data/training/v3/")

import json
import boto3
import os

# This Lambda function triggers when new data lands in S3.
# It kicks off an ML preprocessing pipeline via Step Functions.

def lambda_handler(event, context):
    """Triggered by S3 ObjectCreated event on the raw data bucket.
    
    Expected event structure:
    {
      "Records": [
        {
          "s3": {
            "bucket": {"name": "ml-datasets-raw"},
            "object": {"key": "uploads/batch-2026-02-07/data.csv", "size": 1073741824}
          }
        }
      ]
    }
    """
    sfn_client = boto3.client("stepfunctions", region_name="ap-south-1")
    state_machine_arn = os.environ["PREPROCESSING_STATE_MACHINE_ARN"]

    for record in event["Records"]:
        bucket = record["s3"]["bucket"]["name"]
        key = record["s3"]["object"]["key"]
        size = record["s3"]["object"]["size"]

        # Only trigger for files in the uploads/ prefix
        if not key.startswith("uploads/"):
            continue

        # Only trigger for files larger than 1 MB (skip tiny/empty uploads)
        if size < 1_000_000:
            print(f"Skipping small file: {key} ({size} bytes)")
            continue

        pipeline_input = {
            "source_bucket": bucket,
            "source_key": key,
            "file_size_bytes": size,
            "output_bucket": "ml-datasets-processed",
            "output_prefix": f"processed/{key.split('/')[-1].replace('.csv', '')}/",
            "pipeline_config": {
                "format": "parquet",
                "compression": "snappy",
                "shard_size_mb": 128,
                "validation_enabled": True,
            },
        }

        response = sfn_client.start_execution(
            stateMachineArn=state_machine_arn,
            name=f"preprocess-{key.replace('/', '-')}-{context.aws_request_id[:8]}",
            input=json.dumps(pipeline_input),
        )

        print(
            f"Started pipeline for s3://{bucket}/{key} -> "
            f"execution: {response['executionArn']}"
        )

    return {"statusCode": 200, "body": "Pipeline triggered"}

# MinIO configuration for on-premise ML object storage (docker-compose.yml)
version: '3.8'
services:
  minio:
    image: minio/minio:latest
    command: server /data --console-address ":9001"
    ports:
      - "9000:9000"    # S3 API
      - "9001:9001"    # Web console
    environment:
      MINIO_ROOT_USER: minio-admin
      MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD}
      MINIO_REGION: ap-south-1
    volumes:
      - minio_data:/data
    healthcheck:
      test: ["CMD", "mc", "ready", "local"]
      interval: 30s
      timeout: 10s
      retries: 3

  createbuckets:
    image: minio/mc:latest
    depends_on:
      minio:
        condition: service_healthy
    entrypoint: >
      /bin/sh -c "
      mc alias set local http://minio:9000 minio-admin $${MINIO_ROOT_PASSWORD};
      mc mb --ignore-existing local/ml-datasets-raw;
      mc mb --ignore-existing local/ml-datasets-processed;
      mc mb --ignore-existing local/ml-models;
      mc mb --ignore-existing local/ml-checkpoints;
      mc version enable local/ml-models;
      mc version enable local/ml-checkpoints;
      mc ilm rule add local/ml-checkpoints --transition-days 14 --storage-class STANDARD_IA;
      mc ilm rule add local/ml-checkpoints --expire-delete-marker --noncurrent-expire-days 30;
      echo 'Buckets created and configured.';
      "

volumes:
  minio_data:
    driver: local