Repository: replicate/cog

Issue:
Include training interfaces in OpenAPI schema
Cog has experimental support for fine-tuning. See https://github.com/replicate/cog/pull/976

This involves adding a `train()` function, much like `predict()`:

```py
from cog import BasePredictor, File
import io

def train(param: str) -> File:
    return io.StringIO("hello " + param)
```

We should update Cog to parse that training function signature into a JSON schema describing its inputs and outputs, the way we do for predictors.


Discussion:

Resolving patch:
diff --git a/python/cog/predictor.py b/python/cog/predictor.py
--- a/python/cog/predictor.py
+++ b/python/cog/predictor.py
@@ -224,12 +224,6 @@ def cleanup(self) -> None:
                     pass
 
 
-def get_predict(predictor: Any) -> Callable[..., Any]:
-    if hasattr(predictor, "predict"):
-        return predictor.predict
-    return predictor
-
-
 def validate_input_type(type: Type[Any], name: str) -> None:
     if type is inspect.Signature.empty:
         raise TypeError(
@@ -247,22 +241,7 @@ def validate_input_type(type: Type[Any], name: str) -> None:
             )
 
 
-def get_input_type(predictor: BasePredictor) -> Type[BaseInput]:
-    """
-    Creates a Pydantic Input model from the arguments of a Predictor's predict() method.
-
-    class Predictor(BasePredictor):
-        def predict(self, text: str):
-            ...
-
-    programmatically creates a model like this:
-
-    class Input(BaseModel):
-        text: str
-    """
-
-    predict = get_predict(predictor)
-    signature = inspect.signature(predict)
+def get_input_create_model_kwargs(signature: inspect.Signature) -> Dict[str, Any]:
     create_model_kwargs = {}
 
     order = 0
@@ -308,13 +287,39 @@ class StringEnum(str, enum.Enum):
 
         create_model_kwargs[name] = (InputType, default)
 
+    return create_model_kwargs
+
+
+def get_predict(predictor: Any) -> Callable[..., Any]:
+    if hasattr(predictor, "predict"):
+        return predictor.predict
+    return predictor
+
+
+def get_input_type(predictor: BasePredictor) -> Type[BaseInput]:
+    """
+    Creates a Pydantic Input model from the arguments of a Predictor's predict() method.
+
+    class Predictor(BasePredictor):
+        def predict(self, text: str):
+            ...
+
+    programmatically creates a model like this:
+
+    class Input(BaseModel):
+        text: str
+    """
+
+    predict = get_predict(predictor)
+    signature = inspect.signature(predict)
+
     return create_model(
         "Input",
         __config__=None,
         __base__=BaseInput,
         __module__=__name__,
         __validators__=None,
-        **create_model_kwargs,
+        **get_input_create_model_kwargs(signature),
     )  # type: ignore
 
 
@@ -322,6 +327,7 @@ def get_output_type(predictor: BasePredictor) -> Type[BaseModel]:
     """
     Creates a Pydantic Output model from the return type annotation of a Predictor's predict() method.
     """
+
     predict = get_predict(predictor)
     signature = inspect.signature(predict)
     OutputType: Type[BaseModel]
