dubious upload fix

technillogue · technillogue · commit 6996db3cee81 · 2024-03-12T15:44:38.000-04:00
Signed-off-by: technillogue &lt;technillogue@gmail.com&gt;
diff --git a/python/cog/server/clients.py b/python/cog/server/clients.py
@@ -156,6 +156,7 @@ async def upload_file(self, fh: io.IOBase, url: Optional[str]) -> str:
         # in that case we need to return data uris
         if url is None:
             return file_to_data_uri(fh, content_type)
+        assert url
 
         # ensure trailing slash
         url_with_trailing_slash = url if url.endswith("/") else url + "/"
@@ -169,10 +170,20 @@ async def chunk_file_reader() -> AsyncIterator[bytes]:
                     break
                 yield chunk
 
+        url = url_with_trailing_slash + filename
+        if url and "internal" in url:
+            resp1 = await self.file_client.put(
+                url,
+                content=b"",
+                headers={"Content-Type": content_type},
+                follow_redirects=False,
+            )
+            if resp1.status_code == 307:
+                url = resp1.headers["Location"]
         resp = await self.file_client.put(
-            url_with_trailing_slash + filename,
+            url,
             content=chunk_file_reader(),
-            headers={"Content-type": content_type},
+            headers={"Content-Type": content_type},
         )
         resp.raise_for_status()
 
diff --git a/python/cog/server/runner.py b/python/cog/server/runner.py
@@ -10,7 +10,7 @@
 import typing  # TypeAlias, py3.10
 from datetime import datetime, timezone
 from enum import Enum, auto, unique
-from typing import Any, AsyncIterator, Iterator, Optional, Union, TypeVar
+from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union
 
 import httpx
 import structlog
@@ -483,13 +483,16 @@ async def handle_event_stream(
                 break
         return self.response
 
+    async def noop(self) -> None:
+        pass
+
     def event_to_handle_future(self, event: PublicEventType) -> Awaitable[None]:
         if isinstance(event, Heartbeat):
             # Heartbeat events exist solely to ensure that we have a
             # regular opportunity to check for cancelation and
             # timeouts.
             # We don't need to do anything with them.
-            return
+            return self.noop()
         if isinstance(event, Log):
             return self.append_logs(event.message)
 
@@ -500,9 +503,9 @@ def event_to_handle_future(self, event: PublicEventType) -> Awaitable[None]:
             if self._output_type.multi:
                 return self.set_output([])
         if isinstance(event, PredictionOutput):
-            if output_type is None:
+            if self._output_type is None:
                 return self.failed(error="Predictor returned unexpected output")
-            if output_type.multi:
+            if self._output_type.multi:
                 return self.append_output(event.payload)
             return self.set_output(event.payload)
         if isinstance(event, Done):  # pyright: ignore reportUnnecessaryIsinstance
