pyodide fixes

Author: mhochsteger

webgpu/canvas.py

@@ -6,7 +6,7 @@
 import pathlib
 
 from . import platform
-from .utils import get_device, read_texture
+from .utils import get_device, read_texture, Lock
 from .webgpu_api import *
 
 _TARGET_FPS = 60
@@ -20,9 +20,6 @@ class _DebounceData:
 
 def debounce(arg=None):
     def decorator(func):
-        if platform.is_pyodide:
-            return arg
-
         # Render only once every 1/_TARGET_FPS seconds
         @functools.wraps(func)
         def debounced(obj, *args, **kwargs):
@@ -35,16 +32,25 @@ def debounced(obj, *args, **kwargs):
 
             data = obj._debounce_data[fname]
 
-            if data.timer is not None:
-                # we already have a render scheduled, so do nothing
-                return
+            # check if we already have a render scheduled
+            if platform.is_pyodide:
+                if data.timer is not None and not data.timer.done():
+                    return
+            else:
+                if data.timer is not None:
+                    return
 
             def f():
                 # clear the timer, so we can schedule a new one with the next function call
                 t = time.time()
                 data.timer = None
-                func(obj, *args, **kwargs)
-                data.t_last = t
+                if platform.is_pyodide:
+                    # due to async nature, we need to update t_last before calling func
+                    data.t_last = t
+                    func(obj, *args, **kwargs)
+                else:
+                    func(obj, *args, **kwargs)
+                    data.t_last = t
 
             if data.t_last is None:
                 # first call -> just call the function immediately
@@ -53,8 +59,16 @@ def f():
                 return
 
             t_wait = max(1 / target_fps - (time.time() - data.t_last), 0)
-            data.timer = threading.Timer(t_wait, f)
-            data.timer.start()
+            if platform.is_pyodide:
+                import asyncio
+                async def _runner():
+                    if t_wait > 0:
+                        await asyncio.sleep(t_wait)
+                    f()
+                data.timer = asyncio.create_task(_runner())
+            else:
+                data.timer = threading.Timer(t_wait, f)
+                data.timer.start()
 
         return debounced
 
@@ -91,7 +105,7 @@ class Canvas:
     _on_update_html_canvas: list[Callable]
 
     def __init__(self, device, canvas, multisample_count=4):
-        self._update_mutex = threading.RLock()
+        self._update_mutex = Lock()
         self.target_texture = None
 
         self._on_resize_callbacks = []

webgpu/input_handler.py

@@ -1,5 +1,5 @@
-import threading
 from typing import Callable
+from .utils import Lock
 
 
 class InputHandler:
@@ -24,7 +24,7 @@ def get_set(self):
             return s
 
     def __init__(self):
-        self._mutex = threading.Lock()
+        self._mutex = Lock(True)
         self._callbacks = {}
         self._js_handlers = {}
         self._is_mousedown = False

webgpu/link/base.py

@@ -2,6 +2,7 @@
 import base64
 import itertools
 import json
+import time
 import threading
 from collections.abc import Mapping
 from typing import Callable
@@ -96,6 +97,8 @@ def __init__(self):
         self._objects = {}
         self._cache = {}
 
+        next(self._request_id) # make sure first id is 1, in case 0 is interpreted as None
+
     def _call_data(self, id, prop, args, ignore_result=False):
         buffer = []
         args = self._dump_data(args, buffer)
@@ -313,7 +316,11 @@ async def _on_message_async(self, message: str):
                     self._requests[request_id] = self._load_data(data.get("value", None))
                     if key and data.get("cache", False):
                         self._cache[key] = self._requests[request_id]
-                    event.set()
+                    
+                    if isinstance(event, asyncio.Future):
+                        event.set_result(self._requests[request_id])
+                    else:
+                        event.set()
                     return
 
                 case "call":
@@ -407,11 +414,10 @@ def _on_message(self, message: str):
 
 
 class PyodideLink(LinkBase):
-    def __init__(self, send_message, size_buffer, result_buffer):
+    def __init__(self, send_message):
         super().__init__()
         self._send_message = send_message
-        self._size_buffer = size_buffer
-        self._result_buffer = result_buffer
+        self._requests = {}
 
     def create_proxy(self, func, ignore_return_value=False):
         id_ = id(func)
@@ -424,26 +430,23 @@ def create_proxy(self, func, ignore_return_value=False):
         }
 
     def _send_data(self, metadata, data, key=None):
