Repository: replicate/cog

Issue:
Cog Redis needs to be able to put files to API's file upload endpoint
It'll be a redirect to a signed GCP URL.

We already have a thing for the HTTP API to PUT files. Might be able to reuse same thing so Redis and HTTP are consistent.

This: 

https://github.com/replicate/cog/blob/12dfcaa4d47adb586b75a60842e22f0e5d3a41f5/python/cog/files.py#L12-L17

Except it needs to put the file as the request body rather than a multipart form. And needs to be able to follow a 307 redirect.


Discussion:

Resolving patch:
diff --git a/python/cog/files.py b/python/cog/files.py
--- a/python/cog/files.py
+++ b/python/cog/files.py
@@ -28,3 +28,9 @@ def upload_file(fh: io.IOBase, output_file_prefix: str = None) -> str:
         mime_type = "application/octet-stream"
     s = encoded_body.decode("utf-8")
     return f"data:{mime_type};base64,{s}"
+
+
+def guess_filename(obj: io.IOBase) -> str:
+    """Tries to guess the filename of the given object."""
+    name = getattr(obj, "name", "file")
+    return os.path.basename(name)
diff --git a/python/cog/server/redis_queue.py b/python/cog/server/redis_queue.py
--- a/python/cog/server/redis_queue.py
+++ b/python/cog/server/redis_queue.py
@@ -2,8 +2,10 @@
 import io
 import json
 import os
+from mimetypes import guess_type
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Tuple
+from urllib.parse import urlparse
 import signal
 import sys
 import traceback
@@ -19,6 +21,7 @@
 from ..json import upload_files
 from ..response import Status
 from .runner import PredictionRunner
+from ..files import guess_filename
 
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter  # type: ignore
@@ -361,9 +364,20 @@ def setter(response: Any) -> None:
 
     def upload_files(self, obj: Any) -> Any:
         def upload_file(fh: io.IOBase) -> str:
-            resp = requests.put(self.upload_url, files={"file": fh})
+            filename = guess_filename(fh)
+            content_type, _ = guess_type(filename)
+
+            resp = requests.put(
+                ensure_trailing_slash(self.upload_url) + filename,
+                fh,  # type: ignore
+                headers={"Content-type": content_type},
+            )
             resp.raise_for_status()
-            return resp.json()["url"]
+
+            # strip any signing gubbins from the URL
+            final_url = urlparse(resp.url)._replace(query="").geturl()
+
+            return final_url
 
         return upload_files(obj, upload_file)
 
@@ -378,6 +392,7 @@ def calculate_time_in_queue(message_id: str) -> float:
     return now - queue_time
 
 
