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:
kv_router_config=self.router_config.kv_router_config,
)
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(
tokenizer,
......
......@@ -835,11 +835,13 @@ impl Endpoint {
})
}
fn client<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
self.client2(py, RouterMode::RoundRobin)
}
fn client2<'p>(&self, py: Python<'p>, router_mode: RouterMode) -> PyResult<Bound<'p, PyAny>> {
#[pyo3(signature = (router_mode = None))]
fn client<'p>(
&self,
py: Python<'p>,
router_mode: Option<RouterMode>,
) -> PyResult<Bound<'p, PyAny>> {
let router_mode = router_mode.unwrap_or(RouterMode::RoundRobin);
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let client = inner.client().await.map_err(to_pyerr)?;
......
......@@ -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:
"""
Create a `Client` capable of calling served instances of this endpoint, using a specific
router mode (random, round-robin, kv).
By default this uses round-robin routing when `router_mode` is not provided.
"""
...
......
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