Python dask.get() Examples

def test_ephemeral_persisting(zk, dsk2):
    with Persist(zk, name="dsk2", ns="/test/dags", ephemeral=True):
        with Ran() as r:
            assert get(dsk2, 'e') == 10
            assert r.steps == ['e']
        with Ran() as r:
            assert get(dsk2, 's') == 0.4
            assert r.steps == ['f', 's']
        with Ran() as r:
            assert get(dsk2, 's') == 0.4
            assert r.steps == []

        assert loads(zk.get("/test/dags/dsk2/e")[0]) == 10

    with pytest.raises(NoNodeError):
        zk.get("/test/dags/dsk2/e") 

def enable(func):
        """
        Decorator to attach the current MPI communicator to the input
        keyword arguments of ``func``, via the ``comm`` keyword.
        """
        import functools
        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            kwargs.setdefault('comm', None)
            if kwargs['comm'] is None:
                kwargs['comm'] = CurrentMPIComm.get()
            return func(*args, **kwargs)
        return wrapped 

def get(cls):
        """
        Get the default current MPI communicator. The initial value is ``MPI.COMM_WORLD``.
        """
        return cls._stack[-1] 

def get(cls):
        """
        Return the global cache object. The default size is controlled
        by the ``global_cache_size`` global option; see :class:`set_options`.

        Returns
        -------
        cache : :class:`dask.cache.Cache`
            the cache object, as provided by dask
        """
        # if not created, use default cache size
        return _global_cache 

def __init__(self, **kwargs):
        self.old = _global_options.copy()
        for key in sorted(kwargs):
            if key not in _global_options:
                raise KeyError("Option `%s` is not supported" % key)

        _global_options.update(kwargs)

        # resize the global Cache!
        # FIXME: after https://github.com/dask/cachey/pull/12
        if 'global_cache_size' in kwargs:
            cache = GlobalCache.get().cache
            cache.available_bytes = _global_options['global_cache_size']
            cache.shrink() 

def __exit__(self, type, value, traceback):
        _global_options.clear()
        _global_options.update(self.old)

        # resize Cache to original size
        # FIXME: after https://github.com/dask/cachey/pull/12
        cache = GlobalCache.get().cache
        cache.available_bytes = _global_options['global_cache_size']
        cache.shrink() 

def test_persisting(zk, dsk1):
    with Persist(zk, name="dsk1"):
        with Ran() as r:
            assert get(dsk1, 'w') == 6
            assert r.steps == ['z', 'w']
        with Ran() as r:
            assert get(dsk1, 'w') == 6
            assert r.steps == []

        assert loads(zk.get("/epos/dsk1/z")[0]) == 3
        assert loads(zk.get("/epos/dsk1/w")[0]) == 6

    # tests ephemeral=False, znode still exists after context handler
    assert loads(zk.get("/epos/dsk1/w")[0]) == 6 

def test_locking(zk, dsk1):
    with pytest.raises(LockTimeout):
        # two identical locks; the second cannot acquire
        with Lock(zk, name="dsk1", timeout=1), \
                Lock(zk, name="dsk1", timeout=1):
            get(dsk1, 'w')

    # test lock release
    with Lock(zk, name="dsk1"):
        get(dsk1, 'w')
        get(dsk1, 'w') 

def test_ephemeral_locking(zk, dsk2):
    with pytest.raises(LockTimeout):
        with Lock(zk, name="dsk2", timeout=1, ephemeral=True), \
                Lock(zk, name="dsk2", timeout=1, ephemeral=True):
            get(dsk2, 'f')

    with pytest.raises(NoNodeError):
        zk.get("/epos/dsk2") 

def __call__(self, dsk, keys, **kwargs):
        """Compute dask task and keep track of number of times we do so."""
        import dask
        self.total_computes += 1
        if self.total_computes > self.max_computes:
            raise RuntimeError("Too many dask computations were scheduled: "
                               "{}".format(self.total_computes))
        return dask.get(dsk, keys, **kwargs) 

def __call__(self, dsk, keys, **kwargs):
        """Compute dask task and keep track of number of times we do so."""
        import dask
        self.total_computes += 1
        if self.total_computes > self.max_computes:
            raise RuntimeError("Too many dask computations were scheduled: "
                               "{}".format(self.total_computes))
        return dask.get(dsk, keys, **kwargs)