Skip to main content

Crate ts_client

Crate ts_client 

Source
Expand description

A Rust frontend’s connection to TokenSpeed engines.

The client binds the handshake, input and output sockets; the engines (one per attention-DP rank) connect in and register. The client then routes tokenized requests to engines and streams each request’s outputs back. It does not tokenize or detokenize: the frontend owns both, and so also matches stop strings itself. The engine refuses requests that carry stop strings on this wire; send stop token ids.

use std::time::Duration;
use ts_client::{ClientConfig, EngineClient, LoadBalance, Routing};
use ts_wire::{SamplingParams, TokenizedGenerateReqInput};

let client = EngineClient::connect(ClientConfig {
    handshake_address: "tcp://127.0.0.1:30500".into(),
    input_address: "tcp://127.0.0.1:0".into(),
    output_address: "tcp://127.0.0.1:0".into(),
    first_engine_index: 0,
    num_engines: 1,
    routing: Routing::Aggregated(LoadBalance::LeastInFlight),
    handshake_timeout: Duration::from_secs(600),
    send_timeout: Duration::from_secs(5),
})
.await?;
let mut stream = client
    .generate(TokenizedGenerateReqInput {
        input_ids: Some(vec![1, 2, 3]),
        stream: true,
        ..TokenizedGenerateReqInput::new("r1", SamplingParams::default())
    })
    .await?;
while let Some(chunk) = stream.next().await {
    println!("{:?}", chunk.expect("engine alive").output_ids);
}

Structs§

ClientConfig
Where the client listens and which engines it expects.
ClientListener
The client’s sockets, bound and waiting for the engines.
EngineClient
A connection to one engine group: every attention-DP rank’s engine.
EngineInfo
An engine that completed the handshake.
EngineLoad
An engine’s latest scheduler load, as reported with its outputs.
OutputChunk
One step of a request’s output.
RequestStream
The outputs of one request, in order.

Enums§

Error
An error from connecting to the engines or submitting a request.
LoadBalance
Placement policy for unpinned requests in aggregated serving.
Routing
How requests are assigned to engines (one engine per attention-DP rank).
StreamError
Why a request’s output stream ended without a finished output.

Type Aliases§

Result
Result alias for client operations.