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
* add a note about this release

Author: technillogue

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

README.md

@@ -4,6 +4,8 @@ Cog is an open-source tool that lets you package machine learning models in a st
 
 You can deploy your packaged model to your own infrastructure, or to [Replicate](https://replicate.com/).
 
+This commit is a highly experimental version of cog that supports concurrent predictions and backwards-incompatible file handling changes. It is not thoroughly tested, but needs to be released anyway.
+
 ## Highlights
 
 - 📦 **Docker containers without the pain.** Writing your own `Dockerfile` can be a bewildering process. With Cog, you define your environment with a [simple configuration file](#how-it-works) and it generates a Docker image with all the best practices: Nvidia base images, efficient caching of dependencies, installing specific Python versions, sensible environment variable defaults, and so on.

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 {

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

pyproject.toml

@@ -15,6 +15,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",

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

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)

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:

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