Implement multipart upload (#364)

* implement core multipart upload + test

* fixed whitespaces

* refactor

* Fix multipart name mismatch

* Add test_data/ to gitignore to prevent committing large test files

* Fix

Author: sbassam

.gitignore

@@ -127,3 +127,4 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+test_data/

src/together/constants.py

@@ -15,6 +15,20 @@
 DOWNLOAD_BLOCK_SIZE = 10 * 1024 * 1024  # 10 MB
 DISABLE_TQDM = False
 
+# Upload defaults
+MAX_CONCURRENT_PARTS = 4  # Maximum concurrent parts for multipart upload
+
+# Multipart upload constants
+MIN_PART_SIZE_MB = 5  # Minimum part size (S3 requirement)
+TARGET_PART_SIZE_MB = 100  # Target part size for optimal performance
+MAX_MULTIPART_PARTS = 250  # Maximum parts per upload (S3 limit)
+MULTIPART_UPLOAD_TIMEOUT = 300  # Timeout in seconds for uploading each part
+MULTIPART_THRESHOLD_GB = 5.0  # threshold for switching to multipart upload
+
+# maximum number of GB sized files we support finetuning for
+MAX_FILE_SIZE_GB = 25.0
+
+
 # Messages
 MISSING_API_KEY_MESSAGE = """TOGETHER_API_KEY not found.
 Please set it as an environment variable or set it as together.api_key
@@ -26,8 +40,6 @@
 # the number of bytes in a gigabyte, used to convert bytes to GB for readable comparison
 NUM_BYTES_IN_GB = 2**30
 
-# maximum number of GB sized files we support finetuning for
-MAX_FILE_SIZE_GB = 4.9
 
 # expected columns for Parquet files
 PARQUET_EXPECTED_COLUMNS = ["input_ids", "attention_mask", "labels"]

src/together/filemanager.py

@@ -1,28 +1,40 @@
 from __future__ import annotations
 
+import math
 import os
 import shutil
 import stat
 import tempfile
 import uuid
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from functools import partial
 from pathlib import Path
-from typing import Tuple
+from typing import Any, Dict, List, Tuple
 
 import requests
 from filelock import FileLock
 from requests.structures import CaseInsensitiveDict
 from tqdm import tqdm
-from tqdm.utils import CallbackIOWrapper
 
-import together.utils
 from together.abstract import api_requestor
-from together.constants import DISABLE_TQDM, DOWNLOAD_BLOCK_SIZE, MAX_RETRIES
+from together.constants import (
+    DISABLE_TQDM,
+    DOWNLOAD_BLOCK_SIZE,
+    MAX_CONCURRENT_PARTS,
+    MAX_FILE_SIZE_GB,
+    MAX_RETRIES,
+    MIN_PART_SIZE_MB,
+    NUM_BYTES_IN_GB,
+    TARGET_PART_SIZE_MB,
+    MAX_MULTIPART_PARTS,
+    MULTIPART_UPLOAD_TIMEOUT,
+)
 from together.error import (
     APIError,
     AuthenticationError,
     DownloadError,
     FileTypeError,
+    ResponseError,
 )
 from together.together_response import TogetherResponse
 from together.types import (
@@ -32,6 +44,8 @@
     TogetherClient,
     TogetherRequest,
 )
+from tqdm.utils import CallbackIOWrapper
+import together.utils
 
 
 def chmod_and_replace(src: Path, dst: Path) -> None:
@@ -339,7 +353,7 @@ def upload(
                 )
             redirect_url, file_id = self.get_upload_url(url, file, purpose, filetype)
 
-        file_size = os.stat(file.as_posix()).st_size
+        file_size = os.stat(file).st_size
 
         with tqdm(
             total=file_size,
@@ -385,3 +399,214 @@ def upload(
         assert isinstance(response, TogetherResponse)
 
         return FileResponse(**response.data)
+
+
+class MultipartUploadManager:
+    """Handles multipart uploads for large files"""
+
+    def __init__(self, client: TogetherClient) -> None:
+        self._client = client
+        self.max_concurrent_parts = MAX_CONCURRENT_PARTS
+
+    def upload(
+        self,
+        url: str,
+        file: Path,
+        purpose: FilePurpose,
+    ) -> FileResponse:
+        """Upload large file using multipart upload"""
+
+        file_size = os.stat(file).st_size
+
+        file_size_gb = file_size / NUM_BYTES_IN_GB
+        if file_size_gb > MAX_FILE_SIZE_GB:
+            raise FileTypeError(
+                f"File size {file_size_gb:.1f}GB exceeds maximum supported size of {MAX_FILE_SIZE_GB}GB"
+            )
+
+        part_size, num_parts = self._calculate_parts(file_size)
+
+        file_type = self._get_file_type(file)
+        upload_info = None
+
+        try:
+            upload_info = self._initiate_upload(
+                url, file, file_size, num_parts, purpose, file_type
+            )
+
+            completed_parts = self._upload_parts_concurrent(
+                file, upload_info, part_size
+            )
+
+            return self._complete_upload(
+                url, upload_info["upload_id"], upload_info["file_id"], completed_parts
+            )
+
+        except Exception as e:
+            # Cleanup on failure
+            if upload_info is not None:
+                self._abort_upload(
+                    url, upload_info["upload_id"], upload_info["file_id"]
+                )
+            raise e
+
+    def _get_file_type(self, file: Path) -> str:
+        """Get file type from extension, raising ValueError for unsupported extensions"""
+        if file.suffix == ".jsonl":
+            return "jsonl"
+        elif file.suffix == ".parquet":
+            return "parquet"
+        elif file.suffix == ".csv":
+            return "csv"
+        else:
+            raise ValueError(
+                f"Unsupported file extension: '{file.suffix}'. "
+                f"Supported extensions: .jsonl, .parquet, .csv"
+            )
+
+    def _calculate_parts(self, file_size: int) -> tuple[int, int]:
+        """Calculate optimal part size and count"""
+        min_part_size = MIN_PART_SIZE_MB * 1024 * 1024  # 5MB
+        target_part_size = TARGET_PART_SIZE_MB * 1024 * 1024  # 100MB
+
+        if file_size <= target_part_size:
+            return file_size, 1
+
+        num_parts = min(MAX_MULTIPART_PARTS, math.ceil(file_size / target_part_size))
+        part_size = math.ceil(file_size / num_parts)
+
+        if part_size < min_part_size:
+            part_size = min_part_size
+            num_parts = math.ceil(file_size / part_size)
+
+        return part_size, num_parts
+
+    def _initiate_upload(
+        self,
+        url: str,
+        file: Path,
+        file_size: int,
+        num_parts: int,
+        purpose: FilePurpose,
+        file_type: str,
+    ) -> Any:
+        """Initiate multipart upload with backend"""
+
+        requestor = api_requestor.APIRequestor(client=self._client)
+
+        payload = {
+            "file_name": file.name,
+            "file_size": file_size,
+            "num_parts": num_parts,
+            "purpose": purpose.value,
+            "file_type": file_type,
+        }
+
+        response, _, _ = requestor.request(
+            options=TogetherRequest(
+                method="POST",
+                url="files/multipart/initiate",
+                params=payload,
+            ),
+        )
+
+        return response.data
+
+    def _upload_parts_concurrent(
+        self, file: Path, upload_info: Dict[str, Any], part_size: int
+    ) -> List[Dict[str, Any]]:
+        """Upload file parts concurrently with progress tracking"""
+
+        parts = upload_info["parts"]
+        completed_parts = []
+
+        with ThreadPoolExecutor(max_workers=self.max_concurrent_parts) as executor:
+            with tqdm(total=len(parts), desc="Uploading parts", unit="part") as pbar:
+                future_to_part = {}
+
+                with open(file, "rb") as f:
+                    for part_info in parts:
+                        f.seek((part_info["PartNumber"] - 1) * part_size)
+                        part_data = f.read(part_size)
+
+                        future = executor.submit(
+                            self._upload_single_part, part_info, part_data
+                        )
+                        future_to_part[future] = part_info["PartNumber"]
+
+                # Collect results
+                for future in as_completed(future_to_part):
+                    part_number = future_to_part[future]
+                    try:
+                        etag = future.result()
+                        completed_parts.append(
+                            {"part_number": part_number, "etag": etag}
+                        )
+                        pbar.update(1)
+                    except Exception as e:
+                        raise Exception(f"Failed to upload part {part_number}: {e}")
+
+        completed_parts.sort(key=lambda x: x["part_number"])
+        return completed_parts
+
+    def _upload_single_part(self, part_info: Dict[str, Any], part_data: bytes) -> str:
+        """Upload a single part and return ETag"""
+
+        response = requests.put(
+            part_info["URL"],
+            data=part_data,
+            headers=part_info.get("Headers", {}),
+            timeout=MULTIPART_UPLOAD_TIMEOUT,
+        )
+        response.raise_for_status()
+
+        etag = response.headers.get("ETag", "").strip('"')
+        if not etag:
+            raise ResponseError(f"No ETag returned for part {part_info['PartNumber']}")
+
+        return etag
+
+    def _complete_upload(
+        self,
+        url: str,
+        upload_id: str,
+        file_id: str,
+        completed_parts: List[Dict[str, Any]],
+    ) -> FileResponse:
+        """Complete the multipart upload"""
+
+        requestor = api_requestor.APIRequestor(client=self._client)
+
+        payload = {
+            "upload_id": upload_id,
+            "file_id": file_id,
+            "parts": completed_parts,
+        }
+
+        response, _, _ = requestor.request(
+            options=TogetherRequest(
+                method="POST",
+                url="files/multipart/complete",
+                params=payload,
+            ),
+        )
+
+        return FileResponse(**response.data.get("file", response.data))
+
+    def _abort_upload(self, url: str, upload_id: str, file_id: str) -> None:
+        """Abort the multipart upload"""
+
+        requestor = api_requestor.APIRequestor(client=self._client)
+
+        payload = {
+            "upload_id": upload_id,
+            "file_id": file_id,
+        }
+
+        requestor.request(
+            options=TogetherRequest(
+                method="POST",
+                url="files/multipart/abort",
+                params=payload,
+            ),
+        )

src/together/resources/files.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import os
 from pathlib import Path
 from pprint import pformat
 
 from together.abstract import api_requestor
+from together.constants import MULTIPART_THRESHOLD_GB, NUM_BYTES_IN_GB
 from together.error import FileTypeError
-from together.filemanager import DownloadManager, UploadManager
+from together.filemanager import DownloadManager, UploadManager, MultipartUploadManager
 from together.together_response import TogetherResponse
 from together.types import (
     FileDeleteResponse,
@@ -30,7 +32,6 @@ def upload(
         purpose: FilePurpose | str = FilePurpose.FineTune,
         check: bool = True,
     ) -> FileResponse:
-        upload_manager = UploadManager(self._client)
 
         if check and purpose == FilePurpose.FineTune:
             report_dict = check_file(file)
@@ -47,7 +48,15 @@ def upload(
 
         assert isinstance(purpose, FilePurpose)
 
-        return upload_manager.upload("files", file, purpose=purpose, redirect=True)
+        file_size = os.stat(file).st_size
+        file_size_gb = file_size / NUM_BYTES_IN_GB
+
+        if file_size_gb > MULTIPART_THRESHOLD_GB:
+            multipart_manager = MultipartUploadManager(self._client)
+            return multipart_manager.upload("files", file, purpose)
+        else:
+            upload_manager = UploadManager(self._client)
+            return upload_manager.upload("files", file, purpose=purpose, redirect=True)
 
     def list(self) -> FileList:
         requestor = api_requestor.APIRequestor(

src/together/types/files.py

@@ -52,7 +52,7 @@ class FileResponse(BaseModel):
     """
 
     id: str
-    object: Literal[ObjectType.File]
+    object: str
     # created timestamp
     created_at: int | None = None
     type: FileType | None = None

src/together/utils/files.py

@@ -65,7 +65,7 @@ def check_file(
     else:
         report_dict["found"] = True
 
-    file_size = os.stat(file.as_posix()).st_size
+    file_size = os.stat(file).st_size
 
     if file_size > MAX_FILE_SIZE_GB * NUM_BYTES_IN_GB:
         report_dict["message"] = (