Rust with axum

Cover Page

Back-end Page

Install Rust

ssh to your server and install Rust:

server$ sudo apt install gcc          # cargo depends on gcc's linker
server$ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
server$ hash -r
server$ rustup update

If you see:

rustup: Command not found.

and you need help updating your PATH shell environment variable, please see a teaching staff.

The command rustup update is also how you can subsequently update your installation of Rust to a new version.

Cargo.toml: add dependencies

First create and change into a directory where you want to keep your harnessd package:

server$ cd ~/agentic
server$ cargo new harnessd
# output:
     Created binary (application) `harnessd` package

This will create the ~/agentic/harnessd/ directory for you. Change to this directory and edit the file Cargo.toml to list all the 3rd-party libraries (crates in Rust-speak) we will be using.

server$ cd harnessd
server$ vi Cargo.toml

In Cargo.toml, add the following below the [dependencies] tag:

axum = { version = "0.8.9", features = ["macros"] }
axum-server = { version = "0.8.0", features = ["tls-rustls"] }
mimalloc = { version = "0.1.50" }
rand = "0.10.1"
reqwest = { version = "0.13.3", features = ["json"] }
serde_json = "1.0.145"
tokio = { version = "1.48.0", features = ["full"] }
tower-http = { version = "0.6.7", features = ["normalize-path", "trace"] }
tower-layer = "0.3.3"
tracing = "0.1.43"
tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }

[profile.release]
strip = true
update

To see outdated crates and update all crates to the latest compatible version:

  server$ cargo outdated --root-deps-only
  server$ cargo update

main.rs

In ~/agentic/harnessd/src/ the file main.rs has also been created for you. Edit the file:

server$ vi src/main.rs

and replace all existing lines in main.rs with the server and URL routing code below, starting with the following use lines::

#![allow(non_snake_case)]

mod handlers;
mod runners;

use mimalloc::MiMalloc;
#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;
// `Mimalloc` uses memory more efficiently and supports higher concurrency than 
// either the Linux malloc or jemalloc, especially for harness architecture workload. 

use axum::{
    Router, ServiceExt,
    extract::{ FromRef, Request },
    routing::post,
};
use axum_server::tls_rustls::RustlsConfig;
use reqwest::Client;
use runners::LlmRunners;
use std::{ net::SocketAddr, process, sync::Arc };
use tower_http::{
    normalize_path::NormalizePathLayer,
    trace::{ DefaultMakeSpan, DefaultOnFailure, TraceLayer },
};
use tower_layer::Layer;
use tracing::Level;

After listing all our imports, we export our modules handlers and runners, which we will define later.

We create a struct to hold state(s) that we want to pass to each API handler. For llmprompt, the only state we pass to its handler is an instantiation of the reqwest client, with which we will initiate a connection to the LLM runner:

#[derive(FromRef, Clone)]
pub struct AppState {
    client: Client,
    runners: Arc<LlmRunners>,
}

The FromRef macro creates a “getter” for each of the structure’s property that allows the property to be passed as axum’s State.

Create the following main() function and enable logging (tracing):

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_level(false)
        .with_target(false)
        .compact()
        .init();

    // HTTP client
}

Replace // HTTP client to instantiate an AppState that we will pass to each API handler:

    // Create HTTP client to connect to LLM runner and LLM runner selector
    let appState = AppState {
        client: Client::new(),
        runners: Arc::new(LlmRunners::new()),
    };

    // API end-point router

Continuing inside the main() function, we next create a routing table, in the form of a Router structure, to hold the URL routing information needed by axum_server. We define the route to serve the HTTP POST request with URL endpoint /llmprompt. We route the endpoint to the llmprompt() function. Since we only allow HTTP POST request for the /llmprompt endpoint, we pass the llmprompt() handler to post(). We enable the tracing middleware and pass the appState instantiated above to each API handler. Replace // API end-point router:

    let router = Router::new()
        .route("/llmprompt", post(handlers::llmprompt))
        .layer(
            // must be after all handlers to be traced
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
                .on_failure(DefaultOnFailure::new().level(Level::INFO)),
        )
        .with_state(appState);       

    // launch HTTPS server with cert and key

