gpu info upon notebook import (#410)

kushalkolar · web-flow · commit 7921516f3f72 · 2024-02-27T01:14:44.000-05:00
* displaying fpl logo and gpu info works on import in nb

* add face logo for nb and add to MANIFEST

* remove line that snuck in from other branch

* git lfs to pypi build action

* fix import and add easter egg
diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml
@@ -21,7 +21,14 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
+    - name: Install git-lfs
+      run: |
+        sudo apt install --no-install-recommends -y  git-lfs
     - uses: actions/checkout@v3
+    - name: fetch git lfs files
+      run: |
+        git lfs fetch --all
+        git lfs pull
     - name: Set up Python
       uses: actions/setup-python@v3
       with:
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -1,2 +1,4 @@
 recursive-include fastplotlib/utils/colormaps/ *
 include fastplotlib/VERSION
+recursive-include fastplotlib/assets/ *
+
diff --git a/fastplotlib/__init__.py b/fastplotlib/__init__.py
@@ -3,6 +3,7 @@
 from .layouts import Plot, GridPlot
 from .graphics import *
 from .graphics.selectors import *
+from .utils import _notebook_print_banner, config
 
 from wgpu.gui.auto import run
 
@@ -13,6 +14,16 @@
 else:
     from .widgets import ImageWidget
 
+from wgpu.backends.wgpu_native import enumerate_adapters
+
+adapters = [a.request_adapter_info() for a in enumerate_adapters()]
+
+if len(adapters) < 1:
+    raise IndexError(
+        "No WGPU adapters found, fastplotlib will not work."
+    )
+
+_notebook_print_banner()
 
 with open(Path(__file__).parent.joinpath("VERSION"), "r") as f:
     __version__ = f.read().split("\n")[0]
diff --git a/fastplotlib/assets/egg.gif b/fastplotlib/assets/egg.gif
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7251090fd17fa81eae8a5ad176608755b4a68d620b125ebd571acc6d68daa017
+size 4589
diff --git a/fastplotlib/assets/fastplotlib_face_logo.png b/fastplotlib/assets/fastplotlib_face_logo.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84976200545d143f559b6840a332f3121d58f8d7de8dad7cb8a6d027d854c153
+size 10239
diff --git a/fastplotlib/layouts/_frame/_ipywidget_toolbar.py b/fastplotlib/layouts/_frame/_ipywidget_toolbar.py
@@ -3,6 +3,7 @@
 from itertools import product
 from math import copysign
 from functools import partial
+from pathlib import Path
 from typing import *
 
 
@@ -17,10 +18,12 @@
     BoundedIntText,
     Play,
     jslink,
+    Image,
 )
 
 from ...graphics.selectors import PolygonSelector
 from ._toolbar import ToolBar
+from ...utils import config
 
 
 class IpywidgetToolBar(HBox, ToolBar):
@@ -92,6 +95,19 @@ def __init__(self, plot):
             self._record_button,
         ]
 
+        if config.party_parrot:
+            gif_path = Path(__file__).parent.parent.parent.joinpath("assets", "egg.gif")
+            with open(gif_path, "rb") as f:
+                gif = f.read()
+
+            image = Image(
+                value=gif,
+                format="png",
+                width=35,
+                height=25,
+            )
+            widgets.append(image)
+
         if hasattr(self.plot, "_subplots"):
             positions = list(product(range(self.plot.shape[0]), range(self.plot.shape[1])))
             values = list()
diff --git a/fastplotlib/utils/__init__.py b/fastplotlib/utils/__init__.py
@@ -1 +1,15 @@
+from dataclasses import dataclass
+
+
 from .functions import *
+from ._gpu_info import _notebook_print_banner
+
+
+@dataclass
+class _Config:
+    party_parrot: bool
+
+
+config = _Config(
+    party_parrot=False
+)
diff --git a/fastplotlib/utils/_gpu_info.py b/fastplotlib/utils/_gpu_info.py
@@ -0,0 +1,66 @@
+from pathlib import Path
+
+from wgpu.backends.wgpu_native import enumerate_adapters
+from wgpu.utils import get_default_device
+
+try:
+    ip = get_ipython()
+    from ipywidgets import Image
+    from wgpu.gui.jupyter import JupyterWgpuCanvas
+except (NameError, ModuleNotFoundError, ImportError):
+    NOTEBOOK = False
+else:
+    from IPython.display import display
+    if ip.has_trait("kernel") and (JupyterWgpuCanvas is not False):
+        NOTEBOOK = True
+    else:
+        NOTEBOOK = False
+
+
+def _notebook_print_banner():
+    if NOTEBOOK is False:
+        return
+
+    logo_path = Path(__file__).parent.parent.joinpath("assets", "fastplotlib_face_logo.png")
+
+    with open(logo_path, "rb") as f:
+        logo_data = f.read()
+
+    image = Image(
+        value=logo_data,
+        format="png",
+        width=300,
+        height=55
+    )
+
+    display(image)
+
+    # print logo and adapter info
+    adapters = [a for a in enumerate_adapters()]
+    adapters_info = [a.request_adapter_info() for a in adapters]
+
+    ix_default = adapters_info.index(get_default_device().adapter.request_adapter_info())
+
+    if len(adapters) > 0:
+        print("Available devices:")
+
+    for ix, adapter in enumerate(adapters_info):
+        atype = adapter["adapter_type"]
+        backend = adapter["backend_type"]
+        driver = adapter["description"]
+        device = adapter["device"]
+
+        if atype == "DiscreteGPU" and backend != "OpenGL":
+            charactor = chr(0x2705)
+        elif atype == "IntegratedGPU" and backend != "OpenGL":
+            charactor = chr(0x0001FBC4)
+        else:
+            charactor = chr(0x2757)
+
+        if ix == ix_default:
+            default = " (default) "
+        else:
+            default = " "
+
+        output_str = f"{charactor}{default}| {device} | {atype} | {backend} | {driver}"
+        print(output_str)
