Skip to content

Workload Inference and Classification

Jump to bottom
William Smith edited this page Jul 13, 2026 · 1 revision

Workload Inference and Classification

Relevant source files

The following files were used as context for generating this wiki page:

The Workload Inference and Classification system is responsible for transforming natural language task descriptions into structured TaskRequirements. This process bridges the gap between a user's intent (e.g., "run a 70b coding model") and the deterministic hardware constraints (RAM, tools, context window) required by the Placement Engine.

System Architecture

AXIS employs a hybrid approach to classification, combining a Semantic Path (LLM-based) with a Reflex Path (string-matching). The system prioritizes speed and reliability, ensuring that even if a local LLM is slow or unavailable, a deterministic fallback is provided within milliseconds.

The Inference Pipeline

The entry point for requirement derivation is workload.InferRequirements internal/workload/requirements.go:18-39.

graph TD
    UserPrompt["User Prompt / Description"] --> InferReqs["workload.InferRequirements()"]
    
    subgraph "Classification Layer"
        InferReqs --> ResolveMatch["resolveWorkloadMatch()"]
        ResolveMatch --> Semantic["llmrouter.Engine.Classify() (Semantic Path)"]
        ResolveMatch --> Reflex["workload.Match() (Reflex Path)"]
        Semantic -- "Timeout/Error" --> Reflex
    end

    subgraph "Requirement Derivation"
        ResolveMatch --> ApplyProfile["Apply(match, reqs)"]
        ApplyProfile --> Profiles["workload.profileForClass()"]
        ApplyProfile --> Heuristics["Heuristic Refinement (RAM Floors, Tools)"]
    end

    subgraph "Additive Inference"
        Heuristics --> ContextInference["InferContextWindowTokens()"]
        ContextInference --> BackendInference["InferBackends()"]
    end

    BackendInference --> FinalReqs["models.TaskRequirements"]
Loading

Sources: internal/workload/requirements.go:10-39, internal/llmrouter/engine.go:137-150

1. Semantic LLM Router (llmrouter.Engine)

The llmrouter.Engine acts as a "semantic reflex" classifier. It is designed to run locally with a strict latency budget to avoid delaying the placement pipeline internal/llmrouter/engine.go:5-11.

Implementation Details

Data Structures

The IntentSignal captures the output of the semantic classification:

Field Type Description
Class models.WorkloadClass The identified category (e.g., go-build, local-llm-inference).
Confidence float64 Model-reported confidence [0.0, 1.0].
Source ClassifySource Either semantic or reflex.
Signals []string Keywords or features that drove the decision.

Sources: internal/llmrouter/engine.go:52-71, internal/llmrouter/engine.go:168-173

2. Reflex Fallback and String Matching

When the semantic engine fails or is disabled, AXIS uses workload.Match (Reflex Path). This is a deterministic keyword-based system that maps substrings to WorkloadClass internal/llmrouter/engine.go:40-48.

Priority Hierarchy

The reflex matcher handles "promotion" of specific classes to prevent ambiguous matches. For example, if a prompt contains both "ollama" and "128k context", it is promoted from ClassLocalLLMInference to ClassLongContextInference workload/match.go.

Common Mappings:

Sources: internal/workload/requirements.go:229-235, internal/workload/requirements_test.go:139-191

3. Requirement Derivation (RAM & Tools)

Once a WorkloadClass is identified, workload.Apply populates the TaskRequirements using a combination of static profiles and dynamic heuristics internal/workload/requirements.go:43-78.

RAM Floor Heuristics

The system applies "floors" based on model sizes mentioned in the prompt, which can override the base class requirements:

Mentioned Size RAM Floor (MB)
70b 12,288 MB
13b / 32b 8,192 MB
7b / 8b 4,096 MB
Generic llm 6,144 MB

Sources: internal/workload/requirements.go:169-182

Tool and Backend Detection

The system scans for explicit backend preferences and required tools:

4. Context Window Inference

AXIS performs additive inference for long-context tasks. InferContextWindowTokens parses human-readable context sizes and adjusts RAM requirements accordingly internal/workload/requirements.go:110-124.

Keyword Tokens RAM Adjustment
1m tokens, million token 1,000,000 Floor to 12,288 MB
512k 512,000 Floor to 8,192 MB
128k, long context 128,000 Floor to 6,144 MB

Sources: internal/workload/requirements.go:127-140

5. Cloud Fallback Mechanism

For the axis llm command surface, AXIS supports an interactive cloud fallback if the local classifier is unavailable cmd/axis/llm.go:185-200.

sequenceDiagram
    participant CLI as "cmd/axis (llmCmd)"
    participant Eng as "llmrouter.Engine"
    participant Cloud as "llmrouter.SelectCloudFallback"
    participant Prov as "Cloud Provider (Groq/Anthropic)"

    CLI->>Eng: Classify(prompt)
    Eng-->>CLI: SourceReflex (Local LLM Down)
    CLI->>Cloud: Select best provider (cost/latency)
    Cloud->>Prov: Health Check / Cost Estimate
    Prov-->>Cloud: OK (Latency 40ms)
    Cloud-->>CLI: Selected: Anthropic
    CLI->>Prov: Send(Classification Prompt)
    Prov-->>CLI: JSON {class: "go-build", ...}
Loading

Sources: cmd/axis/llm.go:116-141, internal/llmrouter/cloud.go:160-182, internal/llmrouter/cloud_test.go:162-197

Provider Registry

The llmrouter.Registry manages multiple backends (OpenRouter, Groq, Anthropic). It selects the fallback based on the inference.prefer configuration (defaulting to latency) internal/llmrouter/registry.go:10-19, internal/llmrouter/cloud_test.go:162-184.

6. Summary of Key Entities

Code Entity File Path Role
TaskRequirements internal/models/tasks.go Struct containing RAM, Tools, and Context Window.
Engine internal/llmrouter/engine.go:74-79 Orchestrates local LLM classification and reflex fallback.
InferRequirements internal/workload/requirements.go:18 Main logic for parsing strings into requirements.
Apply internal/workload/requirements.go:43 Maps WorkloadClass to hardware/tool profiles.
CloudProvider internal/llmrouter/cloud.go:54-62 Implementation for OpenAI-compatible and Anthropic APIs.

Sources: internal/llmrouter/engine.go, internal/workload/requirements.go, internal/llmrouter/cloud.go

AXIS Wiki

Architecture Artifacts

DeepWiki

Clone this wiki locally