diff --git a/src/quantem/widget/__init__.py b/src/quantem/widget/__init__.py index 6e151a9b..0fea6e5d 100644 --- a/src/quantem/widget/__init__.py +++ b/src/quantem/widget/__init__.py @@ -64,6 +64,7 @@ "SupportsHtmlExport": ("quantem.widget.export", "SupportsHtmlExport"), "supports_html_export": ("quantem.widget.export", "supports_html_export"), "device_info": ("quantem.widget.info", "device_info"), + "profile": ("quantem.widget.info", "profile"), "WidgetProfile": ("quantem.widget._timing", "WidgetProfile"), "format_timing_table": ("quantem.widget._timing", "format_timing_table"), "format_widget_render_timing": ( @@ -101,63 +102,6 @@ def __dir__() -> list[str]: return sorted(set(globals()) | set(_LAZY_EXPORTS)) -def profile() -> None: - """Print the installed QuantEM stack and active compute environment. - - Use this single report in notebooks and bug reports instead of printing - individual package versions. It records the widget, GPU, and core QuantEM - versions together with the active Torch device and Python version. - - Examples - -------- - >>> import quantem.widget as qw - >>> qw.profile() - """ - import platform - - print(f"quantem.widget {__version__}") - try: - print(f"quantem.gpu {version('quantem.gpu')}") - except PackageNotFoundError: - print("quantem.gpu (not installed)") - try: - import quantem - - print(f"quantem {getattr(quantem, '__version__', '?')}") - print(f" loaded from {quantem.__file__}") - except ImportError: - print("quantem (not importable)") - try: - import torch - - if torch.cuda.is_available(): - dev = f"cuda ({torch.cuda.get_device_name(0)})" - elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): - dev = "mps (Apple)" - else: - dev = "cpu" - print(f"torch {torch.__version__} device={dev}") - if torch.cuda.is_available(): - # Show EVERY visible GPU + how many are visible, so the reader knows up front - # whether the next merge / recon fits and on WHICH card - no surprise mid-run. - # torch live-vs-reserved is the leak signal: if "live" climbs across repeated - # calls, refs are still pinned (del them, then free_gpu() returns the pool). - import os - n = torch.cuda.device_count() - print(f"GPUs {n} visible (CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', 'all')})") - for i in range(n): - free, total = torch.cuda.mem_get_info(i) - print(f" GPU{i} {(total - free) / 1e9:5.1f} used / {total / 1e9:.0f} GB ({free / 1e9:.0f} free) <- run free_gpu() if low") - print(f" torch pool {torch.cuda.memory_allocated() / 1e9:.1f} live / {torch.cuda.memory_reserved() / 1e9:.1f} reserved GB") - elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): - cur = torch.mps.current_allocated_memory() / 1e9 if hasattr(torch.mps, "current_allocated_memory") else 0.0 - drv = torch.mps.driver_allocated_memory() / 1e9 if hasattr(torch.mps, "driver_allocated_memory") else 0.0 - print(f"VRAM (MPS) {cur:.1f} live / {drv:.1f} driver GB") - except ImportError: - print("torch (not importable)") - print(f"python {platform.python_version()}") - - def free_gpu(verbose: bool = True) -> float: """Release cached GPU memory back to the driver, on CUDA (torch + cupy pools) or Apple MPS. Call AFTER ``del``-ing your big objects (the merged 4D stack, the widget): this diff --git a/src/quantem/widget/info.py b/src/quantem/widget/info.py index 0aed7a10..34e28f54 100644 --- a/src/quantem/widget/info.py +++ b/src/quantem/widget/info.py @@ -1,50 +1,208 @@ -"""Human-readable device / version info for notebooks - which GPU, which machine. +"""Human-readable environment information for notebooks.""" -``device_info()`` prints (and returns) a tidy one-block summary so a shared -notebook records what it ran on: widget version, date, compute backend (Apple -Metal / CUDA / CPU), the GPU or Mac chip, and memory. Useful at the top of any -demo so results are reproducible. -""" -from __future__ import annotations - -import datetime +import json +import os +import platform import subprocess +from datetime import datetime +from importlib.metadata import PackageNotFoundError, distribution, version +from importlib.util import find_spec +from pathlib import Path +from urllib.parse import unquote, urlparse +from urllib.request import Request, urlopen + +from packaging.version import Version + + +def profile(*, check_updates: bool = False) -> None: + """Print the installed QuantEM stack and active compute environment. + Use this single report in notebooks and bug reports instead of printing + individual package versions. The default report is local and does not + contact package indexes or Git remotes. -def _mac_chip_mem(): - def _sysctl(key): + Parameters + ---------- + check_updates : bool, default False + Compare installed widget and GPU metadata with TestPyPI. This opt-in + check needs network access. + + Examples + -------- + >>> import quantem.widget as qw + >>> qw.profile() + """ + import quantem.widget as qw + + def editable_source(distribution_name: str) -> Path | None: try: - return subprocess.run(["sysctl", "-n", key], capture_output=True, - text=True, timeout=3).stdout.strip() - except Exception: - return "" - chip = _sysctl("machdep.cpu.brand_string") or "Apple Silicon" - mem = _sysctl("hw.memsize") - gb = f"{int(mem) // (1024 ** 3)} GB" if mem.isdigit() else "?" - return chip, gb - - -def device_info(verbose: bool = True) -> dict: - """Return (and by default print) version + backend + hardware for this machine.""" - import quantem.widget + raw = distribution(distribution_name).read_text("direct_url.json") + except (PackageNotFoundError, OSError): + return None + if not raw: + return None + + try: + direct_url = json.loads(raw) + except ValueError: + return None + if not isinstance(direct_url, dict): + return None + directory = direct_url.get("dir_info") + source_url = direct_url.get("url") + if not isinstance(directory, dict) or not directory.get("editable"): + return None + if not isinstance(source_url, str): + return None + + parsed = urlparse(source_url) + if parsed.scheme != "file": + return None + return Path(unquote(parsed.path)).resolve() + + def print_update(distribution_name: str, installed: str) -> None: + package = distribution_name.replace(".", "-") + request = Request( + f"https://test.pypi.org/pypi/{package}/json", + headers={"User-Agent": "quantem.widget profile()"}, + ) + try: + with urlopen(request, timeout=4) as response: + latest = json.load(response)["info"]["version"] + installed_version = Version(installed) + latest_version = Version(latest) + except (KeyError, OSError, TypeError, ValueError): + print(" release update check unavailable") + return + + print(f" TestPyPI latest {latest}") + if installed_version < latest_version: + print(f" WARNING installed metadata {installed} trails {latest}") + elif installed_version > latest_version: + print(" release newer than TestPyPI") + else: + print(" release current") + + def print_distribution_status( + distribution_name: str, + installed: str, + ) -> None: + source = editable_source(distribution_name) + try: + spec = find_spec(distribution_name) + except (ImportError, ValueError): + spec = None + loaded = ( + Path(spec.origin).resolve() + if spec is not None and spec.origin is not None + else None + ) + + if source is None: + print(" install published package") + elif loaded is not None and not loaded.is_relative_to(source): + print(" install source override (differs from installed metadata)") + else: + print(" install editable checkout") + + if check_updates: + print_update(distribution_name, installed) + + print(f"quantem.widget {qw.__version__}") + print_distribution_status("quantem.widget", qw.__version__) + try: + gpu_version = version("quantem.gpu") + print(f"quantem.gpu {gpu_version}") + print_distribution_status("quantem.gpu", gpu_version) + except PackageNotFoundError: + print("quantem.gpu (not installed)") + try: + import quantem + + print(f"quantem {getattr(quantem, '__version__', '?')}") + except ImportError: + print("quantem (not importable)") + try: + import torch + + if torch.cuda.is_available(): + device = f"cuda ({torch.cuda.get_device_name(0)})" + elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): + device = "mps (Apple)" + else: + device = "cpu" + print(f"torch {torch.__version__} device={device}") + if torch.cuda.is_available(): + count = torch.cuda.device_count() + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "all") + print(f"GPUs {count} visible (CUDA_VISIBLE_DEVICES={visible})") + for index in range(count): + free, total = torch.cuda.mem_get_info(index) + print( + f" GPU{index} {(total - free) / 1e9:5.1f} used / " + f"{total / 1e9:.0f} GB ({free / 1e9:.0f} free)" + ) + print( + f" torch pool {torch.cuda.memory_allocated() / 1e9:.1f} live / " + f"{torch.cuda.memory_reserved() / 1e9:.1f} reserved GB" + ) + elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): + current = ( + torch.mps.current_allocated_memory() / 1e9 + if hasattr(torch.mps, "current_allocated_memory") + else 0.0 + ) + driver = ( + torch.mps.driver_allocated_memory() / 1e9 + if hasattr(torch.mps, "driver_allocated_memory") + else 0.0 + ) + print(f"VRAM (MPS) {current:.1f} live / {driver:.1f} driver GB") + except ImportError: + print("torch (not importable)") + print(f"python {platform.python_version()}") + + +def device_info(verbose: bool = True) -> dict[str, str]: + """Return and optionally print the active device information.""" from quantem.gpu.device import detect + import quantem.widget as qw + backend = detect() - info = { - "widget_version": quantem.widget.__version__, - "date": str(datetime.date.today()), + report = { + "widget_version": qw.__version__, + "date": str(datetime.now().astimezone().date()), "backend": backend, + "device": "CPU", } if backend == "mps": - chip, mem = _mac_chip_mem() - info["device"] = f"Apple Metal (MPS) - {chip}, {mem} unified memory" + + def sysctl(key: str) -> str: + try: + result = subprocess.run( + ["sysctl", "-n", key], + check=True, + capture_output=True, + text=True, + timeout=3, + ) + except (OSError, subprocess.SubprocessError): + return "" + return result.stdout.strip() + + chip = sysctl("machdep.cpu.brand_string") or "Apple Silicon" + memory = sysctl("hw.memsize") + memory_gb = f"{int(memory) // (1024**3)} GB" if memory.isdigit() else "?" + report["device"] = f"Apple Metal (MPS) - {chip}, {memory_gb} unified memory" elif backend == "cuda": try: import torch - info["device"] = f"CUDA - {torch.cuda.get_device_name(0)}" - except Exception: - info["device"] = "CUDA" + + report["device"] = f"CUDA - {torch.cuda.get_device_name(0)}" + except (AssertionError, ImportError, RuntimeError): + report["device"] = "CUDA" if verbose: - print(f"quantem.widget {info['widget_version']} | {info['date']}") - print(f"compute: {info['device']}") - return info + print(f"quantem.widget {report['widget_version']} | {report['date']}") + print(f"compute: {report['device']}") + return report diff --git a/tests/test_profile.py b/tests/test_profile.py index 62037849..6eaa19e1 100644 --- a/tests/test_profile.py +++ b/tests/test_profile.py @@ -1,6 +1,7 @@ -from __future__ import annotations - +import io +import json from pathlib import Path +from types import SimpleNamespace def test_profile_reports_the_installed_quantem_stack(capsys) -> None: @@ -15,6 +16,74 @@ def test_profile_reports_the_installed_quantem_stack(capsys) -> None: assert "quantem" in output assert "torch" in output assert "python" in output + assert "install" in output + + +def test_profile_checks_testpypi_only_when_requested(monkeypatch, capsys) -> None: + """A notebook opts into release checks without changing the normal report.""" + import quantem.widget as qw + from quantem.widget import info + + calls = [] + + def response(request, *, timeout): + calls.append(request.full_url) + payload = json.dumps({"info": {"version": "99.0rc1"}}).encode() + return io.BytesIO(payload) + + monkeypatch.setattr(info, "urlopen", response) + + qw.profile() + assert calls == [] + + qw.profile(check_updates=True) + + output = capsys.readouterr().out + assert "TestPyPI latest 99.0rc1" in output + assert "WARNING" in output + assert calls == [ + "https://test.pypi.org/pypi/quantem-widget/json", + "https://test.pypi.org/pypi/quantem-gpu/json", + ] + assert output.count("TestPyPI latest 99.0rc1") == 2 + + +def test_profile_does_not_print_editable_source_paths(monkeypatch, capsys) -> None: + """A shared profile report labels an editable checkout without leaking paths.""" + import quantem.widget as qw + from quantem.widget import info + + source = Path.cwd() / "private-source" / "quantem.widget" + direct_url = json.dumps( + {"url": source.as_uri(), "dir_info": {"editable": True}} + ) + + def installed_distribution(name): + raw = direct_url if name == "quantem.widget" else None + return SimpleNamespace(read_text=lambda filename: raw) + + monkeypatch.setattr(info, "distribution", installed_distribution) + monkeypatch.setattr( + info, + "find_spec", + lambda name: SimpleNamespace( + origin=source / "src/quantem/widget/__init__.py" + ), + ) + qw.profile() + + output = capsys.readouterr().out + assert "editable checkout" in output + assert str(source) not in output + + loaded = Path.cwd() / "another-checkout/src/quantem/widget/__init__.py" + monkeypatch.setattr(info, "find_spec", lambda name: SimpleNamespace(origin=loaded)) + qw.profile() + + output = capsys.readouterr().out + assert "source override" in output + assert str(source) not in output + assert str(loaded) not in output def test_documented_environment_checks_use_widget_profile() -> None: