Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable multi-prompt input for persistent deployments #275

Merged
merged 10 commits into from
Nov 10, 2023
Merged

Enable multi-prompt input for persistent deployments #275

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
be18a00
enable multi-prompt inputs to persistent deployment
mrwyattii Nov 9, 2023
35fa078
Merge branch 'main' into mrwyattii/multiprompt-input
mrwyattii Nov 9, 2023
5e25585
add flush_uid method
mrwyattii Nov 9, 2023
761f500
fix missing self.lock
mrwyattii Nov 9, 2023
2862627
prototyping multi-prompt support with session_id
mrwyattii Nov 9, 2023
7e7301f
remove session_id
mrwyattii Nov 9, 2023
26db3a4
use tid instead of uid in get_response
mrwyattii Nov 9, 2023
9b5343f
add comments to new code
mrwyattii Nov 9, 2023
02875ba
add unit test
mrwyattii Nov 9, 2023
f8a1e72
fix streaming functionality
mrwyattii Nov 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 11 additions & 29 deletions mii/backend/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@
import asyncio
import grpc
import requests
from typing import Dict, Any, Callable
from typing import Dict, Any, Callable, List, Union

from mii.config import MIIConfig
from mii.constants import GRPC_MAX_MSG_SIZE, TaskType
from mii.constants import GRPC_MAX_MSG_SIZE
from mii.grpc_related.proto import modelresponse_pb2, modelresponse_pb2_grpc
from mii.grpc_related.task_methods import TASK_METHODS_DICT

Expand Down Expand Up @@ -55,18 +55,21 @@ async def _request_async_response_stream(self, request_dict, **query_kwargs):
yield task_methods.unpack_response_from_proto(response)

def generate(self,
prompt: str,
prompts: Union[str,
List[str]],
streaming_fn: Callable = None,
**query_kwargs: Dict[str,
Any]):
if not isinstance(prompt, str):
raise RuntimeError(
"MII client only supports a single query string, multi-string will be added soon"
)
request_dict = {"query": prompt}
if isinstance(prompts, str):
prompts = [prompts]
if streaming_fn is not None:
if len(prompts) > 1:
raise RuntimeError(
"MII client streaming only supports a single prompt input.")
request_dict = {"query": prompts}
return self._generate_stream(streaming_fn, request_dict, **query_kwargs)

request_dict = {"query": prompts}
return self.asyncio_loop.run_until_complete(
self._request_async_response(request_dict,
**query_kwargs))
Expand Down Expand Up @@ -100,24 +103,3 @@ def terminate_server(self):
if self.mii_config.enable_restful_api:
requests.get(
f"http://localhost:{self.mii_config.restful_api_port}/terminate")

async def create_session_async(self, session_id):
return await self.stub.CreateSession(
modelresponse_pb2.SessionID(session_id=session_id))

def create_session(self, session_id):
assert (
self.task == TaskType.TEXT_GENERATION
), f"Session creation only available for task '{TaskType.TEXT_GENERATION}'."
return self.asyncio_loop.run_until_complete(
self.create_session_async(session_id))

async def destroy_session_async(self, session_id):
await self.stub.DestroySession(modelresponse_pb2.SessionID(session_id=session_id)
)

def destroy_session(self, session_id):
assert (
self.task == TaskType.TEXT_GENERATION
), f"Session deletion only available for task '{TaskType.TEXT_GENERATION}'."
self.asyncio_loop.run_until_complete(self.destroy_session_async(session_id))
Loading