Function stream_to_worker

gpustack/server/worker_request.py:214–287

Stream a request to a worker and yield response chunks. Yields tuples of (chunk, headers, status). Automatically handles: - URL construction (advertise_address or ip + port) - Proxy selection (based on use_proxy_env_for_url) - Authorization header Args: worker

(
    worker: Worker,
    method: Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"],
    path: str,
    proxy_client: Optional[aiohttp.ClientSession] = None,
    no_proxy_client: Optional[aiohttp.ClientSession] = None,
    params: Optional[Dict] = None,
    data: Optional[Union[bytes, AsyncIterator[bytes], aiohttp.FormData]] = None,
    headers: Optional[Dict[str, str]] = None,
    timeout: Optional[aiohttp.ClientTimeout] = None,
    on_exception: Optional[
        Callable[[Exception, aiohttp.ClientTimeout], Tuple[str, int]]
    ] = None,
    raw: bool = False,
)

async def stream_to_worker(
    worker: Worker,
    method: Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"],
    path: str,
    proxy_client: Optional[aiohttp.ClientSession] = None,
    no_proxy_client: Optional[aiohttp.ClientSession] = None,
    params: Optional[Dict] = None,
    data: Optional[Union[bytes, AsyncIterator[bytes], aiohttp.FormData]] = None,
    headers: Optional[Dict[str, str]] = None,
    timeout: Optional[aiohttp.ClientTimeout] = None,
    on_exception: Optional[
        Callable[[Exception, aiohttp.ClientTimeout], Tuple[str, int]]
    ] = None,
    raw: bool = False,
) -> AsyncGenerator[Tuple[Union[bytes, str], dict, int], None]:
    """
    Stream a request to a worker and yield response chunks.

    Yields tuples of (chunk, headers, status).

    Automatically handles:
    - URL construction (advertise_address or ip + port)
    - Proxy selection (based on use_proxy_env_for_url)
    - Authorization header

    Args:
        worker: Target worker
        method: HTTP method
        path: API path
        proxy_client: HTTP client with proxy
        no_proxy_client: HTTP client without proxy
        params: Query parameters
        data: Bytes, async iterator of bytes, or FormData
        headers: Additional headers
        timeout: Request timeout
        on_exception: Optional callback(exception, timeout) -> (error_msg, status_code).
            Called when an exception occurs during streaming. If not provided,
            the exception is raised.
        raw: If True, yield raw bytes without SSE line formatting (use for log streams).
            If False (default), format each line as SSE (use for OpenAI-compatible streams).
    """
    try:
        async with _request_to_worker(
            worker=worker,
            method=method,
            path=path,
            proxy_client=proxy_client,
            no_proxy_client=no_proxy_client,
            params=params,
            data=data,
            headers=headers,
            timeout=timeout,
            raise_on_error=False,
        ) as resp:
            if resp.status >= 400:
                body = await resp.read()
                yield body, dict(resp.headers), resp.status
                return