| """Check the CUDA/APE runtime needed by Orienter Stage 2.""" |
|
|
| import argparse |
| import importlib |
| import importlib.metadata |
| import json |
| import platform |
| import sys |
|
|
|
|
| HISTORICAL_REFERENCE = { |
| "python": "3.9.18", |
| "torch": "2.2.0+cu121", |
| "torchvision": "0.17.0+cu121", |
| "detectron2": "0.6", |
| "detrex": "0.3.0", |
| } |
|
|
|
|
| def _version(module_name, distribution_name=None): |
| module = importlib.import_module(module_name) |
| value = getattr(module, "__version__", None) |
| if value is not None: |
| return str(value) |
| return importlib.metadata.version(distribution_name or module_name) |
|
|
|
|
| def check_environment(require_cuda=True, require_ape_extension=True): |
| report = { |
| "python": platform.python_version(), |
| "historical_reference": dict(HISTORICAL_REFERENCE), |
| "versions": {}, |
| "cuda_available": False, |
| "ape_extension": False, |
| "errors": [], |
| } |
|
|
| for module_name, distribution_name in ( |
| ("torch", "torch"), |
| ("torchvision", "torchvision"), |
| ("detectron2", "detectron2"), |
| ("detrex", "detrex"), |
| ): |
| try: |
| report["versions"][module_name] = _version(module_name, distribution_name) |
| except Exception as exc: |
| report["errors"].append(f"cannot import {module_name}: {type(exc).__name__}: {exc}") |
|
|
| try: |
| torch = importlib.import_module("torch") |
| report["torch_cuda_runtime"] = torch.version.cuda |
| report["cuda_available"] = bool(torch.cuda.is_available()) |
| if require_cuda and not report["cuda_available"]: |
| report["errors"].append("torch.cuda.is_available() is false") |
| except Exception: |
| pass |
|
|
| if require_ape_extension: |
| try: |
| importlib.import_module("ape._C") |
| report["ape_extension"] = True |
| except Exception as exc: |
| report["errors"].append( |
| f"cannot import compiled ape._C extension: {type(exc).__name__}: {exc}" |
| ) |
|
|
| return report |
|
|
|
|
| def build_parser(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--allow-no-cuda", action="store_true") |
| parser.add_argument("--skip-ape-extension", action="store_true") |
| return parser |
|
|
|
|
| def main(argv=None): |
| args = build_parser().parse_args(argv) |
| report = check_environment( |
| require_cuda=not args.allow_no_cuda, |
| require_ape_extension=not args.skip_ape_extension, |
| ) |
| print(json.dumps(report, indent=2, sort_keys=True)) |
| return 1 if report["errors"] else 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