The function llmprompt() will be implemented in handlers.rs later.

For now, staying in main.rs, in the main() function, set up the axum_server:


We're done with `main.rs`.  Save and exit the file.

### <tt>handlers.rs</tt>

We implement URL path API handlers in `src/handlers.rs`:
```console
server$ vi src/handlers.rs

Start the file with the following use imports:

#![allow(non_snake_case)]
use axum::{
    extract::{ ConnectInfo, Json, State },
    http::{ Response, StatusCode },
};
use crate::AppState;
use serde_json::Value;
use std::net::SocketAddr;

We add a couple of logging functions to print to console results of handling each HTTP request and, in the case of error, return a tuple of HTTP status code and error message to the client:

fn logOk(clientIP: &SocketAddr, runner: &str, model: &str) {
    tracing::info!("{} | {} | {}:{} |", StatusCode::OK, clientIP, runner, model);
}
fn logInfo(clientIP: &SocketAddr, runner: &str, model: &str, errcode: StatusCode, msg: &str) {
    tracing::info!("{} | {} | {}:{} {} |", errcode, clientIP, runner, model, msg);
}
fn logClientErr(clientIP: &SocketAddr, runner: &str, model: &str, errcode: StatusCode, errmsg: String) -> (StatusCode, String) {
    tracing::info!("{} | {} | {}:{} |", errcode, clientIP, runner, model);
    (errcode, errmsg)
}

We next define the handler llmprompt(), which first pick an LLM runner from a ranked list of available runners. Since different runner provides different models, we patch the "model" field of the incoming prompt to the model available at the runner.

pub async fn llmprompt(
    State(appState): State<AppState>,
    ConnectInfo(clientIP): ConnectInfo<SocketAddr>,
    Json(mut body): Json<Value>,
) -> Result<Response<reqwest::Body>, (StatusCode, String)> {
    let model = body["model"].as_str().unwrap_or_default().to_owned();
    for runner in appState.runners.iter() {
        // patch the "model" field with the model available for this runner, 
        // if specified, else restore original model that came with request.
        body["model"] = if runner.model.is_empty() { model.clone().into() } else { runner.model.into() };

        // send prompt and return response
        
    }

    // no runner found
}

Then we assemble together a request addressed to the runner’s OpenAI-compatible API end point, along with the preloaded API key for that runner, and send it to the LLM runner. If we fail to connect to the LLM runner, we log the error and try the next runner on our list. If connection is successful but the runner returns any kind of error, including being rate limited, we again log the error and try the next runner on our list. If the runner completes the prompt, we log the success and forward the completion to the user. Replace //send prompt and return response with:

        let mut req = appState
            .client
            .post(&format!("{}/v1/chat/completions", runner.baseUrl))
            .json(&body);

        if let Some(key) = &runner.apiKey {
            req = req.bearer_auth(key);
        }

        match req.send().await {
            Ok(response) => {
                // If rate limited or server error, 'continue' naturally advances the iterator
                if response.status() != StatusCode::OK {
                    logInfo(&clientIP, &runner.baseUrl, &runner.model, response.status(),
                        "Server error. Trying another one.");
                    continue;
                }

                // Success! Break the loop and return the streaming response
                logOk(&clientIP, &runner.baseUrl, &runner.model);
                return Ok(response.into());
            }
            Err(err) => {
                // failover if network/connection error
                logInfo(&clientIP, &runner.baseUrl, &runner.model, StatusCode::PROCESSING,
                    &format!("Connection error: {}. Trying another one.", err));
                continue;
            }
        }

If we’ve exhausted all the runners on our list but couldn’t find any to complete the prompt, we log the error and inform the client. Replace // no runner found with:

    // If the for-loop naturally finishes without returning, we are exhausted.
    Err(logClientErr(
        &clientIP,
        "None",
        &model,
        StatusCode::TOO_MANY_REQUESTS,
        "All available LLM providers are rate limited or offline.".to_string(),
    ))

We’re done with handlers.rs. Save and exit the file.

runners.rs

Your harness needs to connect to an LLM to perform prompt completion. This part of the tutorial connects the harness to an LLM runner of your choice. If you have a list of available LLM runners, we do a weighted-random round-robin selection to pick a runner to complete each incoming HTTP-request. You can also provide a locally provisioned LLM runner to fall back to in case no cloud-based runners are available. Once you set this feature up, we will be using it for the rest of the term but will not be spending more time on it.

Create src/runners.rs:

server$ vi src/runners.rs

Start the file with the following use imports:

#![allow(non_snake_case)]

use rand::{RngExt, rng};
use std::{cmp::Ordering, collections::HashMap, env};

Next we define a struct, LlmRunner, to hold the base URL and API key of an LLM runner. LLM runners each hosts a different set of models. Put the string that identifies the model you want to use to the LLM runner in the model property. We will patch the "model" field of user’s prompt with this string before sending it to this runner. If you want the option to run multiple models from the same runner, create a separate LlmRunner for each model. The weight field records the priority you want to give to this runner. The value of the weight should be in the range [0.0, 1.0], with 1.0 being the highest priority. If a model hosted by Runner A has a weight of 0.8 and two models hosted by Runner B each has a weight of 0.5, you may see Runner B selected more frequently in your trace. The weights don’t have to add up to 1.0. Add the following code to your file:

#[derive(Clone, Debug)]
pub struct LlmRunner {
    pub baseUrl: &'static str,
    pub apiKey: Option<String>,
    pub model: &'static str,
    pub weight: f32,
}

The weighted-random round-robin selection is implemented as the iter() method of an Iterator trait. If the selected runner cannot complete the prompt, we move on to the next one on the list. Failover runners are also selected in a weighted-random round-robin manner. Add the following code to create an iterable list of LlmRunners:

pub struct LlmRunners {
    pub runners: Vec<LlmRunner>,
    pub local: Option<LlmRunner>,
}
impl LlmRunners {
    // The iterator that powers the `for (const runner of runners)` loop:
    // to iterate through the runner array randomized per HTTP-request to load balance  
    // Implemented using the Schwartzian Transform (Decorate-Sort-Undecorate).
    pub fn iter(&self) -> impl Iterator<Item = &LlmRunner> {
        let mut randgen = rng();

        // Calculate the Efraimidis and Spirakis Algorithm-Refined Sampling (A-Res) 
        // scores based on weights ["Information Processing Letters", 97:5, March 2006].
        // Each runner gets a randomized score that, when sorted, determines
        // the fallback order for each HTTP-request.
        // Each HTTP-request gets a different ordered list, effecting the
        // round-robin selection of runners across HTTP-requests.
        let mut scored: Vec<(f32, &LlmRunner)> = self.runners
            .iter()
            .map(|runner| { // decorate
                // f32:EPSILON: safeguard against 0.0 weight to prevent divide-by-zero panics
                (-(randgen.random_range::<f32,_>(0.0..1.0)).ln() / 
                    if runner.weight > 0.0 { runner.weight } else { f32::EPSILON },
                    runner)
            })
            .collect();
            // We MUST collect here: cannot sort a lazy, infinite stream,
            // all elements must be present in memory to be sorted.

        // Sort ascending by score (unstable is faster but doesn't preserve original
        // ordering of those with the same score, which is irrelevant in this case
        // given how the scores are computed using random numbers).
        scored.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal));

        // undecorate: drop the tuples, extract just the ordered references &LlmRunner; 
        // then dynamically chain the local fallback.
        scored
            .into_iter()
            .map(|(_, runner)| runner)
            .chain(self.local.as_ref())
                // self.local, being Option<T>, natively implements IntoIterator!
                // If Some, it yields 1 item. If None, it yields 0 items.
    }

    // new()
}

Finally, new() initializes the list of runners. LlmRunners::new() must be called before starting the web server, as we did in main() above. The code here lists LLM runners with free tiers. You have to provide your own API key for each runner in an environment variable. If you have subscription to other LLM runners (such as Claude or ChatGPT), you can add them to the list. If you have provisioned for a local LLM as shown in llmprompt-backend, you can assign it to the local variable as your ultimate fallback server. Replace // new() comment with:

    // Call new() from main() BEFORE starting the axum server.
    pub fn new() -> Self {
        let local: Option<LlmRunner> = None;
        // If you have provisioned a local llama-server, replace None above with:
        // Some(LlmRunner { baseUrl: "http://localhost:11433", apiKey: None, model: "", weight: 0.0 });
        
        let apiKeys = HashMap::from([
            ("cerebras", env::var("CEREBRAS_API_KEY").ok()),
            ("groq", env::var("GROQ_API_KEY").ok()),
            ("mistral", env::var("MISTRAL_API_KEY").ok()),
            ("gemini", env::var("GEMINI_API_KEY").ok()),
            ("nvidia", env::var("NVIDIA_API_KEY").ok()),
            //("openrouter", env::var("OPENROUTER_API_KEY").ok()),
        ]);
    
        let mut runners = vec![
            LlmRunner {
                baseUrl: "https://api.cerebras.ai",
                apiKey: apiKeys.get("cerebras").cloned().flatten(),
                model: "zai-glm-4.7",
                weight: 0.3, // 1M tpd permanent
            },
            LlmRunner {
                baseUrl: "https://api.groq.com/openai",
                apiKey: apiKeys.get("groq").cloned().flatten(),
                model: "llama-3.1-8b-instant",
                weight: 0.4, // 14.4K rpd, 100K tpm on Llama 3.3 70B
            },
            LlmRunner {
                baseUrl: "https://api.groq.com/openai",
                apiKey: apiKeys.get("groq").cloned().flatten(),
                model: "meta-llama/llama-4-scout-17b-16e-instruct",
                weight: 0.4, // 14.4K rpd, 100K tpm on Llama 3.3 70B
            },            
            LlmRunner {
                baseUrl: "https://api.mistral.ai",
                apiKey: apiKeys.get("mistral").cloned().flatten(),
                model: "ministral-8b-latest", // mistral-embed
                weight: 0.3, // slow on weekdays, 1 rps (60 rpm) 500K tpm; has embedding (1.0), drop to 0.4 if doing double duty
            },
            LlmRunner {
                baseUrl: "https://api.mistral.ai",
                apiKey: apiKeys.get("mistral").cloned().flatten(),
                model: "open-mistral-nemo", // mistral-embed
                weight: 0.2, // slow on weekdays, 1 rps (60 rpm) 500K tpm; has embedding (1.0), drop to 0.4 if doing double duty
            },
            LlmRunner {
                baseUrl: "https://integrate.api.nvidia.com",
                apiKey: apiKeys.get("nvidia").cloned().flatten(),
                model: "meta/llama-3.1-8b-instruct", // nvidia/NV-Embed-v1 4K dimension, 32K ctx or nvidia/embedqa-e5-v5 for Q&A 1024 dimensions, 8K ctx
                weight: 0.2, // 40 rpm; has embedding (0.67), cannot mix embedding models
                             // no more 1K lifetime cap (https://forums.developer.nvidia.com/t/request-more-4-000-credits-option-on-build-nvidia-com/344567)
            },
            LlmRunner {
                baseUrl: "https://integrate.api.nvidia.com",
                apiKey: apiKeys.get("nvidia").cloned().flatten(),
                model: "deepseek-ai/deepseek-v4-flash",
                weight: 0.2, // 40 rpm; has embedding (0.67), cannot mix embedding models
                             // no more 1K lifetime cap (https://forums.developer.nvidia.com/t/request-more-4-000-credits-option-on-build-nvidia-com/344567)
            },
            LlmRunner {
                baseUrl: "https://integrate.api.nvidia.com",
                apiKey: apiKeys.get("nvidia").cloned().flatten(),
                model: "nvidia/nemotron-nano-12b-v2-vl",
                weight: 0.3, // 40 rpm; has embedding (0.67), cannot mix embedding models
                             // no more 1K lifetime cap (https://forums.developer.nvidia.com/t/request-more-4-000-credits-option-on-build-nvidia-com/344567)
            },
            /* Don't include Open Router due to its low rate limit and variable model availability
            LlmRunner {
                baseUrl: "https://openrouter.ai/api",
                apiKey: apiKeys.get("openrouter").cloned().flatten(),
                model: "openrouter/free", // whichever best model is available and least loaded
                weight: 0.1,
            },
            */
        ];
    
        // Iterate backwards so `swap_remove` doesn't skip elements
        for i in (0..runners.len()).rev() {
            if runners[i].apiKey.as_deref().unwrap_or_default().is_empty() {
                // API Key missing: O(1) removal. The element swapped into `i`
                // came from the end of the vector, so it has already been checked!
                eprintln!("Runner {} INACTIVE. No API KEY.", runners[i].baseUrl);
                runners.swap_remove(i);
            } else {
                eprintln!("Runner {} ACTIVE.", runners[i].baseUrl);
            }
        }
    
        if runners.is_empty() && local.is_none() {
            eprintln!("CRITICAL: No API keys found and no local LLM configured. Terminating process.");
            std::process::exit(1);
        }
    
        runners.shrink_to_fit(); // free unused heap memory
    
        Self { runners, local }
    }

We’re done with runners.rs. Save and exit the file.

Build and test run

For faster subsequent builds:

server$ cargo install sccache
``
then edit your `~/.bashrc` or `~/.zshrc` and add:
```bash
export RUSTC_WRAPPER=sccache

