replace requests with httpx and factor out clients (#1574)

* input downloads, output uploads, and webhooks are now handled by ClientManager, which persists for the lifetime of runner, allowing us to reuse connections, which may significantly help with large uploads.
* although I was originally going to drop output_file_prefix, it's not actually hard to maintain. the behavior is changed now and objects are uploaded as soon as they're outputted rather than after the prediction is completed.
* there's an ugly hack with uploading an empty body to get the redirect instead of making api time out from trying to upload an 140GB file. that can be fixed by implemented an MPU endpoint and/or a "fetch upload url" endpoint.
* the behavior of the non-indempotent endpoint is changed; the id is now randomly generated if it's not provided in the body. this isn't strictly required for this change alone, but is hard to carve out.
* the behavior of Path is changed significantly. see https://www.notion.so/replicate/Cog-Setup-Path-Problem-2fc41d40bcaf47579ccd8b2f4c71ee24

Signed-off-by: technillogue <technillogue@gmail.com>
Co-authored-by: Mattt <mattt@replicate.com>
Signed-off-by: technillogue <technillogue@gmail.com>

Author: technillogue

python/cog/predictor.py

@@ -91,7 +91,7 @@ def get_weights_argument(
     if weights_type is None:
         return None
     weights_url = os.environ.get("COG_WEIGHTS")
-    weights_path = "weights"  # this is the source of a bug isn't it?
+    weights_path = "weights" # this is the source of a bug isn't it?
 
     # TODO: Cog{File,Path}.validate(...) methods accept either "real"
     # paths/files or URLs to those things. In future we can probably tidy this

python/cog/server/clients.py

@@ -109,6 +109,7 @@ def __init__(self) -> None:
         self.log = structlog.get_logger(__name__).bind()
 
     async def aclose(self) -> None:
+        # not used but it's not actually critical to close them
         await self.webhook_client.aclose()
         await self.retry_webhook_client.aclose()
         await self.file_client.aclose()

python/cog/server/runner.py

@@ -89,8 +89,11 @@ def __init__(
         tee_output: bool = True,
     ) -> None:
         self._shutdown_event = shutdown_event  # __main__ waits for this event
+
         self._upload_url = upload_url
-        self._predictions: "dict[str, tuple[schema.PredictionResponse, PredictionTask]]" = {}
+        self._predictions: (
+            "dict[str, tuple[schema.PredictionResponse, PredictionTask]]"
+        ) = {}
         self._predictions_in_flight: "set[str]" = set()
         # it would be lovely to merge these but it's not fully clear how best to handle it
         # since idempotent requests can kinda come whenever?

python/tests/server/test_webhook.py

@@ -2,6 +2,7 @@
 import requests
 import responses
 from cog.schema import WebhookEvent
+
 from responses import registries
 
 pytest.skip(allow_module_level=True)