Hardware-aware agent to orchestrate and optimize NN training workloads.
TrainPilot is a specialized agent designed to intelligently orchestrate and optimize the execution of neural network training workloads across heterogeneous hardware environments, specifically targeting Apple Neural Engine (ANE) and CPU backends. It addresses the complex challenge of routing ML tasks to the most performant accelerator available, a problem currently unaddressed by commercial tooling.
The core functionality revolves around a ReasoningAgent that observes workload characteristics—such as model architecture complexity, dataset size, and real-time hardware utilization—to make informed decisions. It dynamically selects the optimal training backend, aiming for significant speedups over naive or random hardware assignment.
pip install train-pilot# Run end-to-end orchestration
python train_orchestrated.py --model small_cnn --dataset ./test_data_setup
# Exit code 0 = success, 1 = failureOr use the components directly in Python:
from hardware_orchestrator import HardwareOrchestrator
from train_orchestrated import TrainingEnvironment
import numpy as np
import time
# Initialize orchestrator and environment
agent = HardwareOrchestrator()
X_train = np.random.rand(100, 28, 28)
y_train = np.random.randint(0, 10, 100)
environment = TrainingEnvironment(X_train, y_train, {"name": "small_cnn", "layers": 3, "has_conv": True, "has_rnn": False})
# Profile workload
profile = agent.profile_workload({"name": "small_cnn", "layers": 3, "has_conv": True, "has_rnn": False}, 100)
# Observe hardware and decide backend
hw_state = {"ane_available": False, "cpu_utilization": 0.3, "timestamp": time.time()}
agent.observe(hw_state, profile)
backend = agent.think()
# Execute training
result = agent.act(backend, environment)
print(f"Training result: {result['message']}")The agent can analyze a given model and dataset to extract crucial complexity metrics necessary for performance prediction.
from hardware_orchestrator import HardwareOrchestrator
orchestrator = HardwareOrchestrator()
profile = orchestrator.profile_workload(
{"name": "small_cnn", "layers": 3, "has_conv": True, "has_rnn": False},
dataset_size=10000
)
print(f"Model parameters: {profile['params']}, Depth: {profile['depth']}")Based on workload profile and hardware state, TrainPilot estimates training time and selects the best backend (ANE or CPU).
# Assuming profile was gathered previously
hw_state = {"ane_available": True, "cpu_utilization": 0.2, "timestamp": time.time()}
ane_time = orchestrator.predict_performance(profile, "ane")
cpu_time = orchestrator.predict_performance(profile, "cpu")
backend, confidence = orchestrator.select_backend(profile, hw_state)
print(f"ANE: {ane_time:.2f}s, CPU: {cpu_time:.2f}s => Selected: {backend} (confidence: {confidence})")TrainPilot operates around a central HardwareOrchestrator agent housed within hardware_orchestrator.py. This agent acts as the decision-maker, consuming inputs from workload analysis and real-time hardware monitoring.
The core loop follows an Observe → Think → Act pattern:
-
Input: Model config + dataset size
$\rightarrow$ profile_workload() - Profile: Extract metrics (params, layers, conv/rnn architecture) → dict with performance characteristics
-
Observe: Read hardware state (ANE available, CPU utilization) →
observe(hw_state, profile) -
Think: Deterministic rule cascade selects optimal backend →
think()returns "ane" or "cpu" -
Predict: Estimate execution time for each backend →
predict_performance(profile, backend) -
Act: Execute training on selected backend via
TrainingEnvironment→act(backend, environment) -
Execution:
TrainingEnvironmentroutes to ANE trainer or PyTorch CPU backend
┌─────────────────────────────────────────────────────────────┐
│ HardwareOrchestrator (Decision Loop) │
├─────────────────────────────────────────────────────────────┤
│ observe(hw_state, profile) ──→ think() ──→ act(backend) │
│ │
│ Rules: ANE unavailable? → CPU │
│ RNN model? → CPU (0.95 confidence) │
│ Conv + low CPU? → ANE (0.85 confidence) │
│ Default → CPU │
└──────────────────────────┬──────────────────────────────────┘
│
┌───────▼────────┐
│ TrainingEnv │
├────────────────┤
│ execute_train │
│ (backend, cfg) │
└───────┬────────┘
│
┌────────────────┴────────────────┐
│ │
┌─────▼────┐ ┌──────▼──────┐
│ ANE Path │ │ CPU Path │
│ (PyTorch)│ │ (PyTorch) │
└──────────┘ └─────────────┘
Core decision-making agent for backend selection.
Constructor:
__init__(): Initialize orchestrator with empty observation state.
Methods:
-
profile_workload(model_config: dict, dataset_size: int) -> dict: Extracts workload characteristics. Returns dict with keys:params,layers,has_conv,has_rnn,depth,avg_layer_size.model_config = {"name": "small_cnn", "layers": 3, "has_conv": True, "has_rnn": False} profile = orchestrator.profile_workload(model_config, dataset_size=10000) # {'params': 50000, 'layers': 3, 'has_conv': True, 'has_rnn': False, 'depth': 5, 'avg_layer_size': 1000}
-
observe(hardware_state: dict, workload_profile: dict) -> None: Store current observation (called beforethink()).hw_state = {"ane_available": True, "cpu_utilization": 0.3, "timestamp": time.time()} orchestrator.observe(hw_state, profile)
-
think() -> str: Decide backend ("ane" or "cpu") based on last observation.backend = orchestrator.think() # Returns "ane" or "cpu"
-
predict_performance(profile: dict, backend: str) -> float: Estimate training time (seconds).ane_time = orchestrator.predict_performance(profile, "ane") cpu_time = orchestrator.predict_performance(profile, "cpu")
-
select_backend(profile: dict, hardware_state: dict) -> tuple: Returns(backend: str, confidence: float)using deterministic rules.backend, confidence = orchestrator.select_backend(profile, hw_state) # ("cpu", 0.8) or ("ane", 0.85)
-
act(action: str, environment) -> dict: Execute training on selected backend. Returns dict with keys:status,backend,execution_time,message.result = orchestrator.act("cpu", training_environment) # {'status': 'success', 'backend': 'cpu', 'execution_time': 2.34, 'message': '...'}
Manages training execution on ANE or CPU backends.
Constructor:
__init__(X_train: np.ndarray, y_train: np.ndarray, model_config: dict): Initialize with training data and model configuration.
Methods:
execute_training(backend: str, workload_profile: dict) -> dict: Execute training on specified backend. Returns dict with keys:status,backend,execution_time,message.
This project is inspired by the growing need for efficient resource utilization in edge and local ML inference/training, particularly where proprietary hardware accelerators (like ANE) are involved. The underlying concepts draw from dynamic resource scheduling algorithms used in cloud computing, adapted for the constraints of local, heterogeneous silicon.
Tests are maintained in the tests/ directory and cover the core logic of workload profiling and backend selection, ensuring the agent correctly routes workloads for small CNNs, large dense models, and sequential models.
We welcome contributions! Please see our contribution guidelines for details on submitting pull requests, reporting bugs, or suggesting features.
This work is independent, but related research in heterogeneous computing can be found in papers concerning dynamic workload migration and accelerator offloading.
The project is licensed under the MIT License - see the LICENSE file for details.