To build your server:

server$ cargo build --release
server$ ln -s target/release/harnessd harnessd

The first time around cargo build runs, it will take some time to download and build all the 3rd-party crates. Be patient.

Build release version?

We would normally build for development without the --release flag, but due to the limited disk space on AWS virtual hosts, cargo build for debug version often runs out of space. The release version at least doesn’t keep debug symbols around.

Linking error with cargo build?

When running cargo build --release, if you see:

  error: linking with cc failed: exit status: 1
  note: collect2: fatal error: ld terminated with signal 9 [Killed]

below a long list of object files, try running cargo build --release again. It usually works the second time around, when it will have less remaining linking to do. If the error persisted, please talk to the teaching staff.

:point_right:Rust is a compiled language, like C/C++ and unlike Python, which is an interpreted language. This means you must run cargo build each and every time you made changes to your code, for the changes to show up in your executable.

Storing API keys

When we have only one or two environment variables to set, we can do it on the command-line interface (CLI) right before running harnessd, e.g.,

server$ sudo LLM_RUNNER_API_KEY=YOUR_API_KEY ./harnessd

When you have multiple API keys, however, this becomes impractical. A not very secure option is to store your API keys in /root/.bashrc or /root/.zshrc (:warning: make sure you don’t push this file to your git repo!):

