Add is_healthy to truss (#1283)

* first poc

* change logs to be time elapsed

* reset _first_health_check_failure on is_ready success

* is_ready fail fast on load failure

* move logging to after load finished, update log text

* remove comment

* cr

* Add `is_ready` to chains (#1289)

* first pass at is_ready in chains

* revert example chainlet

* bump ctx builder

* address comments + add tests

* add test chain

* fix test

* cr fixes

* fix test

* cr

* fix formatting

* is_ready -> is_healthy

* more refactoring

* fix assert

* fix couple more tests

* fix for extra args

* marius cr

* fix linting

* new loaded endpoint + customize retries

* fix test_e2e

* fix tests

* filter out loaded logs

* move reroute logic into helper

* docstrings + wait 10s max

* bump ctx builder

Author: spal1

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.59rc017"
+version = "0.9.59rc018"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss-chains/tests/custom_health_checks/custom_health_checks.py

@@ -0,0 +1,18 @@
+import truss_chains as chains
+
+
+class CustomHealthChecks(chains.ChainletBase):
+    """Implements custom health checks."""
+
+    def __init__(self):
+        self._should_succeed_health_checks = True
+
+    def is_healthy(self) -> bool:
+        return self._should_succeed_health_checks
+
+    async def run_remote(self, fail: bool) -> str:
+        if fail:
+            self._should_succeed_health_checks = False
+        else:
+            self._should_succeed_health_checks = True
+        return f"health checks will {'succeed' if self._should_succeed_health_checks else 'fail'}"

truss-chains/tests/test_e2e.py

@@ -5,7 +5,10 @@
 
 import pytest
 import requests
-from truss.tests.test_testing_utilities_for_other_tests import ensure_kill_all
+from truss.tests.test_testing_utilities_for_other_tests import (
+    ensure_kill_all,
+    get_container_logs_from_prefix,
+)
 from truss.truss_handle.build import load
 
 from truss_chains import definitions, framework, public_api, utils
@@ -275,3 +278,39 @@ def test_traditional_truss():
         )
         assert response.status_code == 200
         assert response.json() == 5
+
+
+@pytest.mark.integration
+def test_custom_health_checks_chain():
+    with ensure_kill_all():
+        chain_root = TEST_ROOT / "custom_health_checks" / "custom_health_checks.py"
+        with framework.import_target(chain_root, "CustomHealthChecks") as entrypoint:
+            service = deployment_client.push(
+                entrypoint,
+                options=definitions.PushOptionsLocalDocker(
+                    chain_name="integration-test-custom-health-checks",
+                    only_generate_trusses=False,
+                    use_local_chains_src=True,
+                ),
+            )
+
+            assert service is not None
+            health_check_url = service.run_remote_url.split(":predict")[0]
+
+            response = service.run_remote({"fail": False})
+            assert response.status_code == 200
+            response = requests.get(health_check_url)
+            response.status_code == 200
+            container_logs = get_container_logs_from_prefix(entrypoint.name)
+            assert "Health check failed." not in container_logs
+
+            # Start failing health checks
+            response = service.run_remote({"fail": True})
+            response = requests.get(health_check_url)
+            assert response.status_code == 503
+            container_logs = get_container_logs_from_prefix(entrypoint.name)
+            assert container_logs.count("Health check failed.") == 1
+            response = requests.get(health_check_url)
+            assert response.status_code == 503
+            container_logs = get_container_logs_from_prefix(entrypoint.name)
+            assert container_logs.count("Health check failed.") == 2

truss-chains/tests/test_framework.py

@@ -307,7 +307,7 @@ def run_remote(self, arg: "SomeModel") -> None:
 def test_raises_endpoint_no_method():
     match = (
         rf"{TEST_FILE}:\d+ \(StaticMethod\.run_remote\) \[kind: TYPE_ERROR\].*"
-        r"Endpoint must be a method"
+        r"`run_remote` must be a method"
     )
 
     with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
@@ -321,7 +321,7 @@ def run_remote() -> None:
 def test_raises_endpoint_no_method_arg():
     match = (
         rf"{TEST_FILE}:\d+ \(StaticMethod\.run_remote\) \[kind: TYPE_ERROR\].*"
-        r"Endpoint must be a method"
+        r"`run_remote` must be a method"
     )
 
     with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
@@ -579,3 +579,92 @@ def test_raises_iterator_no_arg():
         class IteratorNoArg(chains.ChainletBase):
             async def run_remote(self) -> AsyncIterator:
                 yield "123"
+
+
+def test_raises_is_healthy_not_a_method():
+    match = rf"{TEST_FILE}:\d+ \(IsHealthyNotMethod\) \[kind: TYPE_ERROR\].* `is_healthy` must be a method."
+
+    with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
+
+        class IsHealthyNotMethod(chains.ChainletBase):
+            is_healthy: int = 3
+
+            async def run_remote(self) -> str:
+                return ""
+
+
+def test_raises_is_healthy_no_arg():
+    match = (
+        rf"{TEST_FILE}:\d+ \(IsHealthyNoArg\.is_healthy\) \[kind: TYPE_ERROR\].*"
+        r"`is_healthy` must be a method, i.e. with `self` as first argument. Got function with no arguments."
+    )
+
+    with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
+
+        class IsHealthyNoArg(chains.ChainletBase):
+            async def is_healthy() -> bool:
+                return True
+
+            async def run_remote(self) -> str:
+                return ""
+
+
+def test_raises_is_healthy_first_arg_not_self():
+    match = (
+        rf"{TEST_FILE}:\d+ \(IsHealthyNoSelfArg\.is_healthy\) \[kind: TYPE_ERROR\].*"
+        r"`is_healthy` must be a method, i.e. with `self` as first argument. Got `hi` as first argument."
+    )
+
+    with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
+
+        class IsHealthyNoSelfArg(chains.ChainletBase):
+            def is_healthy(hi) -> bool:
+                return True
+
+            async def run_remote(self) -> str:
+                return ""
+
+
+def test_raises_is_healthy_multiple_args():
+    match = rf"{TEST_FILE}:\d+ \(IsHealthyManyArgs\.is_healthy\) \[kind: TYPE_ERROR\].* `is_healthy` must have only one argument: `self`."
+
+    with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
+
+        class IsHealthyManyArgs(chains.ChainletBase):
+            def is_healthy(self, hi) -> bool:
+                return True
+
+            async def run_remote(self) -> str:
+                return ""
+
+
+def test_raises_is_healthy_not_type_annotated():
+    match = (
+        rf"{TEST_FILE}:\d+ \(IsHealthyNotTyped\.is_healthy\) \[kind: IO_TYPE_ERROR\].*"
+        r"Return value of health check must be type annotated. Got:\n\tis_healthy\(self\) -> !MISSING!"
+    )
+
+    with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
+
+        class IsHealthyNotTyped(chains.ChainletBase):
+            def is_healthy(self):
+                return True
+
+            async def run_remote(self) -> str:
+                return ""
+
+
+def test_raises_is_healthy_not_boolean_typed():
+    match = (
+        rf"{TEST_FILE}:\d+ \(IsHealthyNotBoolTyped\.is_healthy\) \[kind: IO_TYPE_ERROR\].*"
+        r"Return value of health check must be a boolean. Got:\n\tis_healthy\(self\) -> str -> <class 'str'>"
+    )
+
+    with pytest.raises(definitions.ChainsUsageError, match=match), _raise_errors():
+
+        class IsHealthyNotBoolTyped(chains.ChainletBase):
+            def is_healthy(self) -> str:  # type: ignore[misc]
+                return "not ready"
+
+            async def run_remote(self) -> str:
+                return ""

truss-chains/truss_chains/definitions.py

@@ -34,9 +34,10 @@
 GENERATED_CODE_DIR = ".chains_generated"
 DYNAMIC_CHAINLET_CONFIG_KEY = "dynamic_chainlet_config"
 OTEL_TRACE_PARENT_HEADER_KEY = "traceparent"
-# Below arg names must correspond to `definitions.ABCChainlet`.
 RUN_REMOTE_METHOD_NAME = "run_remote"  # Chainlet method name exposed as endpoint.
 MODEL_ENDPOINT_METHOD_NAME = "predict"  # Model method name exposed as endpoint.
+HEALTH_CHECK_METHOD_NAME = "is_healthy"
+# Below arg names must correspond to `definitions.ABCChainlet`.
 CONTEXT_ARG_NAME = "context"  # Referring to Chainlets `__init__` signature.
 SELF_ARG_NAME = "self"
 REMOTE_CONFIG_NAME = "remote_config"
@@ -625,12 +626,18 @@ def display_name(self) -> str:
         return self.chainlet_cls.display_name
 
 
+class HealthCheckAPIDescriptor(SafeModelNonSerializable):
+    name: str = HEALTH_CHECK_METHOD_NAME
+    is_async: bool
+
+
 class ChainletAPIDescriptor(SafeModelNonSerializable):
     chainlet_cls: Type[ABCChainlet]
     src_path: str
     has_context: bool
     dependencies: Mapping[str, DependencyDescriptor]
     endpoint: EndpointAPIDescriptor
+    health_check: Optional[HealthCheckAPIDescriptor]
 
     def __hash__(self) -> int:
         return hash(self.chainlet_cls)

truss-chains/truss_chains/deployment/code_gen.py

@@ -457,6 +457,20 @@ def _gen_load_src(chainlet_descriptor: definitions.ChainletAPIDescriptor) -> _So
     return _Source(src=src, imports=imports)
 
 
+def _gen_health_check_src(
+    health_check: definitions.HealthCheckAPIDescriptor,
+) -> _Source:
+    """Generates AST for the `is_healthy` method of the truss model."""
+    def_str = "async def" if health_check.is_async else "def"
+    maybe_await = "await " if health_check.is_async else ""
+    src = (
+        f"{def_str} is_healthy(self) -> Optional[bool]:\n"
+        f"""{_indent('if hasattr(self, "_chainlet"):')}"""
+        f"""{_indent(f"return {maybe_await}self._chainlet.is_healthy()")}"""
+    )
+    return _Source(src=src)
+
+
 def _gen_predict_src(chainlet_descriptor: definitions.ChainletAPIDescriptor) -> _Source:
     """Generates AST for the `predict` method of the truss model."""
     imports: set[str] = {
@@ -538,6 +552,10 @@ def _gen_truss_chainlet_model(
         libcst.parse_statement(predict_src.src),
     ]
 
+    if chainlet_descriptor.health_check is not None:
+        health_check_src = _gen_health_check_src(chainlet_descriptor.health_check)
+        new_body.extend([libcst.parse_statement(health_check_src.src)])
+
     user_chainlet_ref = _gen_chainlet_import_and_ref(chainlet_descriptor)
     imports.update(user_chainlet_ref.imports)
 

truss-chains/truss_chains/framework.py

@@ -339,24 +339,29 @@ def _validate_streaming_output_type(
     )
 
 
-def _validate_endpoint_params(
-    params: list[inspect.Parameter], location: _ErrorLocation
-) -> list[definitions.InputArg]:
+def _validate_method_signature(
+    method_name: str, location: _ErrorLocation, params: list[inspect.Parameter]
+) -> None:
     if len(params) == 0:
         _collect_error(
-            f"`Endpoint must be a method, i.e. with `{definitions.SELF_ARG_NAME}` as "
+            f"`{method_name}` must be a method, i.e. with `{definitions.SELF_ARG_NAME}` as "
             "first argument. Got function with no arguments.",
             _ErrorKind.TYPE_ERROR,
             location,
         )
-        return []
-    if params[0].name != definitions.SELF_ARG_NAME:
+    elif params[0].name != definitions.SELF_ARG_NAME:
         _collect_error(
-            f"`Endpoint must be a method, i.e. with `{definitions.SELF_ARG_NAME}` as "
+            f"`{method_name}` must be a method, i.e. with `{definitions.SELF_ARG_NAME}` as "
             f"first argument. Got `{params[0].name}` as first argument.",
             _ErrorKind.TYPE_ERROR,
             location,
         )
+
+
+def _validate_endpoint_params(
+    params: list[inspect.Parameter], location: _ErrorLocation
+) -> list[definitions.InputArg]:
+    _validate_method_signature(definitions.RUN_REMOTE_METHOD_NAME, location, params)
     input_args = []
     for param in params[1:]:  # Skip self argument.
         if param.annotation == inspect.Parameter.empty:
@@ -434,7 +439,7 @@ def _validate_and_describe_endpoint(
     ```
 
     * The name must be `run_remote` for Chainlets, or `predict` for Models.
-    * It can be sync or async or def.
+    * It can be sync or async def.
     * The number and names of parameters are arbitrary, both positional and named
       parameters are ok.
     * All parameters and the return value must have type annotations. See
@@ -742,6 +747,63 @@ def _validate_remote_config(
         )
 
 
+def _validate_health_check(
+    cls: Type[definitions.ABCChainlet], location: _ErrorLocation
+) -> Optional[definitions.HealthCheckAPIDescriptor]:
+    """The `is_healthy` method of a Chainlet must have the following signature:
+    ```
+    [async] def is_healthy(self) -> bool:
+    ```
+    * The name must be `is_healthy`.
+    * It can be sync or async def.
+    * Must not define any parameters other than `self`.
+    * Must return a boolean.
+    """
+    if not hasattr(cls, definitions.HEALTH_CHECK_METHOD_NAME):
+        return None
+
+    health_check_method = getattr(cls, definitions.HEALTH_CHECK_METHOD_NAME)
+    if not inspect.isfunction(health_check_method):
+        _collect_error(
+            f"`{definitions.HEALTH_CHECK_METHOD_NAME}` must be a method.",
+            _ErrorKind.TYPE_ERROR,
+            location,
+        )
+        return None
+
+    line = inspect.getsourcelines(health_check_method)[1]
+    location = location.model_copy(
+        update={"line": line, "method_name": definitions.HEALTH_CHECK_METHOD_NAME}
+    )
+    is_async = inspect.iscoroutinefunction(health_check_method)
+    signature = inspect.signature(health_check_method)
+    params = list(signature.parameters.values())
+    _validate_method_signature(definitions.HEALTH_CHECK_METHOD_NAME, location, params)
+    if len(params) > 1:
+        _collect_error(
+            f"`{definitions.HEALTH_CHECK_METHOD_NAME}` must have only one argument: `{definitions.SELF_ARG_NAME}`.",
+            _ErrorKind.TYPE_ERROR,
+            location,
+        )
+    if signature.return_annotation == inspect.Parameter.empty:
+        _collect_error(
+            "Return value of health check must be type annotated. Got:\n"
+            f"\t{location.method_name}{signature} -> !MISSING!",
+            _ErrorKind.IO_TYPE_ERROR,
+            location,
+        )
+        return None
+    if signature.return_annotation is not bool:
+        _collect_error(
+            "Return value of health check must be a boolean. Got:\n"
+            f"\t{location.method_name}{signature} -> {signature.return_annotation}",
+            _ErrorKind.IO_TYPE_ERROR,
+            location,
+        )
+
+    return definitions.HealthCheckAPIDescriptor(is_async=is_async)
+
+
 def validate_and_register_cls(cls: Type[definitions.ABCChainlet]) -> None:
     """Note that validation errors will only be collected, not raised, and Chainlets.
     with issues, are still added to the registry.  Use `raise_validation_errors` to
@@ -759,6 +821,7 @@ def validate_and_register_cls(cls: Type[definitions.ABCChainlet]) -> None:
         has_context=init_validator.has_context,
         endpoint=_validate_and_describe_endpoint(cls, location),
         src_path=src_path,
+        health_check=_validate_health_check(cls, location),
     )
     logging.debug(
         f"Descriptor for {cls}:\n{pprint.pformat(chainlet_descriptor, indent=4)}\n"

truss-chains/truss_chains/remote_chainlet/model_skeleton.py

@@ -47,6 +47,11 @@ def __init__(
     #         side_effect=stub.factory(SideEffectOnlySubclass, self._context),
     #     )
     #
+    # If chainlet implements is_healthy:
+    # def is_healthy(self) -> Optional[bool]:
+    #     if hasattr(self, "_chainlet"):
+    #         return self._chainlet.is_healthy()
+    #
     # def predict(
     #     self, inputs: TextToNumInput, request: starlette.requests.Request
     # ) -> TextToNumOutput: