Skip to content

Commit 3adbad4

Browse files
authored
Use anywidget-based IPC for zarr gets which do not require serving data on localhost (#329)
* WIP: zarr get using anywidget.experimental command/invoke works * Unused zarr js import * Support local store * Lint * Cleanup * Lint * Bump version
1 parent 834b15f commit 3adbad4

File tree

4 files changed

+127
-15
lines changed

4 files changed

+127
-15
lines changed

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
44

55
[project]
66
name = "vitessce"
7-
version = "3.2.5"
7+
version = "3.2.6"
88
authors = [
99
{ name="Mark Keller", email="mark_keller@hms.harvard.edu" },
1010
]
@@ -73,7 +73,7 @@ docs = [
7373
]
7474
all = [
7575
'jupyter-server-proxy>=1.5.2',
76-
'anywidget>=0.9.3',
76+
'anywidget>=0.9.10',
7777
'uvicorn>=0.17.0',
7878
'ujson>=4.0.1',
7979
'starlette==0.14.0',

vitessce/config.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,16 @@ def get_routes(self):
227227

228228
return routes
229229

230+
def get_stores(self, base_url=None):
231+
stores = {}
232+
for obj in self.objs:
233+
stores = {
234+
**stores,
235+
**obj.get_stores(base_url)
236+
}
237+
238+
return stores
239+
230240

231241
class VitessceConfigViewHConcat:
232242
"""
@@ -1489,6 +1499,21 @@ def get_routes(self):
14891499
routes += d.get_routes()
14901500
return routes
14911501

1502+
def get_stores(self, base_url=None):
1503+
"""
1504+
Convert the routes for this view config from the datasets.
1505+
1506+
:returns: A list of server routes.
1507+
:rtype: list[starlette.routing.Route]
1508+
"""
1509+
stores = {}
1510+
for d in self.config["datasets"]:
1511+
stores = {
1512+
**stores,
1513+
**d.get_stores(base_url)
1514+
}
1515+
return stores
1516+
14921517
def to_python(self):
14931518
"""
14941519
Convert the VitessceConfig instance to a one-line Python code snippet that can be used to generate it.

vitessce/widget.py

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
# Widget dependencies
66
import anywidget
7-
from traitlets import Unicode, Dict, Int, Bool
7+
from traitlets import Unicode, Dict, List, Int, Bool
88
import time
99
import uuid
1010

@@ -201,6 +201,7 @@ def get_uid_str(uid):
201201
const customJsUrl = view.model.get('custom_js_url');
202202
const pluginEsm = view.model.get('plugin_esm');
203203
const remountOnUidChange = view.model.get('remount_on_uid_change');
204+
const storeUrls = view.model.get('store_urls');
204205
205206
const pkgName = (jsDevMode ? "@vitessce/dev" : "vitessce");
206207
@@ -224,6 +225,19 @@ def get_uid_str(uid):
224225
let pluginFileTypes;
225226
let pluginJointFileTypes;
226227
228+
const stores = Object.fromEntries(
229+
storeUrls.map(storeUrl => ([
230+
storeUrl,
231+
{
232+
async get(key) {
233+
const [data, buffers] = await view.experimental.invoke("_zarr_get", [storeUrl, key]);
234+
if (!data.success) return undefined;
235+
return buffers[0].buffer;
236+
},
237+
}
238+
])),
239+
);
240+
227241
try {
228242
const pluginEsmUrl = URL.createObjectURL(new Blob([pluginEsm], { type: "text/javascript" }));
229243
const pluginModule = (await import(pluginEsmUrl)).default;
@@ -310,7 +324,7 @@ def get_uid_str(uid):
310324
const vitessceProps = {
311325
height, theme, config, onConfigChange, validateConfig,
312326
pluginViewTypes, pluginCoordinationTypes, pluginFileTypes, pluginJointFileTypes,
313-
remountOnUidChange,
327+
remountOnUidChange, stores,
314328
};
315329
316330
return e('div', { ref: divRef, style: { height: height + 'px' } },
@@ -383,13 +397,15 @@ class VitessceWidget(anywidget.AnyWidget):
383397

384398
next_port = DEFAULT_PORT
385399

386-
js_package_version = Unicode('3.3.7').tag(sync=True)
400+
js_package_version = Unicode('3.3.12').tag(sync=True)
387401
js_dev_mode = Bool(False).tag(sync=True)
388402
custom_js_url = Unicode('').tag(sync=True)
389403
plugin_esm = Unicode(DEFAULT_PLUGIN_ESM).tag(sync=True)
390404
remount_on_uid_change = Bool(True).tag(sync=True)
391405

392-
def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=False, js_package_version='3.3.7', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
406+
store_urls = List(trait=Unicode(''), default_value=[]).tag(sync=True)
407+
408+
def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=False, js_package_version='3.3.12', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
393409
"""
394410
Construct a new Vitessce widget.
395411
@@ -422,13 +438,15 @@ def __init__(self, config, height=600, theme='auto', uid=None, port=None, proxy=
422438
config_dict = config.to_dict(base_url=base_url)
423439
routes = config.get_routes()
424440

441+
self._stores = config.get_stores(base_url=base_url)
442+
425443
uid_str = get_uid_str(uid)
426444

427445
super(VitessceWidget, self).__init__(
428446
config=config_dict, height=height, theme=theme, proxy=proxy,
429447
js_package_version=js_package_version, js_dev_mode=js_dev_mode, custom_js_url=custom_js_url,
430448
plugin_esm=plugin_esm, remount_on_uid_change=remount_on_uid_change,
431-
uid=uid_str,
449+
uid=uid_str, store_urls=list(self._stores.keys())
432450
)
433451

434452
serve_routes(config, routes, use_port)
@@ -460,10 +478,20 @@ def close(self):
460478
self.config_obj.stop_server(self.port)
461479
super().close()
462480

481+
@anywidget.experimental.command
482+
def _zarr_get(self, params, buffers):
483+
[store_url, key] = params
484+
store = self._stores[store_url]
485+
try:
486+
buffers = [store[key.lstrip("/")]]
487+
except KeyError:
488+
buffers = []
489+
return {"success": len(buffers) == 1}, buffers
490+
463491
# Launch Vitessce using plain HTML representation (no ipywidgets)
464492

465493

466-
def ipython_display(config, height=600, theme='auto', base_url=None, host_name=None, uid=None, port=None, proxy=False, js_package_version='3.3.7', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
494+
def ipython_display(config, height=600, theme='auto', base_url=None, host_name=None, uid=None, port=None, proxy=False, js_package_version='3.3.12', js_dev_mode=False, custom_js_url='', plugin_esm=DEFAULT_PLUGIN_ESM, remount_on_uid_change=True):
467495
from IPython.display import display, HTML
468496
uid_str = "vitessce" + get_uid_str(uid)
469497

@@ -485,6 +513,7 @@ def ipython_display(config, height=600, theme='auto', base_url=None, host_name=N
485513
"height": height,
486514
"theme": theme,
487515
"config": config_dict,
516+
"store_urls": [],
488517
}
489518

490519
# We need to clean up the React and DOM state in any case in which

vitessce/wrappers.py

Lines changed: 65 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import tempfile
44
from uuid import uuid4
55
from pathlib import PurePath, PurePosixPath
6+
import zarr
67

78
from .constants import (
89
norm_enum,
@@ -41,9 +42,11 @@ def __init__(self, **kwargs):
4142
self.out_dir = kwargs['out_dir'] if 'out_dir' in kwargs else tempfile.mkdtemp(
4243
)
4344
self.routes = []
44-
self.is_remote = False
45+
self.is_remote = False # TODO: change to needs_localhost_serving for clarity
46+
self.is_store = False # TODO: change to needs_store_registration for clarity
4547
self.file_def_creators = []
4648
self.base_dir = None
49+
self.stores = {}
4750
self._request_init = kwargs['request_init'] if 'request_init' in kwargs else None
4851

4952
def __repr__(self):
@@ -71,6 +74,20 @@ def get_routes(self):
7174
"""
7275
return self.routes
7376

77+
def get_stores(self, base_url):
78+
"""
79+
Obtain the stores that have been created for this wrapper class.
80+
81+
:returns: A dictionary that maps file URLs to Zarr Store objects.
82+
:rtype: dict[str, zarr.Store]
83+
"""
84+
relative_stores = self.stores
85+
absolute_stores = {}
86+
for relative_url, store in relative_stores.items():
87+
absolute_url = base_url + relative_url
88+
absolute_stores[absolute_url] = store
89+
return absolute_stores
90+
7491
def get_file_defs(self, base_url):
7592
"""
7693
Obtain the file definitions for this wrapper class.
@@ -111,6 +128,30 @@ def get_local_dir_url(self, base_url, dataset_uid, obj_i, local_dir_path, local_
111128
return self._get_url_simple(base_url, file_path_to_url_path(local_dir_path, prepend_slash=False))
112129
return self._get_url(base_url, dataset_uid, obj_i, local_dir_uid)
113130

131+
def register_zarr_store(self, dataset_uid, obj_i, store_or_local_dir_path, local_dir_uid):
132+
if not self.is_remote and self.is_store:
133+
# Set up `store` and `local_dir_path` variables.
134+
if isinstance(store_or_local_dir_path, str):
135+
# TODO: use zarr.FSStore if fsspec is installed?
136+
store = zarr.DirectoryStore(store_or_local_dir_path)
137+
local_dir_path = store_or_local_dir_path
138+
else:
139+
# TODO: check that store_or_local_dir_path is a zarr.Store or StoreLike?
140+
store = store_or_local_dir_path
141+
# A store instance was passed directly, so there is no local directory path.
142+
# Instead we just make one up using _get_route_str but it could be any string.
143+
local_dir_path = self._get_route_str(dataset_uid, obj_i, local_dir_uid)
144+
145+
# Register the store on the same route path
146+
# that will be used for the "url" field in the file definition.
147+
if self.base_dir is None:
148+
route_path = self._get_route_str(dataset_uid, obj_i, local_dir_uid)
149+
else:
150+
route_path = file_path_to_url_path(local_dir_path)
151+
local_dir_path = join(self.base_dir, local_dir_path)
152+
153+
self.stores[route_path] = store
154+
114155
def get_local_dir_route(self, dataset_uid, obj_i, local_dir_path, local_dir_uid):
115156
"""
116157
Obtain the Mount for some local directory
@@ -896,12 +937,14 @@ def image_file_def_creator(base_url):
896937

897938

898939
class AnnDataWrapper(AbstractWrapper):
899-
def __init__(self, adata_path=None, adata_url=None, obs_feature_matrix_path=None, feature_filter_path=None, initial_feature_filter_path=None, obs_set_paths=None, obs_set_names=None, obs_locations_path=None, obs_segmentations_path=None, obs_embedding_paths=None, obs_embedding_names=None, obs_embedding_dims=None, obs_spots_path=None, obs_points_path=None, feature_labels_path=None, obs_labels_path=None, convert_to_dense=True, coordination_values=None, obs_labels_paths=None, obs_labels_names=None, **kwargs):
940+
def __init__(self, adata_path=None, adata_url=None, adata_store=None, obs_feature_matrix_path=None, feature_filter_path=None, initial_feature_filter_path=None, obs_set_paths=None, obs_set_names=None, obs_locations_path=None, obs_segmentations_path=None, obs_embedding_paths=None, obs_embedding_names=None, obs_embedding_dims=None, obs_spots_path=None, obs_points_path=None, feature_labels_path=None, obs_labels_path=None, convert_to_dense=True, coordination_values=None, obs_labels_paths=None, obs_labels_names=None, **kwargs):
900941
"""
901942
Wrap an AnnData object by creating an instance of the ``AnnDataWrapper`` class.
902943
903944
:param str adata_path: A path to an AnnData object written to a Zarr store containing single-cell experiment data.
904945
:param str adata_url: A remote url pointing to a zarr-backed AnnData store.
946+
:param adata_store: A path to pass to zarr.FSStore, or an existing store instance.
947+
:type adata_store: str or zarr.Storage
905948
:param str obs_feature_matrix_path: Location of the expression (cell x gene) matrix, like `X` or `obsm/highly_variable_genes_subset`
906949
:param str feature_filter_path: A string like `var/highly_variable` used in conjunction with `obs_feature_matrix_path` if obs_feature_matrix_path points to a subset of `X` of the full `var` list.
907950
:param str initial_feature_filter_path: A string like `var/highly_variable` used in conjunction with `obs_feature_matrix_path` if obs_feature_matrix_path points to a subset of `X` of the full `var` list.
@@ -927,18 +970,30 @@ def __init__(self, adata_path=None, adata_url=None, obs_feature_matrix_path=None
927970
self._repr = make_repr(locals())
928971
self._adata_path = adata_path
929972
self._adata_url = adata_url
930-
if adata_url is not None and (adata_path is not None):
973+
self._adata_store = adata_store
974+
975+
num_inputs = sum([1 for x in [adata_path, adata_url, adata_store] if x is not None])
976+
if num_inputs > 1:
931977
raise ValueError(
932-
"Did not expect adata_url to be provided with adata_path")
933-
if adata_url is None and (adata_path is None):
978+
"Expected only one of adata_path, adata_url, or adata_store to be provided")
979+
if num_inputs == 0:
934980
raise ValueError(
935-
"Expected either adata_url or adata_path to be provided")
981+
"Expected one of adata_path, adata_url, or adata_store to be provided")
982+
936983
if adata_path is not None:
937984
self.is_remote = False
985+
self.is_store = False
938986
self.zarr_folder = 'anndata.zarr'
939-
else:
987+
elif adata_url is not None:
940988
self.is_remote = True
989+
self.is_store = False
941990
self.zarr_folder = None
991+
else:
992+
# Store case
993+
self.is_remote = False
994+
self.is_store = True
995+
self.zarr_folder = None
996+
942997
self.local_dir_uid = make_unique_filename(".adata.zarr")
943998
self._expression_matrix = obs_feature_matrix_path
944999
self._cell_set_obs_names = obs_set_names
@@ -978,6 +1033,9 @@ def convert_and_save(self, dataset_uid, obj_i, base_dir=None):
9781033
def make_anndata_routes(self, dataset_uid, obj_i):
9791034
if self.is_remote:
9801035
return []
1036+
elif self.is_store:
1037+
self.register_zarr_store(dataset_uid, obj_i, self._adata_store, self.local_dir_uid)
1038+
return []
9811039
else:
9821040
return self.get_local_dir_route(dataset_uid, obj_i, self._adata_path, self.local_dir_uid)
9831041

0 commit comments

Comments
 (0)