Skip to main content

Batch Affinity Scheduler for Throttled Inference on Ollama Networks

Project description

BASTION

CI License: MIT Python 3.11+

Batch Affinity Scheduler for Throttled Inference on Ollama Networks

A system-level GPU inference broker that prevents crashes from concurrent Ollama access. BASTION sits as a transparent HTTP proxy between your applications and Ollama, serializing model loads, enforcing VRAM budgets, and eliminating the memory-mapped I/O conditions that cause GPU driver crashes under heavy inference workloads.


Why BASTION?

Running multiple LLM clients against a single Ollama instance is a recipe for hard crashes. Here is why:

  1. Memory-mapped model loading. Ollama memory-maps model files by default. When multiple models are loaded concurrently, the virtual address space balloons far beyond physical VRAM, and the OS page cache competes with GPU memory for the same pages.

  2. Rapid model cycling. With OLLAMA_MAX_LOADED_MODELS=1, every request for a different model triggers a full unload/load cycle. Under concurrent access from multiple clients, this creates dozens of VRAM allocation/deallocation cycles per minute.

  3. No upstream protection. Ollama has no built-in queue, no VRAM budget enforcement, and no cooldown between model transitions.

The result: after roughly 60 rapid model swaps, the GPU driver hits an unrecoverable state -- kernel OOM, display freeze, or full system crash requiring a hard reboot.

BASTION solves this at the system level, sitting transparently on the standard Ollama port so every client benefits with zero configuration changes.

Prerequisites

  • Python 3.11+
  • Linux with NVIDIA GPU and working drivers (nvidia-smi must respond)
  • Ollama installed
  • At least one model pulled (ollama pull llama3.1:8b)

Quick Start

1. Install

git clone https://github.com/cypwin/bastion.git
cd bastion
pip install -e ".[dev]"

2. Move Ollama to port 11435

sudo mkdir -p /etc/systemd/system/ollama.service.d/
sudo tee /etc/systemd/system/ollama.service.d/override.conf > /dev/null << 'EOF'
[Service]
Environment="OLLAMA_HOST=127.0.0.1:11435"
EOF
sudo systemctl daemon-reload
sudo systemctl restart ollama

Or manually: OLLAMA_HOST=127.0.0.1:11435 ollama serve

3. Configure

bastion --init-config       # Generate ~/.config/bastion/broker.yaml
bastion --detect-models     # Discover installed models

4. Validate your setup

bastion --validate

5. Start BASTION

bastion                              # Default config
bastion --config config/broker.yaml  # Custom config
bastion --admin-port 9999            # Two-port mode

6. Verify

curl http://localhost:11434                # "Ollama is running"
curl http://localhost:11434/broker/status  # Broker status
ollama run llama3.1:8b "Hello, world!"    # Transparent proxy

Key Features

  • Transparent HTTP proxy -- drop-in replacement on port 11434; existing clients work without any changes. Streams NDJSON faithfully so ollama run remains responsive.

  • Affinity-based scheduling -- per-model sub-queues drain all pending requests for the currently loaded model before swapping, dramatically reducing GPU model transitions.

  • Priority aging -- four priority tiers (interactive, agent, pipeline, background) with time-based aging to prevent starvation of lower-priority requests.

  • VRAM budget enforcement -- tracks GPU memory via Ollama /api/ps fused with nvidia-smi data. Enforces a configurable budget and blocks model loads that would exceed it.

  • Crash prevention -- BASTION injects use_mmap: false into scheduled inference requests (/api/generate, /api/chat, /api/embed) when the client hasn't specified use_mmap explicitly. This mitigates PCIe power transients that crashed RTX 5090 during mmap-backed model cycling. Passthrough endpoints (/api/pull, /api/show, /api/tags, etc.) forward unchanged. Cooldown periods between model transitions and swap-rate limiting provide additional protection.

  • 19-panel TUI dashboard -- real-time Textual dashboard showing GPU thermals, VRAM usage, queue depth, scheduler state, A2A tasks, leases, audit events, and more.

  • A2A protocol support -- Agent-to-Agent protocol with agent card discovery, task lifecycle, batch inference, and model reservation leases.

  • Prometheus metrics + OpenTelemetry tracing -- optional observability integrations with graceful no-op fallbacks.

  • Circuit breaker -- three-state (closed/open/half-open) protection against cascading failures when the Ollama backend is unresponsive.

  • Per-IP rate limiting -- configurable rate limiting middleware.

  • Tiered audit logging -- JSONL audit logs with content hashing and automatic rotation.

