File size: 5,917 Bytes
323c6ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"
import psutil
import numpy as np
import tensorflow as tf
from tf_agents.environments import py_environment
from tf_agents.specs import array_spec
from tf_agents.trajectories import time_step as ts
from config import Config
try:
import pynvml
pynvml.nvmlInit()
HAS_NVML = True
except ImportError:
HAS_NVML = False
def configure_gpu():
"""Configures TensorFlow memory growth to prevent VRAM allocation errors."""
gpus = tf.config.list_physical_devices('GPU')
if gpus:
try:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
print(f"✅ Found and configured {len(gpus)} GPU(s).")
except RuntimeError as e:
print(e)
else:
print("⚠️ No GPU found. Falling back to CPU.")
def get_system_telemetry():
"""Fetches real-time RAM, CPU, and GPU usage metrics for TensorBoard profiling."""
metrics = {
'system/ram_percent': psutil.virtual_memory().percent,
'system/cpu_percent': psutil.cpu_percent(),
}
if HAS_NVML:
try:
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
metrics['hardware/gpu_utilization'] = float(util.gpu)
metrics['hardware/gpu_temperature_c'] = float(temp)
metrics['hardware/gpu_memory_percent'] = (mem.used / mem.total) * 100.0
except Exception:
pass
return metrics
class PettingZooToTFAgentsWrapper(py_environment.PyEnvironment):
"""Wraps a PettingZoo Parallel environment into a TF-Agents PyEnvironment."""
def __init__(self, pz_env):
super().__init__()
self._env = pz_env
self._current_info = {}
# Read host array shape sizes straight out of Config
total_nodes = Config.TOTAL_HOSTS()
# Red Actions: total_nodes * 4 choices | Blue Actions: total_nodes * 2 choices
self._action_spec = {
'red': array_spec.BoundedArraySpec(shape=(), dtype=np.int32, minimum=0, maximum=(total_nodes * 4) - 1, name='red_action'),
'blue': array_spec.BoundedArraySpec(shape=(), dtype=np.int32, minimum=0, maximum=(total_nodes * 2) - 1, name='blue_action')
}
# Match the complex dictionary nested schema
self._observation_spec = {
agent: {
'node_features': array_spec.BoundedArraySpec(shape=(total_nodes, 2), dtype=np.float32, minimum=0, maximum=1, name='features'),
'adjacency_matrix': array_spec.BoundedArraySpec(shape=(total_nodes, total_nodes), dtype=np.float32, minimum=0, maximum=1, name='topology')
} for agent in ['red', 'blue']
}
def action_spec(self):
return self._action_spec
def observation_spec(self):
return self._observation_spec
def reward_spec(self):
"""Declares multi-agent reward specs to prevent unpacking errors."""
return {
'red': array_spec.ArraySpec(shape=(), dtype=np.float32, name='red_reward'),
'blue': array_spec.ArraySpec(shape=(), dtype=np.float32, name='blue_reward')
}
def get_current_info(self):
"""Safely forwards custom internal metadata down to the trainer."""
return self._current_info
def _reset(self):
obs, infos = self._env.reset()
self._current_info = infos
# Guarantee observations map explicitly to nested dictionary signatures
formatted_obs = {
'red': {
'node_features': np.array(obs['red']['node_features'], dtype=np.float32),
'adjacency_matrix': np.array(obs['red']['adjacency_matrix'], dtype=np.float32)
},
'blue': {
'node_features': np.array(obs['blue']['node_features'], dtype=np.float32),
'adjacency_matrix': np.array(obs['blue']['adjacency_matrix'], dtype=np.float32)
}
}
initial_rewards = {
'red': np.array(0.0, dtype=np.float32),
'blue': np.array(0.0, dtype=np.float32)
}
return ts.TimeStep(
step_type=np.array(ts.StepType.FIRST, dtype=np.int32),
reward=initial_rewards,
discount=np.array(1.0, dtype=np.float32),
observation=formatted_obs
)
def _step(self, action):
# Convert actions into string keys for env.py mapping compliance
pz_actions = {'red': int(action['red']), 'blue': int(action['blue'])}
obs, rewards, terminations, truncations, infos = self._env.step(pz_actions)
self._current_info = infos
# Re-enforce nested array typing explicitly on step updates
formatted_obs = {
'red': {
'node_features': np.array(obs['red']['node_features'], dtype=np.float32),
'adjacency_matrix': np.array(obs['red']['adjacency_matrix'], dtype=np.float32)
},
'blue': {
'node_features': np.array(obs['blue']['node_features'], dtype=np.float32),
'adjacency_matrix': np.array(obs['blue']['adjacency_matrix'], dtype=np.float32)
}
}
formatted_rewards = {
'red': np.array(rewards['red'], dtype=np.float32),
'blue': np.array(rewards['blue'], dtype=np.float32)
}
if terminations['red'] or truncations['red']:
return ts.termination(formatted_obs, reward=formatted_rewards)
return ts.transition(formatted_obs, reward=formatted_rewards, discount=np.array(0.99, dtype=np.float32)) |