Support Optional in Input (#2216)

* Add integration tests for pydantic 2 output

* Fix test case

* Add mode logging to invalid output

* Add support for Optionals in Input

* Remove Optional import from typing_extensions

* Handle python 3.8/3.9

* Fix > 3.9 check

* Fix >= 3.10

* Use None.__class__ to replace NoneType and type(None)

* Check if optional is Union

Author: 8W9aG

python/cog/predictor.py

@@ -7,6 +7,9 @@
 import types
 import uuid
 from collections.abc import Iterable, Iterator
+
+if sys.version_info >= (3, 10):
+    from types import NoneType
 from typing import (
     Any,
     Callable,
@@ -190,8 +193,20 @@ def validate_input_type(
         elif get_origin(type) in (Union, List, list) or (
             hasattr(types, "UnionType") and get_origin(type) is types.UnionType
         ):  # noqa: E721
-            for t in get_args(type):
-                validate_input_type(t, name)
+            args = get_args(type)
+
+            def is_optional() -> bool:
+                if len(args) != 2 or get_origin(type) is not Union:
+                    return False
+                if sys.version_info >= (3, 10):
+                    return args[1] is NoneType
+                return args[1] is None.__class__
+
+            if is_optional():
+                validate_input_type(args[0], name)
+            else:
+                for t in args:
+                    validate_input_type(t, name)
         else:
             if PYDANTIC_V2:
                 # Cog types are exported as `Annotated[Type, ...]`, but `type` is the inner type

python/cog/server/http.py

@@ -469,7 +469,7 @@ async def _predict(
         try:
             _ = response_type(**response_object)
         except ValidationError as e:
-            _log_invalid_output(e)
+            _log_invalid_output(e, mode)
             raise HTTPException(status_code=500, detail=str(e)) from e
 
         response_object["output"] = upload_files(
@@ -520,17 +520,20 @@ def _maybe_shutdown(exc: BaseException, *, status: Health = Health.DEFUNCT) -> N
     return app
 
 
-def _log_invalid_output(error: Any) -> None:
+def _log_invalid_output(error: Any, mode: Mode) -> None:
+    function_name = "predict()"
+    if mode == Mode.TRAIN:
+        function_name = "train()"
     log.error(
         textwrap.dedent(
             f"""\
-            The return value of predict() was not valid:
+            The return value of {function_name} was not valid:
 
             {error}
 
             Check that your predict function is in this form, where `output_type` is the same as the type you are returning (e.g. `str`):
 
-                def predict(...) -> output_type:
+                def {function_name} -> output_type:
                     ...
            """
         )

test-integration/test_integration/fixtures/pydantic2-output/cog.yaml

@@ -0,0 +1,19 @@
+# Configuration for Cog ⚙️
+# Reference: https://cog.run/yaml
+
+build:
+  # set to true if your model requires a GPU
+  gpu: false
+
+  # python version in the form '3.11' or '3.11.4'
+  python_version: "3.12"
+
+  # a list of packages in the format <package-name>==<version>
+  python_packages:
+    - "pydantic==2.10.6"  # The problematic Pydantic version
+
+# predict.py defines how predictions are run on your model
+predict: "predict.py:Predictor"
+
+# train.py defines how training runs on your model
+train: "train.py:train"

test-integration/test_integration/fixtures/pydantic2-output/predict.py

@@ -0,0 +1,24 @@
+# Prediction interface for Cog ⚙️
+# https://cog.run/python
+
+from cog import BasePredictor, Input, Path
+
+
+class Predictor(BasePredictor):
+    def setup(self) -> None:
+        """Load the model into memory to make running multiple predictions efficient"""
+        # self.model = torch.load("./weights.pth")
+        pass
+
+    def predict(
+        self,
+        image: Path = Input(description="Grayscale input image"),
+        scale: float = Input(
+            description="Factor to scale image by", ge=0, le=10, default=1.5
+        ),
+    ) -> Path:
+        """Run a single prediction on the model"""
+        # processed_input = preprocess(image)
+        # output = self.model(processed_image, scale)
+        # return postprocess(output)
+        return Path(".")

test-integration/test_integration/fixtures/pydantic2-output/train.py

@@ -0,0 +1,76 @@
+import os
+from typing import Optional
+from cog import BaseModel, Input, Path as CogPath, Secret
+
+# We return a path to our trained adapter weights
+class TrainingOutput(BaseModel):
+    weights: CogPath
+
+def train(
+    # Basic input
+    some_input: str = Input(
+        description="A basic string input to satisfy minimum requirements.",
+        default="default value",
+    ),
+    # String input with None default (problematic)
+    hf_repo_id: Optional[str] = Input(
+        description="String with None default - this causes issues.",
+        default=None,
+    ),
+    # Secret with None default (problematic)
+    hf_token: Optional[Secret] = Input(
+        description="Secret with None default - this also causes issues.",
+        default=None,
+    ),
+    # String input with empty string default (works)
+    working_repo_id: str = Input(
+        description="String with empty string default - this works.",
+        default="",
+    ),
+    # Secret with empty string default (works)
+    working_token: Secret = Input(
+        description="Secret with empty string default - this works.",
+        default="",
+    ),
+) -> TrainingOutput:
+    """
+    Minimal example to demonstrate issues with Secret inputs.
+    """
+    print("\n=== Minimal Cog Secret Test ===")
+    print(f"cog version: {os.environ.get('COG_VERSION', 'unknown')}")
+    
+    # Inputs with None defaults
+    print("\n-- Inputs with None defaults (problematic) --")
+    print(f"hf_repo_id: {hf_repo_id}")
+    if hf_token:
+        print(f"hf_token: [PROVIDED]")
+        try:
+            value = hf_token.get_secret_value()
+            print("Secret access successful")
+        except Exception as e:
+            print(f"Error accessing secret: {e}")
+    else:
+        print("hf_token: None")
+    
+    # Inputs with empty string defaults
+    print("\n-- Inputs with empty string defaults (works) --")
+    print(f"working_repo_id: {working_repo_id if working_repo_id else '(empty)'}")
+    if working_token and working_token.get_secret_value():
+        print(f"working_token: [PROVIDED]")
+        try:
+            value = working_token.get_secret_value()
+            print("Secret access successful")
+        except Exception as e:
+            print(f"Error accessing secret: {e}")
+    else:
+        print("working_token: (empty)")
+    
+    # Create a dummy output file
+    output_path = "dummy_output.txt"
+    with open(output_path, "w") as f:
+        f.write("This is a dummy output file.")
+    
+    print("\n=== Test Complete ===")
+    
+    # Return the dummy output path
+    return TrainingOutput(weights=CogPath(output_path))

test-integration/test_integration/test_train.py

@@ -19,3 +19,16 @@ def test_train_takes_input_and_produces_weights(tmpdir_factory):
     with open(out_dir / "weights.bin", "rb") as f:
         assert len(f.read()) == 42
     assert "falling back to slow loader" not in str(result.stderr)
+
+
+def test_train_pydantic2(tmpdir_factory):
+    project_dir = Path(__file__).parent / "fixtures/pydantic2-output"
+    out_dir = pathlib.Path(tmpdir_factory.mktemp("project"))
+    shutil.copytree(project_dir, out_dir, dirs_exist_ok=True)
+    result = subprocess.run(
+        ["cog", "train", "--debug", "-i", 'some_input="hello"'],
+        cwd=out_dir,
+        check=False,
+        capture_output=True,
+    )
+    assert result.returncode == 0