Rex Framework Package Overview

Rex Framework enables inference with remotely stored model weights without downloading full model checkpoints to local storage. Only the chunks needed for a given inference pass are fetched; the full model never resides in local memory or on disk.

Package intent:

• Primary: enable end users to run Rex for conversion, serving, and inference workflows.
• Secondary: support validation-oriented usage in CI and application test environments.

This package is intended for:

• Cloud-first inference where model chunks are fetched on demand.
• Memory-bounded environments where full checkpoint residency is undesirable.
• Notebook workflows, including Kaggle and Google Colab.

What You Get In This Package

• Python API for loading Rex manifests and running inference.
• CLI tools for conversion, validation, inspection, serving, benchmarking, and demo runs.
• Optional extras for PyTorch, cloud storage integrations, and benchmark tooling.
• WebSocket streaming is built into rex-serve and the runtime; it is not a separate install-time extra.
• The package is designed for runtime use and ships the support files needed for validation-oriented installs.

Install

Minimal package (no PyTorch, useful for manifest validation and storage testing):

pip install rex-framework

Recommended for real inference workloads:

pip install "rex-framework[pytorch]"

With all optional features (cloud storage backends, benchmarking tools):

pip install "rex-framework[all]"

Available extras:

Python compatibility notes:

• numpy is auto-installed with the package (no separate install needed).
• The pytorch extra supports Python 3.10 to 3.13.
• PyTorch wheels are not published on PyPI for Python 3.14.
• On some platforms (for example, macOS x86_64), Python 3.13 may still lack compatible torch wheels. Use Python 3.11 in that case.

Verify your install:

python -c "import torch, rex; print(torch.__version__); print(rex.__version__)"

How Rex Finds Your Model: Manifest and Chunk Paths

Rex does not load a model from a single checkpoint file. Instead it reads a manifest (a JSON file describing chunk locations, hashes, and metadata) and fetches individual chunks on demand from a base URL. Understanding how to point Rex at your chunk host is essential for it to work.

What a Rex manifest is

A manifest is a JSON file (manifest.json) generated by rex-convert. It contains:

• Model metadata (architecture, dtype, total size).
• A list of chunks: each chunk has a resource_id, a byte_range expressed as ByteRange(start, end), and a chunk_id.
• The expected base URL where chunks are hosted, usually in provenance.base_url.
• An optional storage_backend hint such as http_range or streaming_ws.

Chunks are served separately (e.g., as files in a directory) and fetched via HTTP Range requests. You do not need a special streaming server for the default path, but you do need an HTTP endpoint that supports Range headers (nginx, S3, Cloudflare R2, or rex-serve). If storage_backend is streaming_ws, Rex uses the WebSocket live streaming backend instead of HTTP Range.

Step 1 — Convert your model to Rex format

rex-convert /path/to/model.pt \
  --output ./rex_output \
  --framework pytorch \
  --model-id my-model

This produces:

rex_output/
  manifest.json        ← the manifest you will point load_model at
  weights/
    chunk_000.bin
    chunk_001.bin
    ...

Step 2 — Host the chunk files

Option A — local HTTP server (for testing):

rex-serve --dir ./rex_output/weights --port 8080

Chunks are now reachable at http://localhost:8080/chunk_000.bin, etc.

Option B — any static HTTP host:

Upload rex_output/weights/ to any static host (nginx, S3, Cloudflare R2, GitHub Releases, Google Drive folder with public sharing). Note the base URL.

Step 3 — Point load_model at the manifest

import rex
from rex.api.config import RexConfig

config = RexConfig()
config.storage.base_url = "http://localhost:8080"   # chunk host

runtime = rex.load_model("./rex_output/manifest.json", config=config)

base_url tells Rex where to fetch chunks from. Every chunk path in manifest.json is appended to base_url when fetching. If the manifest includes provenance.base_url, Rex can use that value too. If the manifest advertises storage_backend: streaming_ws, Rex opens the live WebSocket stream during model load.

Remote manifest (manifest itself is also hosted):

config.storage.base_url = "https://my-host.example.com/weights"
runtime = rex.load_model("https://my-host.example.com/weights/manifest.json", config=config)

Environment variable alternative:

export REX_STORAGE_URL=https://my-host.example.com/weights
python your_script.py

Step 4 — Run inference

import rex
import numpy as np
from rex.api.config import RexConfig

input_data = np.random.randn(1, 768).astype(np.float32)
config = RexConfig()
config.storage.base_url = "http://localhost:8080"

output, metrics = rex.run_inference_sync(
    "./rex_output/manifest.json",
    input_data,
    config=config,
)
print(f"Inference time: {metrics.total_time_ms:.1f} ms")

Storage Backends

Rex supports multiple storage backends. Set config.storage.base_url to the appropriate URL scheme:

For constructor-based overrides, RexConfig(storage__backend_type=StorageBackendType.STREAMING_WS, storage__base_url="wss://...") also works. If the manifest includes storage_backend="streaming_ws", Rex will prefer the WebSocket backend automatically.

Authenticated endpoints (e.g., private S3 or token-gated APIs):

config.storage.auth_token = "Bearer YOUR_TOKEN"

Or via environment variable:

export REX_AUTH_TOKEN=Bearer YOUR_TOKEN

Notebook Usage — Kaggle

Kaggle notebooks run on isolated kernels with internet access. The recommended pattern is to convert your model beforehand, host the chunks somewhere reachable (HTTPS URL, Google Drive public folder, or a Kaggle Dataset), then install Rex and load from that URL.

