From 49269b2891b76e6b6650fa0fb933ae98ee60cc9d Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Wed, 29 Apr 2026 15:27:09 -0700 Subject: [PATCH 01/12] Add runtime discovery --- truss-chains/truss_chains/__init__.py | 1 + truss-chains/truss_chains/public_types.py | 57 +++++++++ .../truss_chains/remote_chainlet/utils.py | 55 ++++---- truss-chains/truss_chains/runtime.py | 117 ++++++++++++++++++ 4 files changed, 207 insertions(+), 23 deletions(-) create mode 100644 truss-chains/truss_chains/runtime.py diff --git a/truss-chains/truss_chains/__init__.py b/truss-chains/truss_chains/__init__.py index a4d762cde..ecd92e71e 100644 --- a/truss-chains/truss_chains/__init__.py +++ b/truss-chains/truss_chains/__init__.py @@ -1,4 +1,5 @@ from truss.base.truss_config import WeightsSource +from truss_chains import runtime # noqa: E402, F401 from truss_chains.framework import ChainletBase, EngineBuilderLLMChainlet, ModelBase from truss_chains.public_api import ( depends, diff --git a/truss-chains/truss_chains/public_types.py b/truss-chains/truss_chains/public_types.py index 1972881a0..3b1238c05 100644 --- a/truss-chains/truss_chains/public_types.py +++ b/truss-chains/truss_chains/public_types.py @@ -727,6 +727,17 @@ class LookaheadDecodingConfig(pydantic.BaseModel): lookahead_decoding_config: Optional[LookaheadDecodingConfig] = None +def _to_websocket_scheme(url: str) -> str: + """Rewrite an http(s):// URL to ws(s)://. Raises if the scheme is unknown.""" + if url.startswith("https://"): + return "wss://" + url[len("https://") :] + if url.startswith("http://"): + return "ws://" + url[len("http://") :] + raise ValueError( + f"Cannot convert to WebSocket scheme; expected http(s):// prefix: {url!r}" + ) + + class DeployedServiceDescriptor(custom_types.SafeModel): """Bundles values to establish an RPC session to a dependency chainlet, specifically with ``StubBase``.""" @@ -756,6 +767,52 @@ def check_at_least_one_url( ) return self + @property + def target_url(self) -> str: + """The URL a ``StubBase`` would post to: prefer ``internal_url``'s + ``gateway_run_remote_url`` (cluster-local, lower-latency), fall back to + ``predict_url``. Mirrors the selection logic in + ``BasetenSession.__init__``.""" + if self.internal_url is not None: + return self.internal_url.gateway_run_remote_url + # Validator guarantees predict_url is set when internal_url is None. + assert self.predict_url is not None + return self.predict_url + + @property + def ws_url(self) -> Optional[str]: + """``predict_url`` with the http(s):// scheme rewritten to ws(s)://. + ``None`` when ``predict_url`` is not set.""" + if self.predict_url is None: + return None + return _to_websocket_scheme(self.predict_url) + + @property + def internal_ws_url(self) -> Optional[str]: + """``internal_url.gateway_run_remote_url`` with the http(s):// scheme + rewritten to ws(s)://. ``None`` when ``internal_url`` is not set.""" + if self.internal_url is None: + return None + return _to_websocket_scheme(self.internal_url.gateway_run_remote_url) + + def with_auth_headers(self, api_key: str) -> dict[str, str]: + """Build the headers a ``StubBase`` would send for outbound calls: + + * ``Authorization: Api-Key {api_key}`` — always. + * ``Host: {internal_url.hostname}`` — only when ``internal_url`` is set, + so cluster-local routing matches the chain hostname. + + ``api_key`` is typically sourced from + ``DeploymentContext.get_baseten_api_key()`` inside a ``ChainletBase`` or + from the standard Truss secrets API (``baseten_chain_api_key``) for raw + Trusses. The helper takes the key explicitly so it can be used outside + a chainlet context. + """ + headers = {"Authorization": f"Api-Key {api_key}"} + if self.internal_url is not None: + headers["Host"] = self.internal_url.hostname + return headers + class Environment(custom_types.SafeModel): """The environment the chainlet is deployed in. diff --git a/truss-chains/truss_chains/remote_chainlet/utils.py b/truss-chains/truss_chains/remote_chainlet/utils.py index 17f350a4d..1d8a1d518 100644 --- a/truss-chains/truss_chains/remote_chainlet/utils.py +++ b/truss-chains/truss_chains/remote_chainlet/utils.py @@ -3,7 +3,6 @@ import collections import contextlib import contextvars -import json import logging import statistics import sys @@ -29,8 +28,7 @@ import httpx import pydantic -from truss.templates.shared import dynamic_config_resolver -from truss_chains import private_types, public_types, utils +from truss_chains import private_types, public_types, runtime, utils if TYPE_CHECKING: import aiohttp @@ -119,29 +117,37 @@ class WebSocketState: # type: ignore[no-redef] def populate_chainlet_service_predict_urls( chainlet_to_service: Mapping[str, private_types.ServiceDescriptor], ) -> Mapping[str, public_types.DeployedServiceDescriptor]: - chainlet_to_deployed_service: dict[str, public_types.DeployedServiceDescriptor] = {} + """Resolve URLs for typed chainlet dependencies from the dynamic config. + + Used by the generated ``TrussChainletModel.__init__`` to combine the static + ``ServiceDescriptor`` map (from ``config.yaml``, carrying the framework's + ``name`` / ``display_name`` / ``RPCOptions`` triple) with the + runtime-resolved URLs from ``/etc/b10_dynamic_config/dynamic_chainlet_config``. + + For raw (non-typed) Truss code that just needs sibling URLs, see the public + :mod:`truss_chains.runtime` API instead. + """ if not chainlet_to_service: return {} - dynamic_chainlet_config_str = dynamic_config_resolver.get_dynamic_config_value_sync( - private_types.DYNAMIC_CHAINLET_CONFIG_KEY - ) - if not dynamic_chainlet_config_str: + try: + dynamic_chainlet_config = runtime._load_dynamic_config() + except public_types.MissingDependencyError: + # Preserve the historical error message — existing tests and downstream + # consumers may match on its text. raise public_types.MissingDependencyError( f"No '{private_types.DYNAMIC_CHAINLET_CONFIG_KEY}' " "found. Cannot override Chainlet configs." ) - dynamic_chainlet_config = json.loads(dynamic_chainlet_config_str) - + chainlet_to_deployed_service: dict[str, public_types.DeployedServiceDescriptor] = {} for chainlet_name, service_descriptor in chainlet_to_service.items(): display_name = service_descriptor.display_name - # NOTE: The Chainlet `display_name` in the Truss CLI - # corresponds to Chainlet `name` in the backend. As - # the dynamic Chainlet config is keyed on the backend - # Chainlet name, we have to look up config values by - # using the `display_name` in the service descriptor. + # NOTE: The Chainlet `display_name` in the Truss CLI corresponds to + # Chainlet `name` in the backend. As the dynamic Chainlet config is + # keyed on the backend Chainlet name, we have to look up config values + # by using the `display_name` in the service descriptor. if display_name not in dynamic_chainlet_config: raise public_types.MissingDependencyError( f"Chainlet '{display_name}' not found in " @@ -149,18 +155,21 @@ def populate_chainlet_service_predict_urls( f"Dynamic Chainlet config keys: {list(dynamic_chainlet_config)}." ) - if internal_url := dynamic_chainlet_config[display_name].get("internal_url"): - url = {"internal_url": internal_url} - else: - predict_url = dynamic_chainlet_config[display_name].get("predict_url") - url = {"predict_url": predict_url} - + # Build URLs via the runtime helper, then preserve the typed-chain + # identity (`name` / `display_name`) and `RPCOptions` from the static + # config map. Historical contract: when ``internal_url`` is present it + # is mutually exclusive with ``predict_url`` — keep ``predict_url`` set + # only when no ``internal_url`` is provided. + urls = runtime._descriptor_from_raw( + display_name, dynamic_chainlet_config[display_name] + ) chainlet_to_deployed_service[chainlet_name] = ( public_types.DeployedServiceDescriptor( - display_name=display_name, name=service_descriptor.name, + display_name=display_name, options=service_descriptor.options, - **url, + predict_url=None if urls.internal_url else urls.predict_url, + internal_url=urls.internal_url, ) ) diff --git a/truss-chains/truss_chains/runtime.py b/truss-chains/truss_chains/runtime.py new file mode 100644 index 000000000..198857df7 --- /dev/null +++ b/truss-chains/truss_chains/runtime.py @@ -0,0 +1,117 @@ +"""Public sibling-discovery API for chainlet runtimes. + +Reads ``/etc/b10_dynamic_config/dynamic_chainlet_config`` (the ConfigMap mounted +by the Baseten operator into every chainlet pod) and returns typed +``DeployedServiceDescriptor`` instances. + +This is a **stable public contract**: any Truss running inside a chain — whether +authored as a typed ``ChainletBase`` or as a raw Truss directory wrapped via +Project 2's ``TrussChainlet`` — can use this module to discover sibling chainlet +URLs without depending on framework internals. + +Schema of the dynamic config file (JSON object), with at least one of +``predict_url`` / ``internal_url`` always present per chainlet:: + + { + "": { + "predict_url": "https://chain-.api.baseten.co/.../run_remote", + "internal_url": { + "gateway_run_remote_url": "https://.api.baseten.co/.../run_remote", + "hostname": "chain-.api.baseten.co" + } + }, + ... + } + +Note on RPC options: the dynamic config does *not* carry ``RPCOptions`` (retries, +timeout, concurrency_limit). Those are typed-chain-only and live in the static +``chainlet_to_service`` map embedded in ``config.yaml`` by codegen. Descriptors +returned here use default ``RPCOptions``; callers using their own HTTP/WS client +should bring their own retry, timeout, and concurrency policy. +""" + +import json +from typing import Mapping + +from truss.templates.shared import dynamic_config_resolver +from truss_chains import private_types, public_types + + +def _load_dynamic_config() -> dict: + """Read and parse ``/etc/b10_dynamic_config/dynamic_chainlet_config``. + + Raises: + MissingDependencyError: if the ConfigMap file is absent or empty, + indicating the Truss is not running inside a chain context. + """ + raw = dynamic_config_resolver.get_dynamic_config_value_sync( + private_types.DYNAMIC_CHAINLET_CONFIG_KEY + ) + if not raw: + raise public_types.MissingDependencyError( + f"No '{private_types.DYNAMIC_CHAINLET_CONFIG_KEY}' configmap found at " + f"`{dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR}`. " + "This Truss is not running inside a chain context." + ) + return json.loads(raw) + + +def _descriptor_from_raw( + name: str, raw: dict +) -> public_types.DeployedServiceDescriptor: + """Build a ``DeployedServiceDescriptor`` from one entry of the dynamic config. + + ``name`` is used for both ``name`` and ``display_name`` in the resulting + descriptor (the dynamic config keys *are* display names — see module + docstring). ``options`` defaults to a fresh ``RPCOptions()``. + """ + kwargs: dict = {} + if internal_url := raw.get("internal_url"): + kwargs["internal_url"] = public_types.DeployedServiceDescriptor.InternalURL( + **internal_url + ) + if predict_url := raw.get("predict_url"): + kwargs["predict_url"] = predict_url + return public_types.DeployedServiceDescriptor( + name=name, display_name=name, options=public_types.RPCOptions(), **kwargs + ) + + +def get_service(name: str) -> public_types.DeployedServiceDescriptor: + """Return the descriptor for a sibling chainlet by name. + + Args: + name: The chainlet's display name (the key used in the chain + declaration; matches the keys in the dynamic config map). + + Returns: + ``DeployedServiceDescriptor`` with ``predict_url``, ``internal_url`` + (when available), and helper methods for building auth headers and + WebSocket URLs. + + Raises: + MissingDependencyError: if no chain context is detected (the dynamic + config file is absent), or if ``name`` is not a registered + sibling. The error message includes the available names. + """ + config = _load_dynamic_config() + if name not in config: + raise public_types.MissingDependencyError( + f"No sibling chainlet named '{name}'. Available: {list(config)}." + ) + return _descriptor_from_raw(name, config[name]) + + +def list_services() -> Mapping[str, public_types.DeployedServiceDescriptor]: + """Return all sibling chainlet descriptors keyed by name. + + Returns an empty mapping if no chain context is detected (i.e., the Truss + is running outside a chain). Use this to enumerate available siblings + without raising on absence — useful for code that wants to detect whether + it is running inside a chain at all. + """ + try: + config = _load_dynamic_config() + except public_types.MissingDependencyError: + return {} + return {n: _descriptor_from_raw(n, r) for n, r in config.items()} From 6a166c99aad71d6ca397d1df859407811f894f2e Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Wed, 29 Apr 2026 16:52:29 -0700 Subject: [PATCH 02/12] Add example and tests --- .../examples/bring-your-own-client/README.md | 102 ++++++ .../examples/bring-your-own-client/chain.py | 326 ++++++++++++++++++ .../examples/bring-your-own-client/invoke.py | 44 +++ .../01_basic_descriptor_access/README.md | 15 + .../01_basic_descriptor_access/chain.py | 54 +++ .../test_descriptor_helpers.py | 101 ++++++ .../02_raw_truss_reading_siblings/README.md | 25 ++ .../plain_truss/config.yaml | 10 + .../plain_truss/model/model.py | 32 ++ .../test_raw_truss.py | 65 ++++ .../03_missing_sibling/README.md | 13 + .../test_missing_sibling.py | 38 ++ .../04_dynamic_config_absent/README.md | 18 + .../test_no_context.py | 25 ++ .../README.md | 21 ++ .../test_backward_compat.py | 108 ++++++ .../tests/runtime_discovery/README.md | 43 +++ .../tests/runtime_discovery/run_all_local.sh | 12 + truss-chains/tests/test_runtime.py | 279 +++++++++++++++ 19 files changed, 1331 insertions(+) create mode 100644 truss-chains/examples/bring-your-own-client/README.md create mode 100644 truss-chains/examples/bring-your-own-client/chain.py create mode 100644 truss-chains/examples/bring-your-own-client/invoke.py create mode 100644 truss-chains/tests/runtime_discovery/01_basic_descriptor_access/README.md create mode 100644 truss-chains/tests/runtime_discovery/01_basic_descriptor_access/chain.py create mode 100644 truss-chains/tests/runtime_discovery/01_basic_descriptor_access/test_descriptor_helpers.py create mode 100644 truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md create mode 100644 truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/config.yaml create mode 100644 truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/model/model.py create mode 100644 truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py create mode 100644 truss-chains/tests/runtime_discovery/03_missing_sibling/README.md create mode 100644 truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py create mode 100644 truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md create mode 100644 truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py create mode 100644 truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md create mode 100644 truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py create mode 100644 truss-chains/tests/runtime_discovery/README.md create mode 100755 truss-chains/tests/runtime_discovery/run_all_local.sh create mode 100644 truss-chains/tests/test_runtime.py diff --git a/truss-chains/examples/bring-your-own-client/README.md b/truss-chains/examples/bring-your-own-client/README.md new file mode 100644 index 000000000..59ed9e6b2 --- /dev/null +++ b/truss-chains/examples/bring-your-own-client/README.md @@ -0,0 +1,102 @@ +# Bring-Your-Own-Client — Composable Chains Project 1 Demo + +A minimal, deployable, **CPU-only** chain that demonstrates what Project 1 +(Runtime Discovery Platform Contract) actually unlocks at runtime. + +## What this proves + +Two chainlets: + +- `Echo` — trivial dependency. `run_remote(text: str) -> str` returns the + uppercased input. +- `Caller` — entrypoint. Calls `Echo` **two different ways** in the same + request: + 1. Via the framework's auto-generated `StubBase` (the existing path). + 2. Via raw `httpx.AsyncClient`, using the new `DeployedServiceDescriptor` + helpers `target_url` and `with_auth_headers(api_key)` to source the URL + and auth from the chain context — without touching `StubBase`. + +If Project 1's helpers correctly mirror what `BasetenSession.__init__` would +construct internally, both paths return the same string. The chain returns +both results plus the descriptor metadata so you can see them line up. + +This is the **bring-your-own-client** pattern: keep the chain's UI grouping, +atomic deploy, and per-pod sibling-URL discovery, but plug in any HTTP / WS +client you like (httpx, websockets, grpc.aio, …). + +## Setup + +```sh +# From the truss repo root +pip install -e . +pip install -e truss-chains +truss login # if not already authenticated +``` + +## Push + +```sh +chains push truss-chains/examples/bring-your-own-client/chain.py +``` + +The CLI prints the chain's invoke URL when the deploy completes. Copy the +`run_remote` URL for the next step. + +## Invoke + +```sh +export BASETEN_API_KEY="..." +export CHAIN_URL="https://chain-.api.baseten.co/environments/production/run_remote" +python truss-chains/examples/bring-your-own-client/invoke.py "hello world" +``` + +Expected output: + +```json +{ + "via_stub": "HELLO WORLD", + "via_byo_httpx": "HELLO WORLD", + "match": true, + "target_url": "https://...api.baseten.co/.../chainlet/.../run_remote", + "has_internal_url": true, + "auth_header_keys": ["Authorization", "Host"] +} + +✓ MATCH +``` + +The interesting fields: + +- `match: true` — proves the BYO-httpx call hit the same endpoint and got the + same answer as the framework stub. +- `target_url` — what `desc.target_url` resolved to inside the chainlet pod. + Should be the cluster-internal beefeater URL when `internal_url` is present + (lower latency than going through the public chain hostname). +- `has_internal_url: true` — confirms the operator injected an `internal_url` + alongside the public `predict_url`. +- `auth_header_keys: ["Authorization", "Host"]` — confirms `with_auth_headers` + produced both the API key header and the chain hostname `Host:` override + (the latter only appears when `internal_url` is set). + +## Local smoke + +You can run `chain.py` directly under `run_local` to inspect the descriptor +helpers without deploying: + +```sh +python truss-chains/examples/bring-your-own-client/chain.py +``` + +Prints the four helper outputs against a fake `Echo` descriptor. Useful for +sanity-checking the helpers in development. + +## What this does *not* show + +- The `truss_chains.runtime.get_service(name)` API for **non-`ChainletBase`** + Trusses. That use case is unlocked by Project 2 (`TrussChainlet`); see + `truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/` for + the contract test. +- WebSocket-flavored siblings. The helpers `ws_url` / `internal_ws_url` are + exercised in the runtime_discovery contract tests; a WebSocket-fronted + deployable demo would need a chainlet using + `truss_config.WebsocketOptions` plus a client that handles WS framing. diff --git a/truss-chains/examples/bring-your-own-client/chain.py b/truss-chains/examples/bring-your-own-client/chain.py new file mode 100644 index 000000000..4d338a28e --- /dev/null +++ b/truss-chains/examples/bring-your-own-client/chain.py @@ -0,0 +1,326 @@ +"""Composable Chains — Project 1 deployable demo. + +Exercises **every** new public surface introduced by Project 1: + + Helpers on ``DeployedServiceDescriptor`` (`truss_chains.public_types`): + - ``target_url`` (property) + - ``ws_url`` (property) + - ``internal_ws_url`` (property) + - ``with_auth_headers(api_key)`` (method) + + Public submodule ``truss_chains.runtime``: + - ``get_service(name)`` + - ``list_services()`` + +The chain has two chainlets: + +* ``Echo`` — trivial dependency, returns input uppercased. +* ``Caller`` — entrypoint. For each request it: + 1. calls ``Echo`` via the framework's auto-generated ``StubBase`` + (existing path), + 2. resolves ``Echo``'s descriptor *both* via + ``DeploymentContext.get_service_descriptor`` (typed-chain path) and + via ``truss_chains.runtime.get_service`` (public-runtime path), + 3. inspects every helper on each descriptor, + 4. enumerates all siblings via ``truss_chains.runtime.list_services``, + 5. issues a raw ``httpx`` POST to the URL the helpers resolved to. + +Every step emits ``logging.info`` describing what's being exercised and why. +The pod logs (visible via the Baseten UI / CLI) read like a guided tour of +Project 1's API surface. The HTTP response also surfaces all the helper +outputs so an external invoker can see the same information. + +Both chainlets are CPU-only (1 vCPU / 512Mi each), so this pushes quickly. +""" + +import logging +from typing import Optional + +import httpx +import pydantic + +import truss_chains as chains +from truss_chains import runtime # NEW public submodule from Project 1 + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Echo — trivial dependency +# --------------------------------------------------------------------------- + + +_TRUSS_OVERRIDE = ( + # Override the published `truss_chains` with the Project 1 branch so + # `truss_chains.runtime` (and the new descriptor helpers) are available + # in every chainlet pod. Drop this once Project 1 is released to PyPI. + "truss @ git+https://github.com/basetenlabs/truss.git@matte/composable-chains-1-runtime" +) + + +class Echo(chains.ChainletBase): + """Returns its input uppercased. Exists solely so ``Caller`` has a sibling + to discover and exercise the helpers against.""" + + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi"), + docker_image=chains.DockerImage(pip_requirements=[_TRUSS_OVERRIDE]), + ) + + async def run_remote(self, text: str) -> str: + return text.upper() + + +# --------------------------------------------------------------------------- +# Output models — surface all the helper values to the external invoker +# --------------------------------------------------------------------------- + + +class DescriptorSnapshot(pydantic.BaseModel): + """Snapshot of every value the Project 1 helpers expose for a single + ``DeployedServiceDescriptor``. We deliberately redact the Authorization + header *value* so the response transcript doesn't leak the chain's API + key — the key shape is interesting, the bytes are not.""" + + name: str + display_name: str + predict_url: Optional[str] + internal_url_gateway: Optional[str] + internal_url_hostname: Optional[str] + target_url: str # NEW Project 1 property + ws_url: Optional[str] # NEW Project 1 property + internal_ws_url: Optional[str] # NEW Project 1 property + auth_header_keys: list[str] # keys present in with_auth_headers(...) output + auth_header_host: Optional[str] # Host header (chain hostname; not sensitive) + auth_header_authorization_format: str # e.g. "Api-Key " + + +class CallerOutput(pydantic.BaseModel): + """Side-by-side view of both invocation paths plus a guided tour of the + helpers' return values.""" + + via_stub: str + via_byo_httpx: str + match: bool + + # `context.get_service_descriptor("Echo")` path (typed-chain) + descriptor_via_context: DescriptorSnapshot + # `truss_chains.runtime.get_service("Echo")` path (public-runtime) + descriptor_via_runtime: DescriptorSnapshot + # `truss_chains.runtime.list_services()` — all siblings in the chain pod + runtime_siblings: list[str] + + +def _snapshot( + desc: chains.DeployedServiceDescriptor, api_key: str, source: str +) -> DescriptorSnapshot: + """Build a snapshot, logging each helper invocation as it runs. + + ``source`` is just a label for the log line so the operator can tell + apart the typed-context and public-runtime descriptor sources. + """ + logger.info( + "[%s] descriptor.name=%r display_name=%r", source, desc.name, desc.display_name + ) + logger.info("[%s] descriptor.predict_url=%s", source, desc.predict_url) + logger.info( + "[%s] descriptor.internal_url=%s", + source, + f"{desc.internal_url}" if desc.internal_url else None, + ) + + # NEW Project 1 helper: target_url. Mirrors BasetenSession's selection + # logic — prefer internal_url's gateway URL when present, fall back to + # predict_url. Saves callers from re-implementing the precedence rule. + target = desc.target_url + logger.info( + "[%s] HELPER target_url -> %s (preferring internal_url over predict_url)", + source, + target, + ) + + # NEW Project 1 helper: ws_url. predict_url with the http(s):// + # scheme rewritten to ws(s)://. Returns None when predict_url is unset. + ws = desc.ws_url + logger.info("[%s] HELPER ws_url -> %s (predict_url with https→wss)", source, ws) + + # NEW Project 1 helper: internal_ws_url. Same scheme rewrite, but for + # the cluster-local internal_url path. + iws = desc.internal_ws_url + logger.info( + "[%s] HELPER internal_ws_url -> %s (internal_url with https→wss)", source, iws + ) + + # NEW Project 1 helper: with_auth_headers(api_key). Builds the dict + # BasetenSession would build internally — Authorization always, Host + # when internal_url is set. Saves callers from hand-rolling either. + headers = desc.with_auth_headers(api_key) + logger.info( + "[%s] HELPER with_auth_headers -> keys=%s " + "(Authorization always; Host iff internal_url is set)", + source, + sorted(headers.keys()), + ) + + # Reduce the headers dict to a non-sensitive shape before returning. + auth_value = headers.get("Authorization", "") + auth_scheme = auth_value.split(" ", 1)[0] if auth_value else "" + return DescriptorSnapshot( + name=desc.name, + display_name=desc.display_name, + predict_url=desc.predict_url, + internal_url_gateway=( + desc.internal_url.gateway_run_remote_url if desc.internal_url else None + ), + internal_url_hostname=( + desc.internal_url.hostname if desc.internal_url else None + ), + target_url=target, + ws_url=ws, + internal_ws_url=iws, + auth_header_keys=sorted(headers.keys()), + auth_header_host=headers.get("Host"), + auth_header_authorization_format=f"{auth_scheme} ", + ) + + +# --------------------------------------------------------------------------- +# Caller — entrypoint that exercises every Project 1 surface +# --------------------------------------------------------------------------- + + +@chains.mark_entrypoint("Composable Chains BYO-Client Demo") +class Caller(chains.ChainletBase): + """Exercises every Project 1 helper while comparing the framework stub + path against a raw ``httpx`` BYO-client path.""" + + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi"), + docker_image=chains.DockerImage( + pip_requirements=[_TRUSS_OVERRIDE, "httpx>=0.27"] + ), + ) + + def __init__( + self, + echo: Echo = chains.depends(Echo), + context: chains.DeploymentContext = chains.depends_context(), + ) -> None: + self._echo = echo + self._context = context + + async def run_remote(self, text: str) -> CallerOutput: + logger.info("=" * 72) + logger.info("Caller.run_remote(text=%r) — Project 1 helper tour", text) + logger.info("=" * 72) + + # ----- path 1: framework stub (existing behavior) ------------------ + logger.info("[stub] await self._echo.run_remote(text)") + via_stub = await self._echo.run_remote(text) + logger.info("[stub] result=%r", via_stub) + + # ----- gather descriptors via BOTH access paths -------------------- + api_key = self._context.get_baseten_api_key() + + # Path A: typed-chain access via the DeploymentContext. + # `get_service_descriptor` is what auto-generated stubs use under the + # hood. Project 1 does NOT introduce this method — it pre-existed. + # Including it for contrast against the new public runtime API. + logger.info("[ctx] context.get_service_descriptor('Echo') — typed-chain path") + desc_ctx = self._context.get_service_descriptor("Echo") + snap_ctx = _snapshot(desc_ctx, api_key, source="ctx") + + # Path B: public-runtime access via the new submodule. + # NEW Project 1 entry point: truss_chains.runtime.get_service. Reads + # the same /etc/b10_dynamic_config/dynamic_chainlet_config file as + # the framework, but exposed publicly so any Truss (typed or raw, + # post Project 2) can call it without depending on framework code. + logger.info("[runtime] truss_chains.runtime.get_service('Echo') — public path") + desc_runtime = runtime.get_service("Echo") + snap_runtime = _snapshot(desc_runtime, api_key, source="runtime") + + # Note: the typed path historically clears `predict_url` when an + # `internal_url` is present (mutually exclusive). The new public + # runtime API carries both — so when both URLs are in the dynamic + # config you'll see snap_ctx.predict_url=None but + # snap_runtime.predict_url=. The descriptor_via_* + # comparison surfaces this explicitly. + + # ----- list all siblings via the new public API -------------------- + # NEW Project 1 entry point: truss_chains.runtime.list_services. + # Returns all siblings registered in the chain (including this + # chainlet itself). Returns {} outside a chain context — meaning + # raw Trusses can use this to detect whether they're running + # inside a chain at all. + logger.info("[runtime] truss_chains.runtime.list_services()") + siblings_map = runtime.list_services() + siblings = sorted(siblings_map.keys()) + logger.info("[runtime] siblings discovered: %s", siblings) + + # ----- BYO httpx call using the helpers --------------------------- + # Use the typed-context descriptor for the HTTP call (same wire as + # the stub would use). The helpers' job is to produce a target URL + # and headers equivalent to what BasetenSession.__init__ builds + # internally — so the BYO call should land at the same chainlet + # and return the same answer as the stub. + # + # Re-derive the actual headers here from the descriptor (rather than + # plucking them off `snap_ctx`, which deliberately redacts the + # Authorization value to keep secrets out of the response). + headers = desc_ctx.with_auth_headers(api_key) + logger.info( + "[httpx] POST %s headers=%s body=%s", + snap_ctx.target_url, + sorted(headers.keys()), + {"text": text}, + ) + async with httpx.AsyncClient(timeout=60) as client: + response = await client.post( + snap_ctx.target_url, headers=headers, json={"text": text} + ) + response.raise_for_status() + via_byo_httpx = response.json() + logger.info("[httpx] response=%r", via_byo_httpx) + + match = via_stub == via_byo_httpx + logger.info( + "MATCH=%s (stub=%r vs byo_httpx=%r)", match, via_stub, via_byo_httpx + ) + + return CallerOutput( + via_stub=via_stub, + via_byo_httpx=via_byo_httpx, + match=match, + descriptor_via_context=snap_ctx, + descriptor_via_runtime=snap_runtime, + runtime_siblings=siblings, + ) + + +# --------------------------------------------------------------------------- +# Local smoke (no deployment needed) +# --------------------------------------------------------------------------- + + +if __name__ == "__main__": + # Inspect the helper shapes against a fake descriptor under run_local. + # Useful for development; for real invocation post-deploy, see invoke.py. + logging.basicConfig(level=logging.INFO, format="%(message)s") + + fake_echo = chains.DeployedServiceDescriptor( + name="Echo", + display_name="Echo", + options=chains.RPCOptions(), + predict_url="https://chain-demo.api.baseten.co/.../echo/run_remote", + internal_url=chains.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://wp.api.baseten.co/.../echo/run_remote", + hostname="chain-demo.api.baseten.co", + ), + ) + with chains.run_local( + secrets={chains.public_types.CHAIN_API_KEY_SECRET_NAME: "local-dev-key"}, + chainlet_to_service={"Echo": fake_echo}, + ): + caller = Caller() + desc = caller._context.get_service_descriptor("Echo") + _snapshot(desc, api_key="local-dev-key", source="local") diff --git a/truss-chains/examples/bring-your-own-client/invoke.py b/truss-chains/examples/bring-your-own-client/invoke.py new file mode 100644 index 000000000..d947a7036 --- /dev/null +++ b/truss-chains/examples/bring-your-own-client/invoke.py @@ -0,0 +1,44 @@ +"""Invoke the deployed Composable Chains BYO-Client demo. + +Usage: + BASETEN_API_KEY=... CHAIN_URL=https://chain-.api.baseten.co/environments/production/run_remote \\ + python invoke.py "hello world" + +The chain returns both the stub-path and the raw-httpx-path results plus the +descriptor metadata, so you can verify the bring-your-own-client pattern +matches the framework stub byte-for-byte. +""" + +import json +import os +import sys + +import requests + + +def main() -> None: + api_key = os.environ.get("BASETEN_API_KEY") + chain_url = os.environ.get("CHAIN_URL") + if not api_key or not chain_url: + sys.exit( + "Set BASETEN_API_KEY and CHAIN_URL. CHAIN_URL is the chain's " + "`/run_remote` endpoint shown after `chains push`." + ) + + text = sys.argv[1] if len(sys.argv) > 1 else "hello world" + + response = requests.post( + chain_url, + headers={"Authorization": f"Api-Key {api_key}"}, + json={"text": text}, + timeout=120, + ) + response.raise_for_status() + result = response.json() + print(json.dumps(result, indent=2)) + print() + print("✓ MATCH" if result.get("match") else "✗ MISMATCH — bug in helpers!") + + +if __name__ == "__main__": + main() diff --git a/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/README.md b/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/README.md new file mode 100644 index 000000000..24647850c --- /dev/null +++ b/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/README.md @@ -0,0 +1,15 @@ +# 01 — Basic Descriptor Access + +A two-chainlet typed chain where `Caller` reads the new helper methods on the +`DeployedServiceDescriptor` returned by `context.get_service_descriptor("Echo")` +and reports them through `run_remote`. + +This is the "typed chain that uses the new helpers" path — no raw Truss yet. +Demonstrates that `target_url`, `ws_url`, `internal_ws_url`, and +`with_auth_headers()` produce the expected shapes from a real descriptor. + +## Run + +```sh +uv run pytest test_descriptor_helpers.py -v +``` diff --git a/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/chain.py b/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/chain.py new file mode 100644 index 000000000..4faa3123b --- /dev/null +++ b/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/chain.py @@ -0,0 +1,54 @@ +"""Simple typed chain showing access to the new descriptor helpers from +``context.get_service_descriptor``.""" + +from typing import Optional + +import pydantic + +import truss_chains as chains + + +class CallerOutput(pydantic.BaseModel): + echoed: str + target_url: str + ws_url: Optional[str] + internal_ws_url: Optional[str] + auth_headers: dict[str, str] + + +class Echo(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + async def run_remote(self, text: str) -> str: + return text + + +class Caller(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + def __init__( + self, + echo: Echo = chains.depends(Echo), + context: chains.DeploymentContext = chains.depends_context(), + ) -> None: + self._echo = echo + self._context = context + + async def run_remote(self, text: str) -> CallerOutput: + # Pull the descriptor for the typed dependency and exercise the new + # helpers. In a real chain you would use `target_url` / `ws_url` to + # build your own client; here we just surface them so the test + # harness can assert against them. + desc = self._context.get_service_descriptor("Echo") + api_key = self._context.get_baseten_api_key() + return CallerOutput( + echoed=await self._echo.run_remote(text), + target_url=desc.target_url, + ws_url=desc.ws_url, + internal_ws_url=desc.internal_ws_url, + auth_headers=desc.with_auth_headers(api_key), + ) diff --git a/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/test_descriptor_helpers.py b/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/test_descriptor_helpers.py new file mode 100644 index 000000000..6a6cecbf8 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/test_descriptor_helpers.py @@ -0,0 +1,101 @@ +"""Smoke tests for the descriptor helpers using ``run_local``-style +construction of a ``DeployedServiceDescriptor``. The chain in ``chain.py`` +shows how a chainlet would consume these helpers in production; here we +construct a descriptor directly and assert the helpers' return values.""" + +import pathlib +import sys + +import truss_chains as chains + +# Make ``chain.py`` importable from this test file. +sys.path.insert(0, str(pathlib.Path(__file__).parent)) + + +def _make_descriptor(with_internal_url: bool): + kwargs = { + "name": "Echo", + "display_name": "Echo", + "options": chains.RPCOptions(), + "predict_url": "https://chain-abc.api.baseten.co/.../echo/run_remote", + } + if with_internal_url: + kwargs["internal_url"] = chains.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://wp.api.baseten.co/.../echo/run_remote", + hostname="chain-abc.api.baseten.co", + ) + return chains.DeployedServiceDescriptor(**kwargs) + + +def test_target_url_with_internal_prefers_internal(): + desc = _make_descriptor(with_internal_url=True) + assert desc.target_url == "https://wp.api.baseten.co/.../echo/run_remote" + + +def test_target_url_falls_back_to_predict(): + desc = _make_descriptor(with_internal_url=False) + assert desc.target_url == "https://chain-abc.api.baseten.co/.../echo/run_remote" + + +def test_ws_url_scheme_swap(): + desc = _make_descriptor(with_internal_url=False) + assert desc.ws_url.startswith("wss://") + assert desc.internal_ws_url is None + + +def test_internal_ws_url_scheme_swap(): + desc = _make_descriptor(with_internal_url=True) + assert desc.internal_ws_url.startswith("wss://") + + +def test_auth_headers_no_internal(): + desc = _make_descriptor(with_internal_url=False) + assert desc.with_auth_headers("k") == {"Authorization": "Api-Key k"} + + +def test_auth_headers_with_internal_includes_host(): + desc = _make_descriptor(with_internal_url=True) + assert desc.with_auth_headers("k") == { + "Authorization": "Api-Key k", + "Host": "chain-abc.api.baseten.co", + } + + +def test_chain_runs_locally_with_deployed_echo(): + """run_local with a ``chainlet_to_service`` override that points Echo at + a fake deployed URL — exercises the descriptor helpers via the typed path + (``context.get_service_descriptor``). + + Caller still calls ``self._echo.run_remote(...)`` over HTTP via the stub + rather than running Echo in-process, so this scenario is not what you'd + use in normal local debugging — but it's what makes the descriptor + helpers observable from a run_local test.""" + from chain import Caller + + fake_descriptor = chains.DeployedServiceDescriptor( + name="Echo", + display_name="Echo", + options=chains.RPCOptions(), + predict_url="https://chain-abc.api.baseten.co/.../echo/run_remote", + internal_url=chains.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://wp.api.baseten.co/.../echo/run_remote", + hostname="chain-abc.api.baseten.co", + ), + ) + + with chains.run_local( + secrets={chains.public_types.CHAIN_API_KEY_SECRET_NAME: "test-key"}, + chainlet_to_service={"Echo": fake_descriptor}, + ): + caller = Caller() + + # We don't actually call run_remote (that would fire HTTP at the fake URL). + # Inspect the descriptor through the context directly. + desc = caller._context.get_service_descriptor("Echo") + assert desc.target_url == "https://wp.api.baseten.co/.../echo/run_remote" + assert desc.ws_url == "wss://chain-abc.api.baseten.co/.../echo/run_remote" + assert desc.internal_ws_url == "wss://wp.api.baseten.co/.../echo/run_remote" + assert desc.with_auth_headers("test-key") == { + "Authorization": "Api-Key test-key", + "Host": "chain-abc.api.baseten.co", + } diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md new file mode 100644 index 000000000..822648c3b --- /dev/null +++ b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md @@ -0,0 +1,25 @@ +# 02 — Raw Truss Reading Siblings + +A regular Truss directory (not a `ChainletBase` subclass) that uses +`from truss_chains.runtime import get_service` to discover sibling URLs. +This is the bring-your-own-client pattern that Project 2 will rely on: +the Truss does not depend on the chains framework code, only on the +`truss_chains` package being installed. + +The `plain_truss/` directory shows the full layout: +``` +plain_truss/ +├── config.yaml +└── model/ + └── model.py +``` + +`model.py` calls `get_service("Diarizer")` in `load()` and stashes the +URL in `predict()` — exactly how a TrussChainlet would surface sibling +information to user-managed RPC code. + +## Run + +```sh +uv run pytest test_raw_truss.py -v +``` diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/config.yaml b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/config.yaml new file mode 100644 index 000000000..d3452f772 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/config.yaml @@ -0,0 +1,10 @@ +model_name: PlainTrussReader +model_class_filename: model.py +model_class_name: Model +python_version: py311 +requirements: + - truss_chains +resources: + cpu: "1" + memory: 512Mi + use_gpu: false diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/model/model.py b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/model/model.py new file mode 100644 index 000000000..b6b2bb7f2 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/model/model.py @@ -0,0 +1,32 @@ +"""A regular Truss model that participates in a chain by reading sibling URLs +through the public ``truss_chains.runtime`` API. + +Notice: no `ChainletBase`, no `run_remote`, no `truss_chains` framework imports +beyond the public runtime module. This Truss could be reused outside any chain +context — its dependence on chain siblings is conditional via +``runtime.list_services()``. +""" + +from truss_chains.runtime import get_service, list_services + + +class Model: + def __init__(self, **kwargs) -> None: + self._diarizer_url = None + self._auth_headers = None + + def load(self) -> None: + # Conditional sibling discovery: if running inside a chain, pick up + # the Diarizer URL; if standalone, fall back to no-op behavior. + siblings = list_services() + if "Diarizer" in siblings: + desc = get_service("Diarizer") + # Caller picks the URL flavor it wants — predict_url for public + # routing, internal_url for cluster-local. + self._diarizer_url = desc.target_url + # In a real Truss, source the api_key from the standard Truss + # secrets API. For this demo we leave it as a placeholder. + self._auth_headers = desc.with_auth_headers(api_key="") + + def predict(self, request: dict) -> dict: + return {"diarizer_url": self._diarizer_url, "auth_headers": self._auth_headers} diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py new file mode 100644 index 000000000..dc02c9a1c --- /dev/null +++ b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py @@ -0,0 +1,65 @@ +"""Drive the plain Truss's model.py with a fake dynamic_chainlet_config to +verify it discovers siblings via the public runtime API.""" + +import json +import pathlib +import sys + +import pytest + +# Make the plain_truss/model module importable. +PLAIN_TRUSS_MODEL_DIR = pathlib.Path(__file__).parent / "plain_truss" / "model" +sys.path.insert(0, str(PLAIN_TRUSS_MODEL_DIR)) + + +@pytest.fixture +def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", + str(tmp_path), + ) + yield + + +def _write_dynamic_config(tmp_path, payload): + with (tmp_path / "dynamic_chainlet_config").open("w") as f: + f.write(json.dumps(payload)) + + +def test_plain_truss_picks_up_siblings(tmp_path, dynamic_config_mount_dir): + _write_dynamic_config( + tmp_path, + { + "Diarizer": { + "predict_url": "https://chain-abc.api.baseten.co/.../diarizer/run_remote", + "internal_url": { + "gateway_run_remote_url": "https://wp.api.baseten.co/.../diarizer/run_remote", + "hostname": "chain-abc.api.baseten.co", + }, + } + }, + ) + + from model import Model # noqa: import after monkeypatch + + m = Model() + m.load() + out = m.predict({}) + # internal_url wins (matches BasetenSession's selection rule). + assert out["diarizer_url"] == "https://wp.api.baseten.co/.../diarizer/run_remote" + assert out["auth_headers"] == { + "Authorization": "Api-Key ", + "Host": "chain-abc.api.baseten.co", + } + + +def test_plain_truss_runs_standalone(tmp_path, dynamic_config_mount_dir): + """No dynamic config file present — list_services returns empty, + Model.load() short-circuits, predict returns the standalone shape.""" + from model import Model + + m = Model() + m.load() + out = m.predict({}) + assert out["diarizer_url"] is None + assert out["auth_headers"] is None diff --git a/truss-chains/tests/runtime_discovery/03_missing_sibling/README.md b/truss-chains/tests/runtime_discovery/03_missing_sibling/README.md new file mode 100644 index 000000000..2ebb81807 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/03_missing_sibling/README.md @@ -0,0 +1,13 @@ +# 03 — Missing Sibling (adversarial) + +When user code calls `runtime.get_service("DoesNotExist")`, the framework +must raise `MissingDependencyError` with a message that lists the available +sibling names — not silently return `None` or crash with a `KeyError`. + +This is the canonical failure-mode test for typos in chainlet names. + +## Run + +```sh +uv run pytest test_missing_sibling.py -v +``` diff --git a/truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py b/truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py new file mode 100644 index 000000000..9bbd24bce --- /dev/null +++ b/truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py @@ -0,0 +1,38 @@ +"""Adversarial: `get_service` raises with available-names list when the +requested chainlet is not registered in the chain.""" + +import json + +import pytest + +from truss_chains import public_types, runtime + + +@pytest.fixture +def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", + str(tmp_path), + ) + yield + + +def test_missing_sibling_raises_with_available_names( + tmp_path, dynamic_config_mount_dir +): + config = { + "Whisper": {"predict_url": "https://x.example/whisper"}, + "Diarizer": {"predict_url": "https://x.example/diarizer"}, + } + with (tmp_path / "dynamic_chainlet_config").open("w") as f: + f.write(json.dumps(config)) + + with pytest.raises(public_types.MissingDependencyError) as excinfo: + runtime.get_service("Typo") + + msg = str(excinfo.value) + # The error must name what was attempted. + assert "Typo" in msg + # And the available alternatives — so users can debug typos quickly. + assert "Whisper" in msg + assert "Diarizer" in msg diff --git a/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md b/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md new file mode 100644 index 000000000..86e15b05a --- /dev/null +++ b/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md @@ -0,0 +1,18 @@ +# 04 — Dynamic Config Absent (adversarial) + +When a Truss runs outside any chain context (no +`/etc/b10_dynamic_config/dynamic_chainlet_config` file present): + +- `runtime.list_services()` returns an empty mapping (does NOT raise) — so + caller code can branch on "am I inside a chain?" without try/except. +- `runtime.get_service(name)` raises `MissingDependencyError` with a clear + "not running inside a chain context" message. + +This documents the contract for code that needs to behave gracefully both +inside and outside a chain (e.g., a Truss that's reused across deployments). + +## Run + +```sh +uv run pytest test_no_context.py -v +``` diff --git a/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py b/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py new file mode 100644 index 000000000..f15757753 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py @@ -0,0 +1,25 @@ +"""Adversarial: behavior when no chain context is present.""" + +import pytest + +from truss_chains import public_types, runtime + + +@pytest.fixture +def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", + str(tmp_path), + ) + yield + + +def test_list_services_empty_when_no_context(tmp_path, dynamic_config_mount_dir): + """No file written; list_services returns empty mapping (does NOT raise).""" + assert runtime.list_services() == {} + + +def test_get_service_raises_clear_message(tmp_path, dynamic_config_mount_dir): + with pytest.raises(public_types.MissingDependencyError) as excinfo: + runtime.get_service("Anything") + assert "not running inside a chain context" in str(excinfo.value) diff --git a/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md b/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md new file mode 100644 index 000000000..2d54ece94 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md @@ -0,0 +1,21 @@ +# 05 — Backward Compat: Existing Typed Chainlets (adversarial) + +Project 1 internally refactors `populate_chainlet_service_predict_urls` to +share code with the new `truss_chains.runtime` API. This example pins the +historical contract of that refactor: + +1. The function's signature and return shape are unchanged. +2. Error messages for "config absent" / "name missing" are preserved + byte-for-byte. +3. The mutually-exclusive ``predict_url`` / ``internal_url`` behavior of + the typed path is preserved: when ``internal_url`` is present in the + dynamic config, ``predict_url`` is cleared on the resulting descriptor + (matching the pre-refactor behavior). +4. Same Python type identity for `DeployedServiceDescriptor` regardless of + import path. + +## Run + +```sh +uv run pytest test_backward_compat.py -v +``` diff --git a/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py b/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py new file mode 100644 index 000000000..060061c9c --- /dev/null +++ b/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py @@ -0,0 +1,108 @@ +"""Adversarial: pin the backward-compatible contract of the typed-chain +``populate_chainlet_service_predict_urls`` after the internal refactor.""" + +import json + +import pytest + +from truss_chains import private_types, public_types, runtime +from truss_chains.remote_chainlet.utils import populate_chainlet_service_predict_urls + + +@pytest.fixture +def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", + str(tmp_path), + ) + yield + + +def _write_config(tmp_path, payload): + with (tmp_path / private_types.DYNAMIC_CHAINLET_CONFIG_KEY).open("w") as f: + f.write(json.dumps(payload)) + + +def test_typed_path_clears_predict_url_when_internal_url_present( + tmp_path, dynamic_config_mount_dir +): + """Historical behavior: in the typed-chain path, when both URLs are in + the dynamic config, the resulting descriptor has only ``internal_url`` + set. The new public ``runtime.get_service`` carries both — but the typed + path stays mutually exclusive for backward compatibility.""" + _write_config( + tmp_path, + { + "Hello World!": { + "predict_url": "https://chain-x.api.baseten.co/.../run_remote", + "internal_url": { + "gateway_run_remote_url": "https://wp.api.baseten.co/.../run_remote", + "hostname": "chain-x.api.baseten.co", + }, + } + }, + ) + + chainlet_to_service = { + "HelloWorld": private_types.ServiceDescriptor( + name="HelloWorld", + display_name="Hello World!", + options=public_types.RPCOptions(), + ) + } + out = populate_chainlet_service_predict_urls(chainlet_to_service) + + # Typed path: predict_url cleared. + assert out["HelloWorld"].predict_url is None + assert out["HelloWorld"].internal_url is not None + + # Public runtime API: both URLs preserved (different contract, same data). + desc = runtime.get_service("Hello World!") + assert desc.predict_url == "https://chain-x.api.baseten.co/.../run_remote" + assert desc.internal_url is not None + + +def test_missing_chainlet_error_message_preserved(tmp_path, dynamic_config_mount_dir): + """Existing test in test_utils.py matches on the substring + "Chainlet 'X' not found". Project 1's refactor must keep this exact + wording.""" + _write_config(tmp_path, {"OtherName": {"predict_url": "https://x"}}) + with pytest.raises( + public_types.MissingDependencyError, match="Chainlet 'RandInt' not found" + ): + populate_chainlet_service_predict_urls( + { + "RandInt": private_types.ServiceDescriptor( + name="RandInt", + display_name="RandInt", + options=public_types.RPCOptions(), + ) + } + ) + + +def test_missing_dynamic_config_error_message_preserved( + tmp_path, dynamic_config_mount_dir +): + """No file: typed path raises with the historical "Cannot override + Chainlet configs" wording, distinct from the new runtime API's + "not running inside a chain context" wording.""" + with pytest.raises( + public_types.MissingDependencyError, match="Cannot override Chainlet configs" + ): + populate_chainlet_service_predict_urls( + { + "X": private_types.ServiceDescriptor( + name="X", display_name="X", options=public_types.RPCOptions() + ) + } + ) + + +def test_descriptor_type_identity_across_imports(): + """Same Python type whether reached via public_types or as a re-export + candidate — there's only one canonical class.""" + from truss_chains import DeployedServiceDescriptor as A + from truss_chains.public_types import DeployedServiceDescriptor as B + + assert A is B diff --git a/truss-chains/tests/runtime_discovery/README.md b/truss-chains/tests/runtime_discovery/README.md new file mode 100644 index 000000000..93d77c661 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/README.md @@ -0,0 +1,43 @@ +# Runtime Discovery — Contract Test Scenarios + +CPU-only pytest scenarios that pin the Project 1 (Runtime Discovery Platform +Contract) API. Each scenario doubles as documentation of the expected shape +and behavior of `truss_chains.runtime`. + +These live under `tests/` rather than `examples/` because they are **not +deployable chains** — they exercise the library by writing fake +`/etc/b10_dynamic_config/dynamic_chainlet_config` files via pytest +monkeypatch, not by pushing real chainlets to Baseten. (For deployable +chain examples, see `truss-chains/examples/`.) + +## Setup + +```sh +pip install -e /Users/mattelim/Documents/truss +pip install -e /Users/mattelim/Documents/truss/truss-chains +pip install pytest +``` + +## Run all + +```sh +# Runs each scenario in its own pytest invocation, with descriptive headers. +bash run_all_local.sh + +# Or as one combined pytest run: +uv run pytest truss-chains/tests/runtime_discovery/ +``` + +## Per-scenario + +| Dir | What it demonstrates | Type | +|---|---|---| +| `01_basic_descriptor_access/` | Typed chainlets accessing the new `target_url`, `ws_url`, `internal_ws_url`, `with_auth_headers()` helpers via `context.get_service_descriptor(...)`. | Positive | +| `02_raw_truss_reading_siblings/` | A non-`ChainletBase` Truss using `from truss_chains.runtime import get_service` to discover sibling URLs. | Positive | +| `03_missing_sibling/` | `MissingDependencyError` with helpful "Available: ..." message when looking up a non-existent sibling. | Adversarial | +| `04_dynamic_config_absent/` | `list_services()` returns `{}` (does not raise) and `get_service(...)` raises with a clear "not running inside a chain context" message when no config file is present. | Adversarial | +| `05_backward_compat_existing_chainlets/` | `populate_chainlet_service_predict_urls` (the typed-chain code path) preserves its historical mutually-exclusive `predict_url` / `internal_url` behavior — proves the internal refactor doesn't change the typed path. | Adversarial | + +Each scenario has a `test_.py` that runs under `pytest`. Scenarios that +involve a Truss directory have a `plain_truss/` subdirectory with a hand-written +`config.yaml` and `model/model.py` showing the bring-your-own-client pattern. diff --git a/truss-chains/tests/runtime_discovery/run_all_local.sh b/truss-chains/tests/runtime_discovery/run_all_local.sh new file mode 100755 index 000000000..b07020782 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/run_all_local.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Run every runtime_discovery example's local test. CPU-only. +set -euo pipefail +cd "$(dirname "$0")" +for dir in 0*; do + for test_file in "$dir"/test_*.py; do + if [[ -f "$test_file" ]]; then + echo "=== $dir / $(basename "$test_file") ===" + uv run pytest "$test_file" -v + fi + done +done diff --git a/truss-chains/tests/test_runtime.py b/truss-chains/tests/test_runtime.py new file mode 100644 index 000000000..4c666ec89 --- /dev/null +++ b/truss-chains/tests/test_runtime.py @@ -0,0 +1,279 @@ +"""Tests for ``truss_chains.runtime``: the public sibling-discovery API. + +Covers the positive paths (descriptor access, URL helpers, list_services) and +the adversarial paths (missing config file, unknown sibling name, empty +mapping). + +The fixture ``dynamic_config_mount_dir`` from ``conftest.py`` / shared test +infra monkeypatches ``DYNAMIC_CONFIG_MOUNT_DIR`` to ``tmp_path`` so a test can +write fake JSON to ``tmp_path / "dynamic_chainlet_config"`` and exercise the +runtime exactly as it would behave inside a real chainlet pod. +""" + +import json + +import pytest + +from truss_chains import private_types, public_types, runtime + +# ---- Test fixtures ----------------------------------------------------------- + +PREDICT_URL_ONLY = { + "Whisper": { + "predict_url": "https://chain-abc123.api.baseten.co/deployment/dep/chainlet/cl/run_remote" + } +} + +INTERNAL_AND_PREDICT_URL = { + "Whisper": { + "predict_url": "https://chain-abc123.api.baseten.co/deployment/dep/chainlet/cl/run_remote", + "internal_url": { + "gateway_run_remote_url": "https://aws-us-west-2-ai7.api.baseten.co/deployment/dep/chainlet/cl/run_remote", + "hostname": "chain-abc123.api.baseten.co", + }, + } +} + +INTERNAL_URL_ONLY = { + "Whisper": { + "internal_url": { + "gateway_run_remote_url": "https://aws-us-west-2-ai7.api.baseten.co/deployment/dep/chainlet/cl/run_remote", + "hostname": "chain-abc123.api.baseten.co", + } + } +} + +MULTIPLE_SIBLINGS = { + "Whisper": { + "predict_url": "https://chain-abc.api.baseten.co/.../whisper/run_remote" + }, + "Diarizer": { + "predict_url": "https://chain-abc.api.baseten.co/.../diarizer/run_remote", + "internal_url": { + "gateway_run_remote_url": "https://wp.api.baseten.co/.../diarizer/run_remote", + "hostname": "chain-abc.api.baseten.co", + }, + }, +} + + +@pytest.fixture +def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", + str(tmp_path), + ) + yield + + +def _write_config(tmp_path, payload): + with (tmp_path / private_types.DYNAMIC_CHAINLET_CONFIG_KEY).open("w") as f: + f.write(json.dumps(payload)) + + +# ---- get_service: positive paths -------------------------------------------- + + +def test_get_service_predict_url_only(tmp_path, dynamic_config_mount_dir): + _write_config(tmp_path, PREDICT_URL_ONLY) + desc = runtime.get_service("Whisper") + assert desc.predict_url == PREDICT_URL_ONLY["Whisper"]["predict_url"] + assert desc.internal_url is None + assert desc.name == "Whisper" + assert desc.display_name == "Whisper" + + +def test_get_service_carries_both_urls_when_present(tmp_path, dynamic_config_mount_dir): + """The public runtime API faithfully reflects the dynamic config — when the + config has both predict_url and internal_url, the descriptor carries both. + (The typed-chain ``populate_chainlet_service_predict_urls`` path retains + the historical mutually-exclusive behavior; this is tested separately in + ``test_utils.py``.)""" + _write_config(tmp_path, INTERNAL_AND_PREDICT_URL) + desc = runtime.get_service("Whisper") + assert desc.predict_url == INTERNAL_AND_PREDICT_URL["Whisper"]["predict_url"] + assert desc.internal_url is not None + assert ( + desc.internal_url.gateway_run_remote_url + == INTERNAL_AND_PREDICT_URL["Whisper"]["internal_url"]["gateway_run_remote_url"] + ) + assert ( + desc.internal_url.hostname + == INTERNAL_AND_PREDICT_URL["Whisper"]["internal_url"]["hostname"] + ) + + +def test_get_service_internal_url_only(tmp_path, dynamic_config_mount_dir): + _write_config(tmp_path, INTERNAL_URL_ONLY) + desc = runtime.get_service("Whisper") + assert desc.predict_url is None + assert desc.internal_url is not None + + +# ---- get_service: adversarial paths ----------------------------------------- + + +def test_get_service_missing_name_raises(tmp_path, dynamic_config_mount_dir): + _write_config(tmp_path, MULTIPLE_SIBLINGS) + with pytest.raises( + public_types.MissingDependencyError, match="No sibling chainlet named 'Nope'" + ): + runtime.get_service("Nope") + + +def test_get_service_missing_name_lists_available(tmp_path, dynamic_config_mount_dir): + """Error message must include the list of available chainlet names so + users can debug typos quickly.""" + _write_config(tmp_path, MULTIPLE_SIBLINGS) + with pytest.raises(public_types.MissingDependencyError) as excinfo: + runtime.get_service("Nope") + assert "Whisper" in str(excinfo.value) + assert "Diarizer" in str(excinfo.value) + + +def test_get_service_no_chain_context_raises(tmp_path, dynamic_config_mount_dir): + # No file written. + with pytest.raises( + public_types.MissingDependencyError, match="not running inside a chain context" + ): + runtime.get_service("Whisper") + + +# ---- list_services ---------------------------------------------------------- + + +def test_list_services_empty_when_no_context(tmp_path, dynamic_config_mount_dir): + """list_services must NOT raise outside a chain — returns empty mapping.""" + assert runtime.list_services() == {} + + +def test_list_services_returns_all(tmp_path, dynamic_config_mount_dir): + _write_config(tmp_path, MULTIPLE_SIBLINGS) + services = runtime.list_services() + assert set(services.keys()) == {"Whisper", "Diarizer"} + assert services["Whisper"].predict_url is not None + assert services["Diarizer"].internal_url is not None + + +# ---- DeployedServiceDescriptor helpers -------------------------------------- + + +def test_descriptor_target_url_prefers_internal(): + """target_url mirrors BasetenSession's selection: internal_url wins.""" + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="https://public.example/predict", + internal_url=public_types.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://internal.example/predict", + hostname="public.example", + ), + ) + assert desc.target_url == "https://internal.example/predict" + + +def test_descriptor_target_url_falls_back_to_predict(): + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="https://public.example/predict", + ) + assert desc.target_url == "https://public.example/predict" + + +def test_descriptor_ws_url_https_to_wss(): + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="https://public.example/predict", + ) + assert desc.ws_url == "wss://public.example/predict" + assert desc.internal_ws_url is None + + +def test_descriptor_ws_url_http_to_ws(): + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="http://localhost:8080/predict", + ) + assert desc.ws_url == "ws://localhost:8080/predict" + + +def test_descriptor_internal_ws_url(): + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + internal_url=public_types.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://internal.example/predict", + hostname="public.example", + ), + ) + assert desc.internal_ws_url == "wss://internal.example/predict" + assert desc.ws_url is None + + +def test_descriptor_with_auth_headers_predict_only(): + """Without internal_url, only Authorization header is set.""" + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="https://public.example/predict", + ) + headers = desc.with_auth_headers("my-key") + assert headers == {"Authorization": "Api-Key my-key"} + + +def test_descriptor_with_auth_headers_with_internal_url(): + """With internal_url, both Authorization and Host are set — Host carries + the chain hostname so cluster-local routing matches the chain identity.""" + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + internal_url=public_types.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://internal.example/predict", + hostname="chain-abc.api.baseten.co", + ), + ) + headers = desc.with_auth_headers("my-key") + assert headers == { + "Authorization": "Api-Key my-key", + "Host": "chain-abc.api.baseten.co", + } + + +def test_with_auth_headers_matches_BasetenSession(): + """Mechanically verify ``with_auth_headers`` produces the exact dict + ``BasetenSession.__init__`` would assemble — preventing drift between the + helper and the framework's internal RPC client.""" + from truss_chains.remote_chainlet import stub + + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + internal_url=public_types.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://internal.example/predict", + hostname="chain-abc.api.baseten.co", + ), + ) + session = stub.BasetenSession(service_descriptor=desc, api_key="my-key") + assert session._headers == desc.with_auth_headers("my-key") + + +def test_descriptor_ws_url_unknown_scheme_raises(): + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="ftp://nope.example/predict", + ) + with pytest.raises(ValueError, match="Cannot convert to WebSocket scheme"): + _ = desc.ws_url From f3226fbfed28503f23415bc38256e076b90c4ce3 Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Wed, 29 Apr 2026 20:24:57 -0700 Subject: [PATCH 03/12] Add TrussChainlet to allow truss as chainlets --- truss-chains/truss_chains/__init__.py | 8 +- .../truss_chains/deployment/code_gen.py | 54 ++++++- .../deployment/deployment_client.py | 5 + truss-chains/truss_chains/framework.py | 150 ++++++++++++++++++ truss-chains/truss_chains/private_types.py | 9 ++ 5 files changed, 224 insertions(+), 2 deletions(-) diff --git a/truss-chains/truss_chains/__init__.py b/truss-chains/truss_chains/__init__.py index ecd92e71e..f85592040 100644 --- a/truss-chains/truss_chains/__init__.py +++ b/truss-chains/truss_chains/__init__.py @@ -1,6 +1,11 @@ from truss.base.truss_config import WeightsSource from truss_chains import runtime # noqa: E402, F401 -from truss_chains.framework import ChainletBase, EngineBuilderLLMChainlet, ModelBase +from truss_chains.framework import ( + ChainletBase, + EngineBuilderLLMChainlet, + ModelBase, + TrussChainlet, +) from truss_chains.public_api import ( depends, depends_context, @@ -49,6 +54,7 @@ "RemoteConfig", "RemoteErrorDetail", "StubBase", + "TrussChainlet", "WebSocketProtocol", "WeightsSource", "depends", diff --git a/truss-chains/truss_chains/deployment/code_gen.py b/truss-chains/truss_chains/deployment/code_gen.py index 9fd3c4591..3a7862da3 100644 --- a/truss-chains/truss_chains/deployment/code_gen.py +++ b/truss-chains/truss_chains/deployment/code_gen.py @@ -382,6 +382,11 @@ def _gen_stub_src_for_deps( imports: set[str] = set() src_parts: list[str] = [] for dep in dependencies: + # TrussChainlet deps don't get a typed stub — the user accesses them + # via the injected `DeployedServiceDescriptor` and brings their own + # client. Skip stub generation entirely for these. + if framework.is_truss_chainlet(dep.chainlet_cls): + continue _update_src(_gen_stub_src(dep), src_parts, imports) if not (imports or src_parts): @@ -453,7 +458,15 @@ def _gen_load_src(chainlet_descriptor: private_types.ChainletAPIDescriptor) -> _ stub_args = [] for name, dep in chainlet_descriptor.dependencies.items(): # `dep.name` is the class name, while `name` is the argument name. - stub_args.append(f"{name}=stub.factory({dep.name}, self._context)") + if framework.is_truss_chainlet(dep.chainlet_cls): + # TrussChainlet deps don't have a generated typed stub. Inject the + # runtime `DeployedServiceDescriptor` so the user can call the + # sibling with their own client (httpx / websockets / ...). + stub_args.append( + f"{name}=self._context.get_service_descriptor({dep.name!r})" + ) + else: + stub_args.append(f"{name}=stub.factory({dep.name}, self._context)") if chainlet_descriptor.has_context: if stub_args: @@ -940,6 +953,37 @@ def gen_truss_model( ) +def _prepare_truss_chainlet_artifact( + chainlet_dir: pathlib.Path, + chainlet_descriptor: private_types.ChainletAPIDescriptor, + dep_services: Mapping[str, private_types.ServiceDescriptor], +) -> pathlib.Path: + """Build a TrussChainlet artifact: copy the user's `truss_dir` to + ``chainlet_dir`` and merge ``model_metadata.chains_metadata`` into its + ``config.yaml``. The user's source files (model.py, custom server, etc.) + are preserved byte-for-byte. Returns ``chainlet_dir``.""" + src_truss_dir = chainlet_descriptor.truss_dir + if src_truss_dir is None or not src_truss_dir.is_dir(): + raise public_types.ChainsUsageError( + f"TrussChainlet `{chainlet_descriptor.name}`: truss_dir is not set " + "or does not exist. Validation should have caught this earlier." + ) + # Replace any pre-existing chainlet_dir from a stale prior generation. + if chainlet_dir.exists(): + shutil.rmtree(chainlet_dir) + truss_path.copy_tree_path(src_truss_dir, chainlet_dir) + + config_path = chainlet_dir / serving_image_builder.CONFIG_FILE + config = truss_config.TrussConfig.from_yaml(config_path) + # Preserve any user-set model_metadata keys; only overwrite chains_metadata. + config.model_metadata = dict(config.model_metadata or {}) + config.model_metadata[private_types.TRUSS_CONFIG_CHAINS_KEY] = ( + private_types.TrussMetadata(chainlet_to_service=dep_services).model_dump() + ) + config.write_to_yaml_file(config_path, verbose=True) + return chainlet_dir + + def gen_truss_chainlet( chain_root: pathlib.Path, chain_name: str, @@ -959,6 +1003,14 @@ def gen_truss_chainlet( f"Code generation for {chainlet_descriptor.chainlet_cls.entity_type} `{chainlet_descriptor.name}` " f"in `{chainlet_dir}`." ) + if chainlet_descriptor.is_truss_chainlet: + # TrussChainlet: copy the user's existing Truss directory as-is and + # merge `model_metadata.chains_metadata` into its `config.yaml` so the + # backend recognizes it as a chain member. No `model.py` generation, + # no chain-root copy. + return _prepare_truss_chainlet_artifact( + chainlet_dir, chainlet_descriptor, dep_services + ) if framework.is_engine_builder_chainlet(chainlet_descriptor.chainlet_cls): engine_builder_config = cast( framework.EngineBuilderChainlet, chainlet_descriptor.chainlet_cls diff --git a/truss-chains/truss_chains/deployment/deployment_client.py b/truss-chains/truss_chains/deployment/deployment_client.py index ec226d641..d03a9719a 100644 --- a/truss-chains/truss_chains/deployment/deployment_client.py +++ b/truss-chains/truss_chains/deployment/deployment_client.py @@ -89,6 +89,11 @@ def _collect_external_package_dirs( seen: set[pathlib.Path] = set() result: list[pathlib.Path] = [] for desc in chainlet_descriptors: + if desc.is_truss_chainlet: + # The user's TrussChainlet has its own packaging (its `truss_dir` + # is bundled separately by Truss's `archive_dir`). Chain-level + # `gather_chain` should not touch it. + continue ext_dirs = desc.chainlet_cls.remote_config.docker_image.external_package_dirs if not ext_dirs: continue diff --git a/truss-chains/truss_chains/framework.py b/truss-chains/truss_chains/framework.py index ffad8905b..516344e84 100644 --- a/truss-chains/truss_chains/framework.py +++ b/truss-chains/truss_chains/framework.py @@ -790,10 +790,18 @@ def _validate_dependency_param( # chainlet class, proper type inference is possible even without annotation. # TODO: `Protocol` is not a proper class and this might be version dependent. # Find a better way to inspect this. + # For TrussChainlet deps, the framework injects a `DeployedServiceDescriptor` + # at runtime (not the chainlet class itself), so allow that annotation too. if not ( param.annotation == inspect.Parameter.empty or utils.issubclass_safe(param.annotation, Protocol) # type: ignore[arg-type] or utils.issubclass_safe(chainlet_cls, param.annotation) + or ( + is_truss_chainlet(chainlet_cls) + and utils.issubclass_safe( + param.annotation, public_types.DeployedServiceDescriptor + ) + ) ): _collect_error( f"The type annotation for `{param.name}` must be a class/subclass of the " @@ -974,6 +982,7 @@ def validate_and_register_cls(cls: Type[private_types.ABCChainlet]) -> None: "ModelBase", "EngineBuilderChainlet", "EngineBuilderLLMChainlet", + "TrussChainlet", ] if cls.__name__ in _skip_class_name: logging.debug(f"Skipping chainlet class validation for `{cls}`.") @@ -983,6 +992,23 @@ def validate_and_register_cls(cls: Type[private_types.ABCChainlet]) -> None: line = inspect.getsourcelines(cls)[1] location = _ErrorLocation(src_path=src_path, line=line, chainlet_name=cls.__name__) + if is_truss_chainlet(cls): + # TrussChainlet declarations don't have remote_config / run_remote / + # health_check / dependencies. Validate the truss_dir, build a + # descriptor with a sentinel endpoint, register, and exit early. + _validate_truss_chainlet_cls(cls, src_path, location) + chainlet_descriptor = private_types.ChainletAPIDescriptor( + chainlet_cls=cls, + dependencies={}, + has_context=False, + endpoint=_DUMMY_ENDPOINT_DESCRIPTOR, + src_path=src_path, + health_check=None, + truss_dir=cls._resolved_truss_dir, + ) + _global_chainlet_registry.register_chainlet(chainlet_descriptor) + return + _validate_display_name(cls, location) _validate_config_class_variable( cls, location, private_types.REMOTE_CONFIG_NAME, public_types.RemoteConfig @@ -1015,6 +1041,68 @@ def validate_and_register_cls(cls: Type[private_types.ABCChainlet]) -> None: _global_chainlet_registry.register_chainlet(chainlet_descriptor) +def _validate_truss_chainlet_cls( + cls: Type[private_types.ABCChainlet], src_path: str, location: _ErrorLocation +) -> None: + """Validate a TrussChainlet declaration's class attributes and resolve + the ``truss_dir``. Sets ``cls._resolved_truss_dir`` on success.""" + # 1. truss_dir class attribute is required (follow MRO so subclasses inherit). + truss_dir_raw = getattr(cls, "truss_dir", None) + if not truss_dir_raw: + _collect_error( + f"`TrussChainlet` subclass `{cls.__name__}` must declare " + "a `truss_dir` class attribute pointing at the wrapped Truss directory.", + _ErrorKind.MISSING_API_ERROR, + location, + ) + return + if not isinstance(truss_dir_raw, (str, pathlib.Path)): + _collect_error( + f"`TrussChainlet.truss_dir` must be a str or pathlib.Path, got " + f"`{type(truss_dir_raw).__name__}`.", + _ErrorKind.TYPE_ERROR, + location, + ) + return + + # 2. Resolve relative to the file where the class was declared. + src_dir = pathlib.Path(src_path).parent + truss_dir = pathlib.Path(truss_dir_raw) + if not truss_dir.is_absolute(): + truss_dir = (src_dir / truss_dir).resolve() + cls._resolved_truss_dir = truss_dir + + # 3. Directory exists. + if not truss_dir.is_dir(): + _collect_error( + f"`TrussChainlet.{cls.__name__}.truss_dir` resolved to " + f"`{truss_dir}`, which is not a directory.", + _ErrorKind.INVALID_CONFIG_ERROR, + location, + ) + return + + # 4. config.yaml is a valid Truss config. + config_path = truss_dir / "config.yaml" + if not config_path.is_file(): + _collect_error( + f"`TrussChainlet.{cls.__name__}.truss_dir` ({truss_dir}) is " + "missing `config.yaml` — not a valid Truss directory.", + _ErrorKind.INVALID_CONFIG_ERROR, + location, + ) + return + try: + truss_config.TrussConfig.from_yaml(config_path) + except Exception as exc: + _collect_error( + f"`TrussChainlet.{cls.__name__}.truss_dir` ({truss_dir}) has " + f"an invalid `config.yaml`: {exc}", + _ErrorKind.INVALID_CONFIG_ERROR, + location, + ) + + # Dependency-Injection / Registry ###################################################### @@ -1691,3 +1779,65 @@ async def run_remote( def is_engine_builder_chainlet(cls: Type[private_types.ABCChainlet]): return issubclass(cls, EngineBuilderChainlet) + + +# TrussChainlet ######################################################################## + + +class TrussChainlet(private_types.ABCChainlet, metaclass=abc.ABCMeta): + """Declares an existing Truss directory as a chain member. + + Unlike ``ChainletBase``, the framework does NOT generate a ``model.py`` + or a typed ``StubBase`` for this declaration. The user's existing Truss + directory (``model.py``-flavored or ``docker_server``) is archived as-is; + only ``model_metadata.chains_metadata`` is merged into its ``config.yaml`` + so the operator's runtime URL injection works for it. + + Subclasses declare the wrapped Truss via a single class attribute:: + + class Whisper(chains.TrussChainlet): + truss_dir = "../whisper-truss" # resolved relative to declaring file + + The chain-graph identity is the class name (``Whisper`` here). To use a + different chainlet name, alias the class. + + Used as a dep in another chainlet's ``__init__``:: + + class Caller(chains.ChainletBase): + def __init__( + self, + whisper: chains.DeployedServiceDescriptor = chains.depends(Whisper), + context: chains.DeploymentContext = chains.depends_context(), + ): + self._url = whisper.target_url + self._headers = whisper.with_auth_headers( + context.get_baseten_api_key()) + """ + + truss_dir: ClassVar[str] + + # `_resolved_truss_dir` is set by `__init_subclass__` after path resolution + # and validation. Codegen reads it to locate the user's Truss directory. + _resolved_truss_dir: ClassVar[Optional[pathlib.Path]] = None + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + cls._framework_config = private_types.FrameworkConfig( + entity_type=private_types.EntityType.TRUSS_CHAINLET, + supports_dependencies=False, + endpoint_method_name="", # no run_remote on TrussChainlet + ) + cls.meta_data = private_types.ChainletMetadata() + validate_and_register_cls(cls) + + +def is_truss_chainlet(cls: Type[private_types.ABCChainlet]) -> bool: + return issubclass(cls, TrussChainlet) + + +def should_skip_codegen(descriptor: private_types.ChainletAPIDescriptor) -> bool: + """Predicate covering both engine-builder and TrussChainlet artifacts — + neither gets a generated `model.py`.""" + return is_engine_builder_chainlet(descriptor.chainlet_cls) or is_truss_chainlet( + descriptor.chainlet_cls + ) diff --git a/truss-chains/truss_chains/private_types.py b/truss-chains/truss_chains/private_types.py index 26b9fe551..3668df795 100644 --- a/truss-chains/truss_chains/private_types.py +++ b/truss-chains/truss_chains/private_types.py @@ -68,6 +68,7 @@ class EntityType(utils.StrEnum): CHAINLET = enum.auto() MODEL = enum.auto() ENGINE_BUILDER_MODEL = enum.auto() + TRUSS_CHAINLET = enum.auto() class FrameworkConfig(custom_types.SafeModelNonSerializable): @@ -242,6 +243,10 @@ class ChainletAPIDescriptor(custom_types.SafeModelNonSerializable): dependencies: Mapping[str, DependencyDescriptor] endpoint: EndpointAPIDescriptor health_check: Optional[HealthCheckAPIDescriptor] + # Set for TRUSS_CHAINLET entity type only; absolute path to the user's + # existing Truss directory. The framework copies this dir as-is (only + # `model_metadata.chains_metadata` is added to its `config.yaml`). + truss_dir: Optional[pathlib.Path] = None def __hash__(self) -> int: return hash(self.chainlet_cls) @@ -254,6 +259,10 @@ def name(self) -> str: def display_name(self) -> str: return self.chainlet_cls.display_name + @property + def is_truss_chainlet(self) -> bool: + return self.chainlet_cls.entity_type == EntityType.TRUSS_CHAINLET + ######################################################################################## From 2aac11d39781c341995c2202ec86b6820e3cae08 Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Wed, 29 Apr 2026 20:40:30 -0700 Subject: [PATCH 04/12] Add TrussChainlet example and tests --- .../examples/truss_chainlet_demo/README.md | 100 +++++ .../examples/truss_chainlet_demo/chain.py | 118 ++++++ .../echo_truss/config.yaml | 10 + .../echo_truss/model/model.py | 12 + .../examples/truss_chainlet_demo/invoke.py | 55 +++ truss-chains/tests/test_truss_chainlet.py | 376 ++++++++++++++++++ 6 files changed, 671 insertions(+) create mode 100644 truss-chains/examples/truss_chainlet_demo/README.md create mode 100644 truss-chains/examples/truss_chainlet_demo/chain.py create mode 100644 truss-chains/examples/truss_chainlet_demo/echo_truss/config.yaml create mode 100644 truss-chains/examples/truss_chainlet_demo/echo_truss/model/model.py create mode 100644 truss-chains/examples/truss_chainlet_demo/invoke.py create mode 100644 truss-chains/tests/test_truss_chainlet.py diff --git a/truss-chains/examples/truss_chainlet_demo/README.md b/truss-chains/examples/truss_chainlet_demo/README.md new file mode 100644 index 000000000..54760d8eb --- /dev/null +++ b/truss-chains/examples/truss_chainlet_demo/README.md @@ -0,0 +1,100 @@ +# TrussChainlet Demo — Composable Chains Project 2 + +A deployable, **CPU-only** chain proving that `chains.ChainletBase` and +`chains.TrussChainlet` coexist in one chain. The entrypoint depends on +**both** kinds at once and calls each via its native API in the same +`run_remote`. + +## What this proves + +Three artifacts in one chain push: + +* **`Caller`** — `ChainletBase` entrypoint. Depends on both `Reverser` and + `EchoTruss`. In a single `run_remote` call it: + 1. invokes `Reverser` via the framework's typed stub + (`await self._reverser.run_remote(text)`), and + 2. invokes `EchoTruss` via raw `httpx`, using Project 1's descriptor + helpers (`target_url` + `with_auth_headers`). +* **`Reverser`** — `ChainletBase` dep. Returns the input reversed. + Codegen produces its `model.py` and a typed stub for callers. +* **`EchoTruss`** — `TrussChainlet` dep. Wraps `./echo_truss/` (a plain + Truss directory) without any rewrite. The user's `model.py` is preserved + byte-for-byte; only `model_metadata.chains_metadata` is added to the + copied `config.yaml`. + +If the framework wires both paths correctly, `match`-style verification: + +``` +input = "hello" +via_chainletbase_reverser → "olleh" # typed stub call worked +via_truss_chainlet_echo → "HELLO" # BYO-httpx call worked +``` + +Both inside one response. That's the guarantee of co-existence. + +## Setup + +```sh +# From the truss repo root +pip install -e . +``` + +`Caller` declares `httpx>=0.27` in its `pip_requirements`. The git-pinned +truss override in each chainlet's `pip_requirements` ships the +Composable Chains framework (Projects 1+2) into every pod. Drop those +overrides once Project 2 ships to PyPI. + +## Push + +```sh +truss chains push --remote matte truss-chains/examples/truss_chainlet_demo/chain.py +``` + +…or `--watch` for live patches as you iterate. + +## Invoke + +```sh +export BASETEN_API_KEY="..." +export CHAIN_URL="https://chain-.api.baseten.co/environments/production/run_remote" +python truss-chains/examples/truss_chainlet_demo/invoke.py "hello" +``` + +Expected output: + +```json +{ + "input": "hello", + "via_chainletbase_reverser": "olleh", + "via_truss_chainlet_echo": "HELLO", + "echo_target_url": "https://...api.baseten.co/.../chainlet/.../run_remote" +} +``` + +The two response paths are different mechanisms hitting different chainlets +in the same chain: + +* `via_chainletbase_reverser` came from `Reverser`'s code-gen'd `model.py`, + reached via the framework stub. +* `via_truss_chainlet_echo` came from `echo_truss/model/model.py` (plain + Truss code, *not* generated), reached via `httpx.post` against the URL + from `desc.target_url`. + +Both succeeded — that's the proof of TrussChainlet's BYO-client integration. + +## File layout + +``` +truss_chainlet_demo/ +├── README.md # this file +├── chain.py # Caller + Reverser + EchoTruss declarations +├── echo_truss/ # plain Truss directory — wrapped by EchoTruss +│ ├── config.yaml +│ └── model/model.py # uppercase echo handler +└── invoke.py # post-deploy invocation script +``` + +Note that `echo_truss/` looks exactly like a standalone `truss push`-able +directory. Project 2's promise is: any such directory can become a chain +member by adding **one class definition** in `chain.py`. No rewrite, no +shape change, no framework imports leaking into the Truss's `model.py`. diff --git a/truss-chains/examples/truss_chainlet_demo/chain.py b/truss-chains/examples/truss_chainlet_demo/chain.py new file mode 100644 index 000000000..82521aad9 --- /dev/null +++ b/truss-chains/examples/truss_chainlet_demo/chain.py @@ -0,0 +1,118 @@ +"""Composable Chains — Project 2 deployable demo. + +Demonstrates that ``chains.ChainletBase`` and ``chains.TrussChainlet`` coexist +in one chain: the entrypoint depends on **both** kinds at once and calls each +via its native API in the same ``run_remote``. + +* ``Reverser`` — typed ``ChainletBase`` dep. Caller invokes it via the + framework's auto-generated ``StubBase`` (``await reverser.run_remote(...)``). +* ``EchoTruss`` — ``TrussChainlet`` wrapping ``./echo_truss/`` (a plain Truss + directory, no rewrite). Caller calls it via raw ``httpx`` using Project 1's + descriptor helpers (``target_url`` + ``with_auth_headers``). + +If the framework wires both paths correctly, the chain returns: + +* ``via_chainletbase_reverser``: input reversed (typed-stub path) +* ``via_truss_chainlet_echo``: input uppercased (BYO-client path) + +Both inside a single ``run_remote`` call. +""" + +import httpx +import pydantic + +import truss_chains as chains + + +class CallerOutput(pydantic.BaseModel): + input: str + reversed_text: str # Reverser's output (the typed-stub call) + reversed_then_uppercased: str # Echo's uppercase of Reverser's output (chained) + echo_target_url: str + + +_TRUSS_OVERRIDE = ( + # Override the published `truss_chains` with the Project 2 branch so + # `chains.TrussChainlet` and the descriptor-injection code path are + # available in every chainlet pod. Drop once shipped to PyPI. + "truss @ git+https://github.com/basetenlabs/truss.git@matte/composable-chains-2-trusschainlet" +) + + +# ----- Plain Truss as a chain member ---------------------------------------- + + +class EchoTruss(chains.TrussChainlet): + """The user's existing Truss directory becomes a chain member with one + declarative line. The framework copies ``./echo_truss/`` byte-for-byte + (only ``model_metadata.chains_metadata`` is added to its config.yaml).""" + + truss_dir = "./echo_truss" + + +# ----- Typed ChainletBase as a chain member --------------------------------- + + +class Reverser(chains.ChainletBase): + """Plain typed chainlet — exists to demonstrate that it coexists with + ``EchoTruss`` in the same chain.""" + + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi"), + docker_image=chains.DockerImage(pip_requirements=[_TRUSS_OVERRIDE]), + ) + + async def run_remote(self, text: str) -> str: + return text[::-1] + + +# ----- Entrypoint: depends on BOTH kinds simultaneously --------------------- + + +@chains.mark_entrypoint("Composable Chains TrussChainlet Demo") +class Caller(chains.ChainletBase): + """Demonstrates ChainletBase / TrussChainlet co-existence by calling + each in the same ``run_remote`` via its native API.""" + + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi"), + docker_image=chains.DockerImage( + pip_requirements=[_TRUSS_OVERRIDE, "httpx>=0.27"] + ), + ) + + def __init__( + self, + # Typed ChainletBase dep — framework injects the auto-generated stub. + # User invokes via ``await reverser.run_remote(text)`` (the existing path). + reverser: Reverser = chains.depends(Reverser), + # TrussChainlet dep — framework injects a DeployedServiceDescriptor. + # User invokes via raw httpx using the Project 1 helpers. + echo: chains.DeployedServiceDescriptor = chains.depends(EchoTruss), + context: chains.DeploymentContext = chains.depends_context(), + ) -> None: + self._reverser = reverser + self._echo_url = echo.target_url + self._echo_headers = echo.with_auth_headers(context.get_baseten_api_key()) + + async def run_remote(self, text: str) -> CallerOutput: + # Step 1: typed call to the ChainletBase dep — same shape as today. + reversed_text = await self._reverser.run_remote(text) + + # Step 2: feed Reverser's output into the TrussChainlet dep via raw httpx. + # Demonstrates data flowing through BOTH dep flavors in one chain hop. + async with httpx.AsyncClient(timeout=60) as client: + response = await client.post( + self._echo_url, + headers=self._echo_headers, + json={"text": reversed_text}, + ) + response.raise_for_status() + reversed_then_uppercased = response.json()["out"] + + return CallerOutput( + input=text, + reversed_text=reversed_text, + reversed_then_uppercased=reversed_then_uppercased, + echo_target_url=self._echo_url, + ) diff --git a/truss-chains/examples/truss_chainlet_demo/echo_truss/config.yaml b/truss-chains/examples/truss_chainlet_demo/echo_truss/config.yaml new file mode 100644 index 000000000..f0c6068d8 --- /dev/null +++ b/truss-chains/examples/truss_chainlet_demo/echo_truss/config.yaml @@ -0,0 +1,10 @@ +model_name: EchoTruss +model_class_filename: model.py +model_class_name: Model +python_version: py311 +requirements: + - truss @ git+https://github.com/basetenlabs/truss.git@matte/composable-chains-2-trusschainlet +resources: + cpu: "1" + memory: 512Mi + use_gpu: false diff --git a/truss-chains/examples/truss_chainlet_demo/echo_truss/model/model.py b/truss-chains/examples/truss_chainlet_demo/echo_truss/model/model.py new file mode 100644 index 000000000..497888fe0 --- /dev/null +++ b/truss-chains/examples/truss_chainlet_demo/echo_truss/model/model.py @@ -0,0 +1,12 @@ +"""Plain Truss — uppercase-echoes the input. Used by `truss_chainlet_demo` +to demonstrate that an existing Truss directory can be a chain member via +``chains.TrussChainlet`` without any rewrite as ``ChainletBase``.""" + + +class Model: + def __init__(self, **kwargs) -> None: + pass + + def predict(self, request: dict) -> dict: + text = str(request.get("text", "")) + return {"out": text.upper()} diff --git a/truss-chains/examples/truss_chainlet_demo/invoke.py b/truss-chains/examples/truss_chainlet_demo/invoke.py new file mode 100644 index 000000000..5988c11d3 --- /dev/null +++ b/truss-chains/examples/truss_chainlet_demo/invoke.py @@ -0,0 +1,55 @@ +"""Invoke the deployed Composable Chains TrussChainlet demo. + +Usage: + BASETEN_API_KEY=... CHAIN_URL=https://chain-.api.baseten.co/.../run_remote \\ + python invoke.py "hello" +""" + +import json +import os +import sys + +import requests + + +def main() -> None: + api_key = os.environ.get("BASETEN_API_KEY") + chain_url = os.environ.get("CHAIN_URL") + if not api_key or not chain_url: + sys.exit( + "Set BASETEN_API_KEY and CHAIN_URL. CHAIN_URL is the chain's " + "`/run_remote` endpoint shown after `chains push`." + ) + + text = sys.argv[1] if len(sys.argv) > 1 else "hello" + response = requests.post( + chain_url, + headers={"Authorization": f"Api-Key {api_key}"}, + json={"text": text}, + timeout=120, + ) + response.raise_for_status() + result = response.json() + print(json.dumps(result, indent=2)) + print() + + expected_reversed = text[::-1] + expected_upper = text.upper() + ok = ( + result.get("via_chainletbase_reverser") == expected_reversed + and result.get("via_truss_chainlet_echo") == expected_upper + ) + if ok: + print( + f"✓ ChainletBase ({expected_reversed!r}) and TrussChainlet " + f"({expected_upper!r}) both worked in one chain." + ) + else: + print( + f"✗ Mismatch: expected reversed={expected_reversed!r}, " + f"upper={expected_upper!r}." + ) + + +if __name__ == "__main__": + main() diff --git a/truss-chains/tests/test_truss_chainlet.py b/truss-chains/tests/test_truss_chainlet.py new file mode 100644 index 000000000..bc932b195 --- /dev/null +++ b/truss-chains/tests/test_truss_chainlet.py @@ -0,0 +1,376 @@ +"""Tests for ``chains.TrussChainlet`` — Project 2 of the Composable Chains +effort. Covers declaration validation, dep validation in another chainlet, +codegen branches, backward compatibility, and adversarial cases. + +The validation layer collects errors rather than raising them — tests use +``framework.raise_validation_errors`` to surface them as a single +``ChainsUsageError``. + +Note on fixture naming: the test fixture is `truss_dir_path` (not `truss_dir`) +to avoid shadowing the `truss_dir` class attribute that `TrussChainlet` +subclasses must declare. Pytest fixture name resolution interacts oddly with +class-body assignments using the same name. +""" + +import pathlib + +import pytest +import yaml + +import truss_chains as chains +from truss_chains import framework, private_types, public_types + +# ---- Fixtures --------------------------------------------------------------- + + +VALID_TRUSS_CONFIG_YAML = """\ +model_name: EchoTruss +model_class_filename: model.py +model_class_name: Model +python_version: py311 +resources: + cpu: '1' + memory: 512Mi + use_gpu: false +""" + + +VALID_MODEL_PY = """\ +class Model: + def __init__(self, **kwargs): + pass + + def predict(self, request): + return {"out": str(request.get("text", "")).upper()} +""" + + +@pytest.fixture +def truss_dir_path(tmp_path: pathlib.Path) -> pathlib.Path: + """A minimal valid Truss directory.""" + d = tmp_path / "echo_truss" + d.mkdir() + (d / "config.yaml").write_text(VALID_TRUSS_CONFIG_YAML) + (d / "model").mkdir() + (d / "model" / "model.py").write_text(VALID_MODEL_PY) + return d + + +@pytest.fixture(autouse=True) +def reset_framework_state(): + """Each test starts with a clean error collector and a snapshot-restored + chainlet registry — test-local classes share names like `_Echo`, which + would otherwise collide on the global registry across tests. We save and + restore the registry contents instead of clearing wholesale, so chainlets + registered at module-import time by *other* test files remain available.""" + framework._global_error_collector.clear() + snapshot_chainlets = dict(framework._global_chainlet_registry._chainlets) + snapshot_names = dict(framework._global_chainlet_registry._name_to_cls) + framework._global_chainlet_registry._chainlets.clear() + framework._global_chainlet_registry._name_to_cls.clear() + try: + yield + finally: + framework._global_error_collector.clear() + framework._global_chainlet_registry._chainlets.clear() + framework._global_chainlet_registry._chainlets.update(snapshot_chainlets) + framework._global_chainlet_registry._name_to_cls.clear() + framework._global_chainlet_registry._name_to_cls.update(snapshot_names) + + +# ---- Declaration validation ------------------------------------------------- + + +def test_declares_truss_dir_required(): + """Subclass without `truss_dir` should record a MISSING_API_ERROR.""" + + class _Bad(chains.TrussChainlet): # noqa: N801 + pass + + with pytest.raises(public_types.ChainsUsageError, match="must declare"): + framework.raise_validation_errors() + + +def test_truss_dir_must_resolve_to_directory(tmp_path): + bad_path = str(tmp_path / "does_not_exist") + + class _Bad(chains.TrussChainlet): + truss_dir = bad_path + + with pytest.raises(public_types.ChainsUsageError, match="not a directory"): + framework.raise_validation_errors() + + +def test_truss_dir_missing_config_yaml(tmp_path): + bad = tmp_path / "bad_truss" + bad.mkdir() + bad_path = str(bad) + + class _Bad(chains.TrussChainlet): + truss_dir = bad_path + + with pytest.raises(public_types.ChainsUsageError, match="missing `config.yaml`"): + framework.raise_validation_errors() + + +def test_truss_dir_invalid_config_yaml(tmp_path): + bad = tmp_path / "bad_truss" + bad.mkdir() + (bad / "config.yaml").write_text("not: a valid: truss config { broken") + bad_path = str(bad) + + class _Bad(chains.TrussChainlet): + truss_dir = bad_path + + with pytest.raises(public_types.ChainsUsageError, match="invalid `config.yaml`"): + framework.raise_validation_errors() + + +def test_truss_dir_resolved_relative_to_declaring_file(): + """Relative `truss_dir` should resolve relative to the file declaring + the subclass — not the cwd.""" + test_dir = pathlib.Path(__file__).parent + rel_truss_dir = test_dir / "_relative_test_truss" + rel_truss_dir.mkdir(exist_ok=True) + try: + (rel_truss_dir / "config.yaml").write_text(VALID_TRUSS_CONFIG_YAML) + (rel_truss_dir / "model").mkdir(exist_ok=True) + (rel_truss_dir / "model" / "model.py").write_text(VALID_MODEL_PY) + + class _Echo(chains.TrussChainlet): + truss_dir = "_relative_test_truss" # relative to this test file + + framework.raise_validation_errors() # should not raise + assert _Echo._resolved_truss_dir == rel_truss_dir.resolve() + finally: + import shutil + + shutil.rmtree(rel_truss_dir, ignore_errors=True) + + +def test_valid_truss_chainlet_registers_descriptor(truss_dir_path): + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Echo) + assert descriptor.is_truss_chainlet + assert descriptor.chainlet_cls is _Echo + assert descriptor.truss_dir == truss_dir_path + # Empty deps + dummy endpoint. + assert descriptor.dependencies == {} + assert descriptor.endpoint == framework._DUMMY_ENDPOINT_DESCRIPTOR + + +# ---- Dep validation in another chainlet ------------------------------------- + + +def test_chainletbase_can_depend_on_truss_chainlet(truss_dir_path): + """A typed ChainletBase can declare a TrussChainlet as a dep with the + DeployedServiceDescriptor type annotation.""" + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + class _Caller(chains.ChainletBase): + async def run_remote(self, x: str) -> str: + return x + + def __init__( + self, + echo: chains.DeployedServiceDescriptor = chains.depends(_Echo), + context: chains.DeploymentContext = chains.depends_context(), + ): + self._echo = echo + + framework.raise_validation_errors() + + +def test_mixed_deps_typed_and_truss(truss_dir_path): + """ChainletBase entrypoint with both a typed ChainletBase dep and a + TrussChainlet dep — both validate.""" + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + class _Reverser(chains.ChainletBase): + async def run_remote(self, x: str) -> str: + return x[::-1] + + class _Caller(chains.ChainletBase): + async def run_remote(self, x: str) -> str: + return x + + def __init__( + self, + reverser: _Reverser = chains.depends(_Reverser), + echo: chains.DeployedServiceDescriptor = chains.depends(_Echo), + context: chains.DeploymentContext = chains.depends_context(), + ): + self._reverser = reverser + self._echo = echo + + framework.raise_validation_errors() + + +def test_truss_chainlet_dep_recorded_on_descriptor(truss_dir_path): + """The dep entry on the entrypoint's descriptor records the TrussChainlet + by class — codegen later branches on `is_truss_chainlet(dep.chainlet_cls)`.""" + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + class _Caller(chains.ChainletBase): + async def run_remote(self, x: str) -> str: + return x + + def __init__( + self, + echo: chains.DeployedServiceDescriptor = chains.depends(_Echo), + context: chains.DeploymentContext = chains.depends_context(), + ): + pass + + framework.raise_validation_errors() + caller_desc = framework.get_descriptor(_Caller) + assert "echo" in caller_desc.dependencies + dep = caller_desc.dependencies["echo"] + assert dep.chainlet_cls is _Echo + assert framework.is_truss_chainlet(dep.chainlet_cls) + + +# ---- Codegen branches ------------------------------------------------------- + + +def test_truss_chainlet_artifact_preserves_user_files(truss_dir_path, tmp_path): + """Generating a TrussChainlet artifact must copy the user's truss_dir + byte-for-byte (model.py preserved exactly), and merge `chains_metadata` + into the copied `config.yaml`.""" + from truss_chains.deployment import code_gen + + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Echo) + chainlet_dir = code_gen.gen_truss_chainlet( + chain_root=tmp_path, chain_name="codegen-test", chainlet_descriptor=descriptor + ) + + # User's model.py is preserved byte-for-byte. + src_model_py = (truss_dir_path / "model" / "model.py").read_text() + dst_model_py = (chainlet_dir / "model" / "model.py").read_text() + assert src_model_py == dst_model_py + + # Generated config.yaml has the chains_metadata block. + dst_config = yaml.safe_load((chainlet_dir / "config.yaml").read_text()) + chains_meta = dst_config["model_metadata"][private_types.TRUSS_CONFIG_CHAINS_KEY] + assert "chainlet_to_service" in chains_meta + # Empty for a leaf TrussChainlet (no nested deps). + assert chains_meta["chainlet_to_service"] == {} + + # User's other config keys preserved. + assert dst_config["model_name"] == "EchoTruss" + assert dst_config["python_version"] == "py311" + + +def test_entrypoint_codegen_emits_descriptor_injection_for_truss_dep( + truss_dir_path, tmp_path +): + """The entrypoint's generated `model.py` must call + `self._context.get_service_descriptor(...)` for TrussChainlet deps, + not `stub.factory(...)`.""" + from truss_chains.deployment import code_gen + + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + class _Caller(chains.ChainletBase): + async def run_remote(self, x: str) -> str: + return x + + def __init__( + self, + echo: chains.DeployedServiceDescriptor = chains.depends(_Echo), + context: chains.DeploymentContext = chains.depends_context(), + ): + self._echo = echo + + framework.raise_validation_errors() + + descriptor = framework.get_descriptor(_Caller) + chain_root = pathlib.Path(__file__).parent + chainlet_dir = code_gen.gen_truss_chainlet( + chain_root=chain_root, chain_name="codegen-test", chainlet_descriptor=descriptor + ) + generated_model_py = (chainlet_dir / "model" / "model.py").read_text() + + # Descriptor injection used. + assert ( + "get_service_descriptor('_Echo')" in generated_model_py + or 'get_service_descriptor("_Echo")' in generated_model_py + ) + # No stub.factory(_Echo, ...) — no typed stub for TrussChainlet deps. + assert "stub.factory(_Echo" not in generated_model_py + # No `class _Echo(stub.StubBase)` — no generated stub class. + assert "class _Echo(stub" not in generated_model_py + + +# ---- Backward compatibility ------------------------------------------------- + + +def test_pure_chainletbase_chain_unaffected(): + """A chain made entirely of ChainletBase chainlets validates identically + to before this change.""" + + class _Inner(chains.ChainletBase): + async def run_remote(self, x: str) -> str: + return x + + class _Outer(chains.ChainletBase): + async def run_remote(self, x: str) -> str: + return x + + def __init__( + self, + inner: _Inner = chains.depends(_Inner), + context: chains.DeploymentContext = chains.depends_context(), + ): + self._inner = inner + + framework.raise_validation_errors() # no errors + inner_desc = framework.get_descriptor(_Inner) + outer_desc = framework.get_descriptor(_Outer) + assert not inner_desc.is_truss_chainlet + assert not outer_desc.is_truss_chainlet + assert "inner" in outer_desc.dependencies + assert not framework.is_truss_chainlet( + outer_desc.dependencies["inner"].chainlet_cls + ) + + +# ---- Adversarial ------------------------------------------------------------ + + +def test_depends_on_random_class_rejected(): + class _Random: + pass + + class _Caller(chains.ChainletBase): + async def run_remote(self) -> None: + pass + + def __init__(self, x=chains.depends(_Random)): # type: ignore[arg-type] + pass + + with pytest.raises(public_types.ChainsUsageError): + framework.raise_validation_errors() From 94e1c341f5d43870c3732c83b1dcaddd04a6b89e Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Fri, 1 May 2026 13:04:01 -0700 Subject: [PATCH 05/12] Fix model name validation which causes promotion issues --- truss-chains/tests/test_truss_chainlet.py | 107 +++++++++++++++++- .../truss_chains/deployment/code_gen.py | 21 +++- truss-chains/truss_chains/framework.py | 5 +- 3 files changed, 127 insertions(+), 6 deletions(-) diff --git a/truss-chains/tests/test_truss_chainlet.py b/truss-chains/tests/test_truss_chainlet.py index bc932b195..887f70d9f 100644 --- a/truss-chains/tests/test_truss_chainlet.py +++ b/truss-chains/tests/test_truss_chainlet.py @@ -276,8 +276,11 @@ class _Echo(chains.TrussChainlet): # Empty for a leaf TrussChainlet (no nested deps). assert chains_meta["chainlet_to_service"] == {} - # User's other config keys preserved. - assert dst_config["model_name"] == "EchoTruss" + # `model_name` is overwritten with the chain-uniquified name; user's + # literal `EchoTruss` is preserved in the source `truss_dir` but not in + # the generated artifact (Project 2.5 fix). Other user-set config keys + # remain preserved. + assert dst_config["model_name"] != "EchoTruss" assert dst_config["python_version"] == "py311" @@ -374,3 +377,103 @@ def __init__(self, x=chains.depends(_Random)): # type: ignore[arg-type] with pytest.raises(public_types.ChainsUsageError): framework.raise_validation_errors() + + +# ---- Project 2.5: TrussChainlet artifact parity with ChainletBase ----------- + + +def test_truss_chainlet_artifact_overrides_user_model_name(truss_dir_path, tmp_path): + """The codegen artifact must carry the chain-uniquified model_name, not + the user's literal `truss_dir/config.yaml.model_name`. This avoids + workspace name collisions on `chains push` (Issue 1) and `OracleVersion.name` + regex failures during UI promotion (Issue 5).""" + from truss_chains.deployment import code_gen + + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Echo) + chainlet_dir = code_gen.gen_truss_chainlet( + chain_root=tmp_path, + chain_name="codegen-test", + chainlet_descriptor=descriptor, + model_name="_Echo-abc12345", + ) + + dst_config = yaml.safe_load((chainlet_dir / "config.yaml").read_text()) + assert dst_config["model_name"] == "_Echo-abc12345" + + # User's source file is untouched. + src_config = yaml.safe_load((truss_dir_path / "config.yaml").read_text()) + assert src_config["model_name"] == "EchoTruss" + + +def test_truss_chainlet_artifact_auto_adds_chain_api_key(truss_dir_path, tmp_path): + """The codegen artifact must include `baseten_chain_api_key` in its + `secrets` map so any TrussChainlet that calls a sibling via the chain + RPC URL has the chain-internal key available (Issue 2). The ChainletBase + path auto-adds this; the TrussChainlet path must too.""" + from truss_chains.deployment import code_gen + + truss_dir_str = str(truss_dir_path) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Echo) + chainlet_dir = code_gen.gen_truss_chainlet( + chain_root=tmp_path, + chain_name="codegen-test", + chainlet_descriptor=descriptor, + model_name="_Echo-abc12345", + ) + + dst_config = yaml.safe_load((chainlet_dir / "config.yaml").read_text()) + assert public_types.CHAIN_API_KEY_SECRET_NAME in dst_config["secrets"] + assert ( + dst_config["secrets"][public_types.CHAIN_API_KEY_SECRET_NAME] + == public_types.SECRET_DUMMY + ) + + +def test_truss_chainlet_artifact_chain_api_key_idempotent(tmp_path): + """If the user's truss `config.yaml` already lists `baseten_chain_api_key` + in its `secrets` map, the codegen must not overwrite it — same idempotency + contract as the ChainletBase path.""" + from truss_chains.deployment import code_gen + + # User truss with chain_api_key already declared. + user_config = ( + VALID_TRUSS_CONFIG_YAML + + f"secrets:\n {public_types.CHAIN_API_KEY_SECRET_NAME}: null\n" + ) + truss_dir = tmp_path / "user_truss" + truss_dir.mkdir() + (truss_dir / "config.yaml").write_text(user_config) + (truss_dir / "model").mkdir() + (truss_dir / "model" / "model.py").write_text(VALID_MODEL_PY) + + truss_dir_str = str(truss_dir) + + class _Echo(chains.TrussChainlet): + truss_dir = truss_dir_str + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Echo) + chainlet_dir = code_gen.gen_truss_chainlet( + chain_root=tmp_path, + chain_name="codegen-test", + chainlet_descriptor=descriptor, + model_name="_Echo-abc12345", + ) + + dst_config = yaml.safe_load((chainlet_dir / "config.yaml").read_text()) + # User's literal value (None) preserved — codegen did NOT overwrite to '***'. + assert ( + dst_config["secrets"][public_types.CHAIN_API_KEY_SECRET_NAME] + != public_types.SECRET_DUMMY + ) diff --git a/truss-chains/truss_chains/deployment/code_gen.py b/truss-chains/truss_chains/deployment/code_gen.py index 3a7862da3..ccbdc2b8a 100644 --- a/truss-chains/truss_chains/deployment/code_gen.py +++ b/truss-chains/truss_chains/deployment/code_gen.py @@ -957,11 +957,15 @@ def _prepare_truss_chainlet_artifact( chainlet_dir: pathlib.Path, chainlet_descriptor: private_types.ChainletAPIDescriptor, dep_services: Mapping[str, private_types.ServiceDescriptor], + model_name: str, ) -> pathlib.Path: """Build a TrussChainlet artifact: copy the user's `truss_dir` to ``chainlet_dir`` and merge ``model_metadata.chains_metadata`` into its ``config.yaml``. The user's source files (model.py, custom server, etc.) - are preserved byte-for-byte. Returns ``chainlet_dir``.""" + are preserved byte-for-byte. ``model_name`` (the chain-uniquified name + matching the ChainletBase form ``-``) overwrites + the user's literal ``config.yaml.model_name`` so workspace/promote-time + name uniqueness is satisfied. Returns ``chainlet_dir``.""" src_truss_dir = chainlet_descriptor.truss_dir if src_truss_dir is None or not src_truss_dir.is_dir(): raise public_types.ChainsUsageError( @@ -975,6 +979,16 @@ def _prepare_truss_chainlet_artifact( config_path = chainlet_dir / serving_image_builder.CONFIG_FILE config = truss_config.TrussConfig.from_yaml(config_path) + config.model_name = model_name + if public_types.CHAIN_API_KEY_SECRET_NAME not in config.secrets: + config.secrets[public_types.CHAIN_API_KEY_SECRET_NAME] = ( + public_types.SECRET_DUMMY + ) + else: + logging.info( + f"Chains automatically add {public_types.CHAIN_API_KEY_SECRET_NAME} " + "to secrets - no need to manually add it." + ) # Preserve any user-set model_metadata keys; only overwrite chains_metadata. config.model_metadata = dict(config.model_metadata or {}) config.model_metadata[private_types.TRUSS_CONFIG_CHAINS_KEY] = ( @@ -1009,7 +1023,10 @@ def gen_truss_chainlet( # backend recognizes it as a chain member. No `model.py` generation, # no chain-root copy. return _prepare_truss_chainlet_artifact( - chainlet_dir, chainlet_descriptor, dep_services + chainlet_dir, + chainlet_descriptor, + dep_services, + model_name=model_name or chain_name, ) if framework.is_engine_builder_chainlet(chainlet_descriptor.chainlet_cls): engine_builder_config = cast( diff --git a/truss-chains/truss_chains/framework.py b/truss-chains/truss_chains/framework.py index 516344e84..6bd7cb4f3 100644 --- a/truss-chains/truss_chains/framework.py +++ b/truss-chains/truss_chains/framework.py @@ -994,8 +994,9 @@ def validate_and_register_cls(cls: Type[private_types.ABCChainlet]) -> None: if is_truss_chainlet(cls): # TrussChainlet declarations don't have remote_config / run_remote / - # health_check / dependencies. Validate the truss_dir, build a - # descriptor with a sentinel endpoint, register, and exit early. + # health_check / dependencies. Validate the truss_dir + display_name, + # build a descriptor with a sentinel endpoint, register, and exit early. + _validate_display_name(cls, location) _validate_truss_chainlet_cls(cls, src_path, location) chainlet_descriptor = private_types.ChainletAPIDescriptor( chainlet_cls=cls, From f2ddba2df0e9208a1f50fa48f0b4be102ee578a6 Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Sat, 2 May 2026 00:07:07 -0700 Subject: [PATCH 06/12] Mock voice agent example --- .../examples/voice_agent_mocked/README.md | 69 +++++ .../examples/voice_agent_mocked/chain.py | 248 ++++++++++++++++++ .../examples/voice_agent_mocked/client.py | 42 +++ .../voice_agent_mocked/llm_mock/config.yaml | 8 + .../llm_mock/model/__init__.py | 0 .../llm_mock/model/model.py | 15 ++ .../voice_agent_mocked/stt_mock/config.yaml | 11 + .../stt_mock/model/__init__.py | 0 .../stt_mock/model/model.py | 34 +++ .../voice_agent_mocked/tts_mock/config.yaml | 11 + .../tts_mock/model/__init__.py | 0 .../tts_mock/model/model.py | 35 +++ 12 files changed, 473 insertions(+) create mode 100644 truss-chains/examples/voice_agent_mocked/README.md create mode 100644 truss-chains/examples/voice_agent_mocked/chain.py create mode 100644 truss-chains/examples/voice_agent_mocked/client.py create mode 100644 truss-chains/examples/voice_agent_mocked/llm_mock/config.yaml create mode 100644 truss-chains/examples/voice_agent_mocked/llm_mock/model/__init__.py create mode 100644 truss-chains/examples/voice_agent_mocked/llm_mock/model/model.py create mode 100644 truss-chains/examples/voice_agent_mocked/stt_mock/config.yaml create mode 100644 truss-chains/examples/voice_agent_mocked/stt_mock/model/__init__.py create mode 100644 truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py create mode 100644 truss-chains/examples/voice_agent_mocked/tts_mock/config.yaml create mode 100644 truss-chains/examples/voice_agent_mocked/tts_mock/model/__init__.py create mode 100644 truss-chains/examples/voice_agent_mocked/tts_mock/model/model.py diff --git a/truss-chains/examples/voice_agent_mocked/README.md b/truss-chains/examples/voice_agent_mocked/README.md new file mode 100644 index 000000000..965897516 --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/README.md @@ -0,0 +1,69 @@ +# voice_agent_mocked — Composable Chains centerpiece example + +A CPU-only mock of the FDE voice-agent chain, designed for **fast iteration on chain-internal routing behavior** without paying for GPU model deploys. Same shape as `/Users/mattelim/Documents/fde/e2e-voice/`, but every chainlet is a few lines of Python returning canned data. + +## What it shows + +| Chainlet | Kind | Transport | Body | +|---|---|---|---| +| `STTMock` | `chains.TrussChainlet` | WebSocket | bytes → `{"text": "text-{N}"}` | +| `LLMMock` | `chains.TrussChainlet` | HTTP `predict` | `{"prompt"}` → `{"completion": "echo: ..."}` | +| `TTSMock` | `chains.TrussChainlet` | WebSocket | text → `text.encode()` | +| `Orchestrator` | `chains.ChainletBase` | WebSocket entrypoint | bytes in → drives STT→LLM→TTS → bytes out | + +The orchestrator uses the same patterns as the FDE chain: + +- GraphQL lookup of sibling oracle IDs (workaround until `dynamic_chainlet_config` exposes them — see `COMPOSABLE_CHAINS_PROPOSAL.md` "Discovered gap"). +- Strip `HTTP_PROXY` / `HTTPS_PROXY` env vars at startup so outbound WS connections don't get intercepted. +- Build per-sibling URLs via `model-.api.baseten.co//`. +- Use `httpx` (HTTP) and `websockets` (WS) with `Authorization: Api-Key`. + +## Push + +```sh +truss chains push --remote matte chain.py --watch +``` + +`--watch` keeps the orchestrator hot-patchable; it's the right mode for iterating on `chain.py` against a running chain. + +A workspace secret named **`baseten_api_key`** must exist in your Baseten workspace; the chain reads it (declared via `Assets(secret_keys=["baseten_api_key"])`) for the GraphQL lookup. Set it once via the Baseten UI: . + +## Invoke + +After the push reports the entrypoint URL: + +```sh +export BASETEN_API_KEY=... +export CHAIN_URL="wss://chain-.api.baseten.co/development/websocket" +python client.py +``` + +Expected output: + +``` +reply (bytes, 15 bytes): +b'echo: text-1024' + +✓ Match: b'echo: text-1024' +``` + +End-to-end the chain has done: send 1024 bytes → STT says `"text-1024"` → LLM says `"echo: text-1024"` → TTS encodes that as 15 bytes → returned to client. + +## Why this exists + +When iterating on **baseten-local** (chain gateway routing, operator VirtualService registration, dynamic config injection, promotion flow), the FDE chain is too slow a reproducer because its STT/LLM/TTS take 10–20 minutes per redeploy on H100s. This example reproduces the same chain shape — including the WS-fronted non-entrypoint chainlet pattern that surfaces the platform routing gap — in ≈2 minutes total deploy time. + +Specifically, it's a fast probe for: + +- **Issue 5/5b parity** (auto-uniquified `model_name` reaches the artifact, no name collisions on publish). +- **Issue 2 parity** (chain-internal `baseten_chain_api_key` auto-added to TrussChainlet artifacts). +- **The non-entrypoint WS routing gap** — does `wss://model-.api.baseten.co/development/websocket` route to a non-entrypoint WS-fronted chainlet? If yes, the orchestrator's pipeline succeeds end-to-end. If no, the WS connect fails before LLM is reached. +- **Promotion flow** — `chains push --environment ` and Baseten-UI promotion against a tiny chain that completes in seconds. + +## Cleanup + +The chain stays cheap (CPU-only, scales to zero), so it's fine to leave running between iterations. To delete: + +```sh +# Via the Baseten dashboard chain page → "Delete chain" +``` diff --git a/truss-chains/examples/voice_agent_mocked/chain.py b/truss-chains/examples/voice_agent_mocked/chain.py new file mode 100644 index 000000000..bf821a30f --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/chain.py @@ -0,0 +1,248 @@ +"""Voice Agent (Mocked) — Composable Chains centerpiece example. + +A CPU-only mock of the FDE voice-agent chain, designed to validate routing +behavior end-to-end without the cost of GPU model deploys. Same shape as +``/Users/mattelim/Documents/fde/e2e-voice``: + +- ``STTMock`` (``TrussChainlet``, WS-fronted): bytes ``→`` ``{"text": "text-{N}"}`` +- ``LLMMock`` (``TrussChainlet``, HTTP): ``{"prompt"}`` ``→`` ``{"completion": "echo: ..."}`` +- ``TTSMock`` (``TrussChainlet``, WS-fronted): text ``→`` ``text.encode()`` +- ``Orchestrator`` (``ChainletBase``, WS entrypoint): a single round-trip wires + STT ``→`` LLM ``→`` TTS using the same descriptor / runtime patterns the FDE + chain uses (GraphQL oracle-id resolution + proxy-env strip). + +Iteration loop: + + truss chains push --remote matte chain.py --watch + +The orchestrator's ``run_remote`` reads bytes from the entrypoint WS, drives +the pipeline, and sends back the synthesized bytes. Useful as a regression +fixture for any baseten-local change that touches chain-internal routing. +""" + +import json +import logging +from urllib.parse import urlparse + +import requests + +import truss_chains as chains + +log = logging.getLogger("voice-agent-mocked") + +_TRUSS_OVERRIDE = ( + "truss @ git+https://github.com/basetenlabs/truss.git@matte/composable-chains" +) + + +# ----- Mock TrussChainlet members -------------------------------------------- + + +class STTMock(chains.TrussChainlet): + truss_dir = "./stt_mock" + + +class LLMMock(chains.TrussChainlet): + truss_dir = "./llm_mock" + + +class TTSMock(chains.TrussChainlet): + truss_dir = "./tts_mock" + + +# ----- URL resolution (workaround for the platform routing gap) -------------- +# +# Same approach as the FDE chain: query the dashboard GraphQL with the +# user-supplied baseten_api_key to map chainlet name -> oracle.id, then build +# `model-.api.baseten.co/...` URLs. Replace with the descriptor's +# native fields once the platform fix lands. + +_GRAPHQL_QUERY = ( + "query Chain($id: String!) {\n" + " chain(id: $id) {\n" + " deployments {\n" + " id\n" + " chainlets {\n" + " name oracle { id } oracle_version { id transport_kind }\n" + " }\n" + " }\n" + " }\n" + "}" +) + + +def _chain_id_from_descriptor(desc: chains.DeployedServiceDescriptor) -> str: + """Extract chain id from ``internal_url.hostname`` (``chain-...``).""" + hostname = desc.internal_url.hostname + first = hostname.split(".", 1)[0] + if not first.startswith("chain-"): + raise RuntimeError(f"Unexpected internal_url hostname shape: {hostname!r}") + return first[len("chain-") :] + + +def _chain_deployment_id_from_descriptor(desc: chains.DeployedServiceDescriptor) -> str: + """Extract chain deployment id from the descriptor's gateway URL path.""" + parts = urlparse(desc.internal_url.gateway_run_remote_url).path.split("/") + # Path: /deployment//chainlet//run_remote + return parts[parts.index("deployment") + 1] + + +def _resolve_oracles( + api_key: str, chain_id: str, chain_deployment_id: str +) -> dict[str, dict]: + r = requests.post( + "https://app.baseten.co/graphql/", + headers={ + "Authorization": f"Api-Key {api_key}", + "Content-Type": "application/json", + }, + json={"query": _GRAPHQL_QUERY, "variables": {"id": chain_id}}, + timeout=10, + ) + r.raise_for_status() + payload = r.json() + chain_data = (payload.get("data") or {}).get("chain") or {} + deployments = chain_data.get("deployments") or [] + target = next((d for d in deployments if d["id"] == chain_deployment_id), None) + if target is None: + raise RuntimeError( + f"GraphQL has no deployment {chain_deployment_id} on chain {chain_id}" + ) + return { + c["name"]: { + "oracle_id": c["oracle"]["id"], + "oracle_version_id": c["oracle_version"]["id"], + "transport": c["oracle_version"]["transport_kind"], + } + for c in target["chainlets"] + } + + +def _model_url(oracle_id: str, oracle_version_id: str, scheme: str, path: str) -> str: + """Build a per-chainlet sibling URL using the published-deployment shape. + + ``://model-.api.baseten.co/deployment//`` + + Cluster-internally this shape works for both draft and published chainlets + (verified empirically on chain ``2328913l``); the equivalent + ``/development/`` alias is unnecessary. Using the deployment shape + uniformly keeps URLs stable across draft → published lifecycles. + """ + return ( + f"{scheme}://model-{oracle_id}.api.baseten.co" + f"/deployment/{oracle_version_id}/{path}" + ) + + +# ----- Orchestrator: WebSocket entrypoint ------------------------------------ + + +@chains.mark_entrypoint("VoiceAgentMocked") +class Orchestrator(chains.ChainletBase): + """Single-round-trip WS entrypoint: bytes in -> bytes out, exercising the + same STT/LLM/TTS sibling-call patterns as the FDE chain.""" + + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi"), + docker_image=chains.DockerImage( + base_image=chains.BasetenImage.PY311, + pip_requirements=[_TRUSS_OVERRIDE, "httpx>=0.27", "websockets>=12"], + ), + assets=chains.Assets(secret_keys=["baseten_api_key"]), + ) + + def __init__( + self, + stt: chains.DeployedServiceDescriptor = chains.depends(STTMock), + llm: chains.DeployedServiceDescriptor = chains.depends(LLMMock), + tts: chains.DeployedServiceDescriptor = chains.depends(TTSMock), + context: chains.DeploymentContext = chains.depends_context(), + ) -> None: + # Strip Baseten's internal HTTP_PROXY env vars — they intercept outbound + # WS connections from inside chainlet pods (same mitigation FDE applies). + import os + + for var in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy"): + os.environ.pop(var, None) + os.environ["NO_PROXY"] = "*" + + try: + user_api_key = context.secrets["baseten_api_key"] + except KeyError as e: + raise RuntimeError( + "baseten_api_key secret is missing. Set it in the workspace and" + " re-deploy. Used to resolve sibling oracle IDs via dashboard" + " GraphQL until the platform exposes them via dynamic_chainlet_config." + ) from e + if not user_api_key or user_api_key == "***": + raise RuntimeError("baseten_api_key secret is empty/placeholder.") + + chain_id = _chain_id_from_descriptor(stt) + chain_deployment_id = _chain_deployment_id_from_descriptor(stt) + oracles = _resolve_oracles(user_api_key, chain_id, chain_deployment_id) + log.warning( + "Resolved oracle map for chain=%s deployment=%s: %s", + chain_id, + chain_deployment_id, + oracles, + ) + + self._stt_url = _model_url( + oracles["STTMock"]["oracle_id"], + oracles["STTMock"]["oracle_version_id"], + "wss", + "websocket", + ) + self._llm_url = _model_url( + oracles["LLMMock"]["oracle_id"], + oracles["LLMMock"]["oracle_version_id"], + "https", + "predict", + ) + self._tts_url = _model_url( + oracles["TTSMock"]["oracle_id"], + oracles["TTSMock"]["oracle_version_id"], + "wss", + "websocket", + ) + self._headers = {"Authorization": f"Api-Key {user_api_key}"} + + async def run_remote(self, websocket: chains.WebSocketProtocol) -> None: + import httpx + import websockets as ws_client + + try: + audio_in = await websocket.receive_bytes() + + # 1. STT: bytes -> text + async with ws_client.connect( + self._stt_url, additional_headers=self._headers + ) as sws: + await sws.send(audio_in) + resp = await sws.recv() + text = json.loads(resp)["text"] + log.warning("STTMock returned: %s", text) + + # 2. LLM: text -> completion + async with httpx.AsyncClient(timeout=30, headers=self._headers) as http: + r = await http.post(self._llm_url, json={"prompt": text}) + r.raise_for_status() + completion = r.json()["completion"] + log.warning("LLMMock returned: %s", completion) + + # 3. TTS: text -> bytes + async with ws_client.connect( + self._tts_url, additional_headers=self._headers + ) as tws: + await tws.send(json.dumps({})) # mock config frame, ignored + await tws.send(completion) + audio_out = await tws.recv() + log.warning("TTSMock returned %d bytes", len(audio_out)) + + await websocket.send_bytes(audio_out) + except Exception as e: + log.exception("Orchestrator pipeline failed: %s", e) + try: + await websocket.send_text(json.dumps({"error": str(e)})) + except Exception: + pass diff --git a/truss-chains/examples/voice_agent_mocked/client.py b/truss-chains/examples/voice_agent_mocked/client.py new file mode 100644 index 000000000..846eb835f --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/client.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +"""Local invocation script for the deployed voice_agent_mocked chain. + +Usage: + + export BASETEN_API_KEY=... + export CHAIN_URL="wss://chain-.api.baseten.co/development/websocket" + python client.py + +Sends a 1024-byte buffer to the entrypoint and prints what comes back. Expected: +the round-trip ``Whisper → LLM → TTS`` produces ``b"echo: text-1024"``. +""" + +import asyncio +import os +import sys + +import websockets + + +async def main() -> None: + url = os.environ.get("CHAIN_URL") + api_key = os.environ.get("BASETEN_API_KEY") + if not url or not api_key: + print("Error: CHAIN_URL and BASETEN_API_KEY must be set.", file=sys.stderr) + sys.exit(1) + + headers = {"Authorization": f"Api-Key {api_key}"} + async with websockets.connect(url, additional_headers=headers) as ws: + await ws.send(b"a" * 1024) + reply = await ws.recv() + print(f"reply ({type(reply).__name__}, {len(reply)} bytes):") + print(repr(reply)) + expected = b"echo: text-1024" # STT-LLM-TTS mock pipeline output + if reply == expected: + print(f"\n✓ Match: {expected!r}") + else: + print(f"\n✗ Mismatch — expected {expected!r}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/truss-chains/examples/voice_agent_mocked/llm_mock/config.yaml b/truss-chains/examples/voice_agent_mocked/llm_mock/config.yaml new file mode 100644 index 000000000..ba3eb959d --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/llm_mock/config.yaml @@ -0,0 +1,8 @@ +model_name: LLMMock +model_class_filename: model.py +model_class_name: Model +python_version: py311 +resources: + cpu: "1" + memory: 256Mi + use_gpu: false diff --git a/truss-chains/examples/voice_agent_mocked/llm_mock/model/__init__.py b/truss-chains/examples/voice_agent_mocked/llm_mock/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/examples/voice_agent_mocked/llm_mock/model/model.py b/truss-chains/examples/voice_agent_mocked/llm_mock/model/model.py new file mode 100644 index 000000000..19d5af11e --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/llm_mock/model/model.py @@ -0,0 +1,15 @@ +"""LLMMock — minimal CPU-only HTTP-fronted Truss that mocks an LLM completion. + +POST /predict with `{"prompt": "..."}` → returns `{"completion": "echo: ..."}`. +Stands in for the real OpenAI-compatible LLM in the FDE voice agent without +the OpenAI SDK / TRT-LLM weight overhead. +""" + + +class Model: + def __init__(self, **kwargs) -> None: + pass + + def predict(self, request: dict) -> dict: + prompt = str(request.get("prompt", "")) + return {"completion": f"echo: {prompt}"} diff --git a/truss-chains/examples/voice_agent_mocked/stt_mock/config.yaml b/truss-chains/examples/voice_agent_mocked/stt_mock/config.yaml new file mode 100644 index 000000000..a2b22a59b --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/stt_mock/config.yaml @@ -0,0 +1,11 @@ +model_name: STTMock +model_class_filename: model.py +model_class_name: Model +python_version: py311 +resources: + cpu: "1" + memory: 256Mi + use_gpu: false +runtime: + transport: + kind: websocket diff --git a/truss-chains/examples/voice_agent_mocked/stt_mock/model/__init__.py b/truss-chains/examples/voice_agent_mocked/stt_mock/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py b/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py new file mode 100644 index 000000000..c1dcd9300 --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py @@ -0,0 +1,34 @@ +"""STTMock — minimal CPU-only WS-fronted Truss that mocks streaming STT. + +Protocol (mirrors Baseten Whisper STT shape, but trivial): +- Client connects WS, optionally sends a JSON metadata frame (ignored). +- Client sends audio chunks (bytes). +- For each chunk, server replies with `{"text": "text-{N}"}` where N is the + byte length of the chunk. +""" + +import json + + +class Model: + def __init__(self, **kwargs) -> None: + pass + + async def websocket(self, ws): + try: + while True: + # fastapi.WebSocket.receive() returns a dict like + # {"type": "websocket.receive", "bytes": ...} or + # {"type": "websocket.receive", "text": ...} — + # use the typed helpers instead. + msg = await ws.receive() + if msg.get("type") == "websocket.disconnect": + return + if "bytes" in msg and msg["bytes"] is not None: + audio = msg["bytes"] + await ws.send_text( + json.dumps({"text": f"text-{len(audio)}"}) + ) + # Text frames (e.g. JSON metadata handshake) are silently consumed. + except Exception: + return diff --git a/truss-chains/examples/voice_agent_mocked/tts_mock/config.yaml b/truss-chains/examples/voice_agent_mocked/tts_mock/config.yaml new file mode 100644 index 000000000..9d570664a --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/tts_mock/config.yaml @@ -0,0 +1,11 @@ +model_name: TTSMock +model_class_filename: model.py +model_class_name: Model +python_version: py311 +resources: + cpu: "1" + memory: 256Mi + use_gpu: false +runtime: + transport: + kind: websocket diff --git a/truss-chains/examples/voice_agent_mocked/tts_mock/model/__init__.py b/truss-chains/examples/voice_agent_mocked/tts_mock/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/examples/voice_agent_mocked/tts_mock/model/model.py b/truss-chains/examples/voice_agent_mocked/tts_mock/model/model.py new file mode 100644 index 000000000..abc9bf270 --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/tts_mock/model/model.py @@ -0,0 +1,35 @@ +"""TTSMock — minimal CPU-only WS-fronted Truss that mocks streaming TTS. + +Protocol (mirrors Baseten Orpheus TTS shape, but trivial): +- Client connects WS. +- Client sends a JSON config frame (ignored). +- Client sends a text frame containing the text to synthesize. +- Server replies with `.encode()` as a single bytes frame, then closes. +""" + +import json + + +class Model: + def __init__(self, **kwargs) -> None: + pass + + async def websocket(self, ws): + try: + while True: + msg = await ws.receive() + if msg.get("type") == "websocket.disconnect": + return + text = msg.get("text") + if text is None: + continue + # If the text frame is a JSON config blob, swallow and wait for + # the actual text-to-synthesize frame. + try: + json.loads(text) + continue + except (json.JSONDecodeError, ValueError): + pass + await ws.send_bytes(text.encode()) + except Exception: + return From 6ffb21d60f8f1041ae0b7a93a4891cf38a8f6df6 Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Mon, 4 May 2026 15:49:55 -0700 Subject: [PATCH 07/12] Fix pre-commit retroactively --- .../README.md | 8 ++--- .../bring_your_own_client/__init__.py | 0 .../chain.py | 0 .../invoke.py | 2 +- .../examples/truss_chainlet_demo/__init__.py | 0 .../examples/truss_chainlet_demo/chain.py | 4 +-- .../examples/voice_agent_mocked/__init__.py | 0 .../examples/voice_agent_mocked/chain.py | 9 ++++++ .../stt_mock/model/model.py | 4 +-- .../tests/numpy_and_binary/__init__.py | 0 .../02_raw_truss_reading_siblings/README.md | 2 +- .../tests/runtime_discovery/README.md | 4 +-- .../README.md | 0 .../basic_descriptor_access/__init__.py | 0 .../chain.py | 0 .../test_descriptor_helpers.py | 0 truss-chains/truss_chains/framework.py | 8 +++-- truss-chains/truss_chains/public_api.py | 30 ++++++++++++++++++- 18 files changed, 53 insertions(+), 18 deletions(-) rename truss-chains/examples/{bring-your-own-client => bring_your_own_client}/README.md (92%) create mode 100644 truss-chains/examples/bring_your_own_client/__init__.py rename truss-chains/examples/{bring-your-own-client => bring_your_own_client}/chain.py (100%) rename truss-chains/examples/{bring-your-own-client => bring_your_own_client}/invoke.py (94%) create mode 100644 truss-chains/examples/truss_chainlet_demo/__init__.py create mode 100644 truss-chains/examples/voice_agent_mocked/__init__.py create mode 100644 truss-chains/tests/numpy_and_binary/__init__.py rename truss-chains/tests/runtime_discovery/{01_basic_descriptor_access => basic_descriptor_access}/README.md (100%) create mode 100644 truss-chains/tests/runtime_discovery/basic_descriptor_access/__init__.py rename truss-chains/tests/runtime_discovery/{01_basic_descriptor_access => basic_descriptor_access}/chain.py (100%) rename truss-chains/tests/runtime_discovery/{01_basic_descriptor_access => basic_descriptor_access}/test_descriptor_helpers.py (100%) diff --git a/truss-chains/examples/bring-your-own-client/README.md b/truss-chains/examples/bring_your_own_client/README.md similarity index 92% rename from truss-chains/examples/bring-your-own-client/README.md rename to truss-chains/examples/bring_your_own_client/README.md index 59ed9e6b2..661dea255 100644 --- a/truss-chains/examples/bring-your-own-client/README.md +++ b/truss-chains/examples/bring_your_own_client/README.md @@ -20,7 +20,7 @@ If Project 1's helpers correctly mirror what `BasetenSession.__init__` would construct internally, both paths return the same string. The chain returns both results plus the descriptor metadata so you can see them line up. -This is the **bring-your-own-client** pattern: keep the chain's UI grouping, +This is the **bring_your_own_client** pattern: keep the chain's UI grouping, atomic deploy, and per-pod sibling-URL discovery, but plug in any HTTP / WS client you like (httpx, websockets, grpc.aio, …). @@ -36,7 +36,7 @@ truss login # if not already authenticated ## Push ```sh -chains push truss-chains/examples/bring-your-own-client/chain.py +chains push truss-chains/examples/bring_your_own_client/chain.py ``` The CLI prints the chain's invoke URL when the deploy completes. Copy the @@ -47,7 +47,7 @@ The CLI prints the chain's invoke URL when the deploy completes. Copy the ```sh export BASETEN_API_KEY="..." export CHAIN_URL="https://chain-.api.baseten.co/environments/production/run_remote" -python truss-chains/examples/bring-your-own-client/invoke.py "hello world" +python truss-chains/examples/bring_your_own_client/invoke.py "hello world" ``` Expected output: @@ -84,7 +84,7 @@ You can run `chain.py` directly under `run_local` to inspect the descriptor helpers without deploying: ```sh -python truss-chains/examples/bring-your-own-client/chain.py +python truss-chains/examples/bring_your_own_client/chain.py ``` Prints the four helper outputs against a fake `Echo` descriptor. Useful for diff --git a/truss-chains/examples/bring_your_own_client/__init__.py b/truss-chains/examples/bring_your_own_client/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/examples/bring-your-own-client/chain.py b/truss-chains/examples/bring_your_own_client/chain.py similarity index 100% rename from truss-chains/examples/bring-your-own-client/chain.py rename to truss-chains/examples/bring_your_own_client/chain.py diff --git a/truss-chains/examples/bring-your-own-client/invoke.py b/truss-chains/examples/bring_your_own_client/invoke.py similarity index 94% rename from truss-chains/examples/bring-your-own-client/invoke.py rename to truss-chains/examples/bring_your_own_client/invoke.py index d947a7036..861560ce1 100644 --- a/truss-chains/examples/bring-your-own-client/invoke.py +++ b/truss-chains/examples/bring_your_own_client/invoke.py @@ -5,7 +5,7 @@ python invoke.py "hello world" The chain returns both the stub-path and the raw-httpx-path results plus the -descriptor metadata, so you can verify the bring-your-own-client pattern +descriptor metadata, so you can verify the bring_your_own_client pattern matches the framework stub byte-for-byte. """ diff --git a/truss-chains/examples/truss_chainlet_demo/__init__.py b/truss-chains/examples/truss_chainlet_demo/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/examples/truss_chainlet_demo/chain.py b/truss-chains/examples/truss_chainlet_demo/chain.py index 82521aad9..278c8f1dc 100644 --- a/truss-chains/examples/truss_chainlet_demo/chain.py +++ b/truss-chains/examples/truss_chainlet_demo/chain.py @@ -103,9 +103,7 @@ async def run_remote(self, text: str) -> CallerOutput: # Demonstrates data flowing through BOTH dep flavors in one chain hop. async with httpx.AsyncClient(timeout=60) as client: response = await client.post( - self._echo_url, - headers=self._echo_headers, - json={"text": reversed_text}, + self._echo_url, headers=self._echo_headers, json={"text": reversed_text} ) response.raise_for_status() reversed_then_uppercased = response.json()["out"] diff --git a/truss-chains/examples/voice_agent_mocked/__init__.py b/truss-chains/examples/voice_agent_mocked/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/examples/voice_agent_mocked/chain.py b/truss-chains/examples/voice_agent_mocked/chain.py index bf821a30f..35547a71c 100644 --- a/truss-chains/examples/voice_agent_mocked/chain.py +++ b/truss-chains/examples/voice_agent_mocked/chain.py @@ -73,6 +73,9 @@ class TTSMock(chains.TrussChainlet): def _chain_id_from_descriptor(desc: chains.DeployedServiceDescriptor) -> str: """Extract chain id from ``internal_url.hostname`` (``chain-...``).""" + assert desc.internal_url is not None, ( + "internal_url is always populated for sibling DeployedServiceDescriptor at runtime" + ) hostname = desc.internal_url.hostname first = hostname.split(".", 1)[0] if not first.startswith("chain-"): @@ -82,6 +85,9 @@ def _chain_id_from_descriptor(desc: chains.DeployedServiceDescriptor) -> str: def _chain_deployment_id_from_descriptor(desc: chains.DeployedServiceDescriptor) -> str: """Extract chain deployment id from the descriptor's gateway URL path.""" + assert desc.internal_url is not None, ( + "internal_url is always populated for sibling DeployedServiceDescriptor at runtime" + ) parts = urlparse(desc.internal_url.gateway_run_remote_url).path.split("/") # Path: /deployment//chainlet//run_remote return parts[parts.index("deployment") + 1] @@ -237,6 +243,9 @@ async def run_remote(self, websocket: chains.WebSocketProtocol) -> None: await tws.send(json.dumps({})) # mock config frame, ignored await tws.send(completion) audio_out = await tws.recv() + assert isinstance(audio_out, bytes), ( + f"TTSMock should return bytes, got {type(audio_out).__name__}" + ) log.warning("TTSMock returned %d bytes", len(audio_out)) await websocket.send_bytes(audio_out) diff --git a/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py b/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py index c1dcd9300..e4c5ee82c 100644 --- a/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py +++ b/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py @@ -26,9 +26,7 @@ async def websocket(self, ws): return if "bytes" in msg and msg["bytes"] is not None: audio = msg["bytes"] - await ws.send_text( - json.dumps({"text": f"text-{len(audio)}"}) - ) + await ws.send_text(json.dumps({"text": f"text-{len(audio)}"})) # Text frames (e.g. JSON metadata handshake) are silently consumed. except Exception: return diff --git a/truss-chains/tests/numpy_and_binary/__init__.py b/truss-chains/tests/numpy_and_binary/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md index 822648c3b..ae45f1f38 100644 --- a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md +++ b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md @@ -2,7 +2,7 @@ A regular Truss directory (not a `ChainletBase` subclass) that uses `from truss_chains.runtime import get_service` to discover sibling URLs. -This is the bring-your-own-client pattern that Project 2 will rely on: +This is the bring_your_own_client pattern that Project 2 will rely on: the Truss does not depend on the chains framework code, only on the `truss_chains` package being installed. diff --git a/truss-chains/tests/runtime_discovery/README.md b/truss-chains/tests/runtime_discovery/README.md index 93d77c661..15d5fb7d1 100644 --- a/truss-chains/tests/runtime_discovery/README.md +++ b/truss-chains/tests/runtime_discovery/README.md @@ -32,7 +32,7 @@ uv run pytest truss-chains/tests/runtime_discovery/ | Dir | What it demonstrates | Type | |---|---|---| -| `01_basic_descriptor_access/` | Typed chainlets accessing the new `target_url`, `ws_url`, `internal_ws_url`, `with_auth_headers()` helpers via `context.get_service_descriptor(...)`. | Positive | +| `basic_descriptor_access/` | Typed chainlets accessing the new `target_url`, `ws_url`, `internal_ws_url`, `with_auth_headers()` helpers via `context.get_service_descriptor(...)`. | Positive | | `02_raw_truss_reading_siblings/` | A non-`ChainletBase` Truss using `from truss_chains.runtime import get_service` to discover sibling URLs. | Positive | | `03_missing_sibling/` | `MissingDependencyError` with helpful "Available: ..." message when looking up a non-existent sibling. | Adversarial | | `04_dynamic_config_absent/` | `list_services()` returns `{}` (does not raise) and `get_service(...)` raises with a clear "not running inside a chain context" message when no config file is present. | Adversarial | @@ -40,4 +40,4 @@ uv run pytest truss-chains/tests/runtime_discovery/ Each scenario has a `test_.py` that runs under `pytest`. Scenarios that involve a Truss directory have a `plain_truss/` subdirectory with a hand-written -`config.yaml` and `model/model.py` showing the bring-your-own-client pattern. +`config.yaml` and `model/model.py` showing the bring_your_own_client pattern. diff --git a/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/README.md b/truss-chains/tests/runtime_discovery/basic_descriptor_access/README.md similarity index 100% rename from truss-chains/tests/runtime_discovery/01_basic_descriptor_access/README.md rename to truss-chains/tests/runtime_discovery/basic_descriptor_access/README.md diff --git a/truss-chains/tests/runtime_discovery/basic_descriptor_access/__init__.py b/truss-chains/tests/runtime_discovery/basic_descriptor_access/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/chain.py b/truss-chains/tests/runtime_discovery/basic_descriptor_access/chain.py similarity index 100% rename from truss-chains/tests/runtime_discovery/01_basic_descriptor_access/chain.py rename to truss-chains/tests/runtime_discovery/basic_descriptor_access/chain.py diff --git a/truss-chains/tests/runtime_discovery/01_basic_descriptor_access/test_descriptor_helpers.py b/truss-chains/tests/runtime_discovery/basic_descriptor_access/test_descriptor_helpers.py similarity index 100% rename from truss-chains/tests/runtime_discovery/01_basic_descriptor_access/test_descriptor_helpers.py rename to truss-chains/tests/runtime_discovery/basic_descriptor_access/test_descriptor_helpers.py diff --git a/truss-chains/truss_chains/framework.py b/truss-chains/truss_chains/framework.py index 6bd7cb4f3..5538502ec 100644 --- a/truss-chains/truss_chains/framework.py +++ b/truss-chains/truss_chains/framework.py @@ -37,7 +37,7 @@ ) import pydantic -from typing_extensions import ParamSpec +from typing_extensions import ParamSpec, TypeGuard from truss.base import custom_types, trt_llm_config, truss_config from truss_chains import private_types, public_types, utils @@ -1043,7 +1043,7 @@ def validate_and_register_cls(cls: Type[private_types.ABCChainlet]) -> None: def _validate_truss_chainlet_cls( - cls: Type[private_types.ABCChainlet], src_path: str, location: _ErrorLocation + cls: Type["TrussChainlet"], src_path: str, location: _ErrorLocation ) -> None: """Validate a TrussChainlet declaration's class attributes and resolve the ``truss_dir``. Sets ``cls._resolved_truss_dir`` on success.""" @@ -1832,7 +1832,9 @@ def __init_subclass__(cls, **kwargs) -> None: validate_and_register_cls(cls) -def is_truss_chainlet(cls: Type[private_types.ABCChainlet]) -> bool: +def is_truss_chainlet( + cls: Type[private_types.ABCChainlet], +) -> TypeGuard[Type["TrussChainlet"]]: return issubclass(cls, TrussChainlet) diff --git a/truss-chains/truss_chains/public_api.py b/truss-chains/truss_chains/public_api.py index 18e66d630..e367096ad 100644 --- a/truss-chains/truss_chains/public_api.py +++ b/truss-chains/truss_chains/public_api.py @@ -43,13 +43,41 @@ def depends_context() -> public_types.DeploymentContext: return framework.ContextDependencyMarker() # type: ignore +# `TrussChainlet` deps don't have a typed ``run_remote`` surface; the framework +# injects a ``DeployedServiceDescriptor`` at runtime, so the caller's annotation +# should be ``DeployedServiceDescriptor`` rather than the wrapped Truss class. +# The overload below makes that assignment type-clean. The overlap-overlap +# warning is silenced because it's intentional — ``Type[TrussChainlet]`` is a +# subtype of ``Type[ChainletT]``, and we want the more specific overload to win. +@overload +def depends( # type: ignore[overload-overlap] + chainlet_cls: Type[framework.TrussChainlet], + retries: int = ..., + timeout_sec: float = ..., + use_binary: bool = ..., + concurrency_limit: int = ..., +) -> public_types.DeployedServiceDescriptor: ... + + +# ``ChainletBase`` deps preserve the wrapped class's typed surface so callers +# get IDE code-completion on dep methods (e.g. ``Reverser.run_remote``). +@overload +def depends( + chainlet_cls: Type[framework.ChainletT], + retries: int = ..., + timeout_sec: float = ..., + use_binary: bool = ..., + concurrency_limit: int = ..., +) -> framework.ChainletT: ... + + def depends( chainlet_cls: Type[framework.ChainletT], retries: int = 1, timeout_sec: float = public_types.DEFAULT_TIMEOUT_SEC, use_binary: bool = False, concurrency_limit: int = public_types.DEFAULT_CONCURRENCY_LIMIT, -) -> framework.ChainletT: +) -> Union[public_types.DeployedServiceDescriptor, framework.ChainletT]: """Sets a "symbolic marker" to indicate to the framework that a chainlet is a dependency of another chainlet. The return value of ``depends`` is intended to be used as a default argument in a chainlet's ``__init__``-method. From b68e7c525c5d749f0d3f385cf07afbf069ca8e60 Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Mon, 4 May 2026 16:03:09 -0700 Subject: [PATCH 08/12] Refactor tests to follow existing convention --- .../examples/bring_your_own_client/README.md | 5 +- .../02_raw_truss_reading_siblings/README.md | 25 ---- .../test_raw_truss.py | 65 ---------- .../03_missing_sibling/README.md | 13 -- .../test_missing_sibling.py | 38 ------ .../04_dynamic_config_absent/README.md | 18 --- .../test_no_context.py | 25 ---- .../README.md | 21 ---- .../test_backward_compat.py | 108 ---------------- .../tests/runtime_discovery/README.md | 43 ------- .../basic_descriptor_access/README.md | 15 --- .../basic_descriptor_access/__init__.py | 0 .../basic_descriptor_access/chain.py | 54 -------- .../test_descriptor_helpers.py | 101 --------------- .../plain_truss/config.yaml | 0 .../plain_truss/model/model.py | 0 .../tests/runtime_discovery/run_all_local.sh | 12 -- truss-chains/tests/test_runtime.py | 119 ++++++++++++++++++ 18 files changed, 122 insertions(+), 540 deletions(-) delete mode 100644 truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md delete mode 100644 truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py delete mode 100644 truss-chains/tests/runtime_discovery/03_missing_sibling/README.md delete mode 100644 truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py delete mode 100644 truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md delete mode 100644 truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py delete mode 100644 truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md delete mode 100644 truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py delete mode 100644 truss-chains/tests/runtime_discovery/README.md delete mode 100644 truss-chains/tests/runtime_discovery/basic_descriptor_access/README.md delete mode 100644 truss-chains/tests/runtime_discovery/basic_descriptor_access/__init__.py delete mode 100644 truss-chains/tests/runtime_discovery/basic_descriptor_access/chain.py delete mode 100644 truss-chains/tests/runtime_discovery/basic_descriptor_access/test_descriptor_helpers.py rename truss-chains/tests/runtime_discovery/{02_raw_truss_reading_siblings => }/plain_truss/config.yaml (100%) rename truss-chains/tests/runtime_discovery/{02_raw_truss_reading_siblings => }/plain_truss/model/model.py (100%) delete mode 100755 truss-chains/tests/runtime_discovery/run_all_local.sh diff --git a/truss-chains/examples/bring_your_own_client/README.md b/truss-chains/examples/bring_your_own_client/README.md index 661dea255..48340ac21 100644 --- a/truss-chains/examples/bring_your_own_client/README.md +++ b/truss-chains/examples/bring_your_own_client/README.md @@ -94,8 +94,9 @@ sanity-checking the helpers in development. - The `truss_chains.runtime.get_service(name)` API for **non-`ChainletBase`** Trusses. That use case is unlocked by Project 2 (`TrussChainlet`); see - `truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/` for - the contract test. + `test_plain_truss_picks_up_siblings` in `truss-chains/tests/test_runtime.py` + (with the fixture under `tests/runtime_discovery/plain_truss/`) for the + contract test. - WebSocket-flavored siblings. The helpers `ws_url` / `internal_ws_url` are exercised in the runtime_discovery contract tests; a WebSocket-fronted deployable demo would need a chainlet using diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md deleted file mode 100644 index ae45f1f38..000000000 --- a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/README.md +++ /dev/null @@ -1,25 +0,0 @@ -# 02 — Raw Truss Reading Siblings - -A regular Truss directory (not a `ChainletBase` subclass) that uses -`from truss_chains.runtime import get_service` to discover sibling URLs. -This is the bring_your_own_client pattern that Project 2 will rely on: -the Truss does not depend on the chains framework code, only on the -`truss_chains` package being installed. - -The `plain_truss/` directory shows the full layout: -``` -plain_truss/ -├── config.yaml -└── model/ - └── model.py -``` - -`model.py` calls `get_service("Diarizer")` in `load()` and stashes the -URL in `predict()` — exactly how a TrussChainlet would surface sibling -information to user-managed RPC code. - -## Run - -```sh -uv run pytest test_raw_truss.py -v -``` diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py b/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py deleted file mode 100644 index dc02c9a1c..000000000 --- a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/test_raw_truss.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Drive the plain Truss's model.py with a fake dynamic_chainlet_config to -verify it discovers siblings via the public runtime API.""" - -import json -import pathlib -import sys - -import pytest - -# Make the plain_truss/model module importable. -PLAIN_TRUSS_MODEL_DIR = pathlib.Path(__file__).parent / "plain_truss" / "model" -sys.path.insert(0, str(PLAIN_TRUSS_MODEL_DIR)) - - -@pytest.fixture -def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", - str(tmp_path), - ) - yield - - -def _write_dynamic_config(tmp_path, payload): - with (tmp_path / "dynamic_chainlet_config").open("w") as f: - f.write(json.dumps(payload)) - - -def test_plain_truss_picks_up_siblings(tmp_path, dynamic_config_mount_dir): - _write_dynamic_config( - tmp_path, - { - "Diarizer": { - "predict_url": "https://chain-abc.api.baseten.co/.../diarizer/run_remote", - "internal_url": { - "gateway_run_remote_url": "https://wp.api.baseten.co/.../diarizer/run_remote", - "hostname": "chain-abc.api.baseten.co", - }, - } - }, - ) - - from model import Model # noqa: import after monkeypatch - - m = Model() - m.load() - out = m.predict({}) - # internal_url wins (matches BasetenSession's selection rule). - assert out["diarizer_url"] == "https://wp.api.baseten.co/.../diarizer/run_remote" - assert out["auth_headers"] == { - "Authorization": "Api-Key ", - "Host": "chain-abc.api.baseten.co", - } - - -def test_plain_truss_runs_standalone(tmp_path, dynamic_config_mount_dir): - """No dynamic config file present — list_services returns empty, - Model.load() short-circuits, predict returns the standalone shape.""" - from model import Model - - m = Model() - m.load() - out = m.predict({}) - assert out["diarizer_url"] is None - assert out["auth_headers"] is None diff --git a/truss-chains/tests/runtime_discovery/03_missing_sibling/README.md b/truss-chains/tests/runtime_discovery/03_missing_sibling/README.md deleted file mode 100644 index 2ebb81807..000000000 --- a/truss-chains/tests/runtime_discovery/03_missing_sibling/README.md +++ /dev/null @@ -1,13 +0,0 @@ -# 03 — Missing Sibling (adversarial) - -When user code calls `runtime.get_service("DoesNotExist")`, the framework -must raise `MissingDependencyError` with a message that lists the available -sibling names — not silently return `None` or crash with a `KeyError`. - -This is the canonical failure-mode test for typos in chainlet names. - -## Run - -```sh -uv run pytest test_missing_sibling.py -v -``` diff --git a/truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py b/truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py deleted file mode 100644 index 9bbd24bce..000000000 --- a/truss-chains/tests/runtime_discovery/03_missing_sibling/test_missing_sibling.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Adversarial: `get_service` raises with available-names list when the -requested chainlet is not registered in the chain.""" - -import json - -import pytest - -from truss_chains import public_types, runtime - - -@pytest.fixture -def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", - str(tmp_path), - ) - yield - - -def test_missing_sibling_raises_with_available_names( - tmp_path, dynamic_config_mount_dir -): - config = { - "Whisper": {"predict_url": "https://x.example/whisper"}, - "Diarizer": {"predict_url": "https://x.example/diarizer"}, - } - with (tmp_path / "dynamic_chainlet_config").open("w") as f: - f.write(json.dumps(config)) - - with pytest.raises(public_types.MissingDependencyError) as excinfo: - runtime.get_service("Typo") - - msg = str(excinfo.value) - # The error must name what was attempted. - assert "Typo" in msg - # And the available alternatives — so users can debug typos quickly. - assert "Whisper" in msg - assert "Diarizer" in msg diff --git a/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md b/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md deleted file mode 100644 index 86e15b05a..000000000 --- a/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/README.md +++ /dev/null @@ -1,18 +0,0 @@ -# 04 — Dynamic Config Absent (adversarial) - -When a Truss runs outside any chain context (no -`/etc/b10_dynamic_config/dynamic_chainlet_config` file present): - -- `runtime.list_services()` returns an empty mapping (does NOT raise) — so - caller code can branch on "am I inside a chain?" without try/except. -- `runtime.get_service(name)` raises `MissingDependencyError` with a clear - "not running inside a chain context" message. - -This documents the contract for code that needs to behave gracefully both -inside and outside a chain (e.g., a Truss that's reused across deployments). - -## Run - -```sh -uv run pytest test_no_context.py -v -``` diff --git a/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py b/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py deleted file mode 100644 index f15757753..000000000 --- a/truss-chains/tests/runtime_discovery/04_dynamic_config_absent/test_no_context.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Adversarial: behavior when no chain context is present.""" - -import pytest - -from truss_chains import public_types, runtime - - -@pytest.fixture -def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", - str(tmp_path), - ) - yield - - -def test_list_services_empty_when_no_context(tmp_path, dynamic_config_mount_dir): - """No file written; list_services returns empty mapping (does NOT raise).""" - assert runtime.list_services() == {} - - -def test_get_service_raises_clear_message(tmp_path, dynamic_config_mount_dir): - with pytest.raises(public_types.MissingDependencyError) as excinfo: - runtime.get_service("Anything") - assert "not running inside a chain context" in str(excinfo.value) diff --git a/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md b/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md deleted file mode 100644 index 2d54ece94..000000000 --- a/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/README.md +++ /dev/null @@ -1,21 +0,0 @@ -# 05 — Backward Compat: Existing Typed Chainlets (adversarial) - -Project 1 internally refactors `populate_chainlet_service_predict_urls` to -share code with the new `truss_chains.runtime` API. This example pins the -historical contract of that refactor: - -1. The function's signature and return shape are unchanged. -2. Error messages for "config absent" / "name missing" are preserved - byte-for-byte. -3. The mutually-exclusive ``predict_url`` / ``internal_url`` behavior of - the typed path is preserved: when ``internal_url`` is present in the - dynamic config, ``predict_url`` is cleared on the resulting descriptor - (matching the pre-refactor behavior). -4. Same Python type identity for `DeployedServiceDescriptor` regardless of - import path. - -## Run - -```sh -uv run pytest test_backward_compat.py -v -``` diff --git a/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py b/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py deleted file mode 100644 index 060061c9c..000000000 --- a/truss-chains/tests/runtime_discovery/05_backward_compat_existing_chainlets/test_backward_compat.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Adversarial: pin the backward-compatible contract of the typed-chain -``populate_chainlet_service_predict_urls`` after the internal refactor.""" - -import json - -import pytest - -from truss_chains import private_types, public_types, runtime -from truss_chains.remote_chainlet.utils import populate_chainlet_service_predict_urls - - -@pytest.fixture -def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR", - str(tmp_path), - ) - yield - - -def _write_config(tmp_path, payload): - with (tmp_path / private_types.DYNAMIC_CHAINLET_CONFIG_KEY).open("w") as f: - f.write(json.dumps(payload)) - - -def test_typed_path_clears_predict_url_when_internal_url_present( - tmp_path, dynamic_config_mount_dir -): - """Historical behavior: in the typed-chain path, when both URLs are in - the dynamic config, the resulting descriptor has only ``internal_url`` - set. The new public ``runtime.get_service`` carries both — but the typed - path stays mutually exclusive for backward compatibility.""" - _write_config( - tmp_path, - { - "Hello World!": { - "predict_url": "https://chain-x.api.baseten.co/.../run_remote", - "internal_url": { - "gateway_run_remote_url": "https://wp.api.baseten.co/.../run_remote", - "hostname": "chain-x.api.baseten.co", - }, - } - }, - ) - - chainlet_to_service = { - "HelloWorld": private_types.ServiceDescriptor( - name="HelloWorld", - display_name="Hello World!", - options=public_types.RPCOptions(), - ) - } - out = populate_chainlet_service_predict_urls(chainlet_to_service) - - # Typed path: predict_url cleared. - assert out["HelloWorld"].predict_url is None - assert out["HelloWorld"].internal_url is not None - - # Public runtime API: both URLs preserved (different contract, same data). - desc = runtime.get_service("Hello World!") - assert desc.predict_url == "https://chain-x.api.baseten.co/.../run_remote" - assert desc.internal_url is not None - - -def test_missing_chainlet_error_message_preserved(tmp_path, dynamic_config_mount_dir): - """Existing test in test_utils.py matches on the substring - "Chainlet 'X' not found". Project 1's refactor must keep this exact - wording.""" - _write_config(tmp_path, {"OtherName": {"predict_url": "https://x"}}) - with pytest.raises( - public_types.MissingDependencyError, match="Chainlet 'RandInt' not found" - ): - populate_chainlet_service_predict_urls( - { - "RandInt": private_types.ServiceDescriptor( - name="RandInt", - display_name="RandInt", - options=public_types.RPCOptions(), - ) - } - ) - - -def test_missing_dynamic_config_error_message_preserved( - tmp_path, dynamic_config_mount_dir -): - """No file: typed path raises with the historical "Cannot override - Chainlet configs" wording, distinct from the new runtime API's - "not running inside a chain context" wording.""" - with pytest.raises( - public_types.MissingDependencyError, match="Cannot override Chainlet configs" - ): - populate_chainlet_service_predict_urls( - { - "X": private_types.ServiceDescriptor( - name="X", display_name="X", options=public_types.RPCOptions() - ) - } - ) - - -def test_descriptor_type_identity_across_imports(): - """Same Python type whether reached via public_types or as a re-export - candidate — there's only one canonical class.""" - from truss_chains import DeployedServiceDescriptor as A - from truss_chains.public_types import DeployedServiceDescriptor as B - - assert A is B diff --git a/truss-chains/tests/runtime_discovery/README.md b/truss-chains/tests/runtime_discovery/README.md deleted file mode 100644 index 15d5fb7d1..000000000 --- a/truss-chains/tests/runtime_discovery/README.md +++ /dev/null @@ -1,43 +0,0 @@ -# Runtime Discovery — Contract Test Scenarios - -CPU-only pytest scenarios that pin the Project 1 (Runtime Discovery Platform -Contract) API. Each scenario doubles as documentation of the expected shape -and behavior of `truss_chains.runtime`. - -These live under `tests/` rather than `examples/` because they are **not -deployable chains** — they exercise the library by writing fake -`/etc/b10_dynamic_config/dynamic_chainlet_config` files via pytest -monkeypatch, not by pushing real chainlets to Baseten. (For deployable -chain examples, see `truss-chains/examples/`.) - -## Setup - -```sh -pip install -e /Users/mattelim/Documents/truss -pip install -e /Users/mattelim/Documents/truss/truss-chains -pip install pytest -``` - -## Run all - -```sh -# Runs each scenario in its own pytest invocation, with descriptive headers. -bash run_all_local.sh - -# Or as one combined pytest run: -uv run pytest truss-chains/tests/runtime_discovery/ -``` - -## Per-scenario - -| Dir | What it demonstrates | Type | -|---|---|---| -| `basic_descriptor_access/` | Typed chainlets accessing the new `target_url`, `ws_url`, `internal_ws_url`, `with_auth_headers()` helpers via `context.get_service_descriptor(...)`. | Positive | -| `02_raw_truss_reading_siblings/` | A non-`ChainletBase` Truss using `from truss_chains.runtime import get_service` to discover sibling URLs. | Positive | -| `03_missing_sibling/` | `MissingDependencyError` with helpful "Available: ..." message when looking up a non-existent sibling. | Adversarial | -| `04_dynamic_config_absent/` | `list_services()` returns `{}` (does not raise) and `get_service(...)` raises with a clear "not running inside a chain context" message when no config file is present. | Adversarial | -| `05_backward_compat_existing_chainlets/` | `populate_chainlet_service_predict_urls` (the typed-chain code path) preserves its historical mutually-exclusive `predict_url` / `internal_url` behavior — proves the internal refactor doesn't change the typed path. | Adversarial | - -Each scenario has a `test_.py` that runs under `pytest`. Scenarios that -involve a Truss directory have a `plain_truss/` subdirectory with a hand-written -`config.yaml` and `model/model.py` showing the bring_your_own_client pattern. diff --git a/truss-chains/tests/runtime_discovery/basic_descriptor_access/README.md b/truss-chains/tests/runtime_discovery/basic_descriptor_access/README.md deleted file mode 100644 index 24647850c..000000000 --- a/truss-chains/tests/runtime_discovery/basic_descriptor_access/README.md +++ /dev/null @@ -1,15 +0,0 @@ -# 01 — Basic Descriptor Access - -A two-chainlet typed chain where `Caller` reads the new helper methods on the -`DeployedServiceDescriptor` returned by `context.get_service_descriptor("Echo")` -and reports them through `run_remote`. - -This is the "typed chain that uses the new helpers" path — no raw Truss yet. -Demonstrates that `target_url`, `ws_url`, `internal_ws_url`, and -`with_auth_headers()` produce the expected shapes from a real descriptor. - -## Run - -```sh -uv run pytest test_descriptor_helpers.py -v -``` diff --git a/truss-chains/tests/runtime_discovery/basic_descriptor_access/__init__.py b/truss-chains/tests/runtime_discovery/basic_descriptor_access/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/truss-chains/tests/runtime_discovery/basic_descriptor_access/chain.py b/truss-chains/tests/runtime_discovery/basic_descriptor_access/chain.py deleted file mode 100644 index 4faa3123b..000000000 --- a/truss-chains/tests/runtime_discovery/basic_descriptor_access/chain.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Simple typed chain showing access to the new descriptor helpers from -``context.get_service_descriptor``.""" - -from typing import Optional - -import pydantic - -import truss_chains as chains - - -class CallerOutput(pydantic.BaseModel): - echoed: str - target_url: str - ws_url: Optional[str] - internal_ws_url: Optional[str] - auth_headers: dict[str, str] - - -class Echo(chains.ChainletBase): - remote_config = chains.RemoteConfig( - compute=chains.Compute(cpu_count=1, memory="512Mi") - ) - - async def run_remote(self, text: str) -> str: - return text - - -class Caller(chains.ChainletBase): - remote_config = chains.RemoteConfig( - compute=chains.Compute(cpu_count=1, memory="512Mi") - ) - - def __init__( - self, - echo: Echo = chains.depends(Echo), - context: chains.DeploymentContext = chains.depends_context(), - ) -> None: - self._echo = echo - self._context = context - - async def run_remote(self, text: str) -> CallerOutput: - # Pull the descriptor for the typed dependency and exercise the new - # helpers. In a real chain you would use `target_url` / `ws_url` to - # build your own client; here we just surface them so the test - # harness can assert against them. - desc = self._context.get_service_descriptor("Echo") - api_key = self._context.get_baseten_api_key() - return CallerOutput( - echoed=await self._echo.run_remote(text), - target_url=desc.target_url, - ws_url=desc.ws_url, - internal_ws_url=desc.internal_ws_url, - auth_headers=desc.with_auth_headers(api_key), - ) diff --git a/truss-chains/tests/runtime_discovery/basic_descriptor_access/test_descriptor_helpers.py b/truss-chains/tests/runtime_discovery/basic_descriptor_access/test_descriptor_helpers.py deleted file mode 100644 index 6a6cecbf8..000000000 --- a/truss-chains/tests/runtime_discovery/basic_descriptor_access/test_descriptor_helpers.py +++ /dev/null @@ -1,101 +0,0 @@ -"""Smoke tests for the descriptor helpers using ``run_local``-style -construction of a ``DeployedServiceDescriptor``. The chain in ``chain.py`` -shows how a chainlet would consume these helpers in production; here we -construct a descriptor directly and assert the helpers' return values.""" - -import pathlib -import sys - -import truss_chains as chains - -# Make ``chain.py`` importable from this test file. -sys.path.insert(0, str(pathlib.Path(__file__).parent)) - - -def _make_descriptor(with_internal_url: bool): - kwargs = { - "name": "Echo", - "display_name": "Echo", - "options": chains.RPCOptions(), - "predict_url": "https://chain-abc.api.baseten.co/.../echo/run_remote", - } - if with_internal_url: - kwargs["internal_url"] = chains.DeployedServiceDescriptor.InternalURL( - gateway_run_remote_url="https://wp.api.baseten.co/.../echo/run_remote", - hostname="chain-abc.api.baseten.co", - ) - return chains.DeployedServiceDescriptor(**kwargs) - - -def test_target_url_with_internal_prefers_internal(): - desc = _make_descriptor(with_internal_url=True) - assert desc.target_url == "https://wp.api.baseten.co/.../echo/run_remote" - - -def test_target_url_falls_back_to_predict(): - desc = _make_descriptor(with_internal_url=False) - assert desc.target_url == "https://chain-abc.api.baseten.co/.../echo/run_remote" - - -def test_ws_url_scheme_swap(): - desc = _make_descriptor(with_internal_url=False) - assert desc.ws_url.startswith("wss://") - assert desc.internal_ws_url is None - - -def test_internal_ws_url_scheme_swap(): - desc = _make_descriptor(with_internal_url=True) - assert desc.internal_ws_url.startswith("wss://") - - -def test_auth_headers_no_internal(): - desc = _make_descriptor(with_internal_url=False) - assert desc.with_auth_headers("k") == {"Authorization": "Api-Key k"} - - -def test_auth_headers_with_internal_includes_host(): - desc = _make_descriptor(with_internal_url=True) - assert desc.with_auth_headers("k") == { - "Authorization": "Api-Key k", - "Host": "chain-abc.api.baseten.co", - } - - -def test_chain_runs_locally_with_deployed_echo(): - """run_local with a ``chainlet_to_service`` override that points Echo at - a fake deployed URL — exercises the descriptor helpers via the typed path - (``context.get_service_descriptor``). - - Caller still calls ``self._echo.run_remote(...)`` over HTTP via the stub - rather than running Echo in-process, so this scenario is not what you'd - use in normal local debugging — but it's what makes the descriptor - helpers observable from a run_local test.""" - from chain import Caller - - fake_descriptor = chains.DeployedServiceDescriptor( - name="Echo", - display_name="Echo", - options=chains.RPCOptions(), - predict_url="https://chain-abc.api.baseten.co/.../echo/run_remote", - internal_url=chains.DeployedServiceDescriptor.InternalURL( - gateway_run_remote_url="https://wp.api.baseten.co/.../echo/run_remote", - hostname="chain-abc.api.baseten.co", - ), - ) - - with chains.run_local( - secrets={chains.public_types.CHAIN_API_KEY_SECRET_NAME: "test-key"}, - chainlet_to_service={"Echo": fake_descriptor}, - ): - caller = Caller() - - # We don't actually call run_remote (that would fire HTTP at the fake URL). - # Inspect the descriptor through the context directly. - desc = caller._context.get_service_descriptor("Echo") - assert desc.target_url == "https://wp.api.baseten.co/.../echo/run_remote" - assert desc.ws_url == "wss://chain-abc.api.baseten.co/.../echo/run_remote" - assert desc.internal_ws_url == "wss://wp.api.baseten.co/.../echo/run_remote" - assert desc.with_auth_headers("test-key") == { - "Authorization": "Api-Key test-key", - "Host": "chain-abc.api.baseten.co", - } diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/config.yaml b/truss-chains/tests/runtime_discovery/plain_truss/config.yaml similarity index 100% rename from truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/config.yaml rename to truss-chains/tests/runtime_discovery/plain_truss/config.yaml diff --git a/truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/model/model.py b/truss-chains/tests/runtime_discovery/plain_truss/model/model.py similarity index 100% rename from truss-chains/tests/runtime_discovery/02_raw_truss_reading_siblings/plain_truss/model/model.py rename to truss-chains/tests/runtime_discovery/plain_truss/model/model.py diff --git a/truss-chains/tests/runtime_discovery/run_all_local.sh b/truss-chains/tests/runtime_discovery/run_all_local.sh deleted file mode 100755 index b07020782..000000000 --- a/truss-chains/tests/runtime_discovery/run_all_local.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/usr/bin/env bash -# Run every runtime_discovery example's local test. CPU-only. -set -euo pipefail -cd "$(dirname "$0")" -for dir in 0*; do - for test_file in "$dir"/test_*.py; do - if [[ -f "$test_file" ]]; then - echo "=== $dir / $(basename "$test_file") ===" - uv run pytest "$test_file" -v - fi - done -done diff --git a/truss-chains/tests/test_runtime.py b/truss-chains/tests/test_runtime.py index 4c666ec89..f8c812566 100644 --- a/truss-chains/tests/test_runtime.py +++ b/truss-chains/tests/test_runtime.py @@ -11,11 +11,20 @@ """ import json +import pathlib +import sys import pytest +import truss_chains as chains from truss_chains import private_types, public_types, runtime +# Make the plain Truss model fixture importable. +_PLAIN_TRUSS_MODEL_DIR = ( + pathlib.Path(__file__).parent / "runtime_discovery" / "plain_truss" / "model" +) +sys.path.insert(0, str(_PLAIN_TRUSS_MODEL_DIR)) + # ---- Test fixtures ----------------------------------------------------------- PREDICT_URL_ONLY = { @@ -277,3 +286,113 @@ def test_descriptor_ws_url_unknown_scheme_raises(): ) with pytest.raises(ValueError, match="Cannot convert to WebSocket scheme"): _ = desc.ws_url + + +def test_descriptor_type_identity_across_imports(): + """One canonical class, regardless of import path.""" + from truss_chains import DeployedServiceDescriptor as A + from truss_chains.public_types import DeployedServiceDescriptor as B + + assert A is B + + +# ---- run_local: typed dependency exposes descriptor helpers ------------------ + + +class _Echo(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + async def run_remote(self, text: str) -> str: + return text + + +class _Caller(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + def __init__( + self, + echo: _Echo = chains.depends(_Echo), + context: chains.DeploymentContext = chains.depends_context(), + ) -> None: + self._echo = echo + self._context = context + + async def run_remote(self, text: str) -> str: + return await self._echo.run_remote(text) + + +def test_chain_runs_locally_with_deployed_dep(): + """``run_local`` with a ``chainlet_to_service`` override exposes the + descriptor helpers via the typed path (``context.get_service_descriptor``).""" + fake_descriptor = chains.DeployedServiceDescriptor( + name="_Echo", + display_name="_Echo", + options=chains.RPCOptions(), + predict_url="https://chain-abc.api.baseten.co/.../echo/run_remote", + internal_url=chains.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://wp.api.baseten.co/.../echo/run_remote", + hostname="chain-abc.api.baseten.co", + ), + ) + + with chains.run_local( + secrets={public_types.CHAIN_API_KEY_SECRET_NAME: "test-key"}, + chainlet_to_service={"_Echo": fake_descriptor}, + ): + caller = _Caller() + + desc = caller._context.get_service_descriptor("_Echo") + assert desc.target_url == "https://wp.api.baseten.co/.../echo/run_remote" + assert desc.ws_url == "wss://chain-abc.api.baseten.co/.../echo/run_remote" + assert desc.internal_ws_url == "wss://wp.api.baseten.co/.../echo/run_remote" + assert desc.with_auth_headers("test-key") == { + "Authorization": "Api-Key test-key", + "Host": "chain-abc.api.baseten.co", + } + + +# ---- Plain Truss bring-your-own-client ------------------------------------- + + +def test_plain_truss_picks_up_siblings(tmp_path, dynamic_config_mount_dir): + """A non-ChainletBase Truss reads sibling URLs via ``runtime.list_services`` + and ``runtime.get_service`` — the bring-your-own-client pattern.""" + _write_config( + tmp_path, + { + "Diarizer": { + "predict_url": "https://chain-abc.api.baseten.co/.../diarizer/run_remote", + "internal_url": { + "gateway_run_remote_url": "https://wp.api.baseten.co/.../diarizer/run_remote", + "hostname": "chain-abc.api.baseten.co", + }, + } + }, + ) + + from model import Model # type: ignore[import-not-found] + + m = Model() + m.load() + out = m.predict({}) + assert out["diarizer_url"] == "https://wp.api.baseten.co/.../diarizer/run_remote" + assert out["auth_headers"] == { + "Authorization": "Api-Key ", + "Host": "chain-abc.api.baseten.co", + } + + +def test_plain_truss_runs_standalone(tmp_path, dynamic_config_mount_dir): + """No dynamic config file present — ``list_services`` returns empty, + Model.load short-circuits, predict returns the standalone shape.""" + from model import Model # type: ignore[import-not-found] + + m = Model() + m.load() + out = m.predict({}) + assert out["diarizer_url"] is None + assert out["auth_headers"] is None From de048a0cb19e5d2b1ca5ed686cf4f33f013c1e67 Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Mon, 4 May 2026 16:04:28 -0700 Subject: [PATCH 09/12] Defer TrussChainlet sub dir validation so that selective watch when unselected dirs are missing. --- truss-chains/tests/test_truss_chainlet.py | 44 +++++++++++++++-- .../truss_chains/deployment/code_gen.py | 17 +++++-- truss-chains/truss_chains/framework.py | 48 ++++++------------- 3 files changed, 70 insertions(+), 39 deletions(-) diff --git a/truss-chains/tests/test_truss_chainlet.py b/truss-chains/tests/test_truss_chainlet.py index 887f70d9f..a577cc400 100644 --- a/truss-chains/tests/test_truss_chainlet.py +++ b/truss-chains/tests/test_truss_chainlet.py @@ -92,16 +92,36 @@ class _Bad(chains.TrussChainlet): # noqa: N801 def test_truss_dir_must_resolve_to_directory(tmp_path): + """A missing `truss_dir` is OK at parse time (matches ChainletBase, which + has no filesystem-state checks at parse). The error surfaces at codegen + time only for chainlets that are actually being built/deployed — which + keeps `truss chains watch --experimental-chainlet-names ` working + when non-targeted chainlets' truss_dirs are absent locally.""" + from truss_chains.deployment import code_gen + bad_path = str(tmp_path / "does_not_exist") class _Bad(chains.TrussChainlet): truss_dir = bad_path + # Parse-time validation now succeeds (source-only invariants pass). + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Bad) + + # Codegen surfaces the error at the natural phase. with pytest.raises(public_types.ChainsUsageError, match="not a directory"): - framework.raise_validation_errors() + code_gen.gen_truss_chainlet( + chain_root=tmp_path, + chain_name="codegen-test", + chainlet_descriptor=descriptor, + ) def test_truss_dir_missing_config_yaml(tmp_path): + """An existing directory without a `config.yaml` is OK at parse time; the + error surfaces at codegen.""" + from truss_chains.deployment import code_gen + bad = tmp_path / "bad_truss" bad.mkdir() bad_path = str(bad) @@ -109,11 +129,22 @@ def test_truss_dir_missing_config_yaml(tmp_path): class _Bad(chains.TrussChainlet): truss_dir = bad_path + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Bad) + with pytest.raises(public_types.ChainsUsageError, match="missing `config.yaml`"): - framework.raise_validation_errors() + code_gen.gen_truss_chainlet( + chain_root=tmp_path, + chain_name="codegen-test", + chainlet_descriptor=descriptor, + ) def test_truss_dir_invalid_config_yaml(tmp_path): + """An unparseable `config.yaml` is OK at parse time; the error surfaces + at codegen.""" + from truss_chains.deployment import code_gen + bad = tmp_path / "bad_truss" bad.mkdir() (bad / "config.yaml").write_text("not: a valid: truss config { broken") @@ -122,8 +153,15 @@ def test_truss_dir_invalid_config_yaml(tmp_path): class _Bad(chains.TrussChainlet): truss_dir = bad_path + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Bad) + with pytest.raises(public_types.ChainsUsageError, match="invalid `config.yaml`"): - framework.raise_validation_errors() + code_gen.gen_truss_chainlet( + chain_root=tmp_path, + chain_name="codegen-test", + chainlet_descriptor=descriptor, + ) def test_truss_dir_resolved_relative_to_declaring_file(): diff --git a/truss-chains/truss_chains/deployment/code_gen.py b/truss-chains/truss_chains/deployment/code_gen.py index ccbdc2b8a..1f51d70fe 100644 --- a/truss-chains/truss_chains/deployment/code_gen.py +++ b/truss-chains/truss_chains/deployment/code_gen.py @@ -969,8 +969,8 @@ def _prepare_truss_chainlet_artifact( src_truss_dir = chainlet_descriptor.truss_dir if src_truss_dir is None or not src_truss_dir.is_dir(): raise public_types.ChainsUsageError( - f"TrussChainlet `{chainlet_descriptor.name}`: truss_dir is not set " - "or does not exist. Validation should have caught this earlier." + f"`TrussChainlet.{chainlet_descriptor.name}.truss_dir` resolved to " + f"`{src_truss_dir}`, which is not a directory." ) # Replace any pre-existing chainlet_dir from a stale prior generation. if chainlet_dir.exists(): @@ -978,7 +978,18 @@ def _prepare_truss_chainlet_artifact( truss_path.copy_tree_path(src_truss_dir, chainlet_dir) config_path = chainlet_dir / serving_image_builder.CONFIG_FILE - config = truss_config.TrussConfig.from_yaml(config_path) + if not config_path.is_file(): + raise public_types.ChainsUsageError( + f"`TrussChainlet.{chainlet_descriptor.name}.truss_dir` ({src_truss_dir}) is " + "missing `config.yaml` — not a valid Truss directory." + ) + try: + config = truss_config.TrussConfig.from_yaml(config_path) + except Exception as exc: + raise public_types.ChainsUsageError( + f"`TrussChainlet.{chainlet_descriptor.name}.truss_dir` ({src_truss_dir}) has " + f"an invalid `config.yaml`: {exc}" + ) from exc config.model_name = model_name if public_types.CHAIN_API_KEY_SECRET_NAME not in config.secrets: config.secrets[public_types.CHAIN_API_KEY_SECRET_NAME] = ( diff --git a/truss-chains/truss_chains/framework.py b/truss-chains/truss_chains/framework.py index 5538502ec..c5ba91035 100644 --- a/truss-chains/truss_chains/framework.py +++ b/truss-chains/truss_chains/framework.py @@ -1045,8 +1045,18 @@ def validate_and_register_cls(cls: Type[private_types.ABCChainlet]) -> None: def _validate_truss_chainlet_cls( cls: Type["TrussChainlet"], src_path: str, location: _ErrorLocation ) -> None: - """Validate a TrussChainlet declaration's class attributes and resolve - the ``truss_dir``. Sets ``cls._resolved_truss_dir`` on success.""" + """Validate a TrussChainlet declaration's source-level class attributes and + resolve the ``truss_dir`` path (relative to the declaring file). Sets + ``cls._resolved_truss_dir`` on success. + + Filesystem-state checks — ``truss_dir`` actually existing on disk and + containing a parseable ``config.yaml`` — are deferred to codegen. They are + performed by ``_prepare_truss_chainlet_artifact`` in + ``truss_chains/deployment/code_gen.py`` only for chainlets that are + actually being built/deployed, which respects ``--experimental-chainlet-names`` + filtering during ``truss chains watch``. Mirrors the ChainletBase pattern + where parse-time validation only checks source-level invariants. + """ # 1. truss_dir class attribute is required (follow MRO so subclasses inherit). truss_dir_raw = getattr(cls, "truss_dir", None) if not truss_dir_raw: @@ -1066,43 +1076,15 @@ def _validate_truss_chainlet_cls( ) return - # 2. Resolve relative to the file where the class was declared. + # 2. Resolve relative to the file where the class was declared. Note that + # ``Path.resolve()`` succeeds even when the target doesn't exist; existence + # is intentionally not checked here (see docstring). src_dir = pathlib.Path(src_path).parent truss_dir = pathlib.Path(truss_dir_raw) if not truss_dir.is_absolute(): truss_dir = (src_dir / truss_dir).resolve() cls._resolved_truss_dir = truss_dir - # 3. Directory exists. - if not truss_dir.is_dir(): - _collect_error( - f"`TrussChainlet.{cls.__name__}.truss_dir` resolved to " - f"`{truss_dir}`, which is not a directory.", - _ErrorKind.INVALID_CONFIG_ERROR, - location, - ) - return - - # 4. config.yaml is a valid Truss config. - config_path = truss_dir / "config.yaml" - if not config_path.is_file(): - _collect_error( - f"`TrussChainlet.{cls.__name__}.truss_dir` ({truss_dir}) is " - "missing `config.yaml` — not a valid Truss directory.", - _ErrorKind.INVALID_CONFIG_ERROR, - location, - ) - return - try: - truss_config.TrussConfig.from_yaml(config_path) - except Exception as exc: - _collect_error( - f"`TrussChainlet.{cls.__name__}.truss_dir` ({truss_dir}) has " - f"an invalid `config.yaml`: {exc}", - _ErrorKind.INVALID_CONFIG_ERROR, - location, - ) - # Dependency-Injection / Registry ###################################################### From c14247ea2f10266c5ef27a69e500f9afa953a61e Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Fri, 8 May 2026 02:15:50 -0700 Subject: [PATCH 10/12] Fix websocket url conversion --- truss-chains/tests/test_runtime.py | 82 +++++++++++++++++++++-- truss-chains/truss_chains/public_types.py | 75 ++++++++++++++------- 2 files changed, 127 insertions(+), 30 deletions(-) diff --git a/truss-chains/tests/test_runtime.py b/truss-chains/tests/test_runtime.py index f8c812566..e4a08807c 100644 --- a/truss-chains/tests/test_runtime.py +++ b/truss-chains/tests/test_runtime.py @@ -214,16 +214,19 @@ def test_descriptor_ws_url_http_to_ws(): def test_descriptor_internal_ws_url(): + """``internal_ws_url`` netloc is the chain hostname, not the gateway host + from ``gateway_run_remote_url`` — WS can't use a Host override to smuggle + the chain host (websockets lib always emits Host from URL netloc).""" desc = public_types.DeployedServiceDescriptor( name="X", display_name="X", options=public_types.RPCOptions(), internal_url=public_types.DeployedServiceDescriptor.InternalURL( gateway_run_remote_url="https://internal.example/predict", - hostname="public.example", + hostname="chain-abc.api.baseten.co", ), ) - assert desc.internal_ws_url == "wss://internal.example/predict" + assert desc.internal_ws_url == "wss://chain-abc.api.baseten.co/predict" assert desc.ws_url is None @@ -258,6 +261,32 @@ def test_descriptor_with_auth_headers_with_internal_url(): } +def test_descriptor_with_ws_auth_headers_omits_host(): + """Authorization-only even when internal_url is set — explicit Host breaks + the WS handshake (websockets lib always emits its own Host from URL).""" + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + internal_url=public_types.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url="https://internal.example/predict", + hostname="chain-abc.api.baseten.co", + ), + ) + headers = desc.with_ws_auth_headers("my-key") + assert headers == {"Authorization": "Api-Key my-key"} + + +def test_descriptor_with_ws_auth_headers_predict_only(): + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="https://public.example/predict", + ) + assert desc.with_ws_auth_headers("my-key") == {"Authorization": "Api-Key my-key"} + + def test_with_auth_headers_matches_BasetenSession(): """Mechanically verify ``with_auth_headers`` produces the exact dict ``BasetenSession.__init__`` would assemble — preventing drift between the @@ -288,6 +317,47 @@ def test_descriptor_ws_url_unknown_scheme_raises(): _ = desc.ws_url +def test_descriptor_ws_url_rewrites_run_remote_path_to_websocket(): + """Chainlet predict URLs end in ``/run_remote``; api-gateway's WS handler + expects ``/websocket``. ``ws_url`` swaps both scheme and terminal path.""" + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url=( + "https://chain-abc.api.baseten.co/deployment/dep/chainlet/cl/run_remote" + ), + internal_url=public_types.DeployedServiceDescriptor.InternalURL( + gateway_run_remote_url=( + "https://wp.api.baseten.co/deployment/dep/chainlet/cl/run_remote" + ), + hostname="chain-abc.api.baseten.co", + ), + ) + # Both URLs use the chain hostname; internal_ws_url drops the gateway host + # from gateway_run_remote_url (see internal_ws_url docstring for why). + assert ( + desc.ws_url + == "wss://chain-abc.api.baseten.co/deployment/dep/chainlet/cl/websocket" + ) + assert ( + desc.internal_ws_url + == "wss://chain-abc.api.baseten.co/deployment/dep/chainlet/cl/websocket" + ) + + +def test_descriptor_ws_url_preserves_non_run_remote_paths(): + """When the predict path doesn't end in ``/run_remote`` (e.g. + standalone-model URLs), only the scheme is swapped.""" + desc = public_types.DeployedServiceDescriptor( + name="X", + display_name="X", + options=public_types.RPCOptions(), + predict_url="https://model.example/predict", + ) + assert desc.ws_url == "wss://model.example/predict" + + def test_descriptor_type_identity_across_imports(): """One canonical class, regardless of import path.""" from truss_chains import DeployedServiceDescriptor as A @@ -347,12 +417,16 @@ def test_chain_runs_locally_with_deployed_dep(): desc = caller._context.get_service_descriptor("_Echo") assert desc.target_url == "https://wp.api.baseten.co/.../echo/run_remote" - assert desc.ws_url == "wss://chain-abc.api.baseten.co/.../echo/run_remote" - assert desc.internal_ws_url == "wss://wp.api.baseten.co/.../echo/run_remote" + # ws_url / internal_ws_url: chain hostname netloc + /websocket path. + assert desc.ws_url == "wss://chain-abc.api.baseten.co/.../echo/websocket" + assert desc.internal_ws_url == "wss://chain-abc.api.baseten.co/.../echo/websocket" assert desc.with_auth_headers("test-key") == { "Authorization": "Api-Key test-key", "Host": "chain-abc.api.baseten.co", } + assert desc.with_ws_auth_headers("test-key") == { + "Authorization": "Api-Key test-key" + } # ---- Plain Truss bring-your-own-client ------------------------------------- diff --git a/truss-chains/truss_chains/public_types.py b/truss-chains/truss_chains/public_types.py index 3b1238c05..450ebbbfd 100644 --- a/truss-chains/truss_chains/public_types.py +++ b/truss-chains/truss_chains/public_types.py @@ -3,6 +3,7 @@ import logging import pathlib import traceback +import urllib.parse from collections.abc import AsyncIterator from typing import ( Any, @@ -727,15 +728,33 @@ class LookaheadDecodingConfig(pydantic.BaseModel): lookahead_decoding_config: Optional[LookaheadDecodingConfig] = None -def _to_websocket_scheme(url: str) -> str: - """Rewrite an http(s):// URL to ws(s)://. Raises if the scheme is unknown.""" +def _to_websocket_url(url: str) -> str: + """Rewrite an http(s):// predict URL to its ws(s):// WS counterpart. + + Two transformations: + + 1. Scheme: ``http`` → ``ws``, ``https`` → ``wss``. + 2. Path: terminal ``/run_remote`` → ``/websocket`` — chains' WS endpoints + are routed by api-gateway via that path suffix, while the matching HTTP + predict path uses ``/run_remote``. Other path shapes pass through + unchanged so this is also useful for standalone-model URLs whose paths + differ. + + Raises ``ValueError`` if the scheme is unknown. + """ if url.startswith("https://"): - return "wss://" + url[len("https://") :] - if url.startswith("http://"): - return "ws://" + url[len("http://") :] - raise ValueError( - f"Cannot convert to WebSocket scheme; expected http(s):// prefix: {url!r}" - ) + ws_url = "wss://" + url[len("https://") :] + elif url.startswith("http://"): + ws_url = "ws://" + url[len("http://") :] + else: + raise ValueError( + f"Cannot convert to WebSocket scheme; expected http(s):// prefix: {url!r}" + ) + + if ws_url.endswith("/run_remote"): + ws_url = ws_url[: -len("/run_remote")] + "/websocket" + + return ws_url class DeployedServiceDescriptor(custom_types.SafeModel): @@ -781,38 +800,42 @@ def target_url(self) -> str: @property def ws_url(self) -> Optional[str]: - """``predict_url`` with the http(s):// scheme rewritten to ws(s)://. - ``None`` when ``predict_url`` is not set.""" + """WebSocket URL derived from ``predict_url``: scheme swapped to + ws(s):// and terminal ``/run_remote`` rewritten to ``/websocket`` + (the path api-gateway routes to its WS handler). ``None`` when + ``predict_url`` is not set. See ``_to_websocket_url``.""" if self.predict_url is None: return None - return _to_websocket_scheme(self.predict_url) + return _to_websocket_url(self.predict_url) @property def internal_ws_url(self) -> Optional[str]: - """``internal_url.gateway_run_remote_url`` with the http(s):// scheme - rewritten to ws(s)://. ``None`` when ``internal_url`` is not set.""" + """Cluster-local WS URL; ``None`` if ``internal_url`` unset. Uses the + chain hostname as netloc (not ``gateway_run_remote_url``'s host) — the + ``websockets`` lib emits Host from the URL netloc, so a Host override + can't be used to smuggle chain-host through like HTTP does.""" if self.internal_url is None: return None - return _to_websocket_scheme(self.internal_url.gateway_run_remote_url) + rewritten = _to_websocket_url(self.internal_url.gateway_run_remote_url) + parsed = urllib.parse.urlparse(rewritten) + return parsed._replace(netloc=self.internal_url.hostname).geturl() def with_auth_headers(self, api_key: str) -> dict[str, str]: - """Build the headers a ``StubBase`` would send for outbound calls: - - * ``Authorization: Api-Key {api_key}`` — always. - * ``Host: {internal_url.hostname}`` — only when ``internal_url`` is set, - so cluster-local routing matches the chain hostname. - - ``api_key`` is typically sourced from - ``DeploymentContext.get_baseten_api_key()`` inside a ``ChainletBase`` or - from the standard Truss secrets API (``baseten_chain_api_key``) for raw - Trusses. The helper takes the key explicitly so it can be used outside - a chainlet context. - """ + """Headers for *HTTP* sibling calls (against ``target_url``): + Authorization, plus ``Host: internal_url.hostname`` when set. + For WS, use :py:meth:`with_ws_auth_headers` — the Host override breaks + the WS handshake.""" headers = {"Authorization": f"Api-Key {api_key}"} if self.internal_url is not None: headers["Host"] = self.internal_url.hostname return headers + def with_ws_auth_headers(self, api_key: str) -> dict[str, str]: + """Headers for *WebSocket* sibling calls (against ``internal_ws_url`` / + ``ws_url``). Authorization-only; an explicit Host would corrupt the + handshake.""" + return {"Authorization": f"Api-Key {api_key}"} + class Environment(custom_types.SafeModel): """The environment the chainlet is deployed in. From 041eb2f28abbac316918231ba5e1ff570ea96af7 Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Mon, 11 May 2026 22:13:54 -0700 Subject: [PATCH 11/12] Enable TrussChainlets to be entrypoints --- truss-chains/tests/test_truss_chainlet.py | 220 ++++++++++++++++++++++ truss-chains/truss_chains/framework.py | 128 ++++++++++++- 2 files changed, 343 insertions(+), 5 deletions(-) diff --git a/truss-chains/tests/test_truss_chainlet.py b/truss-chains/tests/test_truss_chainlet.py index a577cc400..c11e5ef42 100644 --- a/truss-chains/tests/test_truss_chainlet.py +++ b/truss-chains/tests/test_truss_chainlet.py @@ -515,3 +515,223 @@ class _Echo(chains.TrussChainlet): dst_config["secrets"][public_types.CHAIN_API_KEY_SECRET_NAME] != public_types.SECRET_DUMMY ) + + +# ---- Project 3: TrussChainlet as the chain entrypoint ----------------------- + + +def test_truss_chainlet_can_be_entrypoint(truss_dir_path): + """A `TrussChainlet` decorated with `@chains.mark_entrypoint` registers + as a valid entrypoint candidate. The `ChainletImporter._is_target_cls` + override is what unlocks this; without it the importer skips + TrussChainlets at discovery time.""" + truss_dir_str = str(truss_dir_path) + + @chains.mark_entrypoint("Polyglot Front Door") + class _Entry(chains.TrussChainlet): + truss_dir = truss_dir_str + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Entry) + assert descriptor.chainlet_cls.meta_data.is_entrypoint is True + assert descriptor.chainlet_cls.meta_data.chain_name == "Polyglot Front Door" + assert descriptor.is_truss_chainlet is True + # `ChainletImporter._is_target_cls` accepts this class. + assert framework.ChainletImporter._is_target_cls(_Entry) + + +def test_truss_chainlet_deps_class_attr_parses(truss_dir_path): + """A `TrussChainlet` with `deps = [SomeDep]` populates the + `dependencies` map on its descriptor — same shape ChainletBase produces + from `__init__` deps.""" + truss_dir_str = str(truss_dir_path) + + class _Sibling(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + async def run_remote(self, text: str) -> str: + return text + + @chains.mark_entrypoint("Entry") + class _Entry(chains.TrussChainlet): + truss_dir = truss_dir_str + deps = [_Sibling] + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Entry) + assert "_Sibling" in descriptor.dependencies + assert descriptor.dependencies["_Sibling"].chainlet_cls is _Sibling + + +def test_truss_chainlet_entrypoint_with_mixed_deps(truss_dir_path, tmp_path): + """Entrypoint TrussChainlet can declare deps that are a mix of + `ChainletBase` and other `TrussChainlet` types — both shapes are valid + sibling targets.""" + cb_sibling_truss_dir = str(truss_dir_path) + tc_sibling_truss_dir = tmp_path / "tc_sibling" + tc_sibling_truss_dir.mkdir() + (tc_sibling_truss_dir / "config.yaml").write_text(VALID_TRUSS_CONFIG_YAML) + (tc_sibling_truss_dir / "model").mkdir() + (tc_sibling_truss_dir / "model" / "model.py").write_text(VALID_MODEL_PY) + + class _CBSibling(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + async def run_remote(self, text: str) -> str: + return text + + class _TCSibling(chains.TrussChainlet): + truss_dir = str(tc_sibling_truss_dir) + + @chains.mark_entrypoint("Entry") + class _Entry(chains.TrussChainlet): + truss_dir = cb_sibling_truss_dir + deps = [_CBSibling, _TCSibling] + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Entry) + assert set(descriptor.dependencies.keys()) == {"_CBSibling", "_TCSibling"} + assert descriptor.dependencies["_CBSibling"].chainlet_cls is _CBSibling + assert descriptor.dependencies["_TCSibling"].chainlet_cls is _TCSibling + + +def test_truss_chainlet_entrypoint_duplicate_dep_rejected(truss_dir_path): + truss_dir_str = str(truss_dir_path) + + class _Sibling(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + async def run_remote(self, text: str) -> str: + return text + + @chains.mark_entrypoint("Entry") + class _Entry(chains.TrussChainlet): + truss_dir = truss_dir_str + deps = [_Sibling, _Sibling] + + with pytest.raises(public_types.ChainsUsageError, match="duplicate"): + framework.raise_validation_errors() + + +def test_truss_chainlet_entrypoint_self_dep_rejected(truss_dir_path): + truss_dir_str = str(truss_dir_path) + + # Self-reference via a forward declaration: a TrussChainlet that lists + # itself as a dep. We construct this by patching `deps` post-class. + @chains.mark_entrypoint("Entry") + class _Entry(chains.TrussChainlet): + truss_dir = truss_dir_str + + # Clear what the class declaration registered (since `deps = []` passed), + # then re-validate with self-reference. + framework._global_error_collector.clear() + framework._global_chainlet_registry.unregister_chainlet("_Entry") + _Entry.deps = [_Entry] + framework.validate_and_register_cls(_Entry) + with pytest.raises(public_types.ChainsUsageError, match="cannot reference itself"): + framework.raise_validation_errors() + + +def test_truss_chainlet_entrypoint_non_chainlet_dep_rejected(truss_dir_path): + """A non-chainlet class in `deps` is rejected with a clear error.""" + truss_dir_str = str(truss_dir_path) + + class _NotAChainlet: + pass + + @chains.mark_entrypoint("Entry") + class _Entry(chains.TrussChainlet): + truss_dir = truss_dir_str + deps = [_NotAChainlet] + + with pytest.raises( + public_types.ChainsUsageError, match="not a `ChainletBase` or `TrussChainlet`" + ): + framework.raise_validation_errors() + + +def test_truss_chainlet_entrypoint_deps_must_be_list(truss_dir_path): + """`deps` must be a list or tuple — anything else is a clear type error.""" + truss_dir_str = str(truss_dir_path) + + @chains.mark_entrypoint("Entry") + class _Entry(chains.TrussChainlet): + truss_dir = truss_dir_str + deps = "not a list" # type: ignore[assignment] + + with pytest.raises(public_types.ChainsUsageError, match="must be a list"): + framework.raise_validation_errors() + + +def test_truss_chainlet_entrypoint_codegen_preserves_truss_dir( + truss_dir_path, tmp_path +): + """When a TrussChainlet *is* the entrypoint, codegen still uses the + `_prepare_truss_chainlet_artifact` path (not the framework-generated + `model.py` path). The user's `truss_dir` is copied byte-for-byte, with + only `model_metadata.chains_metadata` injected.""" + from truss_chains.deployment import code_gen + + truss_dir_str = str(truss_dir_path) + + @chains.mark_entrypoint("Polyglot") + class _Entry(chains.TrussChainlet): + truss_dir = truss_dir_str + + framework.raise_validation_errors() + descriptor = framework.get_descriptor(_Entry) + chainlet_dir = code_gen.gen_truss_chainlet( + chain_root=tmp_path, + chain_name="codegen-test", + chainlet_descriptor=descriptor, + model_name="_Entry-abc12345", + ) + + # User's model.py is preserved. + assert (chainlet_dir / "model" / "model.py").read_text() == VALID_MODEL_PY + # config.yaml carries the uniquified model_name + chains_metadata. + dst_config = yaml.safe_load((chainlet_dir / "config.yaml").read_text()) + assert dst_config["model_name"] == "_Entry-abc12345" + assert private_types.TRUSS_CONFIG_CHAINS_KEY in ( + dst_config.get("model_metadata") or {} + ) + + +def test_chainletbase_ws_as_sibling_rejected(): + """WebSocket `ChainletBase` chainlets can only be entrypoints, not + siblings. The validator at `framework.py:_validate_dependencies` checks + `endpoint.is_websocket` on each dep. Locks in that behavior for the + Project 3 work — `TrussChainlet` entrypoints with a `ChainletBase`+WS + dep should also be blocked via the same predicate (we route + `_validate_truss_chainlet_deps` through the same check).""" + + class _WSChainlet(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + async def run_remote(self, websocket: public_types.WebSocketProtocol) -> None: + pass + + class _CallerCB(chains.ChainletBase): + remote_config = chains.RemoteConfig( + compute=chains.Compute(cpu_count=1, memory="512Mi") + ) + + def __init__(self, ws: _WSChainlet = chains.depends(_WSChainlet)) -> None: + self._ws = ws + + async def run_remote(self, text: str) -> str: + return text + + with pytest.raises( + public_types.ChainsUsageError, + match="websockets can only be used in the entrypoint", + ): + framework.raise_validation_errors() diff --git a/truss-chains/truss_chains/framework.py b/truss-chains/truss_chains/framework.py index c5ba91035..c11c0b4f9 100644 --- a/truss-chains/truss_chains/framework.py +++ b/truss-chains/truss_chains/framework.py @@ -994,13 +994,15 @@ def validate_and_register_cls(cls: Type[private_types.ABCChainlet]) -> None: if is_truss_chainlet(cls): # TrussChainlet declarations don't have remote_config / run_remote / - # health_check / dependencies. Validate the truss_dir + display_name, - # build a descriptor with a sentinel endpoint, register, and exit early. + # health_check. Validate the truss_dir + display_name + the optional + # ``deps`` class attribute, build a descriptor with a sentinel + # endpoint, register, and exit early. _validate_display_name(cls, location) _validate_truss_chainlet_cls(cls, src_path, location) + truss_chainlet_deps = _validate_truss_chainlet_deps(cls, location) chainlet_descriptor = private_types.ChainletAPIDescriptor( chainlet_cls=cls, - dependencies={}, + dependencies=truss_chainlet_deps, has_context=False, endpoint=_DUMMY_ENDPOINT_DESCRIPTOR, src_path=src_path, @@ -1086,6 +1088,84 @@ def _validate_truss_chainlet_cls( cls._resolved_truss_dir = truss_dir +def _validate_truss_chainlet_deps( + cls: Type["TrussChainlet"], location: _ErrorLocation +) -> Mapping[str, private_types.DependencyDescriptor]: + """Validate ``cls.deps`` (a class-attribute list of chainlet classes) and + return the resulting ``{display_name: DependencyDescriptor}`` mapping. + """ + deps_attr = getattr(cls, "deps", None) or [] + if not isinstance(deps_attr, (list, tuple)): + _collect_error( + f"`TrussChainlet.{cls.__name__}.deps` must be a list of chainlet " + f"classes, got `{type(deps_attr).__name__}`.", + _ErrorKind.TYPE_ERROR, + location, + ) + return {} + + used: set[Type[private_types.ABCChainlet]] = set() + dependencies: dict[str, private_types.DependencyDescriptor] = {} + for dep_cls in deps_attr: + if not isinstance(dep_cls, type): + _collect_error( + f"`TrussChainlet.{cls.__name__}.deps` entries must be chainlet " + f"classes, got `{dep_cls!r}`.", + _ErrorKind.TYPE_ERROR, + location, + ) + continue + if not ( + utils.issubclass_safe(dep_cls, ChainletBase) + or utils.issubclass_safe(dep_cls, TrussChainlet) + ): + _collect_error( + f"`TrussChainlet.{cls.__name__}.deps` entry `{dep_cls.__name__}` " + "is not a `ChainletBase` or `TrussChainlet` subclass.", + _ErrorKind.TYPE_ERROR, + location, + ) + continue + if dep_cls is cls: + _collect_error( + f"`TrussChainlet.{cls.__name__}.deps` cannot reference itself.", + _ErrorKind.TYPE_ERROR, + location, + ) + continue + if dep_cls in used: + _collect_error( + f"`TrussChainlet.{cls.__name__}.deps` contains duplicate " + f"`{dep_cls.__name__}`.", + _ErrorKind.TYPE_ERROR, + location, + ) + continue + if dep_cls not in _global_chainlet_registry._chainlets: + _collect_error( + f"`TrussChainlet.{cls.__name__}.deps` references `{dep_cls.__name__}` " + "before it was registered. Define the dep before the chainlet " + "that uses it.", + _ErrorKind.TYPE_ERROR, + location, + ) + continue + if get_descriptor(dep_cls).endpoint.is_websocket: + _collect_error( + f"`TrussChainlet.{cls.__name__}.deps` entry `{dep_cls.__name__}` " + "uses a websocket. WebSockets can only be used in the entrypoint, " + "not in 'inner' chainlets.", + _ErrorKind.TYPE_ERROR, + location, + ) + continue + dependencies[dep_cls.__name__] = private_types.DependencyDescriptor( + chainlet_cls=dep_cls, options=public_types.RPCOptions() + ) + used.add(dep_cls) + return dependencies + + # Dependency-Injection / Registry ###################################################### @@ -1492,6 +1572,16 @@ def _multiple_entrypoints_error( def _target_cls_type(cls) -> Type[private_types.ABCChainlet]: pass + @classmethod + def _is_target_cls(cls, candidate: Any) -> bool: + """Whether ``candidate`` is a class this importer should pick up. + + Default: any subclass of ``_target_cls_type()``. Subclasses can override + to widen acceptance to multiple base classes (e.g., ``ChainletImporter`` + accepts both ``ChainletBase`` and ``TrussChainlet``). + """ + return utils.issubclass_safe(candidate, cls._target_cls_type()) + @classmethod def _get_chainlets( cls, symbols @@ -1499,7 +1589,7 @@ def _get_chainlets( set[Type[private_types.ABCChainlet]], set[Type[private_types.ABCChainlet]] ]: chainlets: set[Type[private_types.ABCChainlet]] = { - sym for sym in symbols if utils.issubclass_safe(sym, cls._target_cls_type()) + sym for sym in symbols if cls._is_target_cls(sym) } entrypoints: set[Type[private_types.ABCChainlet]] = { chainlet for chainlet in chainlets if chainlet.meta_data.is_entrypoint @@ -1603,7 +1693,7 @@ def import_target( f"Target class `{target_name}` not found " f"in `{resolved_module_path}`." ) - if not utils.issubclass_safe(target_cls, cls._target_cls_type()): + if not cls._is_target_cls(target_cls): raise TypeError( f"Target `{target_cls}` is not a {cls._target_cls_type()}." ) @@ -1654,6 +1744,15 @@ def _multiple_entrypoints_error( def _target_cls_type(cls) -> Type[private_types.ABCChainlet]: return ChainletBase + @classmethod + def _is_target_cls(cls, candidate: Any) -> bool: + """Chains accept ``ChainletBase`` *or* ``TrussChainlet`` subclasses as + chainlets (and as entrypoints). ``TrussChainlet`` doesn't subclass + ``ChainletBase`` so we have to test both branches explicitly.""" + return utils.issubclass_safe(candidate, ChainletBase) or utils.issubclass_safe( + candidate, TrussChainlet + ) + class ModelImporter(_ABCImporter): @classmethod @@ -1795,6 +1894,22 @@ def __init__( self._url = whisper.target_url self._headers = whisper.with_auth_headers( context.get_baseten_api_key()) + + May also be marked as the chain entrypoint via ``@chains.mark_entrypoint``, + in which case it serves as the polyglot front door. Declare the deps the + entrypoint pod needs to discover via the ``deps`` class attribute:: + + @chains.mark_entrypoint("Polyglot Front Door") + class VoiceAgentEntry(chains.TrussChainlet): + truss_dir = "./voice-agent-rust" + deps = [LLM, TTS] + + The framework deploys ``LLM`` + ``TTS`` and writes their URLs into the + entrypoint pod's ``dynamic_chainlet_config`` ConfigMap. The entrypoint's + custom server reads them at runtime (e.g., via + :mod:`truss_chains.runtime` from Python or a native HTTP client from + other languages) and calls them with its own client. No framework Python + runs in the entrypoint pod's data path. """ truss_dir: ClassVar[str] @@ -1803,6 +1918,9 @@ def __init__( # and validation. Codegen reads it to locate the user's Truss directory. _resolved_truss_dir: ClassVar[Optional[pathlib.Path]] = None + # Deps for a TrussChainlet entrypoint. Bare chainlet class references + deps: ClassVar[list[type[private_types.ABCChainlet]]] = [] + def __init_subclass__(cls, **kwargs) -> None: super().__init_subclass__(**kwargs) cls._framework_config = private_types.FrameworkConfig( From 75f7f2a78c5b92ce9facc409f7165d46c848234d Mon Sep 17 00:00:00 2001 From: Matte Lim Date: Tue, 12 May 2026 21:55:15 -0700 Subject: [PATCH 12/12] Edit VoiceAgentMocked example to use chainlet urls --- .../examples/voice_agent_mocked/chain.py | 172 +++++------------- 1 file changed, 43 insertions(+), 129 deletions(-) diff --git a/truss-chains/examples/voice_agent_mocked/chain.py b/truss-chains/examples/voice_agent_mocked/chain.py index 35547a71c..89c52c736 100644 --- a/truss-chains/examples/voice_agent_mocked/chain.py +++ b/truss-chains/examples/voice_agent_mocked/chain.py @@ -8,8 +8,10 @@ - ``LLMMock`` (``TrussChainlet``, HTTP): ``{"prompt"}`` ``→`` ``{"completion": "echo: ..."}`` - ``TTSMock`` (``TrussChainlet``, WS-fronted): text ``→`` ``text.encode()`` - ``Orchestrator`` (``ChainletBase``, WS entrypoint): a single round-trip wires - STT ``→`` LLM ``→`` TTS using the same descriptor / runtime patterns the FDE - chain uses (GraphQL oracle-id resolution + proxy-env strip). + STT ``→`` LLM ``→`` TTS using the framework-injected ``DeployedServiceDescriptor`` + URLs and helpers — ``internal_ws_url`` / ``target_url`` for the URL, + ``with_ws_auth_headers`` / ``with_auth_headers`` for headers. No GraphQL, no + hardcoded hostnames. Iteration loop: @@ -22,9 +24,6 @@ import json import logging -from urllib.parse import urlparse - -import requests import truss_chains as chains @@ -50,103 +49,20 @@ class TTSMock(chains.TrussChainlet): truss_dir = "./tts_mock" -# ----- URL resolution (workaround for the platform routing gap) -------------- -# -# Same approach as the FDE chain: query the dashboard GraphQL with the -# user-supplied baseten_api_key to map chainlet name -> oracle.id, then build -# `model-.api.baseten.co/...` URLs. Replace with the descriptor's -# native fields once the platform fix lands. - -_GRAPHQL_QUERY = ( - "query Chain($id: String!) {\n" - " chain(id: $id) {\n" - " deployments {\n" - " id\n" - " chainlets {\n" - " name oracle { id } oracle_version { id transport_kind }\n" - " }\n" - " }\n" - " }\n" - "}" -) - - -def _chain_id_from_descriptor(desc: chains.DeployedServiceDescriptor) -> str: - """Extract chain id from ``internal_url.hostname`` (``chain-...``).""" - assert desc.internal_url is not None, ( - "internal_url is always populated for sibling DeployedServiceDescriptor at runtime" - ) - hostname = desc.internal_url.hostname - first = hostname.split(".", 1)[0] - if not first.startswith("chain-"): - raise RuntimeError(f"Unexpected internal_url hostname shape: {hostname!r}") - return first[len("chain-") :] - - -def _chain_deployment_id_from_descriptor(desc: chains.DeployedServiceDescriptor) -> str: - """Extract chain deployment id from the descriptor's gateway URL path.""" - assert desc.internal_url is not None, ( - "internal_url is always populated for sibling DeployedServiceDescriptor at runtime" - ) - parts = urlparse(desc.internal_url.gateway_run_remote_url).path.split("/") - # Path: /deployment//chainlet//run_remote - return parts[parts.index("deployment") + 1] - - -def _resolve_oracles( - api_key: str, chain_id: str, chain_deployment_id: str -) -> dict[str, dict]: - r = requests.post( - "https://app.baseten.co/graphql/", - headers={ - "Authorization": f"Api-Key {api_key}", - "Content-Type": "application/json", - }, - json={"query": _GRAPHQL_QUERY, "variables": {"id": chain_id}}, - timeout=10, - ) - r.raise_for_status() - payload = r.json() - chain_data = (payload.get("data") or {}).get("chain") or {} - deployments = chain_data.get("deployments") or [] - target = next((d for d in deployments if d["id"] == chain_deployment_id), None) - if target is None: - raise RuntimeError( - f"GraphQL has no deployment {chain_deployment_id} on chain {chain_id}" - ) - return { - c["name"]: { - "oracle_id": c["oracle"]["id"], - "oracle_version_id": c["oracle_version"]["id"], - "transport": c["oracle_version"]["transport_kind"], - } - for c in target["chainlets"] - } - - -def _model_url(oracle_id: str, oracle_version_id: str, scheme: str, path: str) -> str: - """Build a per-chainlet sibling URL using the published-deployment shape. - - ``://model-.api.baseten.co/deployment//`` - - Cluster-internally this shape works for both draft and published chainlets - (verified empirically on chain ``2328913l``); the equivalent - ``/development/`` alias is unnecessary. Using the deployment shape - uniformly keeps URLs stable across draft → published lifecycles. - """ - return ( - f"{scheme}://model-{oracle_id}.api.baseten.co" - f"/deployment/{oracle_version_id}/{path}" - ) - - # ----- Orchestrator: WebSocket entrypoint ------------------------------------ @chains.mark_entrypoint("VoiceAgentMocked") class Orchestrator(chains.ChainletBase): """Single-round-trip WS entrypoint: bytes in -> bytes out, exercising the - same STT/LLM/TTS sibling-call patterns as the FDE chain.""" + same STT/LLM/TTS sibling-call patterns as the FDE chain. + + Sibling URLs come straight from the framework-injected ``DeployedServiceDescriptor`` + objects — ``desc.ws_url`` / ``desc.target_url`` are populated by the runtime + layer from the platform's dynamic chainlet config. No GraphQL oracle-id + resolution, no hardcoded ``api.*.baseten.co`` host: the descriptor knows + where to go. + """ remote_config = chains.RemoteConfig( compute=chains.Compute(cpu_count=1, memory="512Mi"), @@ -173,45 +89,43 @@ def __init__( os.environ["NO_PROXY"] = "*" try: - user_api_key = context.secrets["baseten_api_key"] + user_api_key = context.secrets["baseten_chain_api_key"] except KeyError as e: raise RuntimeError( "baseten_api_key secret is missing. Set it in the workspace and" - " re-deploy. Used to resolve sibling oracle IDs via dashboard" - " GraphQL until the platform exposes them via dynamic_chainlet_config." + " re-deploy. Used in the Authorization header for sibling calls." ) from e if not user_api_key or user_api_key == "***": raise RuntimeError("baseten_api_key secret is empty/placeholder.") - chain_id = _chain_id_from_descriptor(stt) - chain_deployment_id = _chain_deployment_id_from_descriptor(stt) - oracles = _resolve_oracles(user_api_key, chain_id, chain_deployment_id) - log.warning( - "Resolved oracle map for chain=%s deployment=%s: %s", - chain_id, - chain_deployment_id, - oracles, - ) + # WS sibling URLs: prefer the cluster-local internal_ws_url (chain + # hostname + cluster routing); fall back to ws_url. WS auth headers + # are Authorization-only — see with_ws_auth_headers docstring for why + # an explicit Host override breaks WS handshakes. + stt_url = stt.internal_ws_url or stt.ws_url + tts_url = tts.internal_ws_url or tts.ws_url + if stt_url is None or tts_url is None: + raise RuntimeError( + "STT/TTS sibling chainlets must expose a WS URL " + f"(stt_url={stt_url!r}, tts_url={tts_url!r})." + ) + self._stt_url = stt_url + self._tts_url = tts_url + self._headers_stt = stt.with_ws_auth_headers(user_api_key) + self._headers_tts = tts.with_ws_auth_headers(user_api_key) + + # HTTP sibling URL: target_url uses the workload-plane gateway hostname + # for cluster-local routing; with_auth_headers adds the Host override + # api-gateway needs to map gateway-host requests onto the chain. + self._llm_url = llm.target_url + self._headers_llm = llm.with_auth_headers(user_api_key) - self._stt_url = _model_url( - oracles["STTMock"]["oracle_id"], - oracles["STTMock"]["oracle_version_id"], - "wss", - "websocket", - ) - self._llm_url = _model_url( - oracles["LLMMock"]["oracle_id"], - oracles["LLMMock"]["oracle_version_id"], - "https", - "predict", - ) - self._tts_url = _model_url( - oracles["TTSMock"]["oracle_id"], - oracles["TTSMock"]["oracle_version_id"], - "wss", - "websocket", + log.warning( + "Resolved sibling URLs: stt=%s llm=%s tts=%s", + self._stt_url, + self._llm_url, + self._tts_url, ) - self._headers = {"Authorization": f"Api-Key {user_api_key}"} async def run_remote(self, websocket: chains.WebSocketProtocol) -> None: import httpx @@ -222,7 +136,7 @@ async def run_remote(self, websocket: chains.WebSocketProtocol) -> None: # 1. STT: bytes -> text async with ws_client.connect( - self._stt_url, additional_headers=self._headers + self._stt_url, additional_headers=self._headers_stt ) as sws: await sws.send(audio_in) resp = await sws.recv() @@ -230,7 +144,7 @@ async def run_remote(self, websocket: chains.WebSocketProtocol) -> None: log.warning("STTMock returned: %s", text) # 2. LLM: text -> completion - async with httpx.AsyncClient(timeout=30, headers=self._headers) as http: + async with httpx.AsyncClient(timeout=30, headers=self._headers_llm) as http: r = await http.post(self._llm_url, json={"prompt": text}) r.raise_for_status() completion = r.json()["completion"] @@ -238,7 +152,7 @@ async def run_remote(self, websocket: chains.WebSocketProtocol) -> None: # 3. TTS: text -> bytes async with ws_client.connect( - self._tts_url, additional_headers=self._headers + self._tts_url, additional_headers=self._headers_tts ) as tws: await tws.send(json.dumps({})) # mock config frame, ignored await tws.send(completion)