fix: Optimize memory and timeout handling for 50GB+ files

Author: sbassam

src/together/filemanager.py

@@ -212,6 +212,7 @@ def download(
                     ),
                     remaining_retries=MAX_RETRIES,
                     stream=True,
+                    request_timeout=3600,
                 )
 
                 try:
@@ -522,29 +523,47 @@ def _upload_parts_concurrent(
 
         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:
+                    future_to_part = {}
+                    part_index = 0
+
+                    # Submit initial batch limited by max_concurrent_parts
+                    for i in range(min(self.max_concurrent_parts, len(parts))):
+                        part_info = parts[part_index]
                         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}")
+                        part_index += 1
+
+                    # Process completions and submit new parts (sliding window)
+                    while future_to_part:
+                        done_future = next(as_completed(future_to_part))
+                        part_number = future_to_part.pop(done_future)
+
+                        try:
+                            etag = done_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}")
+
+                        # Submit next part if available
+                        if part_index < len(parts):
+                            part_info = parts[part_index]
+                            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"]
+                            part_index += 1
 
         completed_parts.sort(key=lambda x: x["part_number"])
         return completed_parts

src/together/utils/files.py

@@ -7,6 +7,7 @@
 from traceback import format_exc
 from typing import Any, Dict, List
 
+from tqdm import tqdm
 
 from together.constants import (
     MAX_FILE_SIZE_GB,
@@ -363,14 +364,20 @@ def _check_utf8(file: Path) -> Dict[str, Any]:
         Dict[str, Any]: A dictionary with the results of the check.
     """
     report_dict: Dict[str, Any] = {}
+
     try:
+        # Stream file in chunks to avoid loading entire file into memory
+        chunk_size = 8192  # 8KB chunks
         with file.open(encoding="utf-8") as f:
-            f.read()
+            for chunk in iter(lambda: f.read(chunk_size), ""):
+                pass  # UTF-8 decoding happens automatically during read
+
         report_dict["utf8"] = True
     except UnicodeDecodeError as e:
         report_dict["utf8"] = False
         report_dict["message"] = f"File is not UTF-8 encoded. Error raised: {e}."
         report_dict["is_check_passed"] = False
+
     return report_dict
 
 
@@ -470,7 +477,7 @@ def _check_jsonl(file: Path, purpose: FilePurpose | str) -> Dict[str, Any]:
     with file.open() as f:
         idx = -1
         try:
-            for idx, line in enumerate(f):
+            for idx, line in tqdm(enumerate(f), desc="Validating file", unit=" lines"):
                 json_line = json.loads(line)
 
                 if not isinstance(json_line, dict):