"launch/vscode:/vscode.git/clone" did not exist on "439e977d9c751ef80f1ed72f03078dc408137a74"
Unverified Commit 73dc3be8 authored by Graham King's avatar Graham King Committed by GitHub
Browse files

chore: Merge bindings client and client2 functions (#6158)


Signed-off-by: default avatarGraham King <grahamk@nvidia.com>
parent f46720c9
...@@ -448,7 +448,9 @@ class EngineFactory: ...@@ -448,7 +448,9 @@ class EngineFactory:
kv_router_config=self.router_config.kv_router_config, kv_router_config=self.router_config.kv_router_config,
) )
else: else:
router = await generate_endpoint.client2(self.router_config.router_mode) router = await generate_endpoint.client(
router_mode=self.router_config.router_mode
)
gen = VllmProcessor( gen = VllmProcessor(
tokenizer, tokenizer,
......
...@@ -835,11 +835,13 @@ impl Endpoint { ...@@ -835,11 +835,13 @@ impl Endpoint {
}) })
} }
fn client<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> { #[pyo3(signature = (router_mode = None))]
self.client2(py, RouterMode::RoundRobin) fn client<'p>(
} &self,
py: Python<'p>,
fn client2<'p>(&self, py: Python<'p>, router_mode: RouterMode) -> PyResult<Bound<'p, PyAny>> { router_mode: Option<RouterMode>,
) -> PyResult<Bound<'p, PyAny>> {
let router_mode = router_mode.unwrap_or(RouterMode::RoundRobin);
let inner = self.inner.clone(); let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move { pyo3_async_runtimes::tokio::future_into_py(py, async move {
let client = inner.client().await.map_err(to_pyerr)?; let client = inner.client().await.map_err(to_pyerr)?;
......
...@@ -164,16 +164,11 @@ class Endpoint: ...@@ -164,16 +164,11 @@ class Endpoint:
""" """
... ...
async def client(self) -> Client: async def client(self, router_mode: Optional[RouterMode] = None) -> Client:
""" """
Create a `Client` capable of calling served instances of this endpoint using round-robin routing. Create a `Client` capable of calling served instances of this endpoint.
"""
...
async def client2(self, router_mode: RouterMode) -> Client: By default this uses round-robin routing when `router_mode` is not provided.
"""
Create a `Client` capable of calling served instances of this endpoint, using a specific
router mode (random, round-robin, kv).
""" """
... ...
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment