Model Streaming and Dynamic Batching

Stream model weights from object storage

When a model is stored remotely and downloading every .safetensors file to the container is undesirable, configure the vLLM or SGLang app environment with stream_model=True. The app environment passes the model input as FLYTE_MODEL_LOADER_REMOTE_MODEL_PATH, disables downloading that input, and sets FLYTE_MODEL_LOADER_STREAM_SAFETENSORS to true. The local model path is still supplied as FLYTE_MODEL_LOADER_LOCAL_MODEL_PATH for the other model artifacts.

For example, a vLLM application can select the streaming loader as follows:

from flyteplugins.vllm._app_environment import VLLMAppEnvironment

app = VLLMAppEnvironment(
    model_path="s3://models/my-model",
    model_id="my-model",
    stream_model=True,
)

VLLMAppEnvironment adds --load-format flyte-vllm-streaming to its generated server arguments. SGLangAppEnvironment uses the same remote-input and environment-variable arrangement, but its generated command remains the SGLang server command; its loader is installed by the SGLang shim when both the remote path and streaming flag are enabled. Both environments require model_id and exactly one of model_path or model_hf_path.

Set stream_model=False when you want the normal download-and-mount behavior. In that mode the environments set the streaming flag to false, mark the model input for download, and mount it at the configured model mount path. The relevant environment classes are plugins/vllm/src/flyteplugins/vllm/_app_environment.py and plugins/sglang/src/flyteplugins/sglang/_app_environment.py.

How SafeTensorsStreamer reads a model

The public model-loader package re-exports SafeTensorsStreamer from flyte.app.extras._model_loader. Direct use has this shape:

from flyte.app.extras._model_loader import SafeTensorsStreamer

streamer = SafeTensorsStreamer(
    "s3://models/my-model",
    "/srv/model",
    chunk_size=16 * 1024 * 1024,
    max_concurrency=32,
)

for name, tensor in streamer.get_tensors():
    print(name, tensor.shape)

The remote_path is resolved through Flyte storage. The constructor calls get_underlying_filesystem, splits the remote path into a bucket and prefix, constructs an ObjectStore, and creates an ObstoreParallelReader with the requested chunk_size and max_concurrency. local_path is retained by the streamer and is the path used by the surrounding model-loader configuration; the tensor reads themselves use the object store. The store_kwargs argument is present in the constructor signature but is not consumed by the constructor body.

get_tensors() is a synchronous generator. Internally it creates the asynchronous _get_tensors_async() generator and drives it with asyncio.Runner. Consume it from synchronous code; embedding this runner-based generator in an already-running event loop can conflict with the runner model. The module imports PyTorch eagerly and raises a targeted ModuleNotFoundError if torch is unavailable.

The data flow is:

remote SafeTensors files
        |
        |  header ranges through obstore
        v
SafeTensorsMetadata(path, data_start, tensors)
        |
        |  TensorMetadata -> Source -> DownloadTask chunks
        v
ObstoreParallelReader._as_completed(...)
        |
        v
(name, torch.Tensor)

SafeTensorsStreamer._parse_safetensors_metadata() first reads the eight-byte little-endian header length, then reads the JSON header range. An optional __metadata__.format value is accepted only when it is "pt"; another value raises ValueError. Each tensor entry is validated into TensorMetadata.

TensorMetadata is a Pydantic model with name, shape, dtype, and data_offsets. Its dtype validator maps SafeTensors names such as F16, BF16, F32, I64, and BOOL to Torch dtypes. Its computed size is the byte-range length (end - start), while length is the product of the shape dimensions; len(metadata) returns length. When converting a downloaded buffer, the streamer uses torch.frombuffer, passes the metadata dtype and element count, and reshapes the result to the metadata shape.

SafeTensorsMetadata contains the file path, the absolute data_start immediately after the header, and its TensorMetadata list. It is an internal Pydantic container, not a public convenience API: unlike SafeTensorsStreamer, it is not re-exported by flyte.app.extras._model_loader.

The streamer creates a Source for each tensor with this byte position:

offset=stm.data_start + tm.data_offsets[0]

It then divides the tensor's tm.size into reader chunks and submits DownloadTask objects. Results are yielded through _as_completed, so tensors are exposed as their downloads complete rather than only after every file has been read.

Indexed and tensor-parallel files

For a non-parallel model (tensor_parallel_size=1), metadata loading first tries model.safetensors.index.json. Its weight_map is inverted into file-to-tensor sets. The loader verifies that every tensor expected at a file is present; a missing tensor raises a detailed ValueError. If the index is absent or cannot be decoded, or its structure raises KeyError, the loader falls back to files matching *.safetensors. Pattern loading removes duplicate tensor names, retaining the first occurrence. Other failures are not swallowed by this fallback.

For tensor parallelism, pass the rank and world size:

streamer = SafeTensorsStreamer(
    "s3://models/my-model",
    "/srv/model",
    rank=rank,
    tensor_parallel_size=tensor_parallel_size,
)

When tensor_parallel_size > 1, the loader deliberately ignores the index and lists files matching model-rank-{rank}-part-*.safetensors. No matching files produces ValueError. This is the path used by the sharded vLLM and SGLang adapters.

The vLLM adapter registers FlyteModelLoader under flyte-vllm-streaming. Its ordinary loading path constructs SafeTensorsStreamer(REMOTE_MODEL_PATH, LOCAL_MODEL_PATH), then yields source.prefix + name with each tensor. If streamer construction raises ValueError, it delegates to vLLM's default iterator. For sharded loading it obtains the distributed rank, initializes the model, streams rank-specific tensors, and copies each tensor into the initialized state. Smaller streamed dimensions are copied into narrowed parameter views; remaining state keys raise ValueError.

SGLang's FlyteModelLoader follows the same ordinary and sharded patterns and falls back to SGLang's default iterator when construction raises ValueError. Its module-level monkeypatch replaces SGLang's default loader only when REMOTE_MODEL_PATH is set and streaming is enabled. Both shims' _get_model_files() functions check that the remote model path exists and call prefetch; with streaming enabled, prefetch excludes SafeTensors files while copying other artifacts.

Model-loader configuration

The values in src/flyte/app/extras/_model_loader/config.py are read from the process environment:

Variable	Default	Effect
FLYTE_MODEL_LOADER_REMOTE_MODEL_PATH	unset	Remote object-store model path; required by the plugin shims for streaming.
FLYTE_MODEL_LOADER_LOCAL_MODEL_PATH	/srv/model	Local model/artifact path retained by the loader and set by the app environments.
FLYTE_MODEL_LOADER_CHUNK_SIZE	16777216	Byte chunk size passed to ObstoreParallelReader.
FLYTE_MODEL_LOADER_MAX_CONCURRENCY	32	Maximum reader concurrency.
FLYTE_MODEL_LOADER_STREAM_SAFETENSORS	false	Enables streaming only for the literal value true, case-insensitively.

Check the following conditions when loading fails:

• Install PyTorch before importing the model-loader module. The loader also requires Pydantic and the obstore/Flyte storage components.
• Ensure the remote path exists. The vLLM and SGLang _get_model_files() functions raise FileNotFoundError when it does not.
• Use the SafeTensors layout expected by the parser: an eight-byte header length, JSON metadata, and .safetensors files. If __metadata__.format exists, it must be pt; unsupported dtype strings raise ValueError.
• For non-parallel models, check model.safetensors.index.json and its weight_map; indexed files with missing listed tensors are rejected. For parallel models, check the rank-specific model-rank-{rank}-part-*.safetensors names instead.
• Match the vLLM or SGLang version required by the corresponding plugin shim. Those shims validate the installed framework version while importing.

Token-aware dynamic batching

For an LLM endpoint, submit prompts independently while allowing one inference call to receive a token-budgeted list. TokenBatcher is the public token-oriented API in flyte.extras; the normal lifecycle deliberately has two awaits: submitting returns a future, and awaiting that future returns the per-record result.

from flyte.extras import Prompt, TokenBatcher

async def inference(batch: list[Prompt]) -> list[str]:
    return [f"answer: {prompt.text}" for prompt in batch]

async def handle(prompt: Prompt) -> str:
    async with TokenBatcher(
        inference_fn=inference,
        target_batch_tokens=4096,
        max_batch_size=64,
        batch_timeout_s=0.05,
    ) as batcher:
        future = await batcher.submit(prompt)
        return await future

inference_fn must be an async callable taking list[RecordT] and returning list[ResultT]. Results must have exactly the same length and order as the input batch. TokenBatcher also accepts the base name process_fn; constructing it without either function raises TypeError.

Prompt is a simple dataclass with one text field. Its estimate_tokens() implementation is intentionally rough:

len(self.text) // 4 + 1

It is not tokenizer-backed. Use an explicit estimate or a custom TokenEstimator for a production token limit.

How records become batches

DynamicBatcher, the base class, maintains two asynchronous loops. The aggregation loop waits for the first queue item, then fills a batch until the accumulated estimated cost is no longer below target_batch_cost, max_batch_size is reached, or batch_timeout_s expires. The processing loop takes assembled batches from a bounded prefetch queue and calls the processing function.

Each submission is wrapped in the internal _Envelope dataclass, which pairs the record and estimated cost with an asyncio.Future. The future is created from the active event loop. _Envelope is an internal symbol from flyte.extras._dynamic_batcher, not part of the normal public batching API.

TokenBatcher maps token names onto the base cost machinery:

• target_batch_tokens becomes target_batch_cost.
• default_token_estimate becomes default_cost.
• token_estimator and cost_estimator select the configured estimator.
• estimated_tokens and estimated_cost are accepted by submit, with estimated_tokens taking precedence.

The token cost precedence is explicit submit override, configured estimator, record.estimate_tokens(), record.estimate_cost(), then the default. TokenEstimator and CostEstimator are runtime-checkable protocols, so a record only needs the corresponding method:

from dataclasses import dataclass
from flyte.extras import TokenBatcher, TokenEstimator

@dataclass
class TokenRecord:
    text: str

    def estimate_tokens(self) -> int:
        return len(self.text)

async def record_process(batch: list[TokenRecord]) -> list[str]:
    return [record.text.upper() for record in batch]

async def handle_record(record: TokenRecord) -> str:
    async with TokenBatcher(inference_fn=record_process) as batcher:
        assert isinstance(record, TokenEstimator)
        future = await batcher.submit(record)
        return await future

For general workloads, import DynamicBatcher and use CostEstimator or a configured callable:

import asyncio
from dataclasses import dataclass
from flyte.extras import DynamicBatcher

@dataclass
class ApiRequest:
    payload: str

    def estimate_cost(self) -> int:
        return len(self.payload)

async def process(batch: list[ApiRequest]) -> list[str]:
    return [request.payload for request in batch]

async def handle_requests() -> list[str]:
    async with DynamicBatcher(
        process_fn=process,
        target_batch_cost=100,
        max_queue_size=500,
    ) as batcher:
        futures = await batcher.submit_batch(
            [ApiRequest("a"), ApiRequest("b")],
            estimated_cost=[1, 1],
        )
        return await asyncio.gather(*futures)

The aggregator's cost budget is a batching criterion, not validation of an individual record. The first record is accepted even when its estimated cost exceeds the target; the fill loop then stops once the accumulated cost is no longer below the target. max_batch_size is an independent hard record-count cap. min_batch_size can delay dispatch while a small batch is re-enqueued, but shutdown bypasses that minimum. A bounded submission queue applies backpressure: when max_queue_size is full, submit() waits for space. prefetch_batches bounds the queue between aggregation and processing.

Lifecycle, failures, and metrics

Call start() before submit(), or use async with as above. Submitting while stopped raises RuntimeError, and calling start() twice raises RuntimeError. stop() is graceful: it sends a sentinel through the aggregation loop, waits for assembled work to be processed, and blocks until pending futures are resolved.

If process_fn returns the wrong number of results, the processing loop raises a ValueError for that batch. Any exception from processing is set on every future in the failed batch, while the processing loop continues. Await the futures rather than assuming that await batcher.submit(...) is the result.

Inspect the live BatchStats object through batcher.stats:

stats = batcher.stats
print(stats.total_submitted, stats.total_completed)
print(stats.total_batches, stats.avg_batch_size)
print(stats.total_batch_cost, stats.avg_batch_cost)
print(stats.utilization)

BatchStats records submitted and completed records, dispatched batches, cumulative and average costs and sizes, and processing versus waiting time. utilization is busy_time_s / (busy_time_s + idle_time_s), or 0.0 when neither time counter has advanced. The class is publicly re-exported from flyte.extras, while _Envelope remains internal.

The dedicated tests/flyte/extras/test_dynamic_batcher.py suite exercises lifecycle, cost and token budgets, timeouts, backpressure, estimator protocols, error propagation, statistics, and Prompt integration. These tests also demonstrate the important operational contract: a batch function receives records in submission order and must return one corresponding result per record.