1
Fork 0

refactor(hpc): ruff checks

This commit is contained in:
Tibo De Peuter 2026-04-05 08:05:13 +02:00
parent ef380d073f
commit fb346a06b2
4 changed files with 66 additions and 54 deletions

View file

@ -21,13 +21,14 @@ except ImportError:
print("Error: Missing dependency. Please run: pip install tensorboard") print("Error: Missing dependency. Please run: pip install tensorboard")
sys.exit(1) sys.exit(1)
def explore_run(log_dir): def explore_run(log_dir):
""" """
Extracts and displays a summary of scalar metrics from a TensorBoard log directory. Extracts and displays a summary of scalar metrics from a TensorBoard log directory.
""" """
print(f"\n{'='*20} Exploring Run {'='*20}") print(f"\n{'=' * 20} Exploring Run {'=' * 20}")
print(f"Directory: {log_dir}") print(f"Directory: {log_dir}")
print(f"{'='*55}\n") print(f"{'=' * 55}\n")
if not os.path.exists(log_dir): if not os.path.exists(log_dir):
print(f"Error: Directory '{log_dir}' does not exist.") print(f"Error: Directory '{log_dir}' does not exist.")
@ -35,17 +36,20 @@ def explore_run(log_dir):
# Initialize EventAccumulator # Initialize EventAccumulator
# size_guidance=0 loads all data points for each tag. # size_guidance=0 loads all data points for each tag.
ea = event_accumulator.EventAccumulator(log_dir, size_guidance={ ea = event_accumulator.EventAccumulator(
event_accumulator.SCALARS: 0, log_dir,
event_accumulator.TENSORS: 0, size_guidance={
}) event_accumulator.SCALARS: 0,
event_accumulator.TENSORS: 0,
},
)
print("Loading event files (this may take a moment for large runs)...") print("Loading event files (this may take a moment for large runs)...")
ea.Reload() ea.Reload()
tags = ea.Tags() tags = ea.Tags()
scalar_tags = tags.get('scalars', []) scalar_tags = tags.get("scalars", [])
if not scalar_tags: if not scalar_tags:
print("No scalar metrics found in this directory.") print("No scalar metrics found in this directory.")
return None return None
@ -60,43 +64,48 @@ def explore_run(log_dir):
events = ea.Scalars(tag) events = ea.Scalars(tag)
if not events: if not events:
continue continue
values = [e.value for e in events] values = [e.value for e in events]
last_event = events[-1] last_event = events[-1]
data[tag] = values data[tag] = values
summary.append({ summary.append(
"Metric": tag, {
"Steps": len(events), "Metric": tag,
"Last Value": f"{last_event.value:.4f}", "Steps": len(events),
"Max": f"{max(values):.4f}", "Last Value": f"{last_event.value:.4f}",
"Min": f"{min(values):.4f}" "Max": f"{max(values):.4f}",
}) "Min": f"{min(values):.4f}",
}
)
# Display summary table formatted manually # Display summary table formatted manually
summary = sorted(summary, key=lambda x: x['Metric']) summary = sorted(summary, key=lambda x: x["Metric"])
print(f"{'Metric':<30} {'Steps':>10} {'Last':>12} {'Max':>12} {'Min':>12}") print(f"{'Metric':<30} {'Steps':>10} {'Last':>12} {'Max':>12} {'Min':>12}")
print("-" * 80) print("-" * 80)
for row in summary: for row in summary:
print(f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} {row['Max']:>12} {row['Min']:>12}") print(
f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} "
f"{row['Max']:>12} {row['Min']:>12}"
)
# Calculate and display global metadata # Calculate and display global metadata
if 'charts/SPS' in data: if "charts/SPS" in data:
sps_events = ea.Scalars('charts/SPS') sps_events = ea.Scalars("charts/SPS")
if len(sps_events) > 1: if len(sps_events) > 1:
total_duration_hours = (sps_events[-1].wall_time - sps_events[0].wall_time) / 3600 total_duration_hours = (sps_events[-1].wall_time - sps_events[0].wall_time) / 3600
print(f"\nTotal Recorded Duration: {total_duration_hours:.2f} hours") print(f"\nTotal Recorded Duration: {total_duration_hours:.2f} hours")
# Estimate completion if total_timesteps is available in hyperparameters # Estimate completion if total_timesteps is available in hyperparameters
try: try:
hp_tags = [t for t in tags.get('tensors', []) if 'hyperparameters' in t] hp_tags = [t for t in tags.get("tensors", []) if "hyperparameters" in t]
if hp_tags: if hp_tags:
hp_event = ea.Tensors(hp_tags[0])[0] hp_event = ea.Tensors(hp_tags[0])[0]
hp_text = hp_event.tensor_proto.string_val[0].decode('utf-8') hp_text = hp_event.tensor_proto.string_val[0].decode("utf-8")
if 'total_timesteps' in hp_text: if "total_timesteps" in hp_text:
for line in hp_text.split('\n'): for line in hp_text.split("\n"):
if 'total_timesteps' in line: if "total_timesteps" in line:
target = int(line.split('|')[2].strip()) target = int(line.split("|")[2].strip())
current = ea.Scalars(scalar_tags[0])[-1].step current = ea.Scalars(scalar_tags[0])[-1].step
percent = (current / target) * 100 percent = (current / target) * 100
print(f"Progress: {current:,} / {target:,} steps ({percent:.1f}%)") print(f"Progress: {current:,} / {target:,} steps ({percent:.1f}%)")
@ -105,31 +114,30 @@ def explore_run(log_dir):
return data return data
def main(): def main():
parser = argparse.ArgumentParser(description="Clean, reproducible TensorBoard exploration tool.") parser = argparse.ArgumentParser(description="Reproducible TensorBoard exploration tool.")
parser.add_argument("log_dir", help="Path to the TensorBoard run directory.") parser.add_argument("log_dir", help="Path to the TensorBoard run directory.")
parser.add_argument("--csv", help="Optional: Path to save all scalar data as a CSV.", default=None) parser.add_argument("--csv", help="Optional: Path to export scalar data to CSV.", default=None)
args = parser.parse_args() args = parser.parse_args()
scalar_data = explore_run(args.log_dir) scalar_data = explore_run(args.log_dir)
if args.csv and scalar_data: if args.csv and scalar_data:
# Reloading for wall_time and steps # Reloading for wall_time and steps
ea = event_accumulator.EventAccumulator(args.log_dir).Reload() ea = event_accumulator.EventAccumulator(args.log_dir).Reload()
with open(args.csv, mode='w', newline='') as f: with open(args.csv, mode="w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=["tag", "step", "value", "wall_time"]) writer = csv.DictWriter(f, fieldnames=["tag", "step", "value", "wall_time"])
writer.writeheader() writer.writeheader()
for tag in scalar_data.keys(): for tag in scalar_data.keys():
for e in ea.Scalars(tag): for e in ea.Scalars(tag):
writer.writerow({ writer.writerow(
"tag": tag, {"tag": tag, "step": e.step, "value": e.value, "wall_time": e.wall_time}
"step": e.step, )
"value": e.value,
"wall_time": e.wall_time
})
print(f"\nData exported to: {args.csv}") print(f"\nData exported to: {args.csv}")
if __name__ == "__main__": if __name__ == "__main__":
main() main()

View file

@ -53,17 +53,17 @@ def main() -> None:
dep_dict: dict[str, str] = {} dep_dict: dict[str, str] = {}
for dep in data.get("project", {}).get("dependencies", []): for dep in data.get("project", {}).get("dependencies", []):
dep_dict[normalise(pkg_name(dep))] = dep dep_dict[normalise(pkg_name(dep))] = dep
# Add cuda extras (takes precedence for HPC) # Add cuda extras (takes precedence for HPC)
optional_deps = data.get("project", {}).get("optional-dependencies", {}) optional_deps = data.get("project", {}).get("optional-dependencies", {})
for group in ["cuda"]: for group in ["cuda"]:
for dep in optional_deps.get(group, []): for dep in optional_deps.get(group, []):
dep_dict[normalise(pkg_name(dep))] = dep dep_dict[normalise(pkg_name(dep))] = dep
deps = list(dep_dict.values()) deps = list(dep_dict.values())
final_deps: list[str] = [] final_deps: list[str] = []
print(f"Checking dependencies against HPC module list...", file=sys.stderr) print("Checking dependencies against HPC module list...", file=sys.stderr)
for dep in deps: for dep in deps:
name = normalise(pkg_name(dep)) name = normalise(pkg_name(dep))
# Smart check: if the package name is a substring of any loaded module name # Smart check: if the package name is a substring of any loaded module name

View file

@ -55,10 +55,11 @@ def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
with open(path, "r") as f: with open(path, "r") as f:
if path.endswith(".yaml") or path.endswith(".yml"): if path.endswith(".yaml") or path.endswith(".yml"):
import yaml import yaml
config_dict = yaml.safe_load(f) config_dict = yaml.safe_load(f)
else: else:
config_dict = json.load(f) config_dict = json.load(f)
morphology = MorphologyConfig(**config_dict.get("morphology", {})) morphology = MorphologyConfig(**config_dict.get("morphology", {}))
arena = ArenaConfig(**config_dict.get("arena", {})) arena = ArenaConfig(**config_dict.get("arena", {}))
env = EnvConfig(**config_dict.get("env", {})) env = EnvConfig(**config_dict.get("env", {}))

View file

@ -52,10 +52,12 @@ def train(args: PPOArgs):
args.batch_size = args.num_envs * args.num_steps args.batch_size = args.num_envs * args.num_steps
args.minibatch_size = args.batch_size // args.num_minibatches args.minibatch_size = args.batch_size // args.num_minibatches
args.num_iterations = args.total_timesteps // args.batch_size args.num_iterations = args.total_timesteps // args.batch_size
# Try to get git short hash # Try to get git short hash
try: try:
git_hash = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip() git_hash = (
subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
)
except Exception: except Exception:
git_hash = "none" git_hash = "none"
@ -64,8 +66,9 @@ def train(args: PPOArgs):
if args.run_dir is None: if args.run_dir is None:
args.run_dir = f"runs/{run_name}" args.run_dir = f"runs/{run_name}"
import os import os
os.makedirs(args.run_dir, exist_ok=True) os.makedirs(args.run_dir, exist_ok=True)
if args.track: if args.track:
@ -348,14 +351,14 @@ def train(args: PPOArgs):
remaining_steps = args.total_timesteps - global_step remaining_steps = args.total_timesteps - global_step
eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
eta_str = str(datetime.timedelta(seconds=eta_seconds)) eta_str = str(datetime.timedelta(seconds=eta_seconds))
print( print(
f"Iteration {iteration}/{args.num_iterations} | " f"Iteration {iteration}/{args.num_iterations} | "
f"Step {global_step}/{args.total_timesteps} | " f"Step {global_step}/{args.total_timesteps} | "
f"SPS {sps} | " f"SPS {sps} | "
f"Return {avg_episodic_return:.4f} | " f"Return {avg_episodic_return:.4f} | "
f"ETA {eta_str}", f"ETA {eta_str}",
flush=True flush=True,
) )
if args.save_model: if args.save_model:
@ -388,7 +391,7 @@ def train(args: PPOArgs):
def main() -> None: def main() -> None:
temp_args = tyro.cli(PPOArgs) temp_args = tyro.cli(PPOArgs)
if temp_args.env_config_path is not None: if temp_args.env_config_path is not None:
with open(temp_args.env_config_path, "r") as f: with open(temp_args.env_config_path, "r") as f:
config = yaml.safe_load(f) config = yaml.safe_load(f)
@ -397,7 +400,7 @@ def main() -> None:
for key, value in config.items(): for key, value in config.items():
if hasattr(temp_args, key): if hasattr(temp_args, key):
setattr(temp_args, key, value) setattr(temp_args, key, value)
# Re-parse CLI to ensure they OVERRIDE the yaml # Re-parse CLI to ensure they OVERRIDE the yaml
args = tyro.cli(PPOArgs, default=temp_args) args = tyro.cli(PPOArgs, default=temp_args)
else: else: