add script to collect information about the env (#69)

* add script to collect information about the env

* fix cpuonly check

.github/workflows/install.yml

@@ -46,10 +46,17 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Collect environment information
+        run: python scripts/collect_env.py
+
       - name: Install torch
         run: ltt install --cpuonly --pytorch-channel=${{ matrix.pytorch-channel }} torch
 
       - name: Check if CPU only
-        run:
-          python -c "import sys, torch; sys.exit(hasattr(torch._C,
-          '_cuda_getDeviceCount'))"
+        shell: python
+        run: |
+          import sys
+          import torch
+
+          compiled_with_cuda = hasattr(torch._C, "_cuda_getDeviceCount")
+          sys.exit(compiled_with_cuda)

scripts/check_available_pytorch_dists.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import itertools
 import json
 

scripts/collect_env.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+
+import itertools
+import platform
+import subprocess
+
+import importlib_metadata
+import pip
+from pip._vendor.packaging.requirements import Requirement
+from pip._vendor.packaging.version import InvalidVersion, Version
+
+
+try:
+    import light_the_torch
+except ModuleNotFoundError:
+    light_the_torch = None
+
+NOT_AVAILABLE = "N/A"
+
+# TODO: somehow merge this with light_the_torch._patch.PYTORCH_DISTRIBUTIONS to avoid
+#  duplication
+PYTORCH_DISTRIBUTIONS = {
+    "torch",
+    "torch_model_archiver",
+    "torch_tb_profiler",
+    "torcharrow",
+    "torchaudio",
+    "torchcsprng",
+    "torchdata",
+    "torchdistx",
+    "torchserve",
+    "torchtext",
+    "torchvision",
+}
+
+
+def main():
+    header("System")
+    print(f"Platform: {platform.platform()}")
+    print(f"Machine: {platform.machine()}")
+    print(f"Python version: {platform.python_version()}")
+    nvidia_driver_version = detect_nvidia_driver_version()
+    print(f"NVIDIA driver version: {nvidia_driver_version or NOT_AVAILABLE}")
+
+    header("Environment")
+    for name, version in itertools.chain(
+        [
+            ("pip", pip.__version__),
+            (
+                "light_the_torch",
+                light_the_torch.__version__ if light_the_torch else NOT_AVAILABLE,
+            ),
+        ],
+        detect_pytorch_or_dependent_packages(),
+    ):
+        print(f"- `{name}=={version}`")
+
+
+# TODO: somehow merge this with light_the_torch._cb._detect_nvidia_driver_version to
+#  avoid duplication
+def detect_nvidia_driver_version():
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=driver_version",
+                "--format=csv",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        return Version(result.stdout.splitlines()[-1])
+    except (FileNotFoundError, subprocess.CalledProcessError, InvalidVersion):
+        return None
+
+
+def detect_pytorch_or_dependent_packages():
+    packages = {
+        (dist.name, dist.version)
+        for dist in importlib_metadata.distributions()
+        if any(
+            name in PYTORCH_DISTRIBUTIONS
+            for name in itertools.chain(
+                [dist.name],
+                [Requirement(req_str).name for req_str in dist.requires]
+                if dist.requires
+                else [],
+            )
+        )
+    }
+    return sorted(
+        packages,
+        key=lambda package: (package[0] not in PYTORCH_DISTRIBUTIONS, package[0]),
+    )
+
+
+def header(name):
+    print()
+    print(f"#### {name}")
+    print()
+
+
+if __name__ == "__main__":
+    main()