File size: 3,829 Bytes
dddcc0f 0c8f912 dddcc0f 685b361 dddcc0f 0c8f912 dddcc0f 685b361 dddcc0f 685b361 dddcc0f 0c8f912 685b361 0c8f912 685b361 0c8f912 dddcc0f 685b361 dddcc0f 0c8f912 dddcc0f 0c8f912 dddcc0f | 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 | #!/usr/bin/env python3
"""Probe local HPC resources and suggest safe EDA concurrency settings.
This script dynamically calculates memory per worker from the EDA config file
(configs/eda_optimized.yaml) using max_memory_gib / max_workers.
"""
from __future__ import annotations
import argparse
import json
import os
import platform
import shutil
import sys
from pathlib import Path
import yaml
def _mem_available_gib() -> float:
meminfo = Path("/proc/meminfo")
if not meminfo.exists():
return 0.0
for line in meminfo.read_text().splitlines():
if line.startswith("MemAvailable:"):
kb = int(line.split()[1])
return kb / (1024 * 1024)
return 0.0
def _mem_total_gib() -> float:
meminfo = Path("/proc/meminfo")
if not meminfo.exists():
return 0.0
for line in meminfo.read_text().splitlines():
if line.startswith("MemTotal:"):
kb = int(line.split()[1])
return kb / (1024 * 1024)
return 0.0
def _recommend_workers(cpu_count: int, mem_available_gib: float, mem_per_worker_gib: float) -> int:
# Fast profile for HPC: use more cores while still leaving headroom.
by_cpu = max(1, int(cpu_count * 0.75))
by_mem = max(1, int(mem_available_gib // max(1.0, mem_per_worker_gib)))
return max(1, min(by_cpu, by_mem))
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--config",
type=Path,
default=Path(__file__).parent.parent / "configs" / "eda_optimized.yaml",
help="Path to YAML configuration file.",
)
parser.add_argument(
"--workdir",
type=Path,
help="Path to check disk usage for (if not specified, uses first input_dir from config).",
)
args = parser.parse_args()
# Read config to calculate dynamic mem_per_worker_gib
config_path = args.config
if not config_path.exists():
print(f"Error: Config file not found: {config_path}", file=sys.stderr)
sys.exit(1)
with open(config_path) as f:
config = yaml.safe_load(f)
max_memory_gib = config['resources']['max_memory_gib']
max_workers = config['resources']['max_workers']
mem_per_worker_gib = max_memory_gib / max_workers
# Use workdir from args, or fall back to first input_dir from config
workdir = args.workdir
if workdir is None:
input_dirs = config.get('paths', {}).get('input_dirs', [])
if input_dirs:
workdir = Path(input_dirs[0]).parent.parent # Go up to project root
else:
workdir = Path.cwd()
cpu_count = os.cpu_count() or 1
mem_total_gib = _mem_total_gib()
mem_available_gib = _mem_available_gib()
disk_total, disk_used, disk_free = shutil.disk_usage(workdir)
recommended_workers = _recommend_workers(
cpu_count=cpu_count,
mem_available_gib=mem_available_gib,
mem_per_worker_gib=mem_per_worker_gib,
)
recommended_shards = max(1, min(8, cpu_count // max(1, recommended_workers)))
report = {
"hostname": platform.node(),
"platform": platform.platform(),
"cpu_count": cpu_count,
"memory_total_gib": round(mem_total_gib, 2),
"memory_available_gib": round(mem_available_gib, 2),
"disk_total_gib": round(disk_total / (1024**3), 2),
"disk_used_gib": round(disk_used / (1024**3), 2),
"disk_free_gib": round(disk_free / (1024**3), 2),
"assumptions": {"mem_per_worker_gib": mem_per_worker_gib},
"recommendation": {
"workers_per_node": recommended_workers,
"num_shards_suggestion": recommended_shards,
"chunk_size_suggestion": 4096,
},
}
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|