Skills › Scientific Agent Skills › timesfm-forecasting › check_system.py
check_system.py
501 lines · skills/timesfm-forecasting/scripts/check_system.py at 5f2bc03
1#!/usr/bin/env python3
2"""TimesFM System Requirements Preflight Checker.
4MANDATORY: Run this script before loading TimesFM for the first time.
5It checks RAM, GPU/VRAM, disk space, Python version, and package
6installation. Thresholds are advisory heuristics, not an OOM guarantee.
8Usage:
9 python check_system.py
10 python check_system.py --model v2.5 # default
11 python check_system.py --model v2.0 # archived 500M model
12 python check_system.py --model v1.0 # archived 200M model
13 python check_system.py --json # machine-readable output
14"""
16from __future__ import annotations
18import argparse
19import importlib
20import json
21import os
22import platform
23import shutil
24import re
25from importlib.metadata import version as package_version, PackageNotFoundError
26import sys
27from dataclasses import dataclass, field
28from pathlib import Path
29from typing import Any
32# ---------------------------------------------------------------------------
33# Model requirement profiles
34# ---------------------------------------------------------------------------
36MODEL_PROFILES: dict[str, dict[str, Any]] = {
37 "v3.0": {
38 "name": "TimesFM 3.0 (330M)", "params": "330M",
39 "min_ram_gb": 4.0, "recommended_ram_gb": 8.0,
40 "min_vram_gb": 4.0, "recommended_vram_gb": 8.0,
41 "disk_gb": 3.0, "hf_repo": "google/timesfm-3.0-pytorch",
42 },
43 "v2.5": {
44 "name": "TimesFM 2.5 (200M)",
45 "params": "200M",
46 "min_ram_gb": 2.0,
47 "recommended_ram_gb": 4.0,
48 "min_vram_gb": 2.0,
49 "recommended_vram_gb": 4.0,
50 "disk_gb": 2.0, # model weights + overhead
51 "hf_repo": "google/timesfm-2.5-200m-pytorch",
52 },
53 "v2.0": {
54 "name": "TimesFM 2.0 (500M)",
55 "params": "500M",
56 "min_ram_gb": 8.0,
57 "recommended_ram_gb": 16.0,
58 "min_vram_gb": 4.0,
59 "recommended_vram_gb": 8.0,
60 "disk_gb": 4.0,
61 "hf_repo": "google/timesfm-2.0-500m-pytorch",
62 },
63 "v1.0": {
64 "name": "TimesFM 1.0 (200M)",
65 "params": "200M",
66 "min_ram_gb": 4.0,
67 "recommended_ram_gb": 8.0,
68 "min_vram_gb": 2.0,
69 "recommended_vram_gb": 4.0,
70 "disk_gb": 2.0,
71 "hf_repo": "google/timesfm-1.0-200m-pytorch",
72 },
73}
76# ---------------------------------------------------------------------------
77# Result dataclass
78# ---------------------------------------------------------------------------
81@dataclass
82class CheckResult:
83 name: str
84 status: str # "pass", "warn", "fail"
85 detail: str
86 value: str = ""
88 @property
89 def icon(self) -> str:
90 return {"pass": "[OK]", "warn": "[WARN]", "fail": "[FAIL]"}.get(self.status, "[?]")
92 def __str__(self) -> str:
93 return f"[{self.name:<10}] {self.value:<40} {self.icon} {self.status.upper()}"
96@dataclass
97class SystemReport:
98 model: str
99 checks: list[CheckResult] = field(default_factory=list)
100 verdict: str = ""
101 verdict_detail: str = ""
102 recommended_batch_size: int = 1
103 mode: str = "cpu" # "cpu", "gpu", "mps"
105 @property
106 def passed(self) -> bool:
107 return all(c.status != "fail" for c in self.checks)
109 def to_dict(self) -> dict[str, Any]:
110 return {
111 "model": self.model,
112 "passed": self.passed,
113 "mode": self.mode,
114 "recommended_batch_size": self.recommended_batch_size,
115 "verdict": self.verdict,
116 "verdict_detail": self.verdict_detail,
117 "checks": [
118 {
119 "name": c.name,
120 "status": c.status,
121 "detail": c.detail,
122 "value": c.value,
123 }
124 for c in self.checks
125 ],
126 }
129# ---------------------------------------------------------------------------
130# Individual checks
131# ---------------------------------------------------------------------------
134def _get_total_ram_gb() -> float:
135 """Return total physical RAM in GB, cross-platform."""
136 try:
137 if sys.platform == "linux":
138 with open("/proc/meminfo") as f:
139 for line in f:
140 if line.startswith("MemTotal"):
141 return int(line.split()[1]) / (1024 * 1024)
142 elif sys.platform == "darwin":
143 import subprocess
145 result = subprocess.run(
146 ["sysctl", "-n", "hw.memsize"],
147 capture_output=True,
148 text=True,
149 check=True,
150 )
151 return int(result.stdout.strip()) / (1024**3)
152 elif sys.platform == "win32":
153 import ctypes
155 kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined]
157 class MEMORYSTATUSEX(ctypes.Structure):
158 _fields_ = [
159 ("dwLength", ctypes.c_ulong),
160 ("dwMemoryLoad", ctypes.c_ulong),
161 ("ullTotalPhys", ctypes.c_ulonglong),
162 ("ullAvailPhys", ctypes.c_ulonglong),
163 ("ullTotalPageFile", ctypes.c_ulonglong),
164 ("ullAvailPageFile", ctypes.c_ulonglong),
165 ("ullTotalVirtual", ctypes.c_ulonglong),
166 ("ullAvailVirtual", ctypes.c_ulonglong),
167 ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
168 ]
170 stat = MEMORYSTATUSEX()
171 stat.dwLength = ctypes.sizeof(stat)
172 kernel32.GlobalMemoryStatusEx(ctypes.byref(stat))
173 return stat.ullTotalPhys / (1024**3)
174 except Exception:
175 pass
177 return 0.0 # Unknown; never infer physical memory from pointer width.
180def _get_available_ram_gb() -> float:
181 """Return available RAM in GB."""
182 try:
183 if sys.platform == "linux":
184 with open("/proc/meminfo") as f:
185 for line in f:
186 if line.startswith("MemAvailable"):
187 return int(line.split()[1]) / (1024 * 1024)
188 elif sys.platform == "darwin":
189 import subprocess
191 # Use vm_stat for available memory on macOS
192 result = subprocess.run(
193 ["vm_stat"], capture_output=True, text=True, check=True
194 )
195 free = 0
196 match = re.search(r"page size of (\d+) bytes", result.stdout)
197 if not match:
198 return 0.0
199 page_size = int(match.group(1))
200 for line in result.stdout.split("\n"):
201 if "Pages free" in line or "Pages inactive" in line:
202 val = line.split(":")[1].strip().rstrip(".")
203 free += int(val) * page_size
204 return free / (1024**3)
205 elif sys.platform == "win32":
206 import ctypes
208 kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined]
210 class MEMORYSTATUSEX(ctypes.Structure):
211 _fields_ = [
212 ("dwLength", ctypes.c_ulong),
213 ("dwMemoryLoad", ctypes.c_ulong),
214 ("ullTotalPhys", ctypes.c_ulonglong),
215 ("ullAvailPhys", ctypes.c_ulonglong),
216 ("ullTotalPageFile", ctypes.c_ulonglong),
217 ("ullAvailPageFile", ctypes.c_ulonglong),
218 ("ullTotalVirtual", ctypes.c_ulonglong),
219 ("ullAvailVirtual", ctypes.c_ulonglong),
220 ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
221 ]
223 stat = MEMORYSTATUSEX()
224 stat.dwLength = ctypes.sizeof(stat)
225 kernel32.GlobalMemoryStatusEx(ctypes.byref(stat))
226 return stat.ullAvailPhys / (1024**3)
227 except Exception:
228 pass
229 return 0.0
232def check_ram(profile: dict[str, Any]) -> CheckResult:
233 """Check if system has enough RAM."""
234 total = _get_total_ram_gb()
235 available = _get_available_ram_gb()
236 min_ram = profile["min_ram_gb"]
237 rec_ram = profile["recommended_ram_gb"]
239 value = f"Total: {total:.1f} GB | Available: {available:.1f} GB"
241 if total <= 0 or available <= 0:
242 return CheckResult("RAM", "warn", "RAM measurement unavailable; check manually.", value)
243 if min(total, available) < min_ram:
244 return CheckResult(
245 name="RAM",
246 status="fail",
247 detail=(
248 f"System has {total:.1f} GB total / {available:.1f} GB available RAM but {profile['name']} uses a preflight threshold of "
249 f"at least {min_ram:.0f} GB. The model will likely fail to load "
250 f"or cause the system to swap heavily and become unresponsive."
251 ),
252 value=value,
253 )
254 elif min(total, available) < rec_ram:
255 return CheckResult(
256 name="RAM",
257 status="warn",
258 detail=(
259 f"System has {total:.1f} GB total / {available:.1f} GB available RAM. Preflight recommends "
260 f"{rec_ram:.0f} GB. It may work with small batch sizes but could "
261 f"be tight. Use per_core_batch_size=4 or lower."
262 ),
263 value=value,
264 )
265 else:
266 return CheckResult(
267 name="RAM",
268 status="pass",
269 detail=f"Available RAM {available:.1f} GB meets the {rec_ram:.0f} GB preflight recommendation.",
270 value=value,
271 )
274def check_gpu() -> CheckResult:
275 """Check GPU availability and VRAM."""
276 # Try CUDA first
277 try:
278 import torch
280 if torch.cuda.is_available():
281 name = torch.cuda.get_device_name(0)
282 vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
283 return CheckResult(
284 name="GPU",
285 status="pass",
286 detail=f"{name} with {vram:.1f} GB VRAM detected.",
287 value=f"{name} | VRAM: {vram:.1f} GB",
288 )
289 elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
290 return CheckResult(
291 name="GPU",
292 status="warn",
293 detail="MPS is available, but the 2.5 loader selects CPU when CUDA is absent.",
294 value="Apple Silicon MPS",
295 )
296 else:
297 return CheckResult(
298 name="GPU",
299 status="warn",
300 detail=(
301 "No GPU detected. TimesFM will run on CPU (slower but functional). "
302 "Install CUDA-enabled PyTorch for GPU acceleration."
303 ),
304 value="None (CPU only)",
305 )
306 except ImportError:
307 return CheckResult(
308 name="GPU",
309 status="warn",
310 detail="PyTorch not installed — cannot check GPU. Install torch first.",
311 value="Unknown (torch not installed)",
312 )
315def check_disk(profile: dict[str, Any]) -> CheckResult:
316 """Check available disk space for model download."""
317 # Check HuggingFace cache dir or home dir
318 hf_home = os.environ.get("HF_HOME", str(Path(os.environ.get("XDG_CACHE_HOME", "~/.cache")).expanduser() / "huggingface"))
319 hf_cache = os.environ.get("HF_HUB_CACHE", str(Path(hf_home).expanduser() / "hub"))
320 check_dir = Path(hf_cache).expanduser().resolve()
321 while not check_dir.exists():
322 check_dir = check_dir.parent
324 usage = shutil.disk_usage(str(check_dir))
325 free_gb = usage.free / (1024**3)
326 required = profile["disk_gb"]
328 value = f"Free: {free_gb:.1f} GB (in {check_dir})"
330 if free_gb < required:
331 return CheckResult(
332 name="Disk",
333 status="fail",
334 detail=(
335 f"Only {free_gb:.1f} GB free in {check_dir}. "
336 f"Need at least {required:.0f} GB for model weights. "
337 f"Free up space or set HF_HOME to a larger volume."
338 ),
339 value=value,
340 )
341 else:
342 return CheckResult(
343 name="Disk",
344 status="pass",
345 detail=f"{free_gb:.1f} GB available, exceeds {required:.0f} GB requirement.",
346 value=value,
347 )
350def check_python() -> CheckResult:
351 """Check Python version >= 3.10."""
352 version = sys.version.split()[0]
353 major, minor = sys.version_info[:2]
355 if (major, minor) < (3, 10):
356 return CheckResult(
357 name="Python",
358 status="fail",
359 detail=f"Python {version} detected. TimesFM requires Python >= 3.10.",
360 value=version,
361 )
362 else:
363 return CheckResult(
364 name="Python",
365 status="pass",
366 detail=f"Python {version} meets >= 3.10 requirement.",
367 value=version,
368 )
371def check_package(pkg_name: str, import_name: str | None = None) -> CheckResult:
372 """Check if a Python package is installed."""
373 import_name = import_name or pkg_name
374 try:
375 mod = importlib.import_module(import_name)
376 try:
377 version = package_version(pkg_name)
378 except PackageNotFoundError:
379 version = getattr(mod, "__version__", "unknown")
380 return CheckResult(
381 name=pkg_name,
382 status="pass",
383 detail=f"{pkg_name} {version} is installed.",
384 value=f"Installed ({version})",
385 )
386 except ImportError:
387 return CheckResult(
388 name=pkg_name,
389 status="warn",
390 detail=f"{pkg_name} is not installed. Run: uv pip install {pkg_name}",
391 value="Not installed",
392 )
395# ---------------------------------------------------------------------------
396# Batch size recommendation
397# ---------------------------------------------------------------------------
400def recommend_batch_size(report: SystemReport) -> int:
401 """Recommend per_core_batch_size based on available resources."""
402 # Context, horizon, compilation and workload dominate peak allocation.
403 # Start one series at a time; measured tuning belongs to the caller.
404 return 1
407# ---------------------------------------------------------------------------
408# Main
409# ---------------------------------------------------------------------------
412def run_checks(model_version: str = "v2.5") -> SystemReport:
413 """Run all system checks and return a report."""
414 profile = MODEL_PROFILES[model_version]
415 report = SystemReport(model=profile["name"])
417 # Run checks
418 report.checks.append(check_ram(profile))
419 report.checks.append(check_gpu())
420 report.checks.append(check_disk(profile))
421 report.checks.append(check_python())
422 report.checks.append(check_package("timesfm"))
423 report.checks.append(check_package("torch"))
425 # Determine mode
426 gpu_check = next((c for c in report.checks if c.name == "GPU"), None)
427 if gpu_check and gpu_check.status == "pass":
428 if "MPS" in gpu_check.value:
429 report.mode = "mps"
430 else:
431 report.mode = "gpu"
432 else:
433 report.mode = "cpu"
435 # Batch size
436 report.recommended_batch_size = recommend_batch_size(report)
438 # Verdict
439 if report.passed:
440 report.verdict = (
441 f"[OK] Preflight thresholds passed for {profile['name']} ({report.mode.upper()} mode)"
442 )
443 report.verdict_detail = (
444 f"Start with per_core_batch_size={report.recommended_batch_size}; "
445 "this does not guarantee available peak memory or package compatibility."
446 )
447 else:
448 failed = [c for c in report.checks if c.status == "fail"]
449 report.verdict = f"[FAIL] System does NOT meet requirements for {profile['name']}"
450 report.verdict_detail = "; ".join(c.detail for c in failed)
452 return report
455def print_report(report: SystemReport) -> None:
456 """Print a human-readable report to stdout."""
457 print(f"\n{'=' * 50}")
458 print(f" TimesFM System Requirements Check")
459 print(f" Model: {report.model}")
460 print(f"{'=' * 50}\n")
462 for check in report.checks:
463 print(f" {check}")
464 print()
466 print(f" VERDICT: {report.verdict}")
467 if report.verdict_detail:
468 print(f" {report.verdict_detail}")
469 print()
472def main() -> None:
473 parser = argparse.ArgumentParser(
474 description="Check system requirements for TimesFM."
475 )
476 parser.add_argument(
477 "--model",
478 choices=list(MODEL_PROFILES.keys()),
479 default="v2.5",
480 help="Model version to check requirements for (default: v2.5)",
481 )
482 parser.add_argument(
483 "--json",
484 action="store_true",
485 help="Output results as JSON (machine-readable)",
486 )
487 args = parser.parse_args()
489 report = run_checks(args.model)
491 if args.json:
492 print(json.dumps(report.to_dict(), indent=2))
493 else:
494 print_report(report)
496 # Exit with non-zero if any check failed
497 sys.exit(0 if report.passed else 1)
500if __name__ == "__main__":
501 main()