-        if type(data) is bytes:
-            self._send_message(data)
-        else:
-            if (
-                metadata.get("request_id", None) is not None
-                and metadata["type"] != "response"
-                and not metadata.get("ignore_return_value", False)
-            ):
-                import js
-
-                js.Atomics.store(self._size_buffer, 0, 0)
-                self._send_message(data)
-                js.Atomics.wait(self._size_buffer, 0, 0, 10000)
-                n = self._size_buffer[0]
-                res = bytes(self._result_buffer.slice(0, n))
-                s = res.decode("utf-8")
-                data = json.loads(s)
-                return self._load_data(data.get("value", None))
-            else:
-                self._send_message(data)
+        """Send data to the remote environment,
+        if request_id is set, (blocking-)wait for the response and return it"""
+        request_id = metadata.get("request_id", None)
+        type = metadata.get("type", None)
+        event = None
+        self._send_message(data)
+        if type != "response" and request_id is not None:
+            # from pyodide.ffi import run_sync
+            import asyncio
+            event = asyncio.Future()
+            self._requests[request_id] = event, key
+            # todo: this shouldn't be necessary
+            # but run_sync(event) gives an error
+            while not event.done():
+                time.sleep(0.001)
+
+            return self._requests.pop(request_id)
 
 
 class LinkBaseAsync(LinkBase):

webgpu/platform.py

@@ -35,7 +35,7 @@ def create_event_handler(
 
 
 try:
-    import js
+    import js as pyodide_js
     import pyodide.ffi
     from pyodide.ffi import JsPromise, JsProxy
     from pyodide.ffi import create_proxy as _create_proxy
@@ -48,7 +48,7 @@ def destroy_proxy(proxy):
 
     is_pyodide = True
     try:
-        is_pyodide_main_thread = bool(js.window.document)
+        is_pyodide_main_thread = bool(pyodide_js.window.document)
     except:
         is_pyodide_main_thread = False
 
@@ -88,7 +88,7 @@ def toJS(value):
         value = _convert(value)
         ret = pyodide.ffi.to_js(
             value,
-            dict_converter=js.Object.fromEntries,
+            dict_converter=pyodide_js.Object.fromEntries,
             default_converter=_default_converter,
             create_pyproxies=False,
         )

webgpu/scene.py

@@ -5,7 +5,7 @@
 from .canvas import Canvas, debounce
 from .input_handler import InputHandler
 from .renderer import BaseRenderer, RenderOptions, SelectEvent
-from .utils import max_bounding_box, read_buffer
+from .utils import max_bounding_box, read_buffer, Lock
 from .platform import is_pyodide, is_pyodide_main_thread
 from .webgpu_api import *
 from .camera import Camera
@@ -88,7 +88,7 @@ def init(self, canvas):
         self.input_handler.set_canvas(canvas.canvas)
         self.options.set_canvas(canvas)
 
-        self._render_mutex = canvas._update_mutex
+        self._render_mutex = Lock(True) if is_pyodide else canvas._update_mutex
 
         with self._render_mutex:
             self.options.timestamp = time.time()
@@ -110,7 +110,6 @@ def init(self, canvas):
             )
             self._select_buffer_valid = False
 
-            canvas._update_mutex = self._render_mutex
             canvas.on_resize(self.render)
 
             canvas.on_update_html_canvas(self.__on_update_html_canvas)

webgpu/utils.py

@@ -5,12 +5,38 @@
 from . import platform
 from .webgpu_api import *
 from .webgpu_api import toJS as to_js
-import threading
 
-_device: Device = None
+try:
+    import pyodide.ffi
+    import asyncio
+    class Lock():
+        def __init__(self, do: bool = False):
+            self._lock = asyncio.Lock()
+            self.do = do
+
+        def __enter__(self):
+            if self.do:
+                pyodide.ffi.run_sync(self._lock.acquire())
+
+        def __exit__(self, exc_type, exc, tb):
+            if self.do:
+                self._lock.release()
+
+except ImportError:
+    from threading import RLock
 
-_lock_init_device = threading.Lock()
+    class Lock():
+        def __init__(self, do: bool = True):
+            self._lock = RLock()
 
+        def __enter__(self):
+            self._lock.acquire()
+
+        def __exit__(self, exc_type, exc, tb):
+            self._lock.release()
+
+_lock_init_device = Lock()
+_device: Device = None
 
 def init_device_sync():
     global _device