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..48340ac21 --- /dev/null +++ b/truss-chains/examples/bring_your_own_client/README.md @@ -0,0 +1,103 @@ +# 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 + `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 + `truss_config.WebsocketOptions` plus a client that handles WS framing. 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 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..861560ce1 --- /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/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/__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 new file mode 100644 index 000000000..278c8f1dc --- /dev/null +++ b/truss-chains/examples/truss_chainlet_demo/chain.py @@ -0,0 +1,116 @@ +"""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/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/__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 new file mode 100644 index 000000000..89c52c736 --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/chain.py @@ -0,0 +1,171 @@ +"""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 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: + + 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 + +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" + + +# ----- 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. + + 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"), + 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_chain_api_key"] + except KeyError as e: + raise RuntimeError( + "baseten_api_key secret is missing. Set it in the workspace and" + " 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.") + + # 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) + + log.warning( + "Resolved sibling URLs: stt=%s llm=%s tts=%s", + self._stt_url, + self._llm_url, + self._tts_url, + ) + + 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_stt + ) 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_llm) 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_tts + ) as tws: + 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) + 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..e4c5ee82c --- /dev/null +++ b/truss-chains/examples/voice_agent_mocked/stt_mock/model/model.py @@ -0,0 +1,32 @@ +"""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 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/plain_truss/config.yaml b/truss-chains/tests/runtime_discovery/plain_truss/config.yaml new file mode 100644 index 000000000..d3452f772 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/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/plain_truss/model/model.py b/truss-chains/tests/runtime_discovery/plain_truss/model/model.py new file mode 100644 index 000000000..b6b2bb7f2 --- /dev/null +++ b/truss-chains/tests/runtime_discovery/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/test_runtime.py b/truss-chains/tests/test_runtime.py new file mode 100644 index 000000000..e4a08807c --- /dev/null +++ b/truss-chains/tests/test_runtime.py @@ -0,0 +1,472 @@ +"""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 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 = { + "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(): + """``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="chain-abc.api.baseten.co", + ), + ) + assert desc.internal_ws_url == "wss://chain-abc.api.baseten.co/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_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 + 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 + + +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 + 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" + # 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 ------------------------------------- + + +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 diff --git a/truss-chains/tests/test_truss_chainlet.py b/truss-chains/tests/test_truss_chainlet.py new file mode 100644 index 000000000..c11e5ef42 --- /dev/null +++ b/truss-chains/tests/test_truss_chainlet.py @@ -0,0 +1,737 @@ +"""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): + """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"): + 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) + + 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`"): + 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") + bad_path = str(bad) + + 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`"): + 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(): + """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"] == {} + + # `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" + + +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() + + +# ---- 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 + ) + + +# ---- 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/__init__.py b/truss-chains/truss_chains/__init__.py index a4d762cde..f85592040 100644 --- a/truss-chains/truss_chains/__init__.py +++ b/truss-chains/truss_chains/__init__.py @@ -1,5 +1,11 @@ from truss.base.truss_config import WeightsSource -from truss_chains.framework import ChainletBase, EngineBuilderLLMChainlet, ModelBase +from truss_chains import runtime # noqa: E402, F401 +from truss_chains.framework import ( + ChainletBase, + EngineBuilderLLMChainlet, + ModelBase, + TrussChainlet, +) from truss_chains.public_api import ( depends, depends_context, @@ -48,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..1f51d70fe 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,62 @@ 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], + 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. ``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( + 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(): + shutil.rmtree(chainlet_dir) + truss_path.copy_tree_path(src_truss_dir, chainlet_dir) + + config_path = chainlet_dir / serving_image_builder.CONFIG_FILE + 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] = ( + 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] = ( + 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 +1028,17 @@ 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, + model_name=model_name or chain_name, + ) 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..c11c0b4f9 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 @@ -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,26 @@ 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. 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=truss_chainlet_deps, + 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 +1044,128 @@ 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["TrussChainlet"], src_path: str, location: _ErrorLocation +) -> None: + """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: + _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. 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 + + +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 ###################################################### @@ -1421,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 @@ -1428,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 @@ -1532,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()}." ) @@ -1583,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 @@ -1691,3 +1861,86 @@ 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()) + + 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] + + # `_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 + + # 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( + 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], +) -> TypeGuard[Type["TrussChainlet"]]: + 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 + ######################################################################################## 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. diff --git a/truss-chains/truss_chains/public_types.py b/truss-chains/truss_chains/public_types.py index 1972881a0..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,6 +728,35 @@ class LookaheadDecodingConfig(pydantic.BaseModel): lookahead_decoding_config: Optional[LookaheadDecodingConfig] = None +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://"): + 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): """Bundles values to establish an RPC session to a dependency chainlet, specifically with ``StubBase``.""" @@ -756,6 +786,56 @@ 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]: + """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_url(self.predict_url) + + @property + def internal_ws_url(self) -> Optional[str]: + """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 + 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]: + """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. 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()}