From 4885f19e9f4aee0680fab68a5acd2ff8c00c7d4e Mon Sep 17 00:00:00 2001 From: Mattt Date: Wed, 25 Sep 2024 11:37:12 -0700 Subject: [PATCH 1/3] Add `wait` parameter to prediction creation methods (#354) Signed-off-by: Mattt Zmuda --- replicate/deployment.py | 5 +++++ replicate/model.py | 5 +++++ replicate/prediction.py | 27 +++++++++++++++++++++++++++ 3 files changed, 37 insertions(+) diff --git a/replicate/deployment.py b/replicate/deployment.py index e17edcbc..1f0fdaba 100644 --- a/replicate/deployment.py +++ b/replicate/deployment.py @@ -8,6 +8,7 @@ from replicate.prediction import ( Prediction, _create_prediction_body, + _create_prediction_headers, _json_to_prediction, ) from replicate.resource import Namespace, Resource @@ -425,12 +426,14 @@ def create( client=self._client, file_encoding_strategy=file_encoding_strategy, ) + headers = _create_prediction_headers(wait=params.pop("wait", None)) body = _create_prediction_body(version=None, input=input, **params) resp = self._client._request( "POST", f"/v1/deployments/{self._deployment.owner}/{self._deployment.name}/predictions", json=body, + headers=headers, ) return _json_to_prediction(self._client, resp.json()) @@ -451,12 +454,14 @@ async def async_create( client=self._client, file_encoding_strategy=file_encoding_strategy, ) + headers = _create_prediction_headers(wait=params.pop("wait", None)) body = _create_prediction_body(version=None, input=input, **params) resp = await self._client._async_request( "POST", f"/v1/deployments/{self._deployment.owner}/{self._deployment.name}/predictions", json=body, + headers=headers, ) return _json_to_prediction(self._client, resp.json()) diff --git a/replicate/model.py b/replicate/model.py index ba5e1113..31f625af 100644 --- a/replicate/model.py +++ b/replicate/model.py @@ -9,6 +9,7 @@ from replicate.prediction import ( Prediction, _create_prediction_body, + _create_prediction_headers, _json_to_prediction, ) from replicate.resource import Namespace, Resource @@ -400,12 +401,14 @@ def create( client=self._client, file_encoding_strategy=file_encoding_strategy, ) + headers = _create_prediction_headers(wait=params.pop("wait", None)) body = _create_prediction_body(version=None, input=input, **params) resp = self._client._request( "POST", url, json=body, + headers=headers, ) return _json_to_prediction(self._client, resp.json()) @@ -429,12 +432,14 @@ async def async_create( client=self._client, file_encoding_strategy=file_encoding_strategy, ) + headers = _create_prediction_headers(wait=params.pop("wait", None)) body = _create_prediction_body(version=None, input=input, **params) resp = await self._client._async_request( "POST", url, json=body, + headers=headers, ) return _json_to_prediction(self._client, resp.json()) diff --git a/replicate/prediction.py b/replicate/prediction.py index 9770029b..d09ef504 100644 --- a/replicate/prediction.py +++ b/replicate/prediction.py @@ -383,6 +383,15 @@ class CreatePredictionParams(TypedDict): stream: NotRequired[bool] """Enable streaming of prediction output.""" + wait: NotRequired[Union[int, bool]] + """ + Wait until the prediction is completed before returning. + + If `True`, wait a predetermined number of seconds until the prediction + is completed before returning. + If an `int`, wait for the specified number of seconds. + """ + file_encoding_strategy: NotRequired[FileEncodingStrategy] """The strategy to use for encoding files in the prediction input.""" @@ -463,6 +472,7 @@ def create( # type: ignore client=self._client, file_encoding_strategy=file_encoding_strategy, ) + headers = _create_prediction_headers(wait=params.pop("wait", None)) body = _create_prediction_body( version, input, @@ -472,6 +482,7 @@ def create( # type: ignore resp = self._client._request( "POST", "/v1/predictions", + headers=headers, json=body, ) @@ -554,6 +565,7 @@ async def async_create( # type: ignore client=self._client, file_encoding_strategy=file_encoding_strategy, ) + headers = _create_prediction_headers(wait=params.pop("wait", None)) body = _create_prediction_body( version, input, @@ -563,6 +575,7 @@ async def async_create( # type: ignore resp = await self._client._async_request( "POST", "/v1/predictions", + headers=headers, json=body, ) @@ -603,6 +616,20 @@ async def async_cancel(self, id: str) -> Prediction: return _json_to_prediction(self._client, resp.json()) +def _create_prediction_headers( + *, + wait: Optional[Union[int, bool]] = None, +) -> Dict[str, Any]: + headers = {} + + if wait: + if isinstance(wait, bool): + headers["Prefer"] = "wait" + elif isinstance(wait, int): + headers["Prefer"] = f"wait={wait}" + return headers + + def _create_prediction_body( # pylint: disable=too-many-arguments version: Optional[Union[Version, str]], input: Optional[Dict[str, Any]], From e53bd02f496db1e039075f0367797d9573fde33b Mon Sep 17 00:00:00 2001 From: Mattt Date: Wed, 25 Sep 2024 11:37:21 -0700 Subject: [PATCH 2/3] Add `use_file_output` to streaming methods (#355) Signed-off-by: Mattt Zmuda --- replicate/client.py | 6 ++++-- replicate/prediction.py | 18 ++++++++++++++---- replicate/stream.py | 34 ++++++++++++++++++++++++++++++---- 3 files changed, 48 insertions(+), 10 deletions(-) diff --git a/replicate/client.py b/replicate/client.py index 3da3cc15..52d07f70 100644 --- a/replicate/client.py +++ b/replicate/client.py @@ -190,25 +190,27 @@ def stream( self, ref: str, input: Optional[Dict[str, Any]] = None, + use_file_output: Optional[bool] = None, **params: Unpack["Predictions.CreatePredictionParams"], ) -> Iterator["ServerSentEvent"]: """ Stream a model's output. """ - return stream(self, ref, input, **params) + return stream(self, ref, input, use_file_output, **params) async def async_stream( self, ref: str, input: Optional[Dict[str, Any]] = None, + use_file_output: Optional[bool] = None, **params: Unpack["Predictions.CreatePredictionParams"], ) -> AsyncIterator["ServerSentEvent"]: """ Stream a model's output asynchronously. """ - return async_stream(self, ref, input, **params) + return async_stream(self, ref, input, use_file_output, **params) # Adapted from https://github.com/encode/httpx/issues/108#issuecomment-1132753155 diff --git a/replicate/prediction.py b/replicate/prediction.py index d09ef504..0e5342a1 100644 --- a/replicate/prediction.py +++ b/replicate/prediction.py @@ -153,7 +153,10 @@ async def async_wait(self) -> None: await asyncio.sleep(self._client.poll_interval) await self.async_reload() - def stream(self) -> Iterator["ServerSentEvent"]: + def stream( + self, + use_file_output: Optional[bool] = None, + ) -> Iterator["ServerSentEvent"]: """ Stream the prediction output. @@ -170,9 +173,14 @@ def stream(self) -> Iterator["ServerSentEvent"]: headers["Cache-Control"] = "no-store" with self._client._client.stream("GET", url, headers=headers) as response: - yield from EventSource(response) + yield from EventSource( + self._client, response, use_file_output=use_file_output + ) - async def async_stream(self) -> AsyncIterator["ServerSentEvent"]: + async def async_stream( + self, + use_file_output: Optional[bool] = None, + ) -> AsyncIterator["ServerSentEvent"]: """ Stream the prediction output asynchronously. @@ -194,7 +202,9 @@ async def async_stream(self) -> AsyncIterator["ServerSentEvent"]: async with self._client._async_client.stream( "GET", url, headers=headers ) as response: - async for event in EventSource(response): + async for event in EventSource( + self._client, response, use_file_output=use_file_output + ): yield event def cancel(self) -> None: diff --git a/replicate/stream.py b/replicate/stream.py index 3472799e..4cf0d156 100644 --- a/replicate/stream.py +++ b/replicate/stream.py @@ -15,6 +15,7 @@ from replicate import identifier from replicate.exceptions import ReplicateError +from replicate.helpers import transform_output try: from pydantic import v1 as pydantic # type: ignore @@ -62,10 +63,19 @@ class EventSource: A server-sent event source. """ + client: "Client" response: "httpx.Response" - - def __init__(self, response: "httpx.Response") -> None: + use_file_output: bool + + def __init__( + self, + client: "Client", + response: "httpx.Response", + use_file_output: Optional[bool] = None, + ) -> None: + self.client = client self.response = response + self.use_file_output = use_file_output or False content_type, _, _ = response.headers["content-type"].partition(";") if content_type != "text/event-stream": raise ValueError( @@ -147,6 +157,12 @@ def __iter__(self) -> Iterator[ServerSentEvent]: if sse.event == ServerSentEvent.EventType.ERROR: raise RuntimeError(sse.data) + if ( + self.use_file_output + and sse.event == ServerSentEvent.EventType.OUTPUT + ): + sse.data = transform_output(sse.data, client=self.client) + yield sse if sse.event == ServerSentEvent.EventType.DONE: @@ -161,6 +177,12 @@ async def __aiter__(self) -> AsyncIterator[ServerSentEvent]: if sse.event == ServerSentEvent.EventType.ERROR: raise RuntimeError(sse.data) + if ( + self.use_file_output + and sse.event == ServerSentEvent.EventType.OUTPUT + ): + sse.data = transform_output(sse.data, client=self.client) + yield sse if sse.event == ServerSentEvent.EventType.DONE: @@ -171,6 +193,7 @@ def stream( client: "Client", ref: Union["Model", "Version", "ModelVersionIdentifier", str], input: Optional[Dict[str, Any]] = None, + use_file_output: Optional[bool] = None, **params: Unpack["Predictions.CreatePredictionParams"], ) -> Iterator[ServerSentEvent]: """ @@ -204,13 +227,14 @@ def stream( headers["Cache-Control"] = "no-store" with client._client.stream("GET", url, headers=headers) as response: - yield from EventSource(response) + yield from EventSource(client, response, use_file_output=use_file_output) async def async_stream( client: "Client", ref: Union["Model", "Version", "ModelVersionIdentifier", str], input: Optional[Dict[str, Any]] = None, + use_file_output: Optional[bool] = None, **params: Unpack["Predictions.CreatePredictionParams"], ) -> AsyncIterator[ServerSentEvent]: """ @@ -244,7 +268,9 @@ async def async_stream( headers["Cache-Control"] = "no-store" async with client._async_client.stream("GET", url, headers=headers) as response: - async for event in EventSource(response): + async for event in EventSource( + client, response, use_file_output=use_file_output + ): yield event From d5f5e274ef679c1483d51f61e197fd0cc4ddfef4 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Wed, 25 Sep 2024 11:37:52 -0700 Subject: [PATCH 3/3] 0.34.0 Signed-off-by: Mattt Zmuda --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 82e90bdc..732fb4fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "replicate" -version = "0.33.0" +version = "0.34.0" description = "Python client for Replicate" readme = "README.md" license = { file = "LICENSE" }