feat: Add exclude parameter to S3 pull functionality

- Add typesafe exclude parameter to filter out directories during S3 sync
- Valid exclude options: "checkpoints", "logs", "trajectories"
- Update pull_model_trajectories helper to exclude checkpoints and logs by default
- Enables pulling only trajectories for analysis without downloading large model weights

This is useful when analyzing trajectory data locally without needing the
model checkpoints, which can be very large and consume significant bandwidth.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: corbt

src/art/local/backend.py

@@ -22,7 +22,7 @@
 from transformers.models.auto.tokenization_auto import AutoTokenizer
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 from tqdm import auto as tqdm
-from typing import AsyncIterator, cast
+from typing import AsyncIterator, cast, Literal
 import wandb
 from wandb.sdk.wandb_run import Run
 import weave
@@ -443,6 +443,7 @@ async def _experimental_pull_from_s3(
         prefix: str | None = None,
         verbose: bool = False,
         delete: bool = False,
+        exclude: list[Literal["checkpoints", "logs", "trajectories"]] | None = None,
     ) -> None:
         """Download the model directory from S3 into local Backend storage. Right now this can be used to pull trajectory logs for processing or model checkpoints.
         Args:
@@ -452,7 +453,19 @@ async def _experimental_pull_from_s3(
             prefix: The prefix to pull from S3. If None, the model name will be used.
             verbose: Whether to print verbose output.
             delete: Whether to delete the local model directory.
+            exclude: List of directories to exclude from sync. Valid options: "checkpoints", "logs", "trajectories".
         """