+
 def format_datetime(timestamp: datetime.datetime) -> str:
     """
     Formats a datetime in ISO8601 with a trailing Z, so it's also RFC3339 for
@@ -386,6 +401,16 @@ def format_datetime(timestamp: datetime.datetime) -> str:
     return timestamp.isoformat() + "Z"
 
 
+def ensure_trailing_slash(url: str) -> str:
+    """
+    Adds a trailing slash to `url` if not already present, and then returns it.
+    """
+    if url.endswith("/"):
+        return url
+    else:
+        return url + "/"
+
+
 def _queue_worker_from_argv(
     predictor: BasePredictor,
     redis_host: str,
diff --git a/test-integration/upload_server/app.py b/test-integration/upload_server/app.py
--- a/test-integration/upload_server/app.py
+++ b/test-integration/upload_server/app.py
@@ -1,7 +1,7 @@
 import os
 
 import flask
-from flask import Flask, jsonify, send_from_directory
+from flask import Flask, redirect, request, send_from_directory
 
 app = Flask("upload_server")
 
@@ -11,15 +11,21 @@ def handle_index():
     return "OK"
 
 
-@app.route("/upload", methods=["PUT"])
-def handle_upload():
-    f = flask.request.files["file"]
-    f.save(os.path.join("/uploads", f.filename))
-    return jsonify({"url": "http://upload-server:5000/download/{}".format(f.filename)})
+@app.route("/upload/<name>", methods=["PUT"])
+def redirect_upload(name):
+    print("I'm in here!", name)
+    return redirect(f"/files/{name}", 307)
 
 
-@app.route("/download/<name>")
-def download(name):
+@app.route("/files/<name>", methods=["PUT"])
+def upload_file(name):
+    with open(os.path.join("/uploads", name), "wb") as f:
+        f.write(request.get_data())
+    return "OK"
+
+
+@app.route("/files/<name>", methods=["GET"])
+def download_file(name):
     return send_from_directory("/uploads", name)
 
 


Test patch:
diff --git a/test-integration/test_integration/test_redis_queue.py b/test-integration/test_integration/test_redis_queue.py
--- a/test-integration/test_integration/test_redis_queue.py
+++ b/test-integration/test_integration/test_redis_queue.py
@@ -59,7 +59,7 @@ def test_queue_worker_files(
                 "id": predict_id,
                 "input": {
                     "text": "baz",
-                    "path": "http://upload-server:5000/download/input.txt",
+                    "path": "http://upload-server:5000/files/input.txt",
                 },
                 "webhook": webhook_url,
                 "logs": [],
@@ -83,11 +83,11 @@ def capture_final_response(request):
                 "id": predict_id,
                 "input": {
                     "text": "baz",
-                    "path": "http://upload-server:5000/download/input.txt",
+                    "path": "http://upload-server:5000/files/input.txt",
                 },
                 "webhook": webhook_url,
                 "logs": [],
-                "output": "http://upload-server:5000/download/output.txt",
+                "output": "http://upload-server:5000/files/output.txt",
                 "status": "succeeded",
                 "started_at": mock.ANY,
                 "completed_at": mock.ANY,
@@ -111,7 +111,7 @@ def capture_final_response(request):
                             "id": predict_id,
                             "input": {
                                 "text": "baz",
-                                "path": "http://upload-server:5000/download/input.txt",
+                                "path": "http://upload-server:5000/files/input.txt",
                             },
                             "webhook": webhook_url,
                         }
@@ -172,7 +172,7 @@ def test_queue_worker_yielding_file(
             json={
                 "id": predict_id,
                 "input": {
-                    "path": "http://upload-server:5000/download/input.txt",
+                    "path": "http://upload-server:5000/files/input.txt",
                 },
                 "webhook": webhook_url,
                 "logs": [],
@@ -188,11 +188,11 @@ def test_queue_worker_yielding_file(
             json={
                 "id": predict_id,
                 "input": {
-                    "path": "http://upload-server:5000/download/input.txt",
+                    "path": "http://upload-server:5000/files/input.txt",
                 },
                 "webhook": webhook_url,
                 "logs": [],
-                "output": ["http://upload-server:5000/download/out-0.txt"],
+                "output": ["http://upload-server:5000/files/out-0.txt"],
                 "status": "processing",
                 "started_at": mock.ANY,
             },
@@ -204,13 +204,13 @@ def test_queue_worker_yielding_file(
             json={
                 "id": predict_id,
                 "input": {
-                    "path": "http://upload-server:5000/download/input.txt",
+                    "path": "http://upload-server:5000/files/input.txt",
                 },
                 "webhook": webhook_url,
                 "logs": [],
                 "output": [
-                    "http://upload-server:5000/download/out-0.txt",
-                    "http://upload-server:5000/download/out-1.txt",
+                    "http://upload-server:5000/files/out-0.txt",
+                    "http://upload-server:5000/files/out-1.txt",
                 ],
                 "status": "processing",
                 "started_at": mock.ANY,
@@ -223,14 +223,14 @@ def test_queue_worker_yielding_file(
             json={
                 "id": predict_id,
                 "input": {
-                    "path": "http://upload-server:5000/download/input.txt",
+                    "path": "http://upload-server:5000/files/input.txt",
                 },
                 "webhook": webhook_url,
                 "logs": [],
                 "output": [
-                    "http://upload-server:5000/download/out-0.txt",
-                    "http://upload-server:5000/download/out-1.txt",
-                    "http://upload-server:5000/download/out-2.txt",
+                    "http://upload-server:5000/files/out-0.txt",
+                    "http://upload-server:5000/files/out-1.txt",
+                    "http://upload-server:5000/files/out-2.txt",
                 ],
                 "status": "processing",
                 "started_at": mock.ANY,
@@ -243,14 +243,14 @@ def test_queue_worker_yielding_file(
             json={
                 "id": predict_id,
                 "input": {
-                    "path": "http://upload-server:5000/download/input.txt",
+                    "path": "http://upload-server:5000/files/input.txt",
                 },
                 "webhook": webhook_url,
                 "logs": [],
                 "output": [
-                    "http://upload-server:5000/download/out-0.txt",
-                    "http://upload-server:5000/download/out-1.txt",
-                    "http://upload-server:5000/download/out-2.txt",
+                    "http://upload-server:5000/files/out-0.txt",
+                    "http://upload-server:5000/files/out-1.txt",
+                    "http://upload-server:5000/files/out-2.txt",
                 ],
                 "status": "succeeded",
                 "started_at": mock.ANY,
@@ -274,7 +274,7 @@ def test_queue_worker_yielding_file(
                         {
                             "id": predict_id,
                             "input": {
-                                "path": "http://upload-server:5000/download/input.txt",
+                                "path": "http://upload-server:5000/files/input.txt",
                             },
                             "webhook": webhook_url,
                         }
@@ -1431,7 +1431,7 @@ def test_queue_worker_yielding_list_of_complex_output(
                 "output": [
                     [
                         {
-                            "file": "http://upload-server:5000/download/file",
+                            "file": "http://upload-server:5000/files/file",
                             "text": "hello",
                         }
                     ]
@@ -1452,7 +1452,7 @@ def test_queue_worker_yielding_list_of_complex_output(
                 "output": [
                     [
                         {
-                            "file": "http://upload-server:5000/download/file",
+                            "file": "http://upload-server:5000/files/file",
                             "text": "hello",
                         }
                     ]