truss 0.11.1rc15__py3-none-any.whl → 0.11.2rc2__py3-none-any.whl

truss/cli/train/deploy_checkpoints/deploy_checkpoints.py

@@ -296,18 +296,6 @@ def _get_checkpoint_ids_to_deploy(
     return checkpoint_ids
 
 
-def _select_single_checkpoint(checkpoint_id_options: List[str]) -> List[str]:
-    """Select a single checkpoint using interactive prompt."""
-    checkpoint_id = inquirer.select(
-        message="Select the checkpoints to deploy:", choices=checkpoint_id_options
-    ).execute()
-
-    if not checkpoint_id:
-        raise click.UsageError("A checkpoint must be selected.")
-
-    return [checkpoint_id]
-
-
 def _select_multiple_checkpoints(checkpoint_id_options: List[str]) -> List[str]:
     """Select multiple checkpoints using interactive checkbox."""
     checkpoint_ids = inquirer.checkbox(

truss/templates/control/requirements.txt

@@ -6,7 +6,7 @@ loguru>=0.7.2
 python-json-logger>=2.0.2
 tenacity>=8.1.0
  # To avoid divergence, this should follow the latest release.
-truss==0.11.1rc15
+truss==0.11.1
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.1rc15.dist-info → truss-0.11.2rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.1rc15
+Version: 0.11.2rc2
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co

{truss-0.11.1rc15.dist-info → truss-0.11.2rc2.dist-info}/RECORD

@@ -24,7 +24,7 @@ truss/cli/train/metrics_watcher.py,sha256=smz-zrEsBj_-wJHI0pAZ-EAPrvfCWzq1eQjGiF
 truss/cli/train/poller.py,sha256=TGRzELxsicga0bEXewSX1ujw6lfPmDnHd6nr8zvOFO8,3550
 truss/cli/train/types.py,sha256=alGtr4Q71GeB65PpGMhsoKygw4k_ncR6MKIP1ioP8rI,951
 truss/cli/train/deploy_checkpoints/__init__.py,sha256=wL-M2yu8PxO2tFvjwshXAfPnB-5TlvsBp2v_bdzimRU,99
-truss/cli/train/deploy_checkpoints/deploy_checkpoints.py,sha256=KUaUl5a2lGy-l5rZZJv-rxIAsCJPRNC8h1909fXVKtE,17763
+truss/cli/train/deploy_checkpoints/deploy_checkpoints.py,sha256=nJHUjR4f4_13mFPNetWKq7ecsqr-cradbv3RBBDj2pk,17364
 truss/cli/train/deploy_checkpoints/deploy_checkpoints_helpers.py,sha256=6x5nS_HnWYtS9vi-Pg8akzrJk9L_agjvFhm5EFh1m6Y,1964
 truss/cli/train/deploy_checkpoints/deploy_full_checkpoints.py,sha256=FYRG5KTMlxEMZS-RA_m2gp1wuqWbSpqt2RhdQfLibhA,3968
 truss/cli/train/deploy_checkpoints/deploy_lora_checkpoints.py,sha256=P91dIAzuhl2GlzmrWwCcYI7uCMT1Lm7C79JQHM_exN4,4442
@@ -72,7 +72,7 @@ truss/templates/cache_requirements.txt,sha256=xoPoJ-OVnf1z6oq_RVM3vCr3ionByyqMLj
 truss/templates/copy_cache_files.Dockerfile.jinja,sha256=Os5zFdYLZ_AfCRGq4RcpVTObOTwL7zvmwYcvOzd_Zqo,126
 truss/templates/docker_server_requirements.txt,sha256=PyhOPKAmKW1N2vLvTfLMwsEtuGpoRrbWuNo7tT6v2Mc,18
 truss/templates/server.Dockerfile.jinja,sha256=CUYnF_hgxPGq2re7__0UPWlwzOHMoFkxp6NVKi3U16s,7071
-truss/templates/control/requirements.txt,sha256=x3hyJl3abQQhOsx7SuBYgk_PHkwbyXJ4N_g7vkXYIu0,254
+truss/templates/control/requirements.txt,sha256=tJGr83WoE0CZm2FrloZ9VScK84q-_FTuVXjDYrexhW0,250
 truss/templates/control/control/application.py,sha256=jYeta6hWe1SkfLL3W4IDmdYjg3ZuKqI_UagWYs5RB_E,3793
 truss/templates/control/control/endpoints.py,sha256=VQ1lvZjFvR091yRkiFdvXw1Q7PiNGXT9rJwY7_sX6yg,11828
 truss/templates/control/control/server.py,sha256=R4Y219i1dcz0kkksN8obLoX-YXWGo9iW1igindyG50c,3128
@@ -348,27 +348,27 @@ truss_chains/__init__.py,sha256=QDw1YwdqMaQpz5Oltu2Eq2vzEX9fDrMoqnhtbeh60i4,1278
 truss_chains/framework.py,sha256=CS7tSegPe2Q8UUT6CDkrtSrB3utr_1QN1jTEPjrj5Ug,67519
 truss_chains/private_types.py,sha256=6CaQEPawFLXjEbJ-01lqfexJtUIekF_q61LNENWegFo,8917
 truss_chains/public_api.py,sha256=0AXV6UdZIFAMycUNG_klgo4aLFmBZeKGfrulZEWzR0M,9532
-truss_chains/public_types.py,sha256=8cBstm5m_9aN_U2WoSHQ9s2z6gTdSEl-nVgXyneU3nY,29166
+truss_chains/public_types.py,sha256=Am1vc5pWuEDs65UQp8Be4iOU05kMXjPG3QrJfNXDmHs,29225
 truss_chains/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/pydantic_numpy.py,sha256=MG8Ji_Inwo_JSfM2n7TPj8B-nbrBlDYsY3SOeBwD8fE,4289
 truss_chains/streaming.py,sha256=DGl2LEAN67YwP7Nn9MK488KmYc4KopWmcHuE6WjyO1Q,12521
 truss_chains/utils.py,sha256=LvpCG2lnN6dqPqyX3PwLH9tyjUzqQN3N4WeEFROMHak,6291
 truss_chains/deployment/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truss_chains/deployment/code_gen.py,sha256=AmAUZ3h1hP3uYkl3J6o096K5RFLuBOP7kOFSnFC_C4U,32568
+truss_chains/deployment/code_gen.py,sha256=IBOYdhsWUyW_sBVhlEQAhvwxKcsGflDjgmR-1HyJJLg,32666
 truss_chains/deployment/deployment_client.py,sha256=haFiVmQek42ewlN_YflBaRDQT4ZYbmT20tvvJOkcUX0,32899
 truss_chains/reference_code/reference_chainlet.py,sha256=5feSeqGtrHDbldkfZCfX2R5YbbW0Uhc35mhaP2pXrHw,1340
 truss_chains/reference_code/reference_model.py,sha256=emH3hb23E_nbP98I37PGp1Xk1hz3g3lQ00tiLo55cSM,322
 truss_chains/remote_chainlet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/remote_chainlet/model_skeleton.py,sha256=8ZReLOO2MLcdg7bNZ61C-6j-e68i2Z-fFlyV3sz0qH8,2376
 truss_chains/remote_chainlet/stub.py,sha256=Y2gDUzMY9WRaQNHIz-o4dfLUfFyYV9dUhIRQcfgrY8g,17209
-truss_chains/remote_chainlet/utils.py,sha256=nnJQG3hXIoGghWEGPIAqKfRaueCLgSdwNI2x9B7BSaY,23080
+truss_chains/remote_chainlet/utils.py,sha256=Zn3GZRvK8f65WUa-qa-8uPFZ2pD7ukRFxbLOvT-BL0Q,24063
 truss_train/__init__.py,sha256=7hE6j6-u6UGzCGaNp3CsCN0kAVjBus1Ekups-Bk0fi4,837
 truss_train/definitions.py,sha256=V985HhY4rdXL10DZxpFEpze9ScxzWErMht4WwaPknGU,6789
 truss_train/deployment.py,sha256=lWWANSuzBWu2M4oK4qD7n-oVR1JKdmw2Pn5BJQHg-Ck,3074
 truss_train/loader.py,sha256=0o66EjBaHc2YY4syxxHVR4ordJWs13lNXnKjKq2wq0U,1630
 truss_train/public_api.py,sha256=9N_NstiUlmBuLUwH_fNG_1x7OhGCytZLNvqKXBlStrM,1220
-truss-0.11.1rc15.dist-info/METADATA,sha256=FEwmol_o_slcTpiH9EWPFnpTPTlVQwPImFV2jCvNhdw,6673
-truss-0.11.1rc15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.1rc15.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.1rc15.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.1rc15.dist-info/RECORD,,
+truss-0.11.2rc2.dist-info/METADATA,sha256=l4V9tHBGTjs14pZeFWIR5VCFsS1ArDKc3er-aKHkGqM,6672
+truss-0.11.2rc2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.2rc2.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.2rc2.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.2rc2.dist-info/RECORD,,

truss_chains/deployment/code_gen.py

@@ -544,7 +544,11 @@ async def websocket(self, websocket: fastapi.WebSocket) -> None:
         )"""
     return _Source(
         src=src,
-        imports={"import fastapi", "from truss_chains.remote_chainlet import utils"},
+        imports={
+            "import fastapi",
+            "from starlette.websockets import WebSocketState",
+            "from truss_chains.remote_chainlet import utils",
+        },
     )
 
 

truss_chains/public_types.py

@@ -482,9 +482,11 @@ class WebSocketProtocol(Protocol):
     async def send_bytes(self, data: bytes) -> None: ...
     async def send_json(self, data: Any) -> None: ...
 
-    def iter_text(self) -> AsyncIterator[str]: ...
-    def iter_bytes(self) -> AsyncIterator[bytes]: ...
-    def iter_json(self) -> AsyncIterator[Any]: ...
+    async def iter_text(self) -> AsyncIterator[str]: ...
+    async def iter_bytes(self) -> AsyncIterator[bytes]: ...
+    async def iter_json(self) -> AsyncIterator[Any]: ...
+
+    def is_connected(self) -> bool: ...
 
 
 class EngineBuilderLLMInput(pydantic.BaseModel):

truss_chains/remote_chainlet/utils.py

@@ -39,7 +39,7 @@ if TYPE_CHECKING:
 try:
     import prometheus_client
 except ImportError:
-    logging.warning("Optional `prometheus_client` is not installed. ")
+    logging.warning("Optional `prometheus_client` is not installed.")
 
     class _NoOpMetric:
         def labels(self, *args: object, **kwargs: object) -> "_NoOpMetric":
@@ -65,6 +65,26 @@ except ImportError:
             return _NoOpMetric()
 
 
+try:
+    from fastapi import WebSocketDisconnect
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketDisconnect, in case local environment doesn't have
+    # fastapi.
+    class WebSocketDisconnect(Exception):  # type: ignore[no-redef]
+        def __init__(self, code: int, reason: Optional[str] = None):
+            super().__init__()
+            self.code = code
+            self.reason = reason
+
+
+try:
+    from starlette.websockets import WebSocketState
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketState, in case local environment doesn't have starlette.
+    class WebSocketState:  # type: ignore[no-redef]
+        CONNECTED = "connected"
+
+
 T = TypeVar("T")
 
 _LockT = TypeVar("_LockT", bound=Union[threading.Lock, asyncio.Lock])
@@ -588,17 +608,12 @@ class WebsocketWrapperFastAPI:
         await self._websocket.close(code=code, reason=reason)
 
     async def receive(self) -> Union[str, bytes]:
-        try:
-            import fastapi
-        except ImportError:
-            raise utils.make_optional_import_error("fastapi")
-
         message = await self._websocket.receive()
 
         if message.get("type") == "websocket.disconnect":
             # NB(nikhil): Mimics FastAPI `_raise_on_disconnect`, since otherwise the user has no
             # way of detecting that the client disconnected.
-            raise fastapi.WebSocketDisconnect(message["code"], message.get("reason"))
+            raise WebSocketDisconnect(message["code"], message.get("reason"))
         elif message.get("text"):
             return typing.cast(str, message["text"])
         else:
@@ -623,13 +638,19 @@ class WebsocketWrapperFastAPI:
         await self._websocket.send_json(data)
 
     async def iter_text(self) -> AsyncIterator[str]:
-        while True:
-            yield await self.receive_text()
+        return self._websocket.iter_text()
 
     async def iter_bytes(self) -> AsyncIterator[bytes]:
-        while True:
-            yield await self.receive_bytes()
+        return self._websocket.iter_bytes()
 
     async def iter_json(self) -> AsyncIterator[Any]:
-        while True:
-            yield await self.receive_json()
+        return self._websocket.iter_json()
+
+    async def is_connected(self) -> bool:
+        # NB(nikhil): This isn't a foolproof mechanism for detecting whether a websocket
+        # connection is actually alive, ping/pong messages are best suited for that. However,
+        # as a heuristic to determine if a message is safe to send, this can do a pretty good job.
+        return (
+            self._websocket.application_state == WebSocketState.CONNECTED
+            and self._websocket.client_state == WebSocketState.CONNECTED
+        )