Add top-level pydantic schema around / truss model I/O.

pyproject.toml

@@ -56,6 +56,7 @@ libcst = "<1.2.0"
 autoflake = "<=2.2"
 pytest-asyncio = "^0.23.6"
 
+
 [tool.poetry.group.builder.dependencies]
 blake3 = "^0.3.3"
 boto3 = "^1.26.157"

truss-chains/examples/text_to_num/main.py

@@ -114,7 +114,7 @@ def run_remote(self, data: str) -> int:
         return number
 
 
-@chains.entrypoint
+@chains.mark_entrypoint
 class ExampleChain(chains.ChainletBase):
     def __init__(
         self,

truss-chains/examples/transcribe/adapter.py

@@ -24,7 +24,7 @@
 
 class EnqueueInput(pydantic.BaseModel):
     # This typo `media_for_transciption` is for backwards compatibility.
-    # Also this would ideally be `list[JobDescriptor]` instead of string...
+    # Also, this would ideally be `list[JobDescriptor]` instead of string...
     media_for_transciption: str
 
 
@@ -254,7 +254,7 @@ async def run_remote(
         # When `Transcribe.run` is called with `async_predict` the async framework will
         # invoke the webhook, but it expects that the webhook does not need
         # authentication (instead will validate via payload signature).
-        # We can't run a chainlet (i.e. truss model) without requiring authentication
+        # We can't run a chainlet (i.e. truss model) without requiring authentication,
         # and we don't want to stand up and manage another service - so we basically
         # ignore async's webhook integration and "manually" call a webhook from here
         # where we can explicitly add the authentication.
@@ -299,7 +299,7 @@ async def run_remote(
         return result
 
 
-@chains.entrypoint
+@chains.mark_entrypoint
 class EnqueueBatch(chains.ChainletBase):
     """Accepts a request with multiple transcription jobs and starts the sub-jobs."""
 

truss-chains/examples/transcribe/transcribe.py

@@ -94,7 +94,7 @@ async def run_remote(
         )
 
 
-@chains.entrypoint
+@chains.mark_entrypoint
 class Transcribe(chains.ChainletBase):
     """Transcribes one file end-to-end and sends results to webhook."""
 

truss-chains/tests/chains_test.py

@@ -36,7 +36,7 @@ def test_chain():
             url, json={"length": 300, "num_partitions": 3}, stream=True
         )
         print(response)
-        error = definitions.RemoteErrorDetail.parse_obj(response.json()["error"])
+        error = definitions.RemoteErrorDetail.model_validate(response.json()["error"])
         error_str = error.format()
         print(error_str)
         assert "ValueError: This input is too long: 100." in error_str
@@ -52,7 +52,7 @@ async def test_chain_local():
             with pytest.raises(ValueError):
                 # First time `SplitTextFailOnce` raises an error and
                 # currently local mode does not have retries.
-                result = await entrypoint().run_remote(length=20, num_partitions=5)
+                await entrypoint().run_remote(length=20, num_partitions=5)
 
             result = await entrypoint().run_remote(length=20, num_partitions=5)
             assert result == (4198, "erodfderodfderodfder", 123)

truss-chains/tests/itest_chain/requirements.txt

@@ -1 +1 @@
-git+https://github.com/basetenlabs/truss.git@f3f164f061c465a59d9c520700660424c0e4f8ab
+git+https://github.com/basetenlabs/truss.git

truss-chains/tests/itest_chain/user_package/shared_chainlet.py

@@ -28,6 +28,8 @@ async def run_remote(
     ) -> Tuple[SplitTextOutput, int]:
         import numpy as np
 
+        print(extra_arg)
+
         self._count += 1
         if self._count == 1:
             raise ValueError("Haha this is a fake error.")

truss-chains/truss_chains/__init__.py

@@ -1,13 +1,16 @@
 import warnings
 
-# These warnings show up if we have pydantic v2 installed, but still use v1-compatible
-# APIs, which are of course deprecated...
-warnings.filterwarnings("ignore", message="Valid config keys have changed in V2:*")
-warnings.filterwarnings(
-    "ignore",
-    message="`pydantic.generics:GenericModel` has been moved to `pydantic.BaseModel`.",
-)
+import pydantic
+
+pydantic_major_version = int(pydantic.VERSION.split(".")[0])
+if pydantic_major_version < 2:
+    warnings.warn(
+        f"Pydantic major version is less than 2 ({pydantic.VERSION}). You use Truss, "
+        "but for using Truss-Chains, you must upgrade to pydantic-v2.",
+        UserWarning,
+    )
 
+del pydantic
 del warnings
 
 # flake8: noqa F401
@@ -26,7 +29,7 @@
     depends,
     depends_context,
     deploy_remotely,
-    entrypoint,
+    mark_entrypoint,
     run_local,
 )
 from truss_chains.stub import StubBase