Schedtuna

A comprehensive reinforcement learning framework for optimizing Linux scheduler parameters across multiple scheduler types and parameter configurations.

Overview

Schedtuna uses reinforcement learning algorithms (PPO, SAC) to automatically discover optimal parameter configurations for various Linux schedulers.

Workload-Specific Optimization

The optimizer now supports workload-specific parameter optimization. You can specify the workload type you want to optimize for, and the system will:

Train an RL agent specifically for that workload type
Save optimized parameters to a JSON file (~/rl_scx_params/optimized_params.json)
Automatically apply these parameters when the scheduler detects the workload type

Supported Workload Types

latency_sensitive: High context switch rate, short execution bursts (e.g., web servers, message passing)
cpu_intensive: Low context switches, long runtime (e.g., computation, compilation)
cache_sensitive: Frequent CPU migrations, benefits from cache locality
gpu_intensive: GPU operations detected, needs big core coordination
mixed: Mixed characteristics, balanced approach

Workload Optimization Examples

from rl_scx_params import optimize_scheduler

# Optimize for latency-sensitive workloads
results = optimize_scheduler(
    workload_type="latency_sensitive",
    optimize_metric="pp_tokens_per_sec",
    episodes=100
)

# Optimize for CPU-intensive workloads
results = optimize_scheduler(
    workload_type="cpu_intensive",
    optimize_metric="tg_tokens_per_sec",
    episodes=150
)

Or use configuration files:

# latency_sensitive_config.yaml
scheduler_name: scx_flashyspark
workload_type: latency_sensitive
optimize_metric: pp_tokens_per_sec
episodes: 100

Integration with scx_flashyspark

The scx_flashyspark scheduler automatically loads optimized parameters from ~/rl_scx_params/optimized_params.json. When it detects a workload type change, it applies ALL corresponding optimized parameters if available, completely overriding default values.

Important: When ML-optimized parameters are available for a workload type, ALL parameters from the optimized set are applied, not just specific ones. This ensures consistent behavior controlled entirely by the ML optimization.

The JSON file structure includes ALL scheduler parameters:

{
  "latency_sensitive": {
    "parameters": {
      // Boolean parameters
      "sticky_cpu": false,
      "direct_dispatch": false,
      "aggressive_gpu_tasks": false,
      "local_pcpu": true,
      "no_wake_sync": false,
      "slice_lag_scaling": true,
      "local_kthreads": true,
      "stay_with_kthread": false,
      "native_priority": false,
      "tickless_sched": false,
      "timer_kick": false,
      
      // Time slice parameters (microseconds)
      "slice_us": 2048,
      "slice_us_min": 128,
      "slice_us_lag": 2048,
      "run_us_lag": 16384,
      
      // Other numeric parameters
      "cpu_busy_thresh": 50,
      "max_avg_nvcsw": 256
    },
    "reward": 125.3,
    "experiment_path": "~/rl_experiments/fs_results/...",
    "timestamp": "2024-01-20T10:30:00"
  },
  "cpu_intensive": { ... }
}

Quick Start

Installation

Setup Environment:

./setup.sh

Activate Virtual Environment:

source rl_scx_params/bin/activate

Using the Python API (Recommended)

The easiest way to use rl_scx_params is through the Python API:

from rl_scx_params import optimize_scheduler

# Quick optimization with default settings
results = optimize_scheduler(
    scheduler_name="scx_flashyspark",
    episodes=50,
    algorithm="ppo"
)

See the API Usage section for more details.

Using the Command Line

You can also use the traditional command-line interface:

Optimize scx_flashyspark:

# Quick test run
python main.py --scheduler scx_flashyspark --episodes 10 --baseline-runs 2

# Full optimization
python main.py --scheduler scx_flashyspark --algorithm ppo --episodes 100

Optimize scx_rusty:

python main.py --scheduler scx_rusty --algorithm sac --episodes 50

Custom Configuration:

python main.py --scheduler scx_flashyspark \
    --algorithm ppo \
    --episodes 200 \
    --optimize-metric tg_tokens_per_sec \
    --learning-rate 1e-4 \
    --experiment-name "gpu_workload_optimization"

Advanced Features

Experience Replay and Pre-training

Leverage historical experiment data to accelerate learning:

# Analyze historical patterns
python experience_replay.py --statistics --param-combinations 15

# Pre-train an agent
python experience_replay.py --pretrain-agent --algorithm ppo --pretrain-epochs 20

# Use pre-trained model
python main.py --scheduler scx_flashyspark \
    --pretrained-model pretrained_ppo_model.pt \
    --episodes 100

Parameter Analysis

Understand which parameters work best:

# Analyze parameter effectiveness
python pretrain.py --scheduler scx_flashyspark --action analyze

# Recompute rewards with new function
python pretrain.py --scheduler scx_flashyspark --action recompute

Custom Benchmarks