@@ -350,14 +356,117 @@ def predict(
         field = Field(**{"x-cog-array-type": "iterator"})  # type: ignore
         OutputType: Type[BaseModel] = Annotated[List[get_args(OutputType)[0]], field]  # type: ignore
 
-    if not hasattr(OutputType, "__name__") or OutputType.__name__ != "Output":
-        # Wrap the type in a model called "Output" so it is a consistent name in the OpenAPI schema
+    name = OutputType.__name__ if hasattr(OutputType, "__name__") else ""
+
+    if name == "Output":
+        return OutputType
+
+    # We wrap the OutputType in an Output class to
+    # ensure consistent naming of the interface in the schema.
+    #
+    # NOTE: If the OutputType.__name__ is "TrainingOutput" then cannot use
+    # `__root__` here because this will create a reference for the Object.
+    # e.g.
+    #   {'title': 'Output', '$ref': '#/definitions/TrainingOutput' ... }
+    #
+    # And this reference may conflict with other objects at which
+    # point the item will be namespaced and break our parsing. e.g.
+    #   {'title': 'Output', '$ref': '#/definitions/predict_TrainingOutput' ... }
+    #
+    # So we work around this by inheriting from the original class rather
+    # than using "__root__".
+    if name == "TrainingOutput":
+
+        class Output(OutputType):  # type: ignore
+            pass
+
+        return Output
+    else:
+
         class Output(BaseModel):
             __root__: OutputType  # type: ignore
 
-        OutputType = Output
+        return Output
+
+
+def get_train(predictor: Any) -> Callable[..., Any]:
+    if hasattr(predictor, "train"):
+        return predictor.train
+    return predictor
+
+
+def get_training_input_type(predictor: BasePredictor) -> Type[BaseInput]:
+    """
+    Creates a Pydantic Input model from the arguments of a Predictor's train() method.
+
+    def train(self, text: str):
+        ...
+
+    programmatically creates a model like this:
+
+    class TrainingInput(BaseModel):
+        text: str
+    """
+
+    train = get_train(predictor)
+    signature = inspect.signature(train)
+
+    return create_model(
+        "TrainingInput",
+        __config__=None,
+        __base__=BaseInput,
+        __module__=__name__,
+        __validators__=None,
+        **get_input_create_model_kwargs(signature),
+    )  # type: ignore
+
+
+def get_training_output_type(predictor: BasePredictor) -> Type[BaseModel]:
+    """
+    Creates a Pydantic Output model from the return type annotation of a train() method.
+    """
+
+    train = get_train(predictor)
+    signature = inspect.signature(train)
+
+    if signature.return_annotation is inspect.Signature.empty:
+        raise TypeError(
+            """You must set an output type. If your model can return multiple output types, you can explicitly set `Any` as the output type.
+
+For example:
+
+    from typing import Any
+
+    def train(
+        self,
+        n: int
+    ) -> Any:
+        ...
+"""
+        )
+    else:
+        TrainingOutputType = signature.return_annotation
+
+    name = (
+        TrainingOutputType.__name__ if hasattr(TrainingOutputType, "__name__") else ""
+    )
+    # We wrap the OutputType in a TrainingOutput class to
+    # ensure consistent naming of the interface in the schema
+    # See comment in get_output_type for more info.
+    if name == "TrainingOutput":
+        return TrainingOutputType
+
+    if name == "Output":
+
+        class TrainingOutput(TrainingOutputType):  # type: ignore
+            pass
+
+        return TrainingOutput
+
+    class TrainingOutput(BaseModel):
+        __root__: TrainingOutputType  # type: ignore
 
-    return OutputType
+    return TrainingOutput
 
 
 def human_readable_type_name(t: Type[Any]) -> str:
diff --git a/python/cog/schema.py b/python/cog/schema.py
--- a/python/cog/schema.py
+++ b/python/cog/schema.py
@@ -84,3 +84,9 @@ def with_types(cls, input_type: t.Type[t.Any], output_type: t.Type[t.Any]) -> t.
             input=(t.Optional[input_type], None),
             output=(output_type, None),
         )
+
+class TrainingRequest(PredictionRequest):
+    pass
+
+class TrainingResponse(PredictionResponse):
+    pass
diff --git a/python/cog/server/http.py b/python/cog/server/http.py
--- a/python/cog/server/http.py
+++ b/python/cog/server/http.py
@@ -43,6 +43,8 @@
     get_input_type,
     get_output_type,
     get_predictor_ref,
+    get_training_input_type,
+    get_training_output_type,
     load_config,
     load_predictor_from_ref,
 )
@@ -154,6 +156,46 @@ async def wrapped(*args: "P.args", **kwargs: "P.kwargs") -> "T":
 
         return wrapped
 