server$ sudo su
# You are now root, note the command-line prompt changed from '$' to '#'.
# You can do a lot of harm with root privileges, so be very careful what you do here.
server# vi /root/.bashrc

Add these to the end of the file and provide your available API keys:

export CEREBRAS_API_KEY=""
export GROQ_API_KEY=""
export GEMINI_API_KEY=""
export MISTRAL_API_KEY=""
export NVIDIA_API_KEY=""
export OPENROUTER_API_KEY=""

After you add these environment variables to your /root/.bashrc, you’d have to do:

server# source /root/.bashrc

to load them into your current shell. Subsequent sudo su will load them automatically.

Then you can just run your harnessd:

server# ./harnessd
# Hit ^C to end the test
server# exit
# So that you're no longer root.
server$
`dotenv`

The dotenv library is often used to read .env file. However, common container, serverless, and managed runtime platforms (such as AWS ECS/Fargate and Lambda, Cloudflare Workers, Heroku, Kubernetes, Vercel, etc.), inject configurations directly into process memory, bypassing physical files. We thus import environment variables directly from the parent (shell) process when you run harnessd from the CLI or access them indirectly through systemd service configuration file.

:warning:Now return to the back-end cover page to complete a TODO item and test your implementation following the instructions in the subsequent “Testing llmPrompt APIs” section. You must complete ithe TODO to get full credit for this tutorial.

References

Prepared by Chenglin Li, Xin Jie ‘Joyce’ Liu, Sugih Jamin Last updated June 23rd, 2026