Use different benchmark tools:

python main.py --scheduler scx_flashyspark \
    --benchmark-cmd "/path/to/custom/benchmark" \
    --model-path "/path/to/model.gguf"

API Usage

The rl_scx_params package provides a clean Python API for programmatic usage:

Basic Usage

from rl_scx_params import optimize_scheduler

# Quick optimization with minimal configuration
results = optimize_scheduler(
    scheduler_name="scx_flashyspark",
    episodes=50,
    algorithm="ppo"
)

Using the RLSchedulerOptimizer Class

from rl_scx_params import RLSchedulerOptimizer

# Create optimizer with dictionary configuration
optimizer = RLSchedulerOptimizer({
    "scheduler_name": "scx_flashyspark",
    "algorithm": "ppo",
    "episodes": 100,
    "learning_rate": 0.001,
    "optimize_metric": "throughput"
})

# Get scheduler information
info = optimizer.get_scheduler_info()
print(f"Optimizing {info['scheduler_name']} with {len(info['parameters'])} parameters")

# Train the agent
results = optimizer.train()

# Test the best parameters
test_results = optimizer.test(results['experiment_path'], test_runs=10)

Using Configuration Files

Create a YAML configuration file (config.yaml):

scheduler_name: scx_flashyspark
algorithm: ppo
episodes: 100
learning_rate: 0.001
optimize_metric: throughput
experiment_name: my_experiment

Then use it in your code:

from rl_scx_params import RLSchedulerOptimizer

# Load from config file
optimizer = RLSchedulerOptimizer("config.yaml")
results = optimizer.train()

You can also use JSON configuration files:

{
  "scheduler_name": "scx_flashyspark",
  "algorithm": "sac",
  "episodes": 200,
  "learning_rate": 0.0001,
  "optimize_metric": "latency"
}

Defining Custom Scheduler Parameters

If you're working with a custom scheduler, you can define your own parameters:

optimizer = RLSchedulerOptimizer()

# Define custom parameters
optimizer.define_scheduler_params({
    'enable_turbo': {
        'type': 'boolean',
        'default': False,
        'description': 'Enable turbo boost mode',
        'command_arg': '--turbo'
    },
    'slice_duration': {
        'type': 'integer',
        'default': 10000,
        'min_value': 1000,
        'max_value': 50000,
        'description': 'Time slice duration in microseconds',
        'command_arg': '--slice-duration'
    },
    'load_threshold': {
        'type': 'float',
        'default': 0.75,
        'min_value': 0.0,
        'max_value': 1.0,
        'description': 'CPU load threshold',
        'command_arg': '--load-threshold'
    },
    'scheduling_mode': {
        'type': 'categorical',
        'default': 'balanced',
        'choices': ['performance', 'balanced', 'powersave'],
        'description': 'Overall scheduling mode',
        'command_arg': '--mode'
    }
})

# Train with custom parameters
results = optimizer.train()

Example Scripts

Check the examples/ directory for complete working examples:

basic_usage.py - Simple optimization example
advanced_usage.py - Full API demonstration
config_file_usage.py - Using configuration files
custom_scheduler_params.py - Custom scheduler parameters

Configuration Options

All configuration options available in OptimizerConfig:

Option	Type	Default	Description
`scheduler_name`	str	"scx_flashyspark"	Scheduler to optimize
`algorithm`	str	"ppo"	RL algorithm ("ppo" or "sac")
`episodes`	int	50	Number of training episodes
`baseline_runs`	int	3	Number of baseline evaluation runs
`learning_rate`	float	1e-3	Learning rate for RL agent
`update_frequency`	int	5	Agent update frequency (episodes)
`save_frequency`	int	25	Checkpoint save frequency
`benchmark_cmd`	str	None	Custom benchmark command
`model_path`	str	None	Path to model file
`timeout`	int	300	Benchmark timeout (seconds)
`optimize_metric`	str	"throughput"	Metric to optimize
`reward_scaling`	float	1.0	Reward scaling factor
`experiment_name`	str	None	Custom experiment name
`use_gpu`	bool	True	Use GPU if available
`pretrained_model_path`	str	None	Path to pretrained model

Command Line Reference

Main Training Script (`main.py`)

Core Options:

--scheduler {scx_flashyspark,scx_rusty}: Scheduler to optimize
--algorithm {ppo,sac}: RL algorithm to use
--episodes N: Number of training episodes
--optimize-metric {pp_tokens_per_sec,tg_tokens_per_sec}: Optimization target

Training Configuration:

--learning-rate FLOAT: Learning rate for RL agent
--update-frequency N: Agent update frequency
--save-frequency N: Checkpoint save frequency
--baseline-runs N: Number of baseline evaluation runs

System Configuration:

--model-path PATH: Path to model file for benchmarking
--benchmark-cmd PATH: Path to benchmark executable
--timeout N: Benchmark timeout in seconds
--no-gpu: Force CPU usage