+    if "train" in config:
+        # TODO: avoid loading trainer code in this process
+        trainer = load_predictor_from_ref(config["train"])
+
+        TrainingInputType = get_training_input_type(trainer)
+        TrainingOutputType = get_training_output_type(trainer)
+
+        class TrainingRequest(
+            schema.TrainingRequest.with_types(input_type=TrainingInputType)
+        ):
+            pass
+
+        TrainingResponse = schema.TrainingResponse.with_types(
+            input_type=TrainingInputType, output_type=TrainingOutputType
+        )
+
+        @app.post(
+            "/trainings",
+            response_model=TrainingResponse,
+            response_model_exclude_unset=True,
+        )
+        def train(request: TrainingRequest = Body(default=None), prefer: Union[str, None] = Header(default=None)) -> Any:  # type: ignore
+            return predict(request, prefer)
+
+        @app.put(
+            "/trainings/{training_id}",
+            response_model=PredictionResponse,
+            response_model_exclude_unset=True,
+        )
+        def train_idempotent(
+            training_id: str = Path(..., title="Training ID"),
+            request: TrainingRequest = Body(..., title="Training Request"),
+            prefer: Union[str, None] = Header(default=None),
+        ) -> Any:
+            return predict_idempotent(training_id, request, prefer)
+
+        @app.post("/trainings/{training_id}/cancel")
+        def cancel_training(training_id: str = Path(..., title="Training ID")) -> Any:
+            return cancel(training_id)
+
     @app.on_event("startup")
     def startup() -> None:
         app.state.setup_task = runner.setup()


Test patch:
diff --git a/python/tests/server/conftest.py b/python/tests/server/conftest.py
--- a/python/tests/server/conftest.py
+++ b/python/tests/server/conftest.py
@@ -33,17 +33,36 @@ def uses_predictor(name):
     )
 
 
+def uses_trainer(name):
+    # HACK: `name` can either be in the form "<name>.py:train" or just "<name>".
+    if ":" not in name:
+        name = f"{name}.py:train"
+    options = {"additional_config": {"train": _fixture_path(name)}}
+
+    return pytest.mark.parametrize(
+        "client", [AppConfig(predictor_fixture=name, options=options)], indirect=True
+    )
+
+
 def uses_predictor_with_client_options(name, **options):
     return pytest.mark.parametrize(
         "client", [AppConfig(predictor_fixture=name, options=options)], indirect=True
     )
 
 
-def make_client(fixture_name: str, upload_url: Optional[str] = None):
+def make_client(
+    fixture_name: str,
+    upload_url: Optional[str] = None,
+    additional_config: Optional[dict] = None,
+):
     """
     Creates a fastapi test client for an app that uses the requested Predictor.
     """
+
     config = {"predict": _fixture_path(fixture_name)}
+    if additional_config:
+        config.update(additional_config)
+
     app = create_app(
         config=config,
         shutdown_event=threading.Event(),
diff --git a/python/tests/server/fixtures/train.py b/python/tests/server/fixtures/train.py
new file mode 100644
--- /dev/null
+++ b/python/tests/server/fixtures/train.py
@@ -0,0 +1,17 @@
+from cog import BaseModel, Input, Path
+
+
+class TrainingOutput(BaseModel):
+    weights: Path
+
+
+def train(
+    n: int = Input(description="Dimension of weights to generate"),
+) -> TrainingOutput:
+    with open("weights.bin", "w") as fh:
+        for _ in range(n):
+            fh.write("a")
+
+    return TrainingOutput(
+        weights=Path("weights.bin"),
+    )
diff --git a/python/tests/server/test_http.py b/python/tests/server/test_http.py
--- a/python/tests/server/test_http.py
+++ b/python/tests/server/test_http.py
@@ -7,7 +7,12 @@
 from PIL import Image
 from responses import matchers
 
-from .conftest import make_client, uses_predictor, uses_predictor_with_client_options
+from .conftest import (
+    make_client,
+    uses_predictor,
+    uses_predictor_with_client_options,
+    uses_trainer,
+)
 
 
 def test_setup_healthcheck():
@@ -349,6 +354,47 @@ def test_yielding_strings_from_async_generator_predictors(client, match):
     )
 
 
