InfraGym

Unified Infrastructure Simulation Environments for LLM Agents

Environments

1. Load Balancing (load_balance)

Optimize job distribution across multiple servers with varying capabilities:

• Goal: Minimize job waiting time and balance server loads
• Actions: Choose which server to assign incoming jobs
• Observations: Server loads, job sizes, service rates, queue states

2. Adaptive Bitrate Streaming (abr_sim)

Manage video streaming quality based on network conditions:

• Goal: Maximize video quality while minimizing stalls and bitrate switches
• Actions: Select video bitrate level (0.3-4.3 Mbps)
• Observations: Buffer size, network throughput, past performance

3. Cache Management (cache)

Decide which objects to cache for optimal hit rates:

• Goal: Maximize cache hit rate while managing storage constraints
• Actions: Accept or reject cache requests
• Observations: Cache state, object metadata, access patterns

Usage Examples

Load Balancing

importinfragym# Create load balancing environmentenv=infragym.make('load_balance', num_servers=3, max_episode_steps=100,
reward_scale=1.0)
obs, info=env.reset()
# Environment provides natural language descriptionsprint("Task:", obs['task_description'])
print("Observation:", obs['text_observation']) print("Available actions:", obs['available_actions'])
# Take action (assign job to server 1)obs, reward, done, truncated, info=env.step(1)
print(f"Assigned job to server 1, reward: {reward:.3f}")
env.close()

Adaptive Bitrate Streaming

importinfragym# Create ABR environmentenv=infragym.make('abr_sim',
max_episode_steps=200,
dummy_mode=True) # Use synthetic network tracesobs, info=env.reset()
# Choose bitrate level (0=lowest, 5=highest)bitrate_action=2# Medium bitrateobs, reward, done, truncated, info=env.step(bitrate_action)
print(f"Selected bitrate level {bitrate_action}")
print(f"Buffer: {info['buffer_size']:.1f}s")
print(f"Stall time: {info['stall_time']:.3f}s")
env.close()

Cache Management

importinfragym# Create cache environment env=infragym.make('cache',
cache_size=1000,
max_episode_steps=500)
obs, info=env.reset()
# Make caching decision (1=cache, 0=reject)cache_decision=1# Accept the objectobs, reward, done, truncated, info=env.step(cache_decision)
print(f"Cache decision: {'Accept'ifcache_decisionelse'Reject'}")
print(f"Cache hit rate: {info.get('hit_rate', 0):.2%}")
env.close()

API Reference

Core Functions

importinfragym# Create environmentenv=infragym.make(env_name, **kwargs)
# Environment discoveryenvironments=infragym.list_environments()
info=infragym.get_env_info(env_name)
# Resource cleanupinfragym.close_all()

Environment Parameters

Load Balance (load_balance)

• num_servers: Number of servers (default: 5)
• max_episode_steps: Episode length (default: 50)
• reward_scale: Reward scaling factor (default: 1.0)

ABR Streaming (abr_sim)

• max_episode_steps: Episode length (default: 100)
• dummy_mode: Use synthetic traces (default: True)
• reward_scale: Reward scaling factor (default: 1.0)

Cache Management (cache)

• cache_size: Cache capacity (default: 1000)
• max_episode_steps: Episode length (default: 200)
• trace_type: Trace data source (default: 'synthetic')

Observation Format

All environments return observations with:

{
'text_observation': str, # Natural language description'structured_observation': np.ndarray, # Numerical features 'available_actions': str, # Action descriptions'task_description': str# Environment-specific instructions
}