+        # Validate exclude options
+        validated_exclude = None
+        if exclude:
+            validated_exclude = []
+            for item in exclude:
+                if item not in ["checkpoints", "logs", "trajectories"]:
+                    raise ValueError(
+                        f"Invalid exclude option: {item}. Valid options are: checkpoints, logs, trajectories"
+                    )
+                validated_exclude.append(item)  # type: ignore
+
         await pull_model_from_s3(
             model_name=model.name,
             project=model.project,
@@ -462,6 +475,7 @@ async def _experimental_pull_from_s3(
             verbose=verbose,
             delete=delete,
             art_path=self._path,
+            exclude=validated_exclude,  # type: ignore
         )
 
     async def _experimental_push_to_s3(

src/art/utils/benchmarking/load_trajectories.py

@@ -238,6 +238,7 @@ async def pull_model_trajectories(model: ArtModel) -> None:
             model,
             s3_bucket=bucket,
             verbose=True,
+            exclude=["checkpoints", "logs"],
         )
 
         print("Finished pulling trajectories.", flush=True)

src/art/utils/s3.py

@@ -4,7 +4,7 @@
 import asyncio
 from asyncio.subprocess import DEVNULL
 import tempfile
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Literal
 import zipfile
 
 from art.errors import ForbiddenBucketCreationError
@@ -17,6 +17,8 @@
 
 __all__: Sequence[str] = ("s3_sync",)
 
+ExcludableOption = Literal["checkpoints", "logs", "trajectories"]
+
 
 class S3SyncError(RuntimeError):
     """Raised when the underlying *aws s3 sync* command exits with a non‑zero status."""
@@ -67,6 +69,7 @@ async def s3_sync(
     profile: Optional[str] = None,
     verbose: bool = False,
     delete: bool = False,
+    exclude: list[ExcludableOption] | None = None,
 ) -> None:
     """Synchronise *source* and *destination* using the AWS CLI.
 
@@ -82,6 +85,7 @@ async def s3_sync(
         profile: Optional AWS profile name to pass to the CLI.
         verbose: When *True*, the output of the AWS CLI is streamed to the
             calling process; otherwise it is suppressed.
+        exclude: List of directories to exclude from sync.
 
     Raises:
         S3SyncError: If the *aws s3 sync* command exits with a non‑zero status.
@@ -100,6 +104,12 @@ async def s3_sync(
 
     if delete:
         cmd.append("--delete")
+    
+    # Add exclude patterns for each excluded directory
+    if exclude:
+        for excluded_dir in exclude:
+            cmd.extend(["--exclude", f"{excluded_dir}/*"])
+    
     cmd += [source, destination]
 
     # Suppress output unless verbose mode is requested.
@@ -156,6 +166,7 @@ async def pull_model_from_s3(
     verbose: bool = False,
     delete: bool = False,
     art_path: str | None = None,
+    exclude: list[ExcludableOption] | None = None,
 ) -> str:
     """Pull a model from S3 to the local directory.
 
@@ -169,6 +180,7 @@ async def pull_model_from_s3(
             calling process; otherwise it is suppressed.
         delete: When *True*, delete the local model directory if it exists.
         art_path: The path to the ART directory.
+        exclude: List of directories to exclude from sync.
 
     Returns:
         The local directory path.
@@ -196,7 +208,7 @@ async def pull_model_from_s3(
         prefix=prefix,
     )
     await ensure_bucket_exists(s3_bucket)
-    await s3_sync(s3_path, local_dir, verbose=verbose, delete=delete)
+    await s3_sync(s3_path, local_dir, verbose=verbose, delete=delete, exclude=exclude)
 
     # After pulling, migrate to new structure if needed
     if step is not None:

test_s3_exclude.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+"""Test script to verify S3 exclude functionality."""
+
+import asyncio
+import os
+from dotenv import load_dotenv
+from art.model import Model
+from art.local import LocalBackend
+from art.utils.benchmarking.load_trajectories import pull_model_trajectories
+
+# Load environment variables
+load_dotenv()
+
+async def test_s3_exclude():
+    """Test pulling only trajectories from S3, excluding checkpoints and logs."""
+    
+    # Create a test model
+    model = Model(
+        name="email-agent-216-4",
+        project="email_agent"
+    )
+    
+    # Test 1: Use the helper function (should exclude checkpoints and logs by default now)
+    print("Test 1: Using pull_model_trajectories helper (excludes checkpoints and logs)")
+    print("-" * 60)
+    
+    try:
+        await pull_model_trajectories(model)
+        print("✓ Helper function completed successfully")
+    except Exception as e:
+        print(f"✗ Helper function failed: {e}")
+    
+    # Test 2: Direct backend call with exclude parameter
+    print("\nTest 2: Direct backend call with exclude=['checkpoints', 'logs']")
+    print("-" * 60)
+    
+    with LocalBackend() as backend:
+        bucket = os.getenv("BACKUP_BUCKET")
+        if not bucket:
+            print("✗ BACKUP_BUCKET environment variable not set")
+            return
+            
+        try:
+            await backend._experimental_pull_from_s3(
+                model,
+                s3_bucket=bucket,
+                verbose=True,
+                exclude=["checkpoints", "logs"],
+            )
+            print("✓ Direct backend call completed successfully")
+        except Exception as e:
+            print(f"✗ Direct backend call failed: {e}")
+    
+    # Test 3: Verify the pulled structure
+    print("\nTest 3: Verifying pulled structure")
+    print("-" * 60)
+    
+    # Check current directory .art folder first
+    art_path = ".art"
+    model_path = f"{art_path}/{model.project}/models/{model.name}"
+    
+    if os.path.exists(model_path):
+        print(f"Model directory exists: {model_path}")
+        
+        # Check what was pulled
+        trajectories_path = os.path.join(model_path, "trajectories")
+        checkpoints_path = os.path.join(model_path, "checkpoints")
+        logs_path = os.path.join(model_path, "logs")
+        
+        trajectories_exists = os.path.exists(trajectories_path)
+        checkpoints_exists = os.path.exists(checkpoints_path)
+        logs_exists = os.path.exists(logs_path)
+        
+        print(f"{'✓' if trajectories_exists else '✗'} Trajectories directory exists: {trajectories_exists}")
+        
+        # Check if excluded directories are empty
+        checkpoints_empty = not checkpoints_exists or len(os.listdir(checkpoints_path)) == 0
+        logs_empty = not logs_exists or len(os.listdir(logs_path)) == 0
+        
+        print(f"{'✓' if checkpoints_empty else '✗'} Checkpoints directory empty: {checkpoints_empty}")
+        print(f"{'✓' if logs_empty else '✗'} Logs directory empty: {logs_empty}")
+        
+        # List contents
+        print("\nDirectory contents:")
+        for item in os.listdir(model_path):
+            item_path = os.path.join(model_path, item)
+            if os.path.isdir(item_path):
+                print(f"  - {item}/ ({len(os.listdir(item_path))} items)")
+            else:
+                print(f"  - {item}")
+    else:
+        print(f"✗ Model directory not found: {model_path}")
+
+if __name__ == "__main__":
+    asyncio.run(test_s3_exclude())