Upgrade/loosen dependencies and add invoke_batched parameter

Author: keller-mark

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
   'numpy>=1.21.2',
   'zarr>=2.5.0,<3',
   'numcodecs>=0.5.7,<0.16.0',
-  'anndata>=0.7.8,<0.11.0',
+  'anndata>=0.7.8',
   # scanpy < 1.10.3 does not support numpy >= 2.0.0 and does not
   # Reference: https://github.com/scverse/scanpy/pull/3115/files
   'scanpy>=1.10.2',
@@ -73,7 +73,7 @@ all = [
   'anywidget>=0.9.10',
   'uvicorn>=0.17.0',
   'ujson>=4.0.1',
-  'starlette==0.14.0',
+  'starlette>=0.14.0',
   'generate-tiff-offsets>=0.1.9',
   'kerchunk>=0.2.6',
   'fsspec',
@@ -87,7 +87,10 @@ all = [
 building = []
 testing = []
 linting = []
-notebook = []
+notebook = [
+  'spatialdata>=0.3.0',
+  'dask[dataframe]==2024.11.1',
+]
 demos = [
   "pulp==2.7.0",
   "snakemake",

src/vitessce/responses.py

@@ -0,0 +1,9 @@
+import ujson
+from starlette.responses import JSONResponse
+
+# References:
+# - https://www.starlette.io/responses/#custom-json-serialization
+# - https://github.com/encode/starlette/releases/tag/0.14.1
+class UJSONResponse(JSONResponse):
+    def render(self, content):
+        return ujson.dumps(content, ensure_ascii=False).encode("utf-8")

src/vitessce/widget.py

@@ -233,6 +233,7 @@ def get_uid_str(uid):
     const remountOnUidChange = view.model.get('remount_on_uid_change');
     const storeUrls = view.model.get('store_urls');
     const invokeTimeout = view.model.get('invoke_timeout');
+    const invokeBatched = view.model.get('invoke_batched');
 
     const pageMode = view.model.get('page_mode');
     const pageEsm = view.model.get('page_esm');
@@ -319,23 +320,25 @@ def get_uid_str(uid):
             storeUrl,
             {
                 async get(key) {
-                    return enqueue([storeUrl, key]);
-                    /*
-                    const [data, buffers] = await view.experimental.invoke("_zarr_get", [storeUrl, key], {
-                        signal: AbortSignal.timeout(invokeTimeout),
-                    });
-                    if (!data.success) return undefined;
-
-                    if (key.includes("spatialdata_attrs") && key.endsWith("0") && !ArrayBuffer.isView(buffers[0].buffer)) {
-                        // For some reason, the Zarrita.js UnicodeStringArray does not seem to work with
-                        // ArrayBuffers (throws a TypeError), so here we convert to Uint8Array if needed.
-                        // This error is occurring specifically for the arr.getChunk call within the AnnDataSource._loadString function.
-                        // TODO: figure out a more long-term solution.
-                        return new Uint8Array(buffers[0].buffer);
-                    }
+                    if (invokeBatched) {
+                        return enqueue([storeUrl, key]);
+                    } else {
+                        // Do not submit zarr gets in batches. Instead, submit individually.
+                        const [data, buffers] = await view.experimental.invoke("_zarr_get", [storeUrl, key], {
+                            signal: AbortSignal.timeout(invokeTimeout),
+                        });
+                        if (!data.success) return undefined;
+
+                        if (key.includes("spatialdata_attrs") && key.endsWith("0") && !ArrayBuffer.isView(buffers[0].buffer)) {
+                            // For some reason, the Zarrita.js UnicodeStringArray does not seem to work with
+                            // ArrayBuffers (throws a TypeError), so here we convert to Uint8Array if needed.
+                            // This error is occurring specifically for the arr.getChunk call within the AnnDataSource._loadString function.
+                            // TODO: figure out a more long-term solution.
+                            return new Uint8Array(buffers[0].buffer);
+                        }
 
-                    return buffers[0].buffer;
-                    */
+                        return buffers[0].buffer;
+                    }
                 },
             }
         ])),
@@ -603,10 +606,11 @@ class VitessceWidget(anywidget.AnyWidget):
     page_mode = Bool(False).tag(sync=True)
     page_esm = Unicode('').tag(sync=True)
     invoke_timeout = Int(300000).tag(sync=True)
+    invoke_batched = Bool(True).tag(sync=True)
 
     store_urls = List(trait=Unicode(''), default_value=[]).tag(sync=True)
 
-    def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=False, js_package_version='3.5.11', js_dev_mode=False, custom_js_url='', plugins=None, remount_on_uid_change=True, prefer_local=True, invoke_timeout=300000, page_mode=False, page_esm=None):
+    def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=False, js_package_version='3.5.11', js_dev_mode=False, custom_js_url='', plugins=None, remount_on_uid_change=True, prefer_local=True, invoke_timeout=300000, invoke_batched=True, page_mode=False, page_esm=None):
         """
         Construct a new Vitessce widget.
 
@@ -623,6 +627,7 @@ def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=
         :param bool remount_on_uid_change: Passed to the remountOnUidChange prop of the <Vitessce/> React component. By default, True.
         :param bool prefer_local: Should local data be preferred (only applies to `*_artifact` data objects)? By default, True.
         :param int invoke_timeout: The timeout in milliseconds for invoking Python functions from JavaScript. By default, 300000.
+        :param bool invoke_batched: Should invocations (Zarr gets) be submitted in batch, or individually? By default, True.
         :param bool page_mode: Whether to render the <Vitessce/> component in grid-mode or page-mode. By default, False.
         :param str page_esm: The ES module string for the page component creation function. Optional.
 
@@ -658,7 +663,7 @@ def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=
             js_package_version=js_package_version, js_dev_mode=js_dev_mode, custom_js_url=custom_js_url,
             plugin_esm=plugin_esm, remount_on_uid_change=remount_on_uid_change,
             page_mode=page_mode, page_esm=('' if page_esm is None else page_esm),
-            invoke_timeout=invoke_timeout,
+            invoke_timeout=invoke_timeout, invoke_batched=invoke_batched,
             uid=uid_str, store_urls=list(self._stores.keys())
         )
 

src/vitessce/wrappers.py

@@ -390,8 +390,8 @@ def make_raster_routes(self, dataset_uid, obj_i):
         else:
             # TODO: Move imports back to top when this is factored out.
             from .routes import range_repsonse, JsonRoute, FileRoute
+            from .responses import UJSONResponse
             from generate_tiff_offsets import get_offsets
-            from starlette.responses import UJSONResponse
 
             offsets = get_offsets(self._img_path)
 
@@ -533,8 +533,8 @@ def make_raster_routes(self, dataset_uid, obj_i):
         else:
             # TODO: Move imports back to top when this is factored out.
             from .routes import range_repsonse, JsonRoute, FileRoute
+            from .responses import UJSONResponse
             from generate_tiff_offsets import get_offsets
-            from starlette.responses import UJSONResponse
 
             offsets = get_offsets(self._img_path)
 
@@ -666,8 +666,8 @@ def make_raster_routes(self, dataset_uid, obj_i):
         else:
             # TODO: Move imports back to top when this is factored out.
             from .routes import range_repsonse, JsonRoute, FileRoute
+            from .responses import UJSONResponse
             from generate_tiff_offsets import get_offsets
-            from starlette.responses import UJSONResponse
 
             offsets = get_offsets(self._img_path)