Install in a Kaggle cell

# Cell 1 — install
!pip install "rex-framework[pytorch]" -q
import rex, torch
print(rex.__version__, torch.__version__)

Load from an HTTPS host

# Cell 2 — configure and load
import rex
from rex.api.config import RexConfig

MANIFEST_URL = "https://your-static-host.com/rex_output/manifest.json"
CHUNKS_BASE_URL = "https://your-static-host.com/rex_output/weights"

config = RexConfig()
config.storage.base_url = CHUNKS_BASE_URL
config.cache.max_memory_cache_bytes = 512 * 1024 * 1024  # 512 MB limit

runtime = rex.load_model(MANIFEST_URL, config=config)

Load from a Kaggle Dataset

Upload your rex_output/ directory as a Kaggle Dataset. Kaggle mounts datasets at /kaggle/input/<dataset-name>/.

# Cell 2 — load from Kaggle Dataset mount
!rex-serve --dir /kaggle/input/my-rex-model/weights --port 8080 &

import rex
from rex.api.config import RexConfig

MANIFEST_PATH = "/kaggle/input/my-rex-model/manifest.json"
CHUNKS_BASE_URL = "http://127.0.0.1:8080"

config = RexConfig()
config.storage.base_url = CHUNKS_BASE_URL

runtime = rex.load_model(MANIFEST_PATH, config=config)

Add Kaggle Secrets for authenticated endpoints

from kaggle_secrets import UserSecretsClient

secrets = UserSecretsClient()
token = secrets.get_secret("REX_AUTH_TOKEN")

config.storage.auth_token = f"Bearer {token}"

Run inference on Kaggle

# Cell 3 — inference
import numpy as np

input_data = np.random.randn(1, 768).astype(np.float32)
output, metrics = rex.run_inference_sync(MANIFEST_PATH, input_data, config=config)
print(f"Output shape: {output.shape}")
print(f"Inference time: {metrics.total_time_ms:.1f} ms")

Notebook Usage — Google Colab

Google Colab provides a transient VM with internet access. The same manifest/chunk remote loading pattern applies. Colab T4 or A100 GPUs can be used if your Rex model targets CUDA.

Install in Colab

# Cell 1 — install
!pip install "rex-framework[pytorch]" -q
import rex, torch
print(rex.__version__, torch.__version__)

Load from an HTTPS host

# Cell 2 — configure and load
import rex
from rex.api.config import RexConfig

MANIFEST_URL = "https://your-static-host.com/rex_output/manifest.json"
CHUNKS_BASE_URL = "https://your-static-host.com/rex_output/weights"

config = RexConfig()
config.storage.base_url = CHUNKS_BASE_URL
config.cache.max_memory_cache_bytes = 1 * 1024 * 1024 * 1024  # 1 GB (Colab has more RAM)
config.scheduler.enable_prefetch = True
config.scheduler.prefetch_window = 4

runtime = rex.load_model(MANIFEST_URL, config=config)

Load from Google Drive in Colab

If you uploaded your rex_output/ to your Google Drive, mount it and serve the chunk directory locally:

# Cell 2a — mount Google Drive
from google.colab import drive
drive.mount("/content/drive")

# Cell 2b — serve the mounted directory locally
!rex-serve --dir /content/drive/MyDrive/rex_output/weights --port 8080 &

import rex
from rex.api.config import RexConfig

MANIFEST_PATH = "/content/drive/MyDrive/rex_output/manifest.json"
CHUNKS_BASE_URL = "http://127.0.0.1:8080"

config = RexConfig()
config.storage.base_url = CHUNKS_BASE_URL

runtime = rex.load_model(MANIFEST_PATH, config=config)

Use Colab Secrets for tokens

from google.colab import userdata

config.storage.auth_token = f"Bearer {userdata.get('REX_AUTH_TOKEN')}"

GPU inference in Colab

Rex will use the available CUDA device automatically when PyTorch detects a GPU. Confirm your runtime type is set to T4 GPU or A100 in Colab's Runtime menu.

import torch
print("CUDA available:", torch.cuda.is_available())
print("Device:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU")

Core Principle

Rex executes with bounded local residency by streaming model chunks from remote storage through HTTP range fetches and cache-aware scheduling. At no point does the full model need to exist locally.

Feature Controls (Quick Reference)

Control Rex behaviour through RexConfig:

from rex.api.config import RexConfig

config = RexConfig()

# How much local memory the cache can use
config.cache.max_memory_cache_bytes = 512 * 1024 * 1024

# Fraction of the full model allowed locally at any time (Rex invariant)
config.cache.max_local_fraction_of_model = 0.4

# Cache eviction policy: lru | lfu | weighted_utility
config.cache.policy = "weighted_utility"

# Prefetch ahead of current execution
config.scheduler.enable_prefetch = True
config.scheduler.prefetch_window = 4

# Execution planning mode: graph | sequential
config.scheduler.scheduler_mode = "graph"

# Storage concurrency
config.storage.max_concurrent_fetches = 4
config.storage.adaptive_concurrency = True

# Logging
config.observability.log_level = "INFO"   # DEBUG | INFO | WARNING | ERROR
config.observability.log_format = "console"  # console | json | quiet

For all available config fields and preset profiles (debug, throughput-oriented), see:

• https://rotsl.github.io/rex-framework/

CLI Quick Reference

Package Guide

For full package documentation, including notebook workflows and live streaming guidance, see:

• https://rotsl.github.io/rex-framework/