Experiment Management:

--experiment-name NAME: Custom experiment name
--pretrained-model PATH: Pre-trained model to load
--reward-scaling FLOAT: Reward scaling factor

Experience Replay Script (`experience_replay.py`)

Data Management:

--results-dir PATH: Directory containing experiment results
--experiment-patterns PATTERN [PATTERN ...]: Experiment name patterns to include
--max-episodes-per-exp N: Maximum episodes per experiment

Analysis Options:

--statistics: Show detailed statistics
--param-combinations N: Number of top parameter combinations to analyze
--recalculate: Recalculate rewards with new function

Pre-training:

--pretrain-agent: Pre-train a new agent
--algorithm {ppo,sac}: Algorithm for pre-training
--pretrain-epochs N: Number of pre-training epochs

Analysis Script (`pretrain.py`)

Actions:

--action {analyze,pretrain,recompute}: Action to perform
--scheduler SCHEDULER: Scheduler to analyze

Output and Results

Experiment Structure

results/
├── scx_flashyspark_ppo_100ep_20241201_120000/
│   ├── experiment.log              # Detailed training logs
│   ├── best_parameters.json        # Optimal parameters found
│   ├── episode_results.json        # Complete episode data
│   ├── episode_results.csv         # Episode data in CSV format
│   ├── baseline_results.json       # Baseline performance
│   ├── performance_plots.png       # Training visualizations
│   └── checkpoint_episode_50.pt    # Model checkpoints
└── scx_rusty_sac_50ep_20241201_130000/
    └── ...

Key Metrics

Performance Metrics:

PP tokens/second: Prompt processing throughput
TG tokens/second: Text generation throughput
Execution time: Benchmark completion time
Success rate: Percentage of successful runs

Training Metrics:

Reward: Combined optimization metric
Policy loss: RL algorithm policy loss
Value loss: Value function loss
Exploration: Current exploration rate

Adding New Schedulers

The framework is designed for easy extension. To add a new scheduler:

Define Scheduler Configuration:

def get_scx_newscheduler_config() -> SchedulerConfig:
    parameters = {
        'enable_feature': ParameterSpec(
            'enable_feature', ParameterType.BOOLEAN, False,
            description="Enable special feature",
            command_arg="--enable-feature"
        ),
        'time_slice': ParameterSpec(
            'time_slice', ParameterType.INTEGER, 10000,
            min_value=1000, max_value=50000,
            description="Time slice in microseconds",
            command_arg="--time-slice"
        ),
        'policy_mode': ParameterSpec(
            'policy_mode', ParameterType.CATEGORICAL, 'balanced',
            choices=['balanced', 'performance', 'efficiency'],
            description="Scheduling policy mode",
            command_arg="--policy"
        )
    }
    
    return SchedulerConfig(
        name="scx_newscheduler",
        binary_path="/path/to/scx_newscheduler",
        parameters=parameters,
        description="New scheduler with custom parameters"
    )

Register Configuration:

# In get_scheduler_config() function
configs = {
    'scx_flashyspark': get_scx_flashyspark_config,
    'scx_rusty': get_scx_rusty_config,
    'scx_newscheduler': get_scx_newscheduler_config,  # Add this line
}

Update Command Line Options:

parser.add_argument(
    "--scheduler", "-S",
    choices=["scx_flashyspark", "scx_rusty", "scx_newscheduler"],  # Add new scheduler
    default="scx_flashyspark",
    help="Scheduler to optimize"
)

System Requirements

Operating System:

Linux with sched_ext support
Kernel 6.12+ recommended

Storage:

2GB+ free space for experiments and models
Additional space for benchmark models

Performance Tips

Start Small: Begin with 25-50 episodes to verify system operation
Use Baselines: Always run baseline evaluation for comparison
Monitor Resources: Watch CPU, memory, and GPU utilization
Leverage History: Use experience replay for faster convergence
Choose Metrics Wisely: Optimize for your specific workload requirements
Checkpoint Frequently: Use appropriate save frequency for long runs

Debug Mode

Enable verbose logging:

python main.py --verbose --scheduler scx_flashyspark --episodes 10

Name		Name	Last commit message	Last commit date
examples		examples
.gitignore		.gitignore
EXPERIENCE_REPLAY_GUIDE.md		EXPERIENCE_REPLAY_GUIDE.md
MANIFEST.in		MANIFEST.in
PARAMETER_OPTIMIZATION_COMPLETE.md		PARAMETER_OPTIMIZATION_COMPLETE.md
README.md		README.md
__init__.py		__init__.py
experience_replay.py		experience_replay.py
main.py		main.py
pretrain.py		pretrain.py
requirements.txt		requirements.txt
rl_optimizer.py		rl_optimizer.py
setup.py		setup.py
setup.sh		setup.sh

Latest commit

History

Folders and files

Repository files navigation