Architecture

                          Clients
    (ollama run, Claude Code, Python scripts, A2A agents, curl)
                             |
                             | :11434 (standard Ollama port)
                             v
    +------------------------------------------------------------+
    |                        BASTION                              |
    |                                                             |
    |  +----------------+  +-----------------+  +--------------+  |
    |  | Ollama Proxy   |  | Admin API       |  | A2A Agent    |  |
    |  | /api/*         |  | /broker/*       |  | /a2a/*       |  |
    |  |                |  |                 |  |              |  |
    |  | - use_mmap     |  | - status/queue  |  | - tasks      |  |
    |  |   injection    |  | - health/vram   |  | - streaming  |  |
    |  | - NDJSON       |  | - preload       |  | - leases     |  |
    |  |   streaming    |  | - unload/drain  |  | - agent card |  |
    |  | - priority     |  | - metrics       |  | - batch      |  |
    |  |   detection    |  |                 |  |   inference  |  |
    |  +-------+--------+  +-----------------+  +--------------+  |
    |          |                                                   |
    |  +-------v-------------------------------------------------+ |
    |  |          Affinity Queue + Scheduler                      | |
    |  |                                                          | |
    |  |  - Per-model sub-queues (minimize GPU model swaps)       | |
    |  |  - Priority tiers: INTERACTIVE > AGENT > PIPELINE > BG  | |
    |  |  - Aging: effective = base + (age_seconds * 2.0)        | |
    |  |  - Cooldown: 2s between model transitions               | |
    |  |  - VRAM ledger (assume/confirm/forget pattern)           | |
    |  |  - GPU health gating (temp, power, utilization)          | |
    |  |  - Concurrent co-resident dispatch (up to 3 models)     | |
    |  +-------+-------------------------------------------------+ |
    +-----------|---------------------------------------------------+
                | :11435
    +-----------v---------------------------------------------------+
    |                     Ollama (backend)                           |
    |              OLLAMA_HOST=127.0.0.1:11435                      |
    +---------------------------------------------------------------+

Dashboard

bastion-dashboard
bastion-dashboard --url http://localhost:11434 --interval 2.0

(Or: python -m bastion.dashboard / python -m bastion.dashboard --url http://localhost:11434 --interval 2.0)

Keyboard shortcuts: p preload model, u unload model, d toggle drain mode, r manual refresh, h help overlay, q quit.

Documentation

Guide Description
Getting Started Full installation walkthrough
Configuration Every config option explained
Hardware Guide GPU compatibility and VRAM requirements
Troubleshooting Common issues and fixes
Operations Monitoring, restart, day-2 ops
Security Auth, TLS, network isolation
Crash Prevention How BASTION prevents GPU crashes
API Reference All endpoints with examples
Deployment Systemd, Docker, desktop launcher
Releasing One-time PyPI/OIDC setup and release cut procedure

Optional Extras

pip install -e ".[metrics]"     # Prometheus metrics export
pip install -e ".[dashboard]"   # TUI dashboard (Textual)
pip install -e ".[a2a]"         # A2A SDK types
pip install -e ".[dev]"         # Testing and linting tools

Testing

pip install -e ".[dev]"
python -m pytest tests/ -v

License

MIT License. See LICENSE for details.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

bastion_broker-0.5.0a1.tar.gz (511.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

bastion_broker-0.5.0a1-py3-none-any.whl (253.9 kB view details)

Uploaded Python 3

File details

Details for the file bastion_broker-0.5.0a1.tar.gz.

File metadata

  • Download URL: bastion_broker-0.5.0a1.tar.gz
  • Upload date:
  • Size: 511.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for bastion_broker-0.5.0a1.tar.gz
Algorithm Hash digest
SHA256 22c70e3187aa7f82bd8eafcc68e64b322abbc3d673a51083bcb29045eaa1ea4d
MD5 6ffb4c33478893f4f9adaadd6cd68b34
BLAKE2b-256 1b3415cadac7e186afcf5b059db75d570ceaeec7bdcc06cc27431eccb468f3c7

See more details on using hashes here.

Provenance

The following attestation bundles were made for bastion_broker-0.5.0a1.tar.gz:

Publisher: release.yml on cypwin/bastion

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file bastion_broker-0.5.0a1-py3-none-any.whl.

File metadata

File hashes

Hashes for bastion_broker-0.5.0a1-py3-none-any.whl
Algorithm Hash digest
SHA256 4657124191287c82ece1e170784d71715ee2b20b4086b63a0f4e60db54846310
MD5 66915bf38c83fcf3afd1790723aedbc7
BLAKE2b-256 37971eb2e6c7aff51099d8a9bc64fac47534e58e2ba3293e25b04a18228e50d1

See more details on using hashes here.

Provenance

The following attestation bundles were made for bastion_broker-0.5.0a1-py3-none-any.whl:

Publisher: release.yml on cypwin/bastion

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page