diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/deepseek-ocr-sagemaker/entry.sh b/docs/sagemaker/notebooks/sagemaker-sdk/deepseek-ocr-sagemaker/entry.sh
new file mode 100644
index 0000000000..6725f23a3e
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/deepseek-ocr-sagemaker/entry.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+# SageMaker entry script - installs uv and runs sm_job_runner.py with dependencies
+
+set -e
+
+echo "Installing uv..."
+curl -LsSf https://astral.sh/uv/install.sh | sh
+export PATH="$HOME/.local/bin:$PATH"
+
+echo "Running pipeline with uv..."
+cd /opt/ml/input/data/code
+uv run sm_job_runner.py
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/deepseek-ocr-sagemaker/sm_job_runner.py b/docs/sagemaker/notebooks/sagemaker-sdk/deepseek-ocr-sagemaker/sm_job_runner.py
new file mode 100644
index 0000000000..e4422acdda
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/deepseek-ocr-sagemaker/sm_job_runner.py
@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "boto3",
+#     "s3fs",
+#     "torch",
+#     "datasets[s3]>=4.0.0",
+#     "pyarrow>=12.0.0",
+#     "numpy",
+#     "pillow",
+#     "requests",
+#     "openai",
+#     "huggingface-hub[hf_transfer]",
+#     "rich",
+# ]
+# ///
+"""
+SageMaker Training Job entry point for the DeepSeek OCR pipeline.
+
+This script is run via entry.sh which installs uv and executes: uv run sm_job_runner.py
+Code is automatically available at /opt/ml/input/data/code via SageMaker SourceCode.
+
+Environment Variables:
+    PIPELINE_STAGE: Stage to run (extract, describe, assemble)
+    
+SageMaker Environment Variables (automatically set):
+    SM_MODEL_DIR: /opt/ml/model
+    SM_OUTPUT_DATA_DIR: /opt/ml/output/data
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+
+
+# Code is at /opt/ml/input/data/code (SageMaker SourceCode)
+CODE_DIR = Path("/opt/ml/input/data/code")
+
+
+def setup_logging():
+    """Configure logging."""
+    level = os.environ.get("LOG_LEVEL", "INFO").upper()
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        force=True,
+    )
+
+
+def load_hyperparameters() -> dict:
+    """Load hyperparameters from SageMaker.
+    
+    SageMaker passes hyperparameters as a JSON file at /opt/ml/input/config/hyperparameters.json
+    All values are strings, so we convert them to environment variables.
+    """
+    hp_file = Path("/opt/ml/input/config/hyperparameters.json")
+    if not hp_file.exists():
+        logging.info("No hyperparameters file found at %s", hp_file)
+        return {}
+    
+    with hp_file.open() as f:
+        hyperparameters = json.load(f)
+    
+    logging.info("Loaded hyperparameters: %s", list(hyperparameters.keys()))
+    
+    # Set hyperparameters as environment variables
+    for key, value in hyperparameters.items():
+        # SageMaker wraps values in quotes, strip them
+        if isinstance(value, str):
+            value = value.strip('"').strip("'")
+        os.environ[key] = str(value)
+    
+    return hyperparameters
+
+
+def write_success_marker():
+    """Write a success marker file for SageMaker."""
+    model_dir = Path(os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    model_dir.mkdir(parents=True, exist_ok=True)
+    
+    success_file = model_dir / "_SUCCESS"
+    success_file.write_text("Pipeline completed successfully\n")
+    logging.info("Wrote success marker to %s", success_file)
+
+
+def main() -> None:
+    """Main entry point for SageMaker training job."""
+    setup_logging()
+    logger = logging.getLogger(__name__)
+    
+    logger.info("Starting SageMaker OCR pipeline job")
+    logger.info("Python version: %s", sys.version)
+    logger.info("Working directory: %s", os.getcwd())
+    
+    # Log SageMaker environment
+    sm_vars = {k: v for k, v in os.environ.items() if k.startswith("SM_")}
+    logger.info("SageMaker environment: %s", json.dumps(sm_vars, indent=2))
+    
+    # Load hyperparameters from SageMaker
+    load_hyperparameters()
+    
+    # Add code directory to path
+    sys.path.insert(0, str(CODE_DIR))
+    logger.info("Code directory: %s", CODE_DIR)
+    
+    # Import and run pipeline
+    try:
+        from llm_ocr.cli import main as pipeline_main
+        pipeline_main()
+        write_success_marker()
+        logger.info("Pipeline completed successfully")
+    except Exception as exc:
+        logger.exception("Pipeline failed: %s", exc)
+        # Write failure info for debugging
+        failure_file = Path("/opt/ml/output/failure")
+        failure_file.parent.mkdir(parents=True, exist_ok=True)
+        failure_file.write_text(str(exc))
+        raise
+
+
+if __name__ == "__main__":
+    main()
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/__init__.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/__init__.py
new file mode 100644
index 0000000000..604fc7e8d5
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/__init__.py
@@ -0,0 +1,39 @@
+"""DeepSeek OCR pipeline package.
+
+This package provides tools for running batch OCR inference with DeepSeek-OCR
+and similar vision-language models across different cloud platforms.
+
+Example usage:
+    from llm_ocr import DeepSeekClient, ExtractSettings, run_stage_extract
+
+    client = DeepSeekClient(base_url="http://localhost:8000/v1")
+    settings = ExtractSettings.from_env(client)
+    run_stage_extract(settings)
+"""
+
+# Stage runners - main entry points for pipeline stages
+from llm_ocr.stages import (
+    run_stage_assemble as run_stage_assemble,
+    run_stage_describe as run_stage_describe,
+    run_stage_extract as run_stage_extract,
+)
+
+# Configuration classes
+from llm_ocr.config import (
+    AssembleSettings as AssembleSettings,
+    DescribeSettings as DescribeSettings,
+    ExtractSettings as ExtractSettings,
+    InferenceSettings as InferenceSettings,
+)
+
+# Inference client
+from llm_ocr.server import DeepSeekClient as DeepSeekClient
+
+# Storage backends
+from llm_ocr.storage import (
+    DatasetStorage as DatasetStorage,
+    GCSStorage as GCSStorage,
+    HFHubStorage as HFHubStorage,
+    S3Storage as S3Storage,
+    get_storage as get_storage,
+)
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/__main__.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/__main__.py
new file mode 100644
index 0000000000..ca135d41f2
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/__main__.py
@@ -0,0 +1,6 @@
+"""Entry point for `python -m llm_ocr`."""
+
+from llm_ocr.cli import main
+
+if __name__ == "__main__":
+    main()
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/cli.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/cli.py
new file mode 100644
index 0000000000..25ff9114a0
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/cli.py
@@ -0,0 +1,111 @@
+"""CLI entrypoint for the DeepSeek OCR pipeline."""
+
+from __future__ import annotations
+
+import logging
+
+from .config import AssembleSettings, DescribeSettings, ExtractSettings, env
+from .server import (
+    DeepSeekClient,
+    base_url_from_env,
+    launch_vllm,
+    should_launch_server,
+    shutdown_server,
+    wait_for_server,
+)
+from .stages import run_stage_assemble, run_stage_describe, run_stage_extract
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _setup_logging() -> None:
+    """Configure logging with optional rich handler."""
+    level = env("LOG_LEVEL", "INFO").upper()
+    try:
+        from rich.console import Console
+        from rich.logging import RichHandler
+
+        console = Console(
+            force_terminal=env("FORCE_COLOR", "").lower() in {"1", "true"}
+        )
+        handler = RichHandler(
+            console=console, show_time=True, show_level=True, rich_tracebacks=True
+        )
+        logging.basicConfig(
+            level=level, format="%(message)s", handlers=[handler], force=True
+        )
+    except ImportError:
+        logging.basicConfig(
+            level=level,
+            format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+            force=True,
+        )
+
+
+def _create_client(
+    max_tokens: int, temperature: float, inference_settings
+) -> DeepSeekClient:
+    """Create DeepSeek client from environment."""
+    return DeepSeekClient(
+        base_url=base_url_from_env(),
+        model_name=env("SERVED_MODEL_NAME", "deepseek-ocr"),
+        max_tokens=max_tokens,
+        temperature=temperature,
+        request_timeout=inference_settings.request_timeout,
+        max_retries=inference_settings.max_retries,
+        retry_backoff_seconds=inference_settings.retry_backoff,
+    )
+
+
+def main() -> None:
+    """Main entry point for the pipeline CLI."""
+    _setup_logging()
+
+    stage = env("PIPELINE_STAGE", "extract").lower()
+    if stage not in {"extract", "describe", "assemble"}:
+        raise ValueError(f"Unsupported stage: {stage}")
+
+    needs_server = stage in {"extract", "describe"}
+    launch_server = should_launch_server() and needs_server
+    server_process = None
+
+    try:
+        if launch_server:
+            server_process = launch_vllm()
+
+        if needs_server:
+            base_url = base_url_from_env()
+            health_url = env("HEALTH_URL", f"{base_url}/health")
+            LOGGER.info("Waiting for server at %s", health_url)
+            if not wait_for_server(health_url):
+                raise RuntimeError("vLLM server did not become ready in time")
+
+        if stage == "extract":
+            from .config import InferenceSettings
+
+            inference = InferenceSettings.from_env("EXTRACT")
+            max_tokens = env("DOC_MAX_TOKENS", 2048, int)
+            temperature = env("DOC_TEMPERATURE", 0.0, float)
+            client = _create_client(max_tokens, temperature, inference)
+            settings = ExtractSettings.from_env(client)
+            settings.inference = inference
+            run_stage_extract(settings)
+
+        elif stage == "describe":
+            from .config import InferenceSettings
+
+            inference = InferenceSettings.from_env("DESCRIBE")
+            max_tokens = env("FIGURE_MAX_TOKENS", 512, int)
+            temperature = env("FIGURE_TEMPERATURE", 0.0, float)
+            client = _create_client(max_tokens, temperature, inference)
+            settings = DescribeSettings.from_env(client)
+            settings.inference = inference
+            run_stage_describe(settings)
+
+        elif stage == "assemble":
+            settings = AssembleSettings.from_env()
+            run_stage_assemble(settings)
+
+    finally:
+        if server_process is not None:
+            shutdown_server(server_process)
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/cloudrun_io.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/cloudrun_io.py
new file mode 100644
index 0000000000..4711a48db5
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/cloudrun_io.py
@@ -0,0 +1,153 @@
+"""Google Cloud Storage utilities for Cloud Run jobs."""
+
+from __future__ import annotations
+
+import logging
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_gcs_client():
+    """Get GCS client."""
+    from google.cloud import storage
+
+    return storage.Client()
+
+
+def parse_gcs_uri(uri: str) -> tuple[str, str]:
+    """Parse gs://bucket/key into (bucket, key)."""
+    if not uri.startswith("gs://"):
+        raise ValueError(f"Invalid GCS URI: {uri}")
+    parts = uri[5:].split("/", 1)
+    bucket = parts[0]
+    key = parts[1] if len(parts) > 1 else ""
+    return bucket, key
+
+
+def upload_files_to_gcs(
+    *,
+    output_dir: Path,
+    gcs_uri: str,
+    path_prefix: str = "",
+) -> None:
+    """Upload local directory contents to GCS."""
+    if not gcs_uri:
+        LOGGER.info("No GCS URI provided; skipping upload.")
+        return
+
+    bucket_name, base_prefix = parse_gcs_uri(gcs_uri)
+
+    full_prefix = base_prefix.rstrip("/")
+    if path_prefix:
+        full_prefix = (
+            f"{full_prefix}/{path_prefix.strip('/')}"
+            if full_prefix
+            else path_prefix.strip("/")
+        )
+
+    client = get_gcs_client()
+    bucket = client.bucket(bucket_name)
+    base = output_dir.resolve()
+
+    files = sorted(p for p in base.rglob("*") if p.is_file())
+    if not files:
+        LOGGER.info("Nothing to upload from %s", output_dir)
+        return
+
+    LOGGER.info(
+        "Uploading %d files to gs://%s/%s", len(files), bucket_name, full_prefix
+    )
+
+    for local_path in files:
+        rel = local_path.relative_to(base).as_posix()
+        gcs_key = f"{full_prefix}/{rel}" if full_prefix else rel
+        try:
+            blob = bucket.blob(gcs_key)
+            blob.upload_from_filename(str(local_path))
+        except Exception as exc:
+            LOGGER.error(
+                "Failed to upload %s to gs://%s/%s: %s",
+                local_path,
+                bucket_name,
+                gcs_key,
+                exc,
+            )
+            raise
+
+
+def save_dataset_to_gcs(
+    dataset,
+    gcs_uri: str,
+    name: str = "dataset",
+) -> str:
+    """Save HF dataset to GCS in Arrow format. Returns the GCS URI."""
+    from datasets import DatasetDict
+
+    # Handle DatasetDict by extracting the first split
+    if isinstance(dataset, DatasetDict):
+        if "train" in dataset:
+            dataset = dataset["train"]
+        else:
+            split_name = list(dataset.keys())[0]
+            dataset = dataset[split_name]
+            LOGGER.info("Using split '%s' from DatasetDict", split_name)
+
+    bucket_name, prefix = parse_gcs_uri(gcs_uri)
+    full_prefix = prefix.rstrip("/")
+
+    # Save to local temp directory using Arrow format
+    local_dir = Path(f"/tmp/{name}_arrow_temp")
+    if local_dir.exists():
+        shutil.rmtree(local_dir)
+
+    LOGGER.info("Saving dataset to Arrow format...")
+    dataset.save_to_disk(str(local_dir))
+
+    # Upload entire directory to GCS
+    gcs_prefix = f"{full_prefix}/{name}" if full_prefix else name
+    upload_files_to_gcs(
+        output_dir=local_dir, gcs_uri=f"gs://{bucket_name}/{gcs_prefix}"
+    )
+
+    # Cleanup
+    shutil.rmtree(local_dir)
+
+    result_uri = f"gs://{bucket_name}/{gcs_prefix}"
+    LOGGER.info("Saved dataset to %s", result_uri)
+    return result_uri
+
+
+def load_dataset_from_gcs(gcs_uri: str, split: str = "train") -> "Dataset":
+    """Load HF dataset from GCS. Downloads locally to avoid gcsfs caching issues."""
+    from datasets import load_from_disk
+    import tempfile
+
+    LOGGER.info("Loading dataset from %s", gcs_uri)
+
+    # Parse GCS URI
+    bucket_name, prefix = parse_gcs_uri(gcs_uri)
+
+    # Download to local temp directory (bypasses gcsfs cache)
+    client = get_gcs_client()
+    bucket = client.bucket(bucket_name)
+    local_dir = tempfile.mkdtemp(prefix="gcs_dataset_")
+
+    blobs = list(bucket.list_blobs(prefix=f"{prefix}/"))
+    for blob in blobs:
+        filename = blob.name.split("/")[-1]
+        if filename:  # Skip directory markers
+            local_path = f"{local_dir}/{filename}"
+            blob.download_to_filename(local_path)
+
+    LOGGER.info("Downloaded %d files to %s", len(blobs), local_dir)
+
+    # Load from local
+    ds = load_from_disk(local_dir)
+
+    return ds
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/config.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/config.py
new file mode 100644
index 0000000000..546eb4cc46
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/config.py
@@ -0,0 +1,161 @@
+"""Configuration dataclasses for pipeline stages."""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, TypeVar
+
+LOGGER = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def env(key: str, default: T = None, cast: Type[T] = str) -> T:
+    """Read environment variable with type casting."""
+    raw = os.environ.get(key)
+    if raw is None:
+        return default
+    try:
+        return cast(raw)
+    except (TypeError, ValueError):
+        LOGGER.warning("Invalid %s=%s, using default=%s", key, raw, default)
+        return default
+
+
+@dataclass
+class FigureMetadata:
+    """Metadata for an extracted figure (image stored in dataset, not as file)."""
+
+    figure_id: str
+    label: str
+    bounding_box_pixels: Dict[str, int]
+    description: Optional[str] = None
+
+
+@dataclass
+class InferenceSettings:
+    """Settings for batch inference."""
+
+    batch_size: int = 4
+    max_concurrency: int = 4
+    request_timeout: int = 120
+    max_retries: int = 3
+    retry_backoff: float = 2.0
+
+    @classmethod
+    def from_env(cls, prefix: str = "") -> InferenceSettings:
+        """Load from environment with optional prefix (e.g., 'EXTRACT_')."""
+        p = f"{prefix}_" if prefix else ""
+        return cls(
+            batch_size=max(1, env(f"{p}BATCH_SIZE", 4, int)),
+            max_concurrency=max(1, env(f"{p}MAX_CONCURRENCY", 4, int)),
+            request_timeout=max(1, env(f"{p}REQUEST_TIMEOUT", 120, int)),
+            max_retries=max(0, env(f"{p}MAX_RETRIES", 3, int)),
+            retry_backoff=max(0.1, env(f"{p}RETRY_BACKOFF", 2.0, float)),
+        )
+
+
+@dataclass
+class HubSettings:
+    """Settings for HuggingFace Hub upload."""
+
+    repo_id: Optional[str] = None
+    path_in_repo: str = ""
+    branch: Optional[str] = None
+    commit_message: Optional[str] = None
+
+    @classmethod
+    def from_env(cls, prefix: str = "HF") -> HubSettings:
+        """Load from environment with prefix (e.g., 'HF_' -> HF_REPO_ID)."""
+        p = f"{prefix}_" if prefix else ""
+        return cls(
+            repo_id=env(f"{p}REPO_ID") or env(f"{p}REPO"),
+            path_in_repo=env(f"{p}PATH_IN_REPO", ""),
+            branch=env(f"{p}BRANCH"),
+            commit_message=env(f"{p}COMMIT_MESSAGE"),
+        )
+
+
+@dataclass
+class ExtractSettings:
+    """Settings for the extract stage."""
+
+    dataset_name: str
+    dataset_config: str
+    dataset_split: str
+    output_dir: Path
+    client: Any  # DeepSeekClient
+    prompt: str = "<image>\n<|grounding|>Convert this document to Markdown."
+    max_tokens: int = 2048
+    temperature: float = 0.0
+    max_samples: Optional[int] = None
+    stream_dataset: bool = True
+    inference: InferenceSettings = field(default_factory=InferenceSettings)
+    hub: HubSettings = field(default_factory=HubSettings)
+
+    @classmethod
+    def from_env(cls, client: Any) -> ExtractSettings:
+        """Load settings from environment variables."""
+        return cls(
+            dataset_name=env("DATASET_NAME", "HuggingFaceM4/FineVision"),
+            dataset_config=env("DATASET_CONFIG", "olmOCR-mix-0225-documents"),
+            dataset_split=env("DATASET_SPLIT", "train"),
+            output_dir=Path(env("OUTPUT_DIR", "./outputs/extract")),
+            client=client,
+            prompt=env(
+                "DOC_PROMPT", "<image>\n<|grounding|>Convert this document to Markdown."
+            ),
+            max_tokens=env("DOC_MAX_TOKENS", 2048, int),
+            temperature=env("DOC_TEMPERATURE", 0.0, float),
+            max_samples=env("MAX_SAMPLES", None, int),
+            stream_dataset=env("STREAM_DATASET", "true").lower() == "true",
+            inference=InferenceSettings.from_env("EXTRACT"),
+            hub=HubSettings.from_env("HF"),
+        )
+
+
+@dataclass
+class DescribeSettings:
+    """Settings for the describe stage."""
+
+    output_dir: Path
+    client: Any  # DeepSeekClient
+    source_repo_id: Optional[str] = None
+    prompt: str = "<image>\nDescribe this image in detail"
+    max_tokens: int = 512
+    temperature: float = 0.0
+    inference: InferenceSettings = field(default_factory=InferenceSettings)
+    hub: HubSettings = field(default_factory=HubSettings)
+
+    @classmethod
+    def from_env(cls, client: Any) -> DescribeSettings:
+        """Load settings from environment variables."""
+        return cls(
+            output_dir=Path(env("OUTPUT_DIR", "./outputs/describe")),
+            client=client,
+            source_repo_id=env("SOURCE_REPO_ID") or env("HF_REPO_ID"),
+            prompt=env("FIGURE_PROMPT", "<image>\nDescribe this image in detail"),
+            max_tokens=env("FIGURE_MAX_TOKENS", 512, int),
+            temperature=env("FIGURE_TEMPERATURE", 0.0, float),
+            inference=InferenceSettings.from_env("DESCRIBE"),
+            hub=HubSettings.from_env("HF"),
+        )
+
+
+@dataclass
+class AssembleSettings:
+    """Settings for the assemble stage."""
+
+    source_repo_id: Optional[str] = None
+    hub: HubSettings = field(default_factory=HubSettings)
+
+    @classmethod
+    def from_env(cls) -> AssembleSettings:
+        """Load settings from environment variables."""
+        return cls(
+            source_repo_id=env("SOURCE_REPO_ID") or env("HF_REPO_ID"),
+            hub=HubSettings.from_env("HF"),
+        )
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/document.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/document.py
new file mode 100644
index 0000000000..63cfa5246e
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/document.py
@@ -0,0 +1,388 @@
+"""Document processing: markdown extraction, figure handling, and caption enrichment."""
+
+from __future__ import annotations
+
+import ast
+import base64
+import json
+import logging
+import re
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+from .config import FigureMetadata
+
+LOGGER = logging.getLogger(__name__)
+
+GROUNDING_PATTERN = re.compile(
+    r"<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>",
+    re.DOTALL,
+)
+
+# Matches both old path format and new figure: URI format
+FIGURE_MARKDOWN_PATTERN = re.compile(
+    r"!\[(?:Figure )?(?P<figure_id>[^\]]+)\]\((?P<path>[^)]+)\)"
+)
+
+
+def encode_image(image: Image.Image) -> str:
+    """Encode a PIL Image to base64 PNG string."""
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def extract_grounding_blocks(text: str) -> List[Dict[str, Any]]:
+    """Extract grounding blocks (ref/det tags) from model response."""
+    matches: List[Dict[str, Any]] = []
+    for match in GROUNDING_PATTERN.finditer(text):
+        label = match.group(1).strip()
+        coords_text = match.group(2).strip()
+        coordinates = None
+        if coords_text:
+            try:
+                coordinates = ast.literal_eval(coords_text)
+            except Exception:
+                coordinates = None
+        matches.append(
+            {
+                "label": label,
+                "coordinates": coordinates,
+                "raw": match.group(0),
+                "span": match.span(),
+            }
+        )
+    return matches
+
+
+def postprocess_markdown(text: str) -> str:
+    """Clean up markdown text from model output."""
+    cleaned = (
+        text.replace("\\coloneqq", ":=")
+        .replace("\\eqqcolon", "=:")
+        .replace("<|image_pad|>", "")
+    )
+    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
+    return cleaned.strip()
+
+
+def apply_replacements(text: str, replacements: List[Tuple[int, int, str]]) -> str:
+    """Apply text replacements at specified spans."""
+    if not replacements:
+        return postprocess_markdown(text)
+    sorted_replacements = sorted(replacements, key=lambda item: item[0])
+    segments: List[str] = []
+    cursor = 0
+    for start, end, replacement in sorted_replacements:
+        segments.append(text[cursor:start])
+        segments.append(replacement)
+        cursor = end
+    segments.append(text[cursor:])
+    return postprocess_markdown("".join(segments))
+
+
+def crop_figure(
+    image: Image.Image,
+    sample_id: str,
+    figure_index: int,
+    pixel_box: List[int],
+    label: str,
+) -> Tuple[FigureMetadata, Image.Image]:
+    """Crop a figure region from the source image.
+
+    Args:
+        pixel_box: [x1, y1, x2, y2] bounding box in pixels
+
+    Returns:
+        (metadata, cropped_image) tuple for embedding in dataset
+    """
+    x1, y1, x2, y2 = pixel_box
+    crop = image.crop((x1, y1, x2, y2)).copy()
+
+    figure_id = f"{sample_id}_fig{figure_index:02d}"
+
+    metadata = FigureMetadata(
+        figure_id=figure_id,
+        label=label,
+        bounding_box_pixels={"x1": x1, "y1": y1, "x2": x2, "y2": y2},
+    )
+
+    return metadata, crop
+
+
+def write_text(path: Path, content: str) -> None:
+    """Write text content to a file."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(content, encoding="utf-8")
+
+
+def write_json(path: Path, payload: Any) -> None:
+    """Write JSON content to a file."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=False)
+
+
+def build_document_markdown(
+    image: Image.Image,
+    response_text: str,
+    sample_id: str,
+) -> Tuple[str, List[FigureMetadata], List[Image.Image], Image.Image]:
+    """Process model response to extract markdown and figures.
+
+    Returns:
+        (markdown, figure_metadata, figure_images, annotated_image) tuple
+    """
+    blocks = extract_grounding_blocks(response_text)
+    replacements: List[Tuple[int, int, str]] = []
+    figures: List[FigureMetadata] = []
+    figure_images: List[Image.Image] = []
+    figure_index = 1
+
+    img_draw = image.copy()
+    draw = ImageDraw.Draw(img_draw)
+    overlay = Image.new("RGBA", img_draw.size, (0, 0, 0, 0))
+    draw_overlay = ImageDraw.Draw(overlay)
+    font = ImageFont.load_default()
+
+    width, height = image.size
+
+    for block in blocks:
+        label = block["label"].lower()
+        start, end = block["span"]
+
+        # Random color for this block
+        color = (
+            np.random.randint(0, 200),
+            np.random.randint(0, 200),
+            np.random.randint(0, 255),
+        )
+        color_alpha = color + (20,)
+
+        # Convert normalized coords to pixels
+        raw_box = block["coordinates"][0]
+        x1 = int(raw_box[0] / 999 * width)
+        y1 = int(raw_box[1] / 999 * height)
+        x2 = int(raw_box[2] / 999 * width)
+        y2 = int(raw_box[3] / 999 * height)
+        pixel_box = (x1, y1, x2, y2)
+
+        # Extract figures (images)
+        if label == "image":
+            metadata, crop = crop_figure(
+                image=image,
+                sample_id=sample_id,
+                figure_index=figure_index,
+                pixel_box=pixel_box,
+                label=block["label"],
+            )
+            figures.append(metadata)
+            figure_images.append(crop)
+            # Use figure:{id} URI format - clearly an identifier, not a file path
+            replacements.append(
+                (
+                    start,
+                    end,
+                    f"![{metadata.figure_id}](figure:{metadata.figure_id})",
+                )
+            )
+            figure_index += 1
+        else:
+            replacements.append((start, end, ""))
+
+        # Draw bounding box
+        box_width = 4 if label == "title" else 2
+        draw.rectangle([x1, y1, x2, y2], outline=color, width=box_width)
+        draw_overlay.rectangle([x1, y1, x2, y2], fill=color_alpha)
+
+        # Draw label
+        text_x, text_y = x1, max(0, y1 - 15)
+        text_bbox = draw.textbbox((0, 0), label, font=font)
+        text_w, text_h = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
+        draw.rectangle(
+            [text_x, text_y, text_x + text_w, text_y + text_h], fill=(255, 255, 255, 30)
+        )
+        draw.text((text_x, text_y), label, font=font, fill=color)
+
+    img_draw.paste(overlay, (0, 0), overlay)
+    markdown = apply_replacements(response_text, replacements)
+    return markdown, figures, figure_images, img_draw
+
+
+def _truncate_for_alt(description: str, max_length: int = 120) -> str:
+    """Create a short alt text from a description (first sentence, truncated)."""
+    # Take first sentence
+    first_sentence = description.split(". ")[0].split(".\n")[0]
+    if len(first_sentence) <= max_length:
+        return first_sentence.strip()
+    # Truncate at word boundary
+    truncated = first_sentence[:max_length].rsplit(" ", 1)[0]
+    return truncated.strip() + "..."
+
+
+def enrich_markdown_with_captions(
+    markdown: str,
+    description_map: Dict[str, Dict[str, Any]],
+) -> str:
+    """Add figure captions to markdown. Alt text is truncated; full description below."""
+    used: set[str] = set()
+
+    def replace(match: re.Match[str]) -> str:
+        alt_text = match.group("figure_id").strip()
+        path = match.group("path").strip()
+
+        # Extract figure_id from figure:{id} URI or from alt text
+        if path.startswith("figure:"):
+            figure_id = path[7:]  # Remove "figure:" prefix
+        else:
+            # Legacy format - figure_id is in alt text after "Figure "
+            figure_id = alt_text.replace("Figure ", "").split(":")[0].strip()
+
+        entry = description_map.get(figure_id)
+        if not entry:
+            return match.group(0)
+
+        description = (entry.get("description") or "").strip()
+        if not description:
+            return match.group(0)
+
+        # Alt text: short summary (first sentence, max 120 chars)
+        short_alt = _truncate_for_alt(description)
+
+        # Image tag with short alt text
+        rendered = f"![{figure_id}: {short_alt}]({path})"
+
+        # Add full caption below (only once per figure)
+        if figure_id not in used:
+            rendered += f"\n\n*Figure {figure_id}: {description}*\n"
+            used.add(figure_id)
+        return rendered
+
+    return FIGURE_MARKDOWN_PATTERN.sub(replace, markdown)
+
+
+def render_markdown_with_images(
+    markdown: str,
+    figure_images: List[Image.Image],
+    figure_metadata: List[Dict[str, Any]],
+) -> str:
+    """Replace figure:{id} URIs in markdown with base64-encoded images."""
+    # Build figure_id -> image mapping
+    id_to_image: Dict[str, Image.Image] = {}
+    for i, meta in enumerate(figure_metadata):
+        fig_id = meta.get("figure_id", "")
+        if fig_id and i < len(figure_images) and figure_images[i] is not None:
+            id_to_image[fig_id] = figure_images[i]
+
+    def replace(match: re.Match[str]) -> str:
+        alt_text = match.group("figure_id").strip()
+        path = match.group("path").strip()
+
+        # Extract figure_id from figure:{id} URI or use alt_text as fallback
+        if path.startswith("figure:"):
+            figure_id = path[7:]  # Remove "figure:" prefix
+        else:
+            # Legacy path format - extract figure_id from alt_text
+            figure_id = alt_text.replace("Figure ", "").split(":")[0].strip()
+
+        img = id_to_image.get(figure_id)
+        if img is None:
+            return match.group(0)  # Keep original if image not found
+
+        # Embed as base64 data URI
+        data_uri = f"data:image/png;base64,{encode_image(img)}"
+        return f"![{alt_text}]({data_uri})"
+
+    return FIGURE_MARKDOWN_PATTERN.sub(replace, markdown)
+
+
+def render_sample_markdown(sample: Dict[str, Any]) -> str:
+    """Render dataset sample's markdown with embedded base64 images."""
+    markdown = (
+        sample.get("document_final_markdown") or sample.get("document_markdown") or ""
+    )
+
+    # Parse metadata
+    raw_metadata = sample.get("extracted_figures_metadata") or []
+    metadata = []
+    for m in raw_metadata:
+        if isinstance(m, str):
+            metadata.append(json.loads(m))
+        else:
+            metadata.append(m)
+
+    images = sample.get("extracted_figures") or []
+
+    return render_markdown_with_images(
+        markdown=markdown,
+        figure_images=images,
+        figure_metadata=metadata,
+    )
+
+
+def display_markdown(sample: Dict[str, Any]) -> None:
+    """Display sample's markdown with embedded images in Jupyter."""
+    from IPython.display import display, Markdown
+
+    rendered = render_sample_markdown(sample)
+    display(Markdown(rendered))
+
+
+def display_samples(dataset, num_samples: int = 2) -> None:
+    """Display samples with source images, markdown, and figure descriptions."""
+    from IPython.display import display
+
+    print(f"Dataset: {len(dataset)} samples")
+    print(f"Columns: {list(dataset.column_names)}")
+    print()
+
+    for i in range(min(num_samples, len(dataset))):
+        sample = dataset[i]
+        print(f"=== Sample {i}: {sample.get('sample_id', i)} ===")
+
+        # Show source image
+        if sample.get("source_image"):
+            print("Source image:")
+            img = sample["source_image"]
+            img.thumbnail((500, 500))  # Resize to max 500px
+            display(img)
+
+        # Show markdown preview
+        md = sample.get("document_markdown") or sample.get("document_markdown_text", "")
+        if md:
+            print(f"\nMarkdown preview ({len(md)} chars):")
+            print(md[:500] + "..." if len(md) > 500 else md)
+
+        # Show final markdown if available
+        final_md = sample.get("document_final_markdown") or sample.get(
+            "document_final_markdown_text", ""
+        )
+        if final_md:
+            print(f"\nFinal markdown preview ({len(final_md)} chars):")
+            print(final_md[:500] + "..." if len(final_md) > 500 else final_md)
+
+        # Show figures and their descriptions
+        figures = sample.get("extracted_figures", [])
+        metadata = sample.get("extracted_figures_metadata", [])
+        if figures:
+            print(f"\nExtracted figures: {len(figures)}")
+            for j, fig in enumerate(figures[:2]):  # Show max 2 figures
+                fig.thumbnail((500, 500))
+                display(fig)
+                # Show figure description if available
+                if j < len(metadata):
+                    try:
+                        meta = (
+                            json.loads(metadata[j])
+                            if isinstance(metadata[j], str)
+                            else metadata[j]
+                        )
+                        if meta.get("description"):
+                            print(f"  📝 Description: {meta['description'][:200]}...")
+                    except Exception:
+                        pass
+        print()
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/server.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/server.py
new file mode 100644
index 0000000000..f4457fdfaf
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/server.py
@@ -0,0 +1,280 @@
+"""vLLM server management and async inference client."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import signal
+import subprocess
+import threading
+import time
+from typing import TYPE_CHECKING, Any, Awaitable, Dict, List, Sequence
+
+import requests
+from openai import AsyncOpenAI
+
+from .document import encode_image
+
+if TYPE_CHECKING:
+    from PIL import Image
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _stream_output(pipe, prefix: str) -> None:
+    """Stream subprocess output to stdout with prefix."""
+    try:
+        for line in iter(pipe.readline, ""):
+            print(f"[{prefix}] {line.rstrip()}", flush=True)
+    finally:
+        pipe.close()
+
+
+def launch_vllm() -> subprocess.Popen:
+    """Launch vLLM server as subprocess."""
+    model_id = os.environ.get("MODEL_ID", "deepseek-ai/DeepSeek-OCR")
+    served_name = os.environ.get("SERVED_MODEL_NAME", "deepseek-ocr")
+    port = os.environ.get("PORT", "8080")
+    host = os.environ.get("HOST", "0.0.0.0")
+
+    cmd: List[str] = [
+        "vllm",
+        "serve",
+        "--model",
+        model_id,
+        "--served-model-name",
+        served_name,
+        "--tensor-parallel-size",
+        os.environ.get("TENSOR_PARALLEL_SIZE", "1"),
+        "--max-model-len",
+        os.environ.get("MAX_MODEL_LEN", "4096"),
+        "--gpu-memory-utilization",
+        os.environ.get("GPU_MEMORY_UTILIZATION", "0.90"),
+        "--port",
+        port,
+        "--host",
+        host,
+        "--trust-remote-code",
+        "--enable-chunked-prefill",
+        "--no-enable-prefix-caching",
+        "--mm-processor-cache-gb",
+        os.environ.get("MM_PROCESSOR_CACHE_GB", "0"),
+        "--logits-processors",
+        os.environ.get(
+            "LOGITS_PROCESSORS",
+            "vllm.model_executor.models.deepseek_ocr:NGramPerReqLogitsProcessor",
+        ),
+    ]
+
+    extra_args = os.environ.get("EXTRA_VLLM_ARGS")
+    if extra_args:
+        cmd.extend(extra_args.split())
+
+    LOGGER.info("Launching vLLM server: %s", " ".join(cmd))
+    process = subprocess.Popen(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1
+    )
+
+    # Start output streaming threads
+    threads = []
+    for name, pipe in [("STDOUT", process.stdout), ("STDERR", process.stderr)]:
+        if pipe:
+            t = threading.Thread(
+                target=_stream_output, args=(pipe, f"vLLM {name}"), daemon=True
+            )
+            t.start()
+            threads.append(t)
+
+    process._log_threads = threads  # type: ignore
+    return process
+
+
+def shutdown_server(server_process: subprocess.Popen) -> None:
+    """Gracefully shutdown vLLM server."""
+    LOGGER.info("Shutting down vLLM server")
+    server_process.send_signal(signal.SIGTERM)
+    try:
+        server_process.wait(timeout=30)
+    except subprocess.TimeoutExpired:
+        LOGGER.warning("Server did not exit in time, sending SIGKILL")
+        server_process.kill()
+
+    for thread in getattr(server_process, "_log_threads", []):
+        thread.join(timeout=1)
+
+
+def _format_duration(seconds: float) -> str:
+    """Format duration as mm:ss."""
+    minutes = int(seconds // 60)
+    secs = int(seconds % 60)
+    return f"{minutes:02d}:{secs:02d}"
+
+
+def wait_for_server(url: str, timeout_s: int = None, interval_s: int = 5) -> bool:
+    """Wait for server health endpoint to respond."""
+    if timeout_s is None:
+        timeout_s = int(os.environ.get("VLLM_STARTUP_TIMEOUT", "600"))  # 10 min default
+
+    start_time = time.time()
+    LOGGER.info("⏳ Waiting for vLLM server to start...")
+
+    deadline = time.time() + timeout_s
+    while time.time() < deadline:
+        try:
+            if requests.get(url, timeout=5).ok:
+                elapsed = time.time() - start_time
+                LOGGER.info("✅ vLLM server ready in %s", _format_duration(elapsed))
+                return True
+        except Exception:
+            pass
+        time.sleep(interval_s)
+
+    elapsed = time.time() - start_time
+    LOGGER.error("❌ vLLM server failed to start after %s", _format_duration(elapsed))
+    return False
+
+
+def should_launch_server() -> bool:
+    """Check if server should be auto-launched."""
+    return os.environ.get("SKIP_SERVER_LAUNCH", "").lower() not in {"1", "true", "yes"}
+
+
+def base_url_from_env() -> str:
+    """Get vLLM base URL from environment."""
+    port = os.environ.get("PORT", "8080")
+    return os.environ.get("BASE_URL", f"http://127.0.0.1:{port}")
+
+
+def _prepare_payload(
+    image: "Image.Image",
+    model_name: str,
+    prompt: str,
+    max_tokens: int,
+    temperature: float,
+) -> Dict[str, Any]:
+    """Prepare OpenAI-compatible chat completion payload."""
+    return {
+        "model": model_name,
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{encode_image(image)}"
+                        },
+                    },
+                ],
+            }
+        ],
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "extra_body": {
+            "skip_special_tokens": False,
+            "vllm_xargs": {
+                "ngram_size": 30,
+                "window_size": 90,
+                "whitelist_token_ids": "[128821,128822]",
+            },
+        },
+    }
+
+
+class DeepSeekClient:
+    """Async batch inference client for DeepSeek OCR via vLLM."""
+
+    def __init__(
+        self,
+        base_url: str,
+        model_name: str,
+        max_tokens: int,
+        temperature: float,
+        *,
+        request_timeout: int = 120,
+        max_retries: int = 3,
+        retry_backoff_seconds: float = 2.0,
+        max_retry_wait_seconds: float = 60.0,
+    ) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.model_name = model_name
+        self.default_max_tokens = max_tokens
+        self.default_temperature = temperature
+        self.default_request_timeout = request_timeout
+        self.max_retries = max(0, max_retries)
+        self.retry_backoff_seconds = max(0.0, retry_backoff_seconds)
+        self.max_retry_wait_seconds = max_retry_wait_seconds
+        self._client = AsyncOpenAI(api_key="vllm", base_url=f"{self.base_url}/v1")
+
+    async def _async_completion(self, payload: Dict[str, Any], timeout: int) -> str:
+        """Execute single async completion request."""
+        try:
+            response = await self._client.chat.completions.create(
+                model=payload["model"],
+                messages=payload["messages"],
+                max_tokens=payload["max_tokens"],
+                temperature=payload["temperature"],
+                timeout=timeout,
+                extra_body=payload.get("extra_body"),
+            )
+        except Exception as exc:
+            LOGGER.error("DeepSeek request failed: %s", exc)
+            raise
+
+        if not response.choices:
+            return ""
+        return getattr(response.choices[0].message, "content", "") or ""
+
+    def infer(self, requests_data: Sequence[Dict[str, Any]]) -> List[str]:
+        """Run batch inference synchronously.
+
+        Args:
+            requests_data: List of dicts with keys: image (PIL.Image), prompt (str),
+                           optional: max_tokens, temperature, request_timeout
+
+        Returns:
+            List of response strings, one per request
+        """
+        if not requests_data:
+            return []
+
+        payloads = []
+        timeouts = []
+        for req in requests_data:
+            payloads.append(
+                _prepare_payload(
+                    image=req["image"],
+                    model_name=self.model_name,
+                    prompt=req.get("prompt", ""),
+                    max_tokens=req.get("max_tokens", self.default_max_tokens),
+                    temperature=req.get("temperature", self.default_temperature),
+                )
+            )
+            timeouts.append(req.get("request_timeout") or self.default_request_timeout)
+
+        return self._run_async(self._async_infer_batch(payloads, timeouts))
+
+    async def _async_infer_batch(
+        self, payloads: Sequence[Dict[str, Any]], timeouts: Sequence[int]
+    ) -> List[str]:
+        """Run batch of async completions concurrently."""
+        tasks = [
+            asyncio.create_task(self._async_completion(p, t))
+            for p, t in zip(payloads, timeouts)
+        ]
+        return await asyncio.gather(*tasks)
+
+    @staticmethod
+    def _run_async(coro: Awaitable[Any]) -> Any:
+        """Run async coroutine in new event loop."""
+        loop = asyncio.new_event_loop()
+        try:
+            asyncio.set_event_loop(loop)
+            result = loop.run_until_complete(coro)
+            loop.run_until_complete(loop.shutdown_asyncgens())
+            return result
+        finally:
+            asyncio.set_event_loop(None)
+            loop.close()
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/sm_io.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/sm_io.py
new file mode 100644
index 0000000000..5d14284d56
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/sm_io.py
@@ -0,0 +1,152 @@
+"""Amazon S3 utilities for SageMaker jobs."""
+
+from __future__ import annotations
+
+import logging
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Tuple
+
+import boto3
+from botocore.config import Config
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_s3_client():
+    """Get S3 client with retry configuration."""
+    config = Config(
+        retries={"max_attempts": 3, "mode": "standard"},
+        max_pool_connections=50,
+    )
+    return boto3.client("s3", config=config)
+
+
+def parse_s3_uri(uri: str) -> Tuple[str, str]:
+    """Parse s3://bucket/key into (bucket, key)."""
+    if not uri.startswith("s3://"):
+        raise ValueError(f"Invalid S3 URI: {uri}")
+    parts = uri[5:].split("/", 1)
+    bucket = parts[0]
+    key = parts[1] if len(parts) > 1 else ""
+    return bucket, key
+
+
+def upload_files_to_s3(
+    *,
+    output_dir: Path,
+    s3_uri: str,
+    path_prefix: str = "",
+) -> None:
+    """Upload local directory contents to S3."""
+    if not s3_uri:
+        LOGGER.info("No S3 URI provided; skipping upload.")
+        return
+
+    bucket, base_prefix = parse_s3_uri(s3_uri)
+
+    full_prefix = base_prefix.rstrip("/")
+    if path_prefix:
+        full_prefix = (
+            f"{full_prefix}/{path_prefix.strip('/')}"
+            if full_prefix
+            else path_prefix.strip("/")
+        )
+
+    s3 = get_s3_client()
+    base = output_dir.resolve()
+
+    files = sorted(p for p in base.rglob("*") if p.is_file())
+    if not files:
+        LOGGER.info("Nothing to upload from %s", output_dir)
+        return
+
+    LOGGER.info("Uploading %d files to s3://%s/%s", len(files), bucket, full_prefix)
+
+    for local_path in files:
+        rel = local_path.relative_to(base).as_posix()
+        s3_key = f"{full_prefix}/{rel}" if full_prefix else rel
+        try:
+            s3.upload_file(str(local_path), bucket, s3_key)
+        except Exception as exc:
+            LOGGER.error(
+                "Failed to upload %s to s3://%s/%s: %s", local_path, bucket, s3_key, exc
+            )
+            raise
+
+
+def save_dataset_to_s3(
+    dataset,
+    s3_uri: str,
+    name: str = "dataset",
+) -> str:
+    """Save HF dataset to S3 in Arrow format. Returns the S3 URI."""
+    from datasets import DatasetDict
+
+    # Handle DatasetDict by extracting the first split
+    if isinstance(dataset, DatasetDict):
+        if "train" in dataset:
+            dataset = dataset["train"]
+        else:
+            split_name = list(dataset.keys())[0]
+            dataset = dataset[split_name]
+            LOGGER.info("Using split '%s' from DatasetDict", split_name)
+
+    bucket, prefix = parse_s3_uri(s3_uri)
+    full_prefix = prefix.rstrip("/")
+
+    # Save to local temp directory using Arrow format
+    local_dir = Path(f"/tmp/{name}_arrow_temp")
+    if local_dir.exists():
+        shutil.rmtree(local_dir)
+
+    LOGGER.info("Saving dataset to Arrow format...")
+    dataset.save_to_disk(str(local_dir))
+
+    # Upload entire directory to S3
+    s3_prefix = f"{full_prefix}/{name}" if full_prefix else name
+    upload_files_to_s3(output_dir=local_dir, s3_uri=f"s3://{bucket}/{s3_prefix}")
+
+    # Cleanup
+    shutil.rmtree(local_dir)
+
+    result_uri = f"s3://{bucket}/{s3_prefix}"
+    LOGGER.info("Saved dataset to %s", result_uri)
+    return result_uri
+
+
+def load_dataset_from_s3(s3_uri: str, split: str = "train") -> "Dataset":
+    """Load HF dataset from S3. Downloads locally to avoid s3fs caching issues."""
+    from datasets import load_from_disk
+    import tempfile
+
+    LOGGER.info("Loading dataset from %s", s3_uri)
+
+    # Parse S3 URI
+    bucket_name, prefix = parse_s3_uri(s3_uri)
+
+    # Download to local temp directory (bypasses s3fs cache)
+    s3 = get_s3_client()
+    local_dir = tempfile.mkdtemp(prefix="s3_dataset_")
+
+    # List and download all objects
+    paginator = s3.get_paginator("list_objects_v2")
+    download_count = 0
+    for page in paginator.paginate(Bucket=bucket_name, Prefix=f"{prefix}/"):
+        for obj in page.get("Contents", []):
+            key = obj["Key"]
+            filename = key.split("/")[-1]
+            if filename:  # Skip directory markers
+                local_path = f"{local_dir}/{filename}"
+                s3.download_file(bucket_name, key, local_path)
+                download_count += 1
+
+    LOGGER.info("Downloaded %d files to %s", download_count, local_dir)
+
+    # Load from local
+    ds = load_from_disk(local_dir)
+
+    return ds
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/stages.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/stages.py
new file mode 100644
index 0000000000..94060ad3e0
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/stages.py
@@ -0,0 +1,461 @@
+"""Pipeline stages: extract, describe, assemble."""
+
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+import time
+from dataclasses import asdict
+from datetime import datetime
+from typing import Any, Dict, List
+
+from datasets import Features, Sequence, Value, load_dataset, Image as HfImage
+
+from .config import AssembleSettings, DescribeSettings, ExtractSettings, env
+from .document import build_document_markdown, enrich_markdown_with_captions, write_json
+from .storage import get_storage, get_source_storage
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _format_duration(seconds: float) -> str:
+    """Format duration as mm:ss."""
+    minutes = int(seconds // 60)
+    secs = int(seconds % 60)
+    return f"{minutes:02d}:{secs:02d}"
+
+
+def _now_iso() -> str:
+    return datetime.utcnow().isoformat() + "Z"
+
+
+def _dataset_features() -> Features:
+    """Dataset schema - all data is embedded, no external file paths."""
+    return Features(
+        {
+            "sample_id": Value("string"),
+            "dataset_index": Value("int64"),
+            "source_image": HfImage(),
+            "document_with_boxes_image": HfImage(),
+            "document_markdown": Value("string"),
+            "extracted_figures": Sequence(HfImage()),
+            "extracted_figures_metadata": Sequence(Value("string")),
+            "document_final_markdown": Value("string"),
+        }
+    )
+
+
+def run_stage_extract(settings: ExtractSettings) -> None:
+    """Run OCR extraction on dataset samples."""
+    stage_start = time.time()
+    LOGGER.info("⏳ Starting EXTRACT stage...")
+
+    # Lazy import torch - only needed for extract stage
+    from torch.utils.data import DataLoader
+
+    dataset = load_dataset(
+        settings.dataset_name,
+        settings.dataset_config,
+        split=settings.dataset_split,
+        streaming=settings.stream_dataset,
+    )
+
+    # Setup iterator with optional DataLoader for streaming
+    if settings.stream_dataset:
+        num_workers = env("DATALOADER_WORKERS", 2, int)
+        prefetch = env("DATALOADER_PREFETCH", 2, int)
+        kwargs = {
+            "batch_size": 1,
+            "num_workers": num_workers,
+            "collate_fn": lambda b: b[0],
+        }
+        if num_workers > 0:
+            kwargs["prefetch_factor"] = prefetch
+        sample_iter = iter(DataLoader(dataset, **kwargs))
+    else:
+        sample_iter = iter(dataset)
+
+    settings.output_dir.mkdir(parents=True, exist_ok=True)
+    batches_dir = settings.output_dir / "document_batches"
+    if batches_dir.exists():
+        shutil.rmtree(batches_dir)
+    batches_dir.mkdir(parents=True, exist_ok=True)
+
+    batch_files: List[str] = []
+    batch_idx = 0
+    doc_count = 0
+    failures: List[Dict[str, Any]] = []
+    chunk_size = settings.inference.batch_size
+
+    LOGGER.info(
+        "Extract | dataset=%s/%s/%s | max_samples=%s | batch=%s",
+        settings.dataset_name,
+        settings.dataset_config,
+        settings.dataset_split,
+        settings.max_samples,
+        chunk_size,
+    )
+
+    contexts: List[Dict[str, Any]] = []
+    requests: List[Dict[str, Any]] = []
+
+    def flush():
+        nonlocal contexts, requests, doc_count, batch_idx
+        if not contexts:
+            return
+
+        try:
+            responses = settings.client.infer(requests)
+        except Exception as exc:
+            LOGGER.exception("Batch inference failed for %d samples", len(contexts))
+            for ctx in contexts:
+                failures.append({"sample_id": ctx["sample_id"], "error": str(exc)})
+                if hasattr(ctx.get("image"), "close"):
+                    ctx["image"].close()
+            contexts, requests = [], []
+            return
+
+        docs: List[Dict[str, Any]] = []
+        for i, ctx in enumerate(contexts):
+            img = ctx.get("image")
+            try:
+                text = responses[i].strip() if i < len(responses) else ""
+                if not text:
+                    raise RuntimeError("Empty response")
+
+                sample_dir = ctx["sample_dir"]
+                sample_id = ctx["sample_id"]
+
+                markdown, figures, figure_images, img_draw = build_document_markdown(
+                    image=img,
+                    response_text=text,
+                    sample_id=sample_id,
+                )
+
+                # Save images locally for dataset loading (HfImage needs file paths)
+                source_path = sample_dir / "source.png"
+                boxes_path = sample_dir / "document_with_boxes.png"
+                img_draw.save(boxes_path)
+
+                # Save figure images for dataset loading
+                figures_dir = sample_dir / "figures"
+                figures_dir.mkdir(parents=True, exist_ok=True)
+                figure_paths = []
+                for fig_meta, fig_img in zip(figures, figure_images):
+                    fig_path = figures_dir / f"{fig_meta.figure_id}.png"
+                    fig_img.save(fig_path)
+                    figure_paths.append(str(fig_path))
+
+                docs.append(
+                    {
+                        "sample_id": sample_id,
+                        "dataset_index": ctx["dataset_index"],
+                        "source_image": str(source_path),
+                        "document_with_boxes_image": str(boxes_path),
+                        "document_markdown": markdown,
+                        "extracted_figures": figure_paths,
+                        "extracted_figures_metadata": [
+                            json.dumps(asdict(f)) for f in figures
+                        ],
+                        "document_final_markdown": "",  # Filled in assemble stage
+                    }
+                )
+            except Exception as exc:
+                LOGGER.exception("Failed sample %s", ctx["sample_id"])
+                failures.append({"sample_id": ctx["sample_id"], "error": str(exc)})
+            finally:
+                if hasattr(img, "close"):
+                    img.close()
+
+        if docs:
+            batch_file = batches_dir / f"batch_{batch_idx:05d}.json"
+            write_json(batch_file, docs)
+            batch_files.append(str(batch_file))
+            batch_idx += 1
+            doc_count += len(docs)
+
+        contexts, requests = [], []
+
+    for idx, sample in enumerate(sample_iter):
+        if settings.max_samples and idx >= settings.max_samples:
+            break
+
+        sample_id = f"sample_{idx:05d}"
+        sample_dir = settings.output_dir / sample_id
+        sample_dir.mkdir(parents=True, exist_ok=True)
+
+        img = sample["images"][0].copy()
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        img.save(sample_dir / "source.png")
+
+        contexts.append(
+            {
+                "sample_id": sample_id,
+                "dataset_index": idx,
+                "sample_dir": sample_dir,
+                "image": img.copy(),
+            }
+        )
+        requests.append(
+            {
+                "image": contexts[-1]["image"],
+                "prompt": settings.prompt,
+                "max_tokens": settings.max_tokens,
+                "temperature": settings.temperature,
+                "request_timeout": settings.inference.request_timeout,
+            }
+        )
+        img.close()
+
+        if len(requests) >= chunk_size:
+            flush()
+
+    flush()
+
+    # Save manifest
+    write_json(
+        settings.output_dir / "manifest.json",
+        {
+            "generated_at": _now_iso(),
+            "stage": "extract",
+            "documents_count": doc_count,
+            "failures": failures,
+        },
+    )
+
+    # Load as HF dataset
+    ds = load_dataset("json", data_files=batch_files, features=_dataset_features())
+    shutil.rmtree(batches_dir)
+
+    # Get storage backend and save
+    storage = get_storage(repo_id=settings.hub.repo_id)
+    storage.save_dataset(ds, "dataset")
+
+    elapsed = time.time() - stage_start
+    LOGGER.info(
+        "✅ Extract complete | docs=%d | failures=%d | duration=%s",
+        doc_count,
+        len(failures),
+        _format_duration(elapsed),
+    )
+
+
+def run_stage_describe(settings: DescribeSettings) -> None:
+    """Describe figures in the dataset that lack descriptions."""
+    stage_start = time.time()
+    LOGGER.info("⏳ Starting DESCRIBE stage...")
+
+    # Get source storage and load dataset
+    source_storage = get_source_storage(
+        source_repo_id=settings.source_repo_id or settings.hub.repo_id
+    )
+    if not source_storage.is_configured:
+        raise ValueError(
+            "No source configured for describe stage (set SOURCE_REPO_ID, HF_REPO_ID, or S3_INPUT_URI)"
+        )
+
+    dataset = source_storage.load_dataset()
+    if dataset is None:
+        raise RuntimeError("Failed to load source dataset for describe stage")
+
+    settings.output_dir.mkdir(parents=True, exist_ok=True)
+    desc_dir = settings.output_dir / "descriptions"
+    if desc_dir.exists():
+        shutil.rmtree(desc_dir)
+    desc_dir.mkdir(parents=True, exist_ok=True)
+
+    chunk_size = settings.inference.batch_size
+    failures: List[Dict[str, Any]] = []
+    contexts: List[Dict[str, Any]] = []
+    requests: List[Dict[str, Any]] = []
+    batch_idx = 0
+    described = 0
+
+    def flush():
+        nonlocal contexts, requests, batch_idx, described
+        if not contexts:
+            return
+
+        results = []
+        try:
+            responses = settings.client.infer(requests)
+            for i, ctx in enumerate(contexts):
+                desc = responses[i].strip() if i < len(responses) else ""
+                if desc:
+                    results.append({"figure_id": ctx["figure_id"], "description": desc})
+                    described += 1
+        except Exception as exc:
+            LOGGER.exception("Describe batch failed")
+            for ctx in contexts:
+                failures.append({"figure_id": ctx.get("figure_id"), "error": str(exc)})
+        finally:
+            for ctx in contexts:
+                if hasattr(ctx.get("image"), "close"):
+                    ctx["image"].close()
+            contexts, requests = [], []
+
+        if results:
+            with (desc_dir / f"batch_{batch_idx:05d}.jsonl").open("w") as f:
+                for r in results:
+                    f.write(json.dumps(r) + "\n")
+            batch_idx += 1
+
+    # Queue figures needing descriptions
+    pending = 0
+    for row in dataset:
+        sample_id = row["sample_id"]
+        metas = row.get("extracted_figures_metadata") or []
+        images = row.get("extracted_figures") or []
+
+        for i, meta_json in enumerate(metas):
+            meta = json.loads(meta_json) if isinstance(meta_json, str) else meta_json
+            if meta.get("description"):
+                continue
+
+            pending += 1
+            fig_id = meta.get("figure_id", "")
+
+            if i >= len(images) or images[i] is None:
+                failures.append(
+                    {
+                        "sample_id": sample_id,
+                        "figure_id": fig_id,
+                        "reason": "missing_image",
+                    }
+                )
+                continue
+
+            fig_img = images[i]
+            contexts.append(
+                {"sample_id": sample_id, "figure_id": fig_id, "image": fig_img}
+            )
+            requests.append(
+                {
+                    "image": fig_img,
+                    "prompt": settings.prompt,
+                    "max_tokens": settings.max_tokens,
+                    "temperature": settings.temperature,
+                    "request_timeout": settings.inference.request_timeout,
+                }
+            )
+
+            if len(requests) >= chunk_size:
+                flush()
+
+    flush()
+
+    if pending == 0:
+        LOGGER.info("No figures need descriptions")
+        return
+
+    # Load descriptions and apply to dataset
+    lookup = {}
+    for f in sorted(desc_dir.glob("batch_*.jsonl")):
+        for line in f.read_text().splitlines():
+            if line.strip():
+                r = json.loads(line)
+                lookup[r["figure_id"]] = r["description"]
+
+    if not lookup:
+        LOGGER.info("No descriptions generated")
+        return
+
+    LOGGER.info(
+        "Lookup has %d descriptions, first few: %s",
+        len(lookup),
+        list(lookup.keys())[:3],
+    )
+
+    def apply(row):
+        metas = row.get("extracted_figures_metadata") or []
+        new_metas = []
+        for m in metas:
+            meta = json.loads(m) if isinstance(m, str) else m
+            if meta.get("figure_id") in lookup:
+                meta["description"] = lookup[meta["figure_id"]]
+            new_metas.append(json.dumps(meta))
+        return {"extracted_figures_metadata": new_metas}
+
+    # Don't pass features - let map infer them
+    updated = dataset.map(apply)
+    shutil.rmtree(desc_dir)
+
+    # Verify before saving
+    test_meta = updated[0]["extracted_figures_metadata"]
+    if test_meta:
+        test_parsed = (
+            json.loads(test_meta[0]) if isinstance(test_meta[0], str) else test_meta[0]
+        )
+        LOGGER.info(
+            "VERIFY before save - description present: %s",
+            test_parsed.get("description") is not None,
+        )
+
+    # Get output storage and save
+    storage = get_storage(repo_id=settings.hub.repo_id)
+    storage.save_dataset(updated, "dataset")
+
+    elapsed = time.time() - stage_start
+    LOGGER.info(
+        "✅ Describe complete | described=%d | failures=%d | duration=%s",
+        described,
+        len(failures),
+        _format_duration(elapsed),
+    )
+
+
+def run_stage_assemble(settings: AssembleSettings) -> None:
+    """Enrich markdown with figure descriptions."""
+    stage_start = time.time()
+    LOGGER.info("⏳ Starting ASSEMBLE stage...")
+
+    # Get source storage and load dataset
+    source_storage = get_source_storage(
+        source_repo_id=settings.source_repo_id or settings.hub.repo_id
+    )
+    if not source_storage.is_configured:
+        raise ValueError(
+            "No source configured for assemble stage (set SOURCE_REPO_ID, HF_REPO_ID, or S3_INPUT_URI)"
+        )
+
+    dataset = source_storage.load_dataset()
+    if dataset is None:
+        raise RuntimeError("Failed to load source dataset for assemble stage")
+
+    assembled_count = 0
+
+    def assemble(row):
+        nonlocal assembled_count
+        markdown = row.get("document_markdown") or ""
+        if not markdown:
+            return row
+
+        # Parse figure metadata for description lookup
+        desc_map = {}
+        for m in row.get("extracted_figures_metadata") or []:
+            meta = json.loads(m) if isinstance(m, str) else m
+            if meta.get("figure_id"):
+                desc_map[meta["figure_id"]] = meta
+
+        # Enrich with captions (keeps figure: URIs, adds descriptions)
+        # Images are rendered on-demand using render_sample_markdown()
+        final_markdown = enrich_markdown_with_captions(markdown, desc_map)
+
+        row["document_final_markdown"] = final_markdown
+        assembled_count += 1
+        return row
+
+    dataset = dataset.map(assemble)
+
+    # Get output storage and save
+    storage = get_storage(repo_id=settings.hub.repo_id)
+    storage.save_dataset(dataset, "dataset")
+
+    elapsed = time.time() - stage_start
+    LOGGER.info(
+        "✅ Assemble complete | assembled=%d | duration=%s",
+        assembled_count,
+        _format_duration(elapsed),
+    )
diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/storage.py b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/storage.py
new file mode 100644
index 0000000000..848d945a4b
--- /dev/null
+++ b/docs/sagemaker/notebooks/sagemaker-sdk/llm_ocr/storage.py
@@ -0,0 +1,267 @@
+"""Unified storage abstraction for dataset I/O.
+
+This module provides a common interface for saving/loading HuggingFace datasets,
+abstracting away whether we're using HuggingFace Hub, S3, or GCS.
+
+Usage:
+    from .storage import get_storage
+
+    storage = get_storage()
+    storage.save_dataset(dataset, "my_dataset")
+    dataset = storage.load_dataset()
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+from .config import env
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+LOGGER = logging.getLogger(__name__)
+
+
+class DatasetStorage(ABC):
+    """Abstract base class for dataset storage backends."""
+
+    @abstractmethod
+    def save_dataset(self, dataset: "Dataset", name: str) -> bool:
+        """Save dataset to storage. Returns True on success."""
+        pass
+
+    @abstractmethod
+    def load_dataset(self, split: str = "train") -> Optional["Dataset"]:
+        """Load dataset from storage. Returns None if unavailable."""
+        pass
+
+    @property
+    @abstractmethod
+    def is_configured(self) -> bool:
+        """Check if this storage backend is configured."""
+        pass
+
+
+class HFHubStorage(DatasetStorage):
+    """HuggingFace Hub storage backend."""
+
+    def __init__(
+        self,
+        repo_id: Optional[str] = None,
+        branch: Optional[str] = None,
+        commit_message: Optional[str] = None,
+    ):
+        self.repo_id = repo_id or env("HF_REPO_ID")
+        self.branch = branch or env("HF_BRANCH")
+        self.commit_message = commit_message or env("HF_COMMIT_MESSAGE")
+        self._token = env("HF_TOKEN")
+
+    @property
+    def is_configured(self) -> bool:
+        return bool(self.repo_id)
+
+    def save_dataset(self, dataset: "Dataset", name: str) -> bool:
+        if not self.is_configured:
+            LOGGER.debug("HF Hub not configured, skipping dataset save")
+            return False
+
+        # Fail early if token is missing or empty
+        if not self._token:
+            raise ValueError(
+                "HF_TOKEN is required to push datasets to the Hub. "
+                "Set the HF_TOKEN environment variable with a token that has write permissions."
+            )
+
+        try:
+            dataset.push_to_hub(
+                self.repo_id,
+                token=self._token,
+                revision=self.branch,
+                commit_message=self.commit_message or f"Add {name}",
+            )
+            LOGGER.info("Pushed dataset to HF Hub: %s", self.repo_id)
+            return True
+        except Exception as exc:
+            LOGGER.exception("HF Hub dataset push failed: %s", exc)
+            raise
+
+    def load_dataset(self, split: str = "train") -> Optional["Dataset"]:
+        if not self.is_configured:
+            LOGGER.debug("HF Hub not configured, cannot load dataset")
+            return None
+
+        try:
+            from datasets import load_dataset
+
+            LOGGER.info("Loading dataset from HF Hub: %s", self.repo_id)
+            return load_dataset(self.repo_id, split=split, token=self._token)
+        except Exception as exc:
+            LOGGER.exception("HF Hub dataset load failed: %s", exc)
+            return None
+
+
+class S3Storage(DatasetStorage):
+    """Amazon S3 storage backend."""
+
+    def __init__(
+        self,
+        output_uri: Optional[str] = None,
+        input_uri: Optional[str] = None,
+    ):
+        self.output_uri = output_uri or env("S3_OUTPUT_URI")
+        self.input_uri = input_uri or env("S3_INPUT_URI")
+
+    @property
+    def is_configured(self) -> bool:
+        return bool(self.output_uri or self.input_uri)
+
+    def save_dataset(self, dataset: "Dataset", name: str) -> bool:
+        if not self.output_uri:
+            LOGGER.debug("S3 output URI not configured, skipping dataset save")
+            return False
+
+        try:
+            from .sm_io import save_dataset_to_s3
+
+            save_dataset_to_s3(dataset, self.output_uri, name)
+            return True
+        except ImportError as exc:
+            LOGGER.warning("S3 save failed (missing dependency): %s", exc)
+            return False
+        except Exception as exc:
+            LOGGER.exception("S3 dataset save failed: %s", exc)
+            return False
+
+    def load_dataset(self, split: str = "train") -> Optional["Dataset"]:
+        if not self.input_uri:
+            LOGGER.debug("S3 input URI not configured, cannot load dataset")
+            return None
+
+        try:
+            from .sm_io import load_dataset_from_s3
+
+            return load_dataset_from_s3(self.input_uri, split=split)
+        except ImportError as exc:
+            LOGGER.warning("S3 load failed (missing dependency): %s", exc)
+            return None
+        except Exception as exc:
+            LOGGER.exception("S3 dataset load failed: %s", exc)
+            return None
+
+
+class GCSStorage(DatasetStorage):
+    """Google Cloud Storage backend."""
+
+    def __init__(
+        self,
+        output_uri: Optional[str] = None,
+        input_uri: Optional[str] = None,
+    ):
+        self.output_uri = output_uri or env("GCS_OUTPUT_URI")
+        self.input_uri = input_uri or env("GCS_INPUT_URI")
+
+    @property
+    def is_configured(self) -> bool:
+        return bool(self.output_uri or self.input_uri)
+
+    def save_dataset(self, dataset: "Dataset", name: str) -> bool:
+        if not self.output_uri:
+            LOGGER.debug("GCS output URI not configured, skipping dataset save")
+            return False
+
+        try:
+            from .cloudrun_io import save_dataset_to_gcs
+
+            save_dataset_to_gcs(dataset, self.output_uri, name)
+            return True
+        except ImportError as exc:
+            LOGGER.warning("GCS save failed (missing dependency): %s", exc)
+            return False
+        except Exception as exc:
+            LOGGER.exception("GCS dataset save failed: %s", exc)
+            return False
+
+    def load_dataset(self, split: str = "train") -> Optional["Dataset"]:
+        if not self.input_uri:
+            LOGGER.debug("GCS input URI not configured, cannot load dataset")
+            return None
+
+        try:
+            from .cloudrun_io import load_dataset_from_gcs
+
+            return load_dataset_from_gcs(self.input_uri, split=split)
+        except ImportError as exc:
+            LOGGER.warning("GCS load failed (missing dependency): %s", exc)
+            return None
+        except Exception as exc:
+            LOGGER.exception("GCS dataset load failed: %s", exc)
+            return None
+
+
+def get_storage(
+    *,
+    repo_id: Optional[str] = None,
+    s3_output_uri: Optional[str] = None,
+    s3_input_uri: Optional[str] = None,
+    gcs_output_uri: Optional[str] = None,
+    gcs_input_uri: Optional[str] = None,
+) -> DatasetStorage:
+    """Get the configured storage backend. Exactly one must be configured."""
+    gcs = GCSStorage(output_uri=gcs_output_uri, input_uri=gcs_input_uri)
+    s3 = S3Storage(output_uri=s3_output_uri, input_uri=s3_input_uri)
+    hf = HFHubStorage(repo_id=repo_id)
+
+    configured = [
+        name
+        for name, backend in [("GCS", gcs), ("S3", s3), ("HF Hub", hf)]
+        if backend.is_configured
+    ]
+
+    if len(configured) > 1:
+        raise ValueError(
+            f"Multiple storage backends configured: {configured}. Configure only one."
+        )
+    if len(configured) == 0:
+        raise ValueError(
+            "No storage backend configured. Set HF_REPO_ID, S3_OUTPUT_URI, or GCS_OUTPUT_URI."
+        )
+
+    if gcs.is_configured:
+        return gcs
+    if s3.is_configured:
+        return s3
+    return hf
+
+
+def get_source_storage(
+    *,
+    source_repo_id: Optional[str] = None,
+) -> DatasetStorage:
+    """Get storage for loading source data. Exactly one must be configured."""
+    gcs_input = env("GCS_INPUT_URI")
+    s3_input = env("S3_INPUT_URI")
+    hf_repo = source_repo_id or env("SOURCE_REPO_ID") or env("HF_REPO_ID")
+
+    configured = [
+        name
+        for name, val in [("GCS", gcs_input), ("S3", s3_input), ("HF Hub", hf_repo)]
+        if val
+    ]
+
+    if len(configured) > 1:
+        raise ValueError(
+            f"Multiple source backends configured: {configured}. Configure only one."
+        )
+    if len(configured) == 0:
+        raise ValueError(
+            "No source storage configured. Set SOURCE_REPO_ID, S3_INPUT_URI, or GCS_INPUT_URI."
+        )
+
+    if gcs_input:
+        return GCSStorage(input_uri=gcs_input)
+    if s3_input:
+        return S3Storage(input_uri=s3_input)
+    return HFHubStorage(repo_id=hf_repo)