def test_cache(comm):
cat = UniformCatalog(nbar=10000, BoxSize=1.0, comm=comm)
cat['test'] = cat['Position']**5
test = cat['test'].compute()
# cache should no longer be empty
cache = GlobalCache.get()
assert cache.cache.total_bytes > 0
# resize
GlobalCache.resize(100)
assert cache.cache.total_bytes < 100
def test_set_options(comm):
with CurrentMPIComm.enter(comm):
with set_options(global_cache_size=5e9, dask_chunk_size=75):
s = UniformCatalog(1000, 1.0)
# check cache size
cache = GlobalCache.get()
assert cache.cache.available_bytes == 5e9
# check chunk size
assert s['Position'].chunks[0][0] == 75
s = UniformCatalog(1000, 1.0)
assert s['Position'].chunks[0][0] == s.size
def test_set_options(comm):
with CurrentMPIComm.enter(comm):
with set_options(global_cache_size=5e9, dask_chunk_size=75):
s = UniformCatalog(1000, 1.0)
# check cache size
cache = GlobalCache.get()
assert cache.cache.available_bytes == 5e9
# check chunk size
assert s['Position'].chunks[0][0] == 75
s = UniformCatalog(1000, 1.0)
assert s['Position'].chunks[0][0] == s.size
def test_cache(comm):
cat = UniformCatalog(nbar=10000, BoxSize=1.0, comm=comm)
cat['test'] = cat['Position'] ** 5
test = cat['test'].compute()
# cache should no longer be empty
cache = GlobalCache.get()
assert cache.cache.total_bytes > 0
# resize
cache.cache.available_bytes = 100
cache.cache.shrink()
assert cache.cache.total_bytes < 100
def compute(self, *args, **kwargs):
"""
Our version of :func:`dask.compute` that computes
multiple delayed dask collections at once.
This should be called on the return value of :func:`read`
to converts any dask arrays to numpy arrays.
This uses the global cache as controlled by
:class:`nbodykit.GlobalCache` to cache dask task computations.
The default size is controlled by the ``global_cache_size`` global
option; see :class:`set_options`. To set the size, see
:func:`nbodykit.GlobalCache.resize`.
.. note::
If the :attr:`base` attribute is set, ``compute()``
will called using :attr:`base` instead of ``self``.
Parameters
-----------
args : object
Any number of objects. If the object is a dask
collection, it's computed and the result is returned.
Otherwise it's passed through unchanged.
Notes
-----
The dask default optimizer induces too many (unnecesarry)
IO calls -- we turn this off feature off by default. Eventually we
want our own optimizer probably.
"""
import dask
# return the base compute if it exists
if self.base is not None:
return self.base.compute(*args, **kwargs)
# do not optimize graph (can lead to slower optimizations)
kwargs.setdefault('optimize_graph', False)
# compute using global cache
with GlobalCache.get():
toret = dask.compute(*args, **kwargs)
# do not return tuples of length one
if len(toret) == 1: toret = toret[0]
return toret
def compute(self, *args, **kwargs):
"""
Our version of :func:`dask.compute` that computes
multiple delayed dask collections at once.
This should be called on the return value of :func:`read`
to converts any dask arrays to numpy arrays.
This uses the global cache as controlled by
:class:`nbodykit.GlobalCache` to cache dask task computations.
The default size is controlled by the ``global_cache_size`` global
option; see :class:`set_options`. To set the size, see
:func:`nbodykit.GlobalCache.resize`.
.. note::
If the :attr:`base` attribute is set, ``compute()``
will called using :attr:`base` instead of ``self``.
Parameters
-----------
args : object
Any number of objects. If the object is a dask
collection, it's computed and the result is returned.
Otherwise it's passed through unchanged.
"""
import dask
# return the base compute if it exists
if self.base is not None:
return self.base.compute(*args, **kwargs)
# compute using global cache
with GlobalCache.get():
toret = dask.compute(*args, **kwargs)
# do not return tuples of length one
if len(toret) == 1: toret = toret[0]
return toret
def test_cache_size():
with set_options(global_cache_size=100):
cache = GlobalCache.get()
assert cache.cache.available_bytes == 100
def test_cache_size():
with set_options(global_cache_size=100):
cache = GlobalCache.get()
assert cache.cache.available_bytes == 100