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

* 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

Co-authored-by: Mattt <[email protected]>

* format
* stick a %s on line 190 clients.py (#1707)
* local upload server can be called cluster.local in addition to .internal (#1714)

Signed-off-by: technillogue <[email protected]>

Author: technillogue

Diff for: pyproject.toml

@@ -10,11 +10,13 @@ authors = [{ name = "Replicate", email = "[email protected]" }]
 license.file = "LICENSE"
 urls."Source" = "https://github.com/replicate/cog"
 
-requires-python = ">=3.7"
+requires-python = ">=3.8"
 dependencies = [
   # intentionally loose. perhaps these should be vendored to not collide with user code?
   "attrs>=20.1,<24",
   "fastapi>=0.75.2,<0.99.0",
+  # we may not need http2
+  "httpx[http2]>=0.21.0,<1",
   "pydantic>=1.9,<2",
   "PyYAML",
   "requests>=2,<3",
@@ -27,9 +29,9 @@ dependencies = [
 optional-dependencies = { "dev" = [
   "black",
   "build",
-  "httpx",
   'hypothesis<6.80.0; python_version < "3.8"',
   'hypothesis; python_version >= "3.8"',
+  "respx",
   'numpy<1.22.0; python_version < "3.8"',
   'numpy; python_version >= "3.8"',
   "pillow",

Diff for: python/cog/files.py

@@ -1,13 +1,10 @@
-import io
 from datetime import datetime
 from enum import Enum
 from types import GeneratorType
-from typing import Any, Callable
+from typing import Any
 
 from pydantic import BaseModel
 
-from .types import Path
-
 
 def make_encodeable(obj: Any) -> Any:
     """
@@ -39,24 +36,3 @@ def make_encodeable(obj: Any) -> Any:
         if isinstance(obj, np.ndarray):
             return obj.tolist()
     return obj
-
-
-def upload_files(obj: Any, upload_file: Callable[[io.IOBase], str]) -> Any:
-    """
-    Iterates through an object from make_encodeable and uploads any files.
-
-    When a file is encountered, it will be passed to upload_file. Any paths will be opened and converted to files.
-    """
-    # skip four isinstance checks for fast text models
-    if type(obj) == str:  # noqa: E721
-        return obj
-    if isinstance(obj, dict):
-        return {key: upload_files(value, upload_file) for key, value in obj.items()}
-    if isinstance(obj, list):
-        return [upload_files(value, upload_file) for value in obj]
-    if isinstance(obj, Path):
-        with obj.open("rb") as f:
-            return upload_file(f)
-    if isinstance(obj, io.IOBase):
-        return upload_file(obj)
-    return obj

Diff for: python/cog/json.py

@@ -86,4 +86,5 @@ def setup_logging(*, log_level: int = logging.NOTSET) -> None:
 
     # Reconfigure log levels for some overly chatty libraries
     logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
+    # FIXME: no more urllib3(?)
     logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)

Diff for: python/cog/logging.py

@@ -1,7 +1,6 @@
 import enum
 import importlib.util
 import inspect
-import io
 import os.path
 import sys
 import types
@@ -11,6 +10,7 @@
 from pathlib import Path
 from typing import (
     Any,
+    Awaitable,
     Callable,
     Dict,
     List,
@@ -22,34 +22,24 @@
 )
 from unittest.mock import patch
 
-import structlog
-
-import cog.code_xforms as code_xforms
-
 try:
     from typing import get_args, get_origin
 except ImportError:  # Python < 3.8
     from typing_compat import get_args, get_origin  # type: ignore
 
+import structlog
 import yaml
 from pydantic import BaseModel, Field, create_model
 from pydantic.fields import FieldInfo
 
 # Added in Python 3.9. Can be from typing if we drop support for <3.9
 from typing_extensions import Annotated
 
+from . import code_xforms
 from .errors import ConfigDoesNotExist, PredictorNotSet
-from .types import (
-    CogConfig,
-    Input,
-    URLPath,
-)
-from .types import (
-    File as CogFile,
-)
-from .types import (
-    Path as CogPath,
-)
+from .types import CogConfig, Input, URLTempFile
+from .types import File as CogFile
+from .types import Path as CogPath
 from .types import Secret as CogSecret
 
 log = structlog.get_logger("cog.server.predictor")
@@ -66,7 +56,7 @@
 
 
 class BasePredictor(ABC):
-    def setup(self, weights: Optional[Union[CogFile, CogPath, str]] = None) -> None:
+    def setup(self, weights: Optional[Union[CogFile, CogPath, str]] = None) -> Optional[Awaitable[None]]:
         """
         An optional method to prepare the model so multiple predictions run efficiently.
         """
@@ -81,51 +71,69 @@ def predict(self, **kwargs: Any) -> Any:
 
 
 def run_setup(predictor: BasePredictor) -> None:
+    weights = get_weights_argument(predictor)
+    if weights:
+        predictor.setup(weights=weights)
+    else:
+        predictor.setup()
+
+
+async def run_setup_async(predictor: BasePredictor) -> None:
+    weights = get_weights_argument(predictor)
+    maybe_coro = predictor.setup(weights=weights) if weights else predictor.setup()
+    if maybe_coro:
+        return await maybe_coro
+
+
+def get_weights_argument(
+    predictor: BasePredictor,
+) -> Union[CogFile, CogPath, str, None]:
+    # by the time we get here we assume predictor has a setup method
     weights_type = get_weights_type(predictor.setup)
 
     # No weights need to be passed, so just run setup() without any arguments.
     if weights_type is None:
         predictor.setup()
         return
 
-    weights: Union[io.IOBase, Path, str, None]
-
     weights_url = os.environ.get("COG_WEIGHTS")
-    weights_path = "weights"
+    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
     # up a little bit.
     # TODO: CogFile/CogPath should have subclasses for each of the subtypes
+
+    # this is a breaking change
+    # previously, CogPath wouldn't be converted in setup(); now it is
+    # essentially everyone needs to switch from Path to str (or a new URL type)
     if weights_url:
         if weights_type == CogFile:
-            weights = cast(CogFile, CogFile.validate(weights_url))
+            return cast(CogFile, CogFile.validate(weights_url))
         elif weights_type == CogPath:
             # TODO: So this can be a url. evil!
-            weights = cast(CogPath, CogPath.validate(weights_url))
+            return cast(CogPath, CogPath.validate(weights_url))
         # allow people to download weights themselves
         elif weights_type == str:  # noqa: E721
-            weights = weights_url
+            return weights_url
         else:
             raise ValueError(
                 f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File, Path and str are supported"
             )
     elif os.path.exists(weights_path):
         if weights_type == CogFile:
-            weights = cast(CogFile, open(weights_path, "rb"))
+            return cast(CogFile, open(weights_path, "rb"))
         elif weights_type == CogPath:
-            weights = CogPath(weights_path)
-        else:
-            raise ValueError(
-                f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File, Path and str are supported"
-            )
-    else:
-        weights = None
-
-    predictor.setup(weights=weights)
+            return CogPath(weights_path)
+        raise ValueError(
+            f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File, Path and str are supported"
+        )
+    return None
 
 
-def get_weights_type(setup_function: Callable[[Any], None]) -> Optional[Any]:
+def get_weights_type(
+    setup_function: Callable[[Any], Optional[Awaitable[None]]],
+) -> Optional[Any]:
     signature = inspect.signature(setup_function)
     if "weights" not in signature.parameters:
         return None
@@ -266,12 +274,19 @@ def cleanup(self) -> None:
         Cleanup any temporary files created by the input.
         """
         for _, value in self:
-            # Handle URLPath objects specially for cleanup.
+            # Handle URLTempFile objects specially for cleanup.
             # Also handle pathlib.Path objects, which cog.Path is a subclass of.
             # A pathlib.Path object shouldn't make its way here,
             # but both have an unlink() method, so we may as well be safe.
-            if isinstance(value, (URLPath, Path)):
-                value.unlink(missing_ok=True)
+            if isinstance(value, (URLTempFile, Path)):
+                try:
+                    value.unlink(missing_ok=True)
+                except FileNotFoundError:
+                    pass
+
+    # if we had a separate method to traverse the input and apply some function to each value
+    # we could have cleanup/get_tempfile/convert functions that operate on a single value
+    # and do it that way. convert is supposed to mutate though, so it's tricky
 
 
 def validate_input_type(type: Type[Any], name: str) -> None:

Diff for: python/cog/predictor.py

@@ -1,6 +1,7 @@
 import importlib.util
 import os
 import os.path
+import secrets
 import sys
 import typing as t
 from datetime import datetime
@@ -43,7 +44,14 @@ class PredictionBaseModel(pydantic.BaseModel, extra=pydantic.Extra.allow):
 
 
 class PredictionRequest(PredictionBaseModel):
-    id: t.Optional[str]
+    # there's a problem here where the idempotent endpoint is supposed to
+    # let you pass id in the route and omit it from the input
+    # however this fills in the default
+    # maybe it should be allowed to be optional without the factory initially
+    # and be filled in later
+    #
+    # actually, this changes the public api so we should really do this differently
+    id: str = pydantic.Field(default_factory=lambda: secrets.token_hex(4))
     created_at: t.Optional[datetime]
 
     # TODO: deprecate this