R2 Signed URL update (#5)

* Temporary R2 redirect

* Preprocess post for R2

* minor fix

* minor fix again

* minor fix again 2

* New R2 endpoint polish. Add preliminary debug logs for upload.

* Write output to stdout

* ruff formatting fix

Author: orangetin

src/together/commands/files.py

@@ -32,6 +32,13 @@ def _add_upload(parser: argparse._SubParsersAction[argparse.ArgumentParser]) ->
         help="File to upload",
         type=str,
     )
+    upload_file_parser.add_argument(
+        "--log-level",
+        default="WARNING",
+        choices=["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "NOTSET"],
+        type=str,
+        help="Set logging level. Defaults to WARNING. DEBUG will show all logs.",
+    )
     upload_file_parser.set_defaults(func=_run_upload)
 
 
@@ -87,7 +94,7 @@ def _run_list(args: argparse.Namespace) -> None:
 
 
 def _run_upload(args: argparse.Namespace) -> None:
-    files = Files(args.endpoint)
+    files = Files(args.endpoint, log_level=args.log_level)
     response = files.upload_file(args.file)
     print(json.dumps(response, indent=4))
 

src/together/files.py

@@ -1,9 +1,10 @@
 import json
+import logging
 import os
 import posixpath
 import sys
 import urllib.parse
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import requests
 from tqdm import tqdm
@@ -12,13 +13,6 @@
 DEFAULT_ENDPOINT = "https://api.together.xyz/"
 
 
-def exception_handler(exception_type: Any, exception: Any, traceback: Any) -> None:
-    # All your trace are belong to us!
-    # your format
-    # print("%s: %s" % (exception_type.__name__, exception))
-    pass
-
-
 class JSONException(Exception):
     pass
 
@@ -39,6 +33,7 @@ class Files:
     def __init__(
         self,
         endpoint_url: Optional[str] = None,
+        log_level: str = "WARNING",
     ) -> None:
         self.together_api_key = os.environ.get("TOGETHER_API_KEY", None)
         if self.together_api_key is None:
@@ -51,6 +46,17 @@ def __init__(
 
         self.endpoint_url = urllib.parse.urljoin(endpoint_url, "/v1/files/")
 
+        self.logger = logging.getLogger(__name__)
+
+        # Setup logging
+        logging.basicConfig(
+            format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+            datefmt="%m/%d/%Y %H:%M:%S",
+            handlers=[logging.StreamHandler(sys.stdout)],
+        )
+
+        self.logger.setLevel(log_level)
+
     def list_files(self) -> Dict[str, List[Dict[str, Union[str, int]]]]:
         headers = {
             "Authorization": f"Bearer {self.together_api_key}",
@@ -72,40 +78,66 @@ def list_files(self) -> Dict[str, List[Dict[str, Union[str, int]]]]:
         return response_json
 
     def upload_file(self, file: str) -> Dict[str, Union[str, int]]:
-        data = {"purpose": "fine-tune"}
-        headers = {
-            "Authorization": f"Bearer {self.together_api_key}",
-        }
+        try:
+            data = {"purpose": "fine-tune", "file_name": os.path.basename(file)}
 
-        if not validate_json(file=file):
-            raise ValueError("Could not load file: invalid .jsonl file detected.")
+            headers = {
+                "Authorization": f"Bearer {self.together_api_key}",
+            }
+
+            if not validate_json(file=file):
+                raise ValueError("Could not load file: invalid .jsonl file detected.")
+
+            session = requests.Session()
+            init_endpoint = self.endpoint_url[:-1]
+
+            self.logger.debug(
+                f"Upload file POST request: data={data}, headers={headers}, URL={init_endpoint}, allow_redirects=False"
+            )
+
+            response = session.post(
+                init_endpoint,
+                data=data,
+                headers=headers,
+                allow_redirects=False,
+            )
+            self.logger.debug(f"Response: {response.text}")
+            r2_signed_url = response.headers["Location"]
+            file_id = response.headers["X-Together-File-Id"]
+
+            self.logger.info(f"R2 Signed URL: {r2_signed_url}")
+            self.logger.info("File-ID")
+
+            self.logger.info("Uploading file...")
+            # print("> Uploading file...")
 
-        sys.excepthook = exception_handler  # send request
-        try:
             with open(file, "rb") as f:
-                response = requests.post(
-                    self.endpoint_url,
-                    headers=headers,
-                    files={"file": f},
-                    data=data,
-                )
+                response = session.put(r2_signed_url, files={"file": f})
 
-        except Exception:
-            print(
-                "ERROR: An exception occurred during file upload, likely due to trying to upload a large file. Please note, that we have a 100MB file upload limit at the moment. Up to 5GB uploads are coming soon, stay tuned!"
+            self.logger.info("> File uploaded.")
+            self.logger.debug(f"status code: {response.status_code}")
+            self.logger.info("> Processing file...")
+            preprocess_url = urllib.parse.urljoin(
+                self.endpoint_url, f"{file_id}/preprocess"
             )
-            sys.exit(0)
-            # print(response.text)
-            # raise ValueError(f"Error raised by files endpoint: {e}")
 
-        try:
-            response_json = dict(response.json())
-        except Exception:
-            raise ValueError(
-                f"JSON Error raised. \nResponse status code: {str(response.status_code)}"
+            response = session.post(
+                preprocess_url,
+                headers=headers,
             )
 
-        return response_json
+            self.logger.info("> File processed")
+            self.logger.debug(f"Status code: {response.status_code}")
+
+        except Exception:
+            self.logger.critical("Response error raised.")
+            sys.exit(1)
+
+        return {
+            "filename": os.path.basename(file),
+            "id": str(file_id),
+            "object": "file",
+        }
 
     def delete_file(self, file_id: str) -> Dict[str, str]:
         delete_url = urllib.parse.urljoin(self.endpoint_url, file_id)
@@ -131,6 +163,7 @@ def delete_file(self, file_id: str) -> Dict[str, str]:
 
     def retrieve_file(self, file_id: str) -> Dict[str, Union[str, int]]:
         retrieve_url = urllib.parse.urljoin(self.endpoint_url, file_id)
+        print(retrieve_url)
 
         headers = {
             "Authorization": f"Bearer {self.together_api_key}",