+@uses_trainer("train.py:train")
+def test_train_openapi_specification(client):
+    resp = client.get("/openapi.json")
+    assert resp.status_code == 200
+
+    schema = resp.json()
+    assert schema["openapi"] == "3.0.2"
+    assert schema["info"] == {"title": "Cog", "version": "0.1.0"}
+
+    assert schema["components"]["schemas"]["TrainingInput"] == {
+        "title": "TrainingInput",
+        "type": "object",
+        "properties": {
+            "n": {
+                "type": "integer",
+                "x-order": 0,
+                "title": "N",
+                "description": "Dimension of weights to generate",
+            },
+        },
+        "required": [
+            "n",
+        ],
+    }
+
+    assert schema["components"]["schemas"]["TrainingOutput"] == {
+        "title": "TrainingOutput",
+        "type": "object",
+        "properties": {
+            "weights": {
+                "type": "string",
+                "format": "uri",
+                "title": "Weights",
+            },
+        },
+        "required": [
+            "weights",
+        ],
+    }
+
+
 @uses_predictor("yield_strings")
 def test_yielding_strings_from_generator_predictors(client, match):
     resp = client.post("/predictions")
diff --git a/python/tests/server/test_worker.py b/python/tests/server/test_worker.py
--- a/python/tests/server/test_worker.py
+++ b/python/tests/server/test_worker.py
@@ -23,7 +23,7 @@
 
 # Set a longer deadline on CI as the instances are a bit slower.
 settings.register_profile("ci", max_examples=100, deadline=1000)
-settings.register_profile("default", max_examples=10, deadline=500)
+settings.register_profile("default", max_examples=10, deadline=1000)
 settings.register_profile("slow", max_examples=10, deadline=2000)
 settings.load_profile(os.getenv("HYPOTHESIS_PROFILE", "default"))
 
diff --git a/test-integration/test_integration/fixtures/train-project/cog.yaml b/test-integration/test_integration/fixtures/train-project/cog.yaml
new file mode 100644
--- /dev/null
+++ b/test-integration/test_integration/fixtures/train-project/cog.yaml
@@ -0,0 +1,4 @@
+build:
+  python_version: "3.8"
+predict: "predict.py:Predictor"
+train: "train.py:train"
diff --git a/test-integration/test_integration/fixtures/train-project/predict.py b/test-integration/test_integration/fixtures/train-project/predict.py
new file mode 100644
--- /dev/null
+++ b/test-integration/test_integration/fixtures/train-project/predict.py
@@ -0,0 +1,6 @@
+from cog import BasePredictor
+
+
+class Predictor(BasePredictor):
+    def predict(self, s: str) -> str:
+        return "hello " + s
diff --git a/test-integration/test_integration/fixtures/train-project/train.py b/test-integration/test_integration/fixtures/train-project/train.py
new file mode 100644
--- /dev/null
+++ b/test-integration/test_integration/fixtures/train-project/train.py
@@ -0,0 +1,15 @@
+from cog import BaseModel, Input, Path
+
+class TrainingOutput(BaseModel):
+    weights: Path
+
+def train(
+    n: int,
+) -> TrainingOutput:
+    with open("weights.bin", "w") as fh:
+        for _ in range(n):
+            fh.write("a")
+
+    return TrainingOutput(
+        weights=Path("weights.bin"),
+    )
diff --git a/test-integration/test_integration/test_train.py b/test-integration/test_integration/test_train.py
new file mode 100644
--- /dev/null
+++ b/test-integration/test_integration/test_train.py
@@ -0,0 +1,19 @@
+import pathlib
+import shutil
+import subprocess
+from pathlib import Path
+
+
+def test_train_takes_input_and_produces_weights(tmpdir_factory):
+    project_dir = Path(__file__).parent / "fixtures/train-project"
+    out_dir = pathlib.Path(tmpdir_factory.mktemp("project"))
+    shutil.copytree(project_dir, out_dir, dirs_exist_ok=True)
+    result = subprocess.run(
+        ["cog", "train", "-i", "n=42"],
+        cwd=out_dir,
+        check=True,
+        capture_output=True,
+    )
+    assert result.stdout == b""
+    with open(out_dir / "weights.bin", "rb") as f:
+        assert len(f.read()) == 42