-
Notifications
You must be signed in to change notification settings - Fork 75
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #1266 from basetenlabs/bump-version-0.9.55
Release 0.9.55
- Loading branch information
Showing
60 changed files
with
2,727 additions
and
1,151 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,92 @@ | ||
| import numpy as np | ||
| import pydantic | ||
|
|
||
| import truss_chains as chains | ||
| from truss_chains import pydantic_numpy | ||
|
|
||
|
|
||
| class DataModel(pydantic.BaseModel): | ||
| msg: str | ||
| np_array: pydantic_numpy.NumpyArrayField | ||
|
|
||
|
|
||
| class SyncChainlet(chains.ChainletBase): | ||
| def run_remote(self, data: DataModel) -> DataModel: | ||
| print(data) | ||
| return data.model_copy(update={"msg": "From sync"}) | ||
|
|
||
|
|
||
| class AsyncChainlet(chains.ChainletBase): | ||
| async def run_remote(self, data: DataModel) -> DataModel: | ||
| print(data) | ||
| return data.model_copy(update={"msg": "From async"}) | ||
|
|
||
|
|
||
| class AsyncChainletNoInput(chains.ChainletBase): | ||
| async def run_remote(self) -> DataModel: | ||
| data = DataModel(msg="From async no input", np_array=np.full((2, 2), 3)) | ||
| print(data) | ||
| return data | ||
|
|
||
|
|
||
| class AsyncChainletNoOutput(chains.ChainletBase): | ||
| async def run_remote(self, data: DataModel) -> None: | ||
| print(data) | ||
|
|
||
|
|
||
| class HostJSON(chains.ChainletBase): | ||
| """Calls various chainlets in JSON mode.""" | ||
|
|
||
| def __init__( | ||
| self, | ||
| sync_chainlet=chains.depends(SyncChainlet, use_binary=False), | ||
| async_chainlet=chains.depends(AsyncChainlet, use_binary=False), | ||
| async_chainlet_no_output=chains.depends( | ||
| AsyncChainletNoOutput, use_binary=False | ||
| ), | ||
| async_chainlet_no_input=chains.depends(AsyncChainletNoInput, use_binary=False), | ||
| ): | ||
| self._sync_chainlet = sync_chainlet | ||
| self._async_chainlet = async_chainlet | ||
| self._async_chainlet_no_output = async_chainlet_no_output | ||
| self._async_chainlet_no_input = async_chainlet_no_input | ||
|
|
||
| async def run_remote(self) -> tuple[DataModel, DataModel, DataModel]: | ||
| a = np.ones((3, 2, 1)) | ||
| data = DataModel(msg="From Host", np_array=a) | ||
| sync_result = self._sync_chainlet.run_remote(data) | ||
| print(sync_result) | ||
| async_result = await self._async_chainlet.run_remote(data) | ||
| print(async_result) | ||
| await self._async_chainlet_no_output.run_remote(data) | ||
| async_no_input = await self._async_chainlet_no_input.run_remote() | ||
| print(async_no_input) | ||
| return sync_result, async_result, async_no_input | ||
|
|
||
|
|
||
| class HostBinary(chains.ChainletBase): | ||
| """Calls various chainlets in binary mode.""" | ||
|
|
||
| def __init__( | ||
| self, | ||
| sync_chainlet=chains.depends(SyncChainlet, use_binary=True), | ||
| async_chainlet=chains.depends(AsyncChainlet, use_binary=True), | ||
| async_chainlet_no_output=chains.depends(AsyncChainletNoOutput, use_binary=True), | ||
| async_chainlet_no_input=chains.depends(AsyncChainletNoInput, use_binary=True), | ||
| ): | ||
| self._sync_chainlet = sync_chainlet | ||
| self._async_chainlet = async_chainlet | ||
| self._async_chainlet_no_output = async_chainlet_no_output | ||
| self._async_chainlet_no_input = async_chainlet_no_input | ||
|
|
||
| async def run_remote(self) -> tuple[DataModel, DataModel, DataModel]: | ||
| a = np.ones((3, 2, 1)) | ||
| data = DataModel(msg="From Host", np_array=a) | ||
| sync_result = self._sync_chainlet.run_remote(data) | ||
| print(sync_result) | ||
| async_result = await self._async_chainlet.run_remote(data) | ||
| print(async_result) | ||
| await self._async_chainlet_no_output.run_remote(data) | ||
| async_no_input = await self._async_chainlet_no_input.run_remote() | ||
| print(async_no_input) | ||
| return sync_result, async_result, async_no_input |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,107 @@ | ||
| import asyncio | ||
| import time | ||
| from typing import AsyncIterator | ||
|
|
||
| import pydantic | ||
|
|
||
| import truss_chains as chains | ||
| from truss_chains import streaming | ||
|
|
||
|
|
||
| class Header(pydantic.BaseModel): | ||
| time: float | ||
| msg: str | ||
|
|
||
|
|
||
| class MyDataChunk(pydantic.BaseModel): | ||
| words: list[str] | ||
|
|
||
|
|
||
| class Footer(pydantic.BaseModel): | ||
| time: float | ||
| duration_sec: float | ||
| msg: str | ||
|
|
||
|
|
||
| class ConsumerOutput(pydantic.BaseModel): | ||
| header: Header | ||
| chunks: list[MyDataChunk] | ||
| footer: Footer | ||
| strings: str | ||
|
|
||
|
|
||
| STREAM_TYPES = streaming.stream_types( | ||
| MyDataChunk, header_type=Header, footer_type=Footer | ||
| ) | ||
|
|
||
|
|
||
| class Generator(chains.ChainletBase): | ||
| """Example that streams fully structured pydantic items with header and footer.""" | ||
|
|
||
| async def run_remote(self) -> AsyncIterator[bytes]: | ||
| print("Entering Generator") | ||
| streamer = streaming.stream_writer(STREAM_TYPES) | ||
| header = Header(time=time.time(), msg="Start.") | ||
| yield streamer.yield_header(header) | ||
| for i in range(1, 5): | ||
| data = MyDataChunk( | ||
| words=[chr(x + 70) * x for x in range(1, i + 1)], | ||
| ) | ||
| print("Yield") | ||
| yield streamer.yield_item(data) | ||
| await asyncio.sleep(0.05) | ||
|
|
||
| end_time = time.time() | ||
| footer = Footer(time=end_time, duration_sec=end_time - header.time, msg="Done.") | ||
| yield streamer.yield_footer(footer) | ||
| print("Exiting Generator") | ||
|
|
||
|
|
||
| class StringGenerator(chains.ChainletBase): | ||
| """Minimal streaming example with strings (e.g. for raw LLM output).""" | ||
|
|
||
| async def run_remote(self) -> AsyncIterator[str]: | ||
| # Note: the "chunk" boundaries are lost, when streaming raw strings. You must | ||
| # add spaces and linebreaks to the items yourself.. | ||
| yield "First " | ||
| yield "second " | ||
| yield "last." | ||
|
|
||
|
|
||
| class Consumer(chains.ChainletBase): | ||
| """Consume that reads the raw streams and parses them.""" | ||
|
|
||
| def __init__( | ||
| self, | ||
| generator=chains.depends(Generator), | ||
| string_generator=chains.depends(StringGenerator), | ||
| ): | ||
| self._generator = generator | ||
| self._string_generator = string_generator | ||
|
|
||
| async def run_remote(self) -> ConsumerOutput: | ||
| print("Entering Consumer") | ||
| reader = streaming.stream_reader(STREAM_TYPES, self._generator.run_remote()) | ||
| print("Consuming...") | ||
| header = await reader.read_header() | ||
| chunks = [] | ||
| async for data in reader.read_items(): | ||
| print(f"Read: {data}") | ||
| chunks.append(data) | ||
|
|
||
| footer = await reader.read_footer() | ||
| strings = [] | ||
| async for part in self._string_generator.run_remote(): | ||
| strings.append(part) | ||
|
|
||
| print("Exiting Consumer") | ||
| return ConsumerOutput( | ||
| header=header, chunks=chunks, footer=footer, strings="".join(strings) | ||
| ) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| with chains.run_local(): | ||
| chain = Consumer() | ||
| result = asyncio.run(chain.run_remote()) | ||
| print(result) |
Oops, something went wrong.