Scary temporary commit for a hemorrhaging-edge release

* add concurrency to config
* this basically works!
* more descriptive names for predict functions
* maybe pass through prediction id and try to make cancelation do both?
* don't cancel from signal handler if a loop is running. expose worker busy state to runner
* move handle_event_stream to PredictionEventHandler
* make setup and canceling work
* drop some checks around cancelation
* try out eager_predict_state_change
* keep track of multiple runner prediction tasks to make idempotent endpoint return the same result and fix tests somewhat
* fix idempotent tests
* fix remaining errors?
* worker predict_generator shouldn't be eager
* wip: make the stuff that handles events and sends webhooks etc async
* drop Runner._result
* drop comments
* inline client code
* get started
* inline webhooks
* move clients into runner, switch to httpx, move create_event_handler into runner
* add some comments
* more notes
* rip out webhooks and most of files and put them in a new ClientManager that handles most of everything. inline upload_files for that
* move create_event_handler into PredictionEventHandler.__init__
* fix one test
* break out Path.validate into value_to_path and inline get_filename and File.validate
* split out URLPath into BackwardsCompatibleDataURLTempFilePath and URLThatCanBeConvertedToPath with the download part of URLFile inlined
* let's make DataURLTempFilePath also use convert and move value_to_path back to Path.validate
* use httpx for downloading input urls and follow redirects
* take get_filename back out for tests
* don't upload in http and delete cog/files.py
* drop should_cancel
* prediction->request
* split up predict/inner/prediction_ctx into enter_predict/exit_predict/prediction_ctx/inner_async_predict/predict/good_predict as one way to do it. however, exposing all of those for runner predict enter/coro exit still sucks, but this is still an improvement
* bigish change: inline predict_and_handle_errors
* inline make_error_handler into setup
* move runner.setup into runner.Runner.setup
* add concurrency to config in go
* try explicitly using prediction_ctx __enter__ and __exit__
* make runner setup more correct and marginally better
* fix a few tests
* notes
* wip ClientManager.convert
* relax setup argument requirement to str
* glom worker into runner
* add logging message
* fix prediction retry and improve logging
* split out handle_event
* use CURL_CA_BUNDLE for file upload
* clean up comments
* dubious upload fix
* small fixes
* attempt to add context logging?
* tweak names
* fix error for predictionOutputType(multi=False)
* improve comments
* fix lints
* skip worker and webhook tests since those were erroring on removed imports. fix or xfail runner tests
* upload in http instead of PredictionEventHandler. this makes tests pass and fixes some problems with validation, but also prevents streaming files and causes new problems. also xfail all the responses tests that need to be replaced with respx
* format
* fix some new-style type signatures and drop 3.8 support
* drop 3.7 in code

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

Author: technillogue

Diff for: .github/workflows/ci.yaml

@@ -53,7 +53,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        # python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
     defaults:
       run:
         shell: bash

Diff for: Makefile

@@ -11,7 +11,7 @@ GO := go
 GOOS := $(shell $(GO) env GOOS)
 GOARCH := $(shell $(GO) env GOARCH)
 
-PYTHON := python
+PYTHON ?= python
 PYTEST := $(PYTHON) -m pytest
 PYRIGHT := $(PYTHON) -m pyright
 RUFF := $(PYTHON) -m ruff

Diff for: pkg/config/config.go

@@ -50,10 +50,11 @@ type Example struct {
 }
 
 type Config struct {
-	Build   *Build `json:"build" yaml:"build"`
-	Image   string `json:"image,omitempty" yaml:"image"`
-	Predict string `json:"predict,omitempty" yaml:"predict"`
-	Train   string `json:"train,omitempty" yaml:"train"`
+	Build       *Build `json:"build" yaml:"build"`
+	Image       string `json:"image,omitempty" yaml:"image"`
+	Predict     string `json:"predict,omitempty" yaml:"predict"`
+	Train       string `json:"train,omitempty" yaml:"train"`
+	Concurrency int    `json:"concurrency,omitempty" yaml:"concurrency"`
 }
 
 func DefaultConfig() *Config {

Diff for: pkg/config/data/config_schema_v1.0.json

@@ -140,6 +140,11 @@
       "$id": "#/properties/train",
       "type": "string",
       "description": "The pointer to the `Predictor` object in your code, which defines how predictions are run on your model."
+    },
+    "concurrency": {
+      "$id": "#/properties/concurrency",
+      "type": "number",
+      "description": "Allowed concurrency."
     }
   },
   "additionalProperties": false

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",
+  # this version specification is pretty arbitrary, and we may not need http2
+  "httpx[http2]>=0.25.0,<0.27",
   "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,21 +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.
