rlx_pkg/
remote.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! HTTP(S) Range reader for remote flat `.rlxp` (`remote` feature).
6
7use anyhow::{Context, Result, bail};
8use std::io::Read;
9
10/// Fetch a byte range from `url` (`Range: bytes=start-end` inclusive).
11#[cfg(feature = "remote")]
12pub fn http_range(url: &str, start: u64, end_inclusive: u64) -> Result<Vec<u8>> {
13    let range = format!("bytes={start}-{end_inclusive}");
14    let resp = ureq::get(url)
15        .set("Range", &range)
16        .call()
17        .with_context(|| format!("GET {url} Range {range}"))?;
18    let status = resp.status();
19    if status != 206 && status != 200 {
20        bail!("unexpected HTTP {status} for Range request");
21    }
22    let mut buf = Vec::new();
23    resp.into_reader()
24        .take(end_inclusive.saturating_sub(start).saturating_add(1) + 64)
25        .read_to_end(&mut buf)
26        .context("read range body")?;
27    Ok(buf)
28}
29
30#[cfg(not(feature = "remote"))]
31pub fn http_range(_url: &str, _start: u64, _end_inclusive: u64) -> Result<Vec<u8>> {
32    bail!("rlx-pkg was built without the `remote` feature")
33}
34
35/// Remote flat-pack handle: TOC in memory; tensor bytes fetched on demand.
36#[cfg(feature = "remote")]
37pub struct RemoteFlat {
38    pub url: String,
39    pub header: crate::flat::FlatHeader,
40    pub toc: crate::flat::FlatToc,
41    pub data_start: u64,
42}
43
44#[cfg(feature = "remote")]
45impl RemoteFlat {
46    /// Fetch header + TOC only (does not pull the data region).
47    pub fn open(url: &str) -> Result<Self> {
48        use crate::flat::{DATA_ALIGN, FLAT_MAGIC, FlatHeader, FlatToc};
49        let prefix = http_range(url, 0, 23)?;
50        if prefix.len() < FlatHeader::SIZE || &prefix[..8] != FLAT_MAGIC {
51            bail!("remote URL is not RLXPFLAT");
52        }
53        let header = FlatHeader::decode(&prefix)?;
54        let toc_end = (FlatHeader::SIZE as u64) + header.toc_len - 1;
55        let hdr_toc = http_range(url, 0, toc_end)?;
56        // The server's response length is untrusted: a short reply (fewer bytes
57        // than the header-declared TOC) would panic the slice below. Bounds-check
58        // it (checked_add guards the usize overflow) and error cleanly instead.
59        let toc_end_idx = match FlatHeader::SIZE.checked_add(header.toc_len as usize) {
60            Some(e) if e <= hdr_toc.len() => e,
61            _ => bail!(
62                "remote TOC truncated: header declares {} TOC bytes but the server returned {}",
63                header.toc_len,
64                hdr_toc.len()
65            ),
66        };
67        let toc_bytes = &hdr_toc[FlatHeader::SIZE..toc_end_idx];
68        let mut toc: FlatToc = if header.is_bincode_toc() {
69            let wire: crate::flat::FlatTocBin = bincode::deserialize(toc_bytes)?;
70            wire.into_toc()?
71        } else {
72            serde_json::from_slice(toc_bytes)?
73        };
74        toc.resolve_names()?;
75        toc.manifest.validate()?;
76        let data_start =
77            (FlatHeader::SIZE as u64 + header.toc_len + DATA_ALIGN - 1) & !(DATA_ALIGN - 1);
78        Ok(Self {
79            url: url.to_string(),
80            header,
81            toc,
82            data_start,
83        })
84    }
85
86    /// Fetch and decode one tensor by name.
87    pub fn tensor_bytes(&self, name: &str) -> Result<Vec<u8>> {
88        use crate::tier::decode_payload;
89        let t = self
90            .toc
91            .tensors
92            .iter()
93            .find(|t| t.name == name)
94            .with_context(|| format!("tensor {name}"))?;
95        let abs0 = self.data_start + t.offset;
96        let abs1 = abs0 + t.length - 1;
97        let stored = http_range(&self.url, abs0, abs1)?;
98        if stored.len() as u64 != t.length {
99            bail!("range length {} != TOC length {}", stored.len(), t.length);
100        }
101        decode_payload(t.codec, &stored)
102    }
103}
rlx_pkg/remote.rs

rlx_pkg/
remote.rs