-    """
-    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

@@ -36,10 +36,7 @@
 from .types import (
     File as CogFile,
 )
-from .types import (
-    Input,
-    URLPath,
-)
+from .types import Input
 from .types import (
     Path as CogPath,
 )
@@ -49,7 +46,7 @@
 
 class BasePredictor(ABC):
     def setup(
-        self, weights: Optional[Union[CogFile, CogPath]] = None
+        self, weights: Optional[Union[CogFile, CogPath, str]] = None
     ) -> Optional[Awaitable[None]]:
         """
         An optional method to prepare the model so multiple predictions run efficiently.
@@ -79,34 +76,40 @@ async def run_setup_async(predictor: BasePredictor) -> None:
         return await maybe_coro
 
 
-def get_weights_argument(predictor: BasePredictor) -> Union[CogFile, CogPath, None]:
+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)
     if weights_type is None:
         return 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; now it is
+    # essentially everyone needs to switch from Path to str (or a new URL type)
     if weights_url:
         if weights_type == CogFile:
             return cast(CogFile, CogFile.validate(weights_url))
         if weights_type == CogPath:
             # TODO: So this can be a url. evil!
             return cast(CogPath, CogPath.validate(weights_url))
+        if weights_type == str:
+            return weights_url
         raise ValueError(
-            f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File and Path are supported"
+            f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File, Path and str are supported"
         )
     if os.path.exists(weights_path):
         if weights_type == CogFile:
             return cast(CogFile, open(weights_path, "rb"))
         if weights_type == CogPath:
             return CogPath(weights_path)
         raise ValueError(
-            f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File and Path are supported"
+            f"Predictor.setup() has an argument 'weights' of type {weights_type}, but only File, Path and str are supported"
         )
     return None
 
@@ -212,17 +215,37 @@ def cleanup(self) -> None:
         Cleanup any temporary files created by the input.
         """
         for _, value in self:
-            # Handle URLPath objects specially for cleanup.
-            if isinstance(value, URLPath):
-                value.unlink()
-            # Note this is pathlib.Path, which cog.Path is a subclass of. A pathlib.Path object shouldn't make its way here,
+            # # Handle URLPath objects specially for cleanup.
+            # if isinstance(value, URLPath):
+            #     value.unlink()
+            # Note this is pathlib.Path, of which cog.Path is a subclass of.
+            # A pathlib.Path object shouldn't make its way here,
             # but both have an unlink() method, so may as well be safe.
-            elif isinstance(value, Path):
+            #
+            # URLTempFile, DataURLTempFilePath, pathlib.Path, doesn't matter
+            # everyone can be unlinked
+            if isinstance(value, Path):
                 try:
                     value.unlink()
                 except FileNotFoundError:
                     pass
 
+    # if we had a separate method to traverse the input and apply some function to each value
+    # we could use something like these functions here
+
+    # def cleanup():
+    #     if isinstance(value, Path):
+    #         value.unlink()
+
+    # def get_tempfile():
+    #     if isinstance(value, URLTempFile):
+    #         return (value.url, value._path)
+
+    # # this one is very annoying because it's supposed to mutate
+    # def convert():
+    #     if isinstance(value, URLTempFile):
+    #         return value.convert()
+
 
 def validate_input_type(type: Type[Any], name: str) -> None:
     if type is inspect.Signature.empty:

Diff for: python/cog/predictor.py

@@ -1,3 +1,4 @@
+import secrets
 import typing as t
 from datetime import datetime
 from enum import Enum
@@ -36,7 +37,15 @@ 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
@@ -85,8 +94,10 @@ def with_types(cls, input_type: t.Type[t.Any], output_type: t.Type[t.Any]) -> t.
             output=(output_type, None),
         )
 
+
 class TrainingRequest(PredictionRequest):
     pass
 
+
 class TrainingResponse(PredictionResponse):
     pass