How to use the elfi.env.client function in elfi

def run_local_object_cache_test(self, local_store):
        sleep_time = .2
        simfn = get_sleep_simulator(sleep_time)
        sim = elfi.Simulator("sim", simfn, observed=0, store=local_store)
        run_cache_test(sim, sleep_time)
        assert local_store._read_data(sim.id, 0)[0] == 1

        # Test that nodes derived from `sim` benefit from the storing
        summ = elfi.Summary("sum", lambda x : x, sim)
        t0 = timeit.default_timer()
        res = summ.acquire(1).compute()
        td = timeit.default_timer() - t0
        assert td < sleep_time
        assert res[0][0] == 1

        elfi.env.client().shutdown()

def clear_elfi_client():
    elfi.env.client().shutdown()
    elfi.env.set_option(client=None)

sim1 = elfi.Simulator('sim', lambda *args, **kwargs: args[0], p1, observed=1)

    for i in range(10):
        y_prev = y
        t_prev = t

        y = sim1.acquire(N, batch_size=bs).compute()
        t = p1.acquire(N, batch_size=bs).compute()

        if y_prev is not None:
            assert np.all(y != y_prev)
            assert np.all(t != t_prev)

        p1.reset()

    elfi.env.client().shutdown()

def test_reset_specific_scheduler_keys():
    """This test fails if keys are not different"""
    elfi.env.client(n_workers=2, threads_per_worker=1)
    N = 20
    bs = 10

    y = None
    t = None

    p1 = elfi.Prior('p', 'Uniform')
    sim1 = elfi.Simulator('sim', lambda *args, **kwargs: args[0], p1, observed=1)

    for i in range(10):
        y_prev = y
        t_prev = t

        y = sim1.acquire(N, batch_size=bs).compute()
        t = p1.acquire(N, batch_size=bs).compute()

y = sim.acquire(N, batch_size=bs).compute()
        t = p.acquire(N, batch_size=bs).compute()

        if y_prev is not None:
            assert np.all(y != y_prev)
            assert np.all(t != t_prev)

        p_id = p.id
        sim_id = sim.id
        if p_prev_id is not None:
            assert p_id != p_prev_id
            assert sim_id != sim_prev_id

        elfi.new_inference_task()

    elfi.env.client().shutdown()

def test_new_inference_task():
    """This test fails if keys that the dask scheduler gets are not different. We
    run the loop 10 times in trying to get key collisions. The collisions occur
    when dask is unable to clear the key of previous computation in time before
    the next computation with the exact same key comes in. This can happen at least
    with the Distributed scheduler.

    """
    elfi.env.client(n_workers=2, threads_per_worker=1)
    N = 20
    bs = 10

    p_id = None
    sim_id = None
    y = None
    t = None

    for i in range(10):
        p_prev_id = p_id
        sim_prev_id = sim_id
        y_prev = y
        t_prev = t

        p = elfi.Prior('p', 'Uniform', i)
        sim = elfi.Simulator('sim', lambda *args, **kwargs: args[0], p, observed=1)

if == ALL_COMPLETED returns when all tasks completed
        Currently supports only FIRST_COMPLETED.

    Returns
    -------
    tuple : (result, index, unfinished_futures)
    """
    if return_when not in (FIRST_COMPLETED, ALL_COMPLETED):
        raise ValueError("Unknown value for 'return_when'." +
                "Expected {} or {}.".format(FIRST_COMPLETED, ALL_COMPLETED) +
                "Received {}.".format(return_when))

    if return_when == ALL_COMPLETED:
        raise NotImplementedError("Support for ALL_COMPLETED not implemented.")

    client = client or elfi_client()
    futures = client.compute(collections)
    f = dc.as_completed(futures).__next__()
    i = futures.index(f)
    del futures[i]
    res = f.result()
    return res, i, futures

def write(self, output, done_callback=None):
        key = output.key
        # Persist key to client
        d = env.client().persist(output)
        self._persisted[key] = d

        future = d.dask[key]
        if done_callback is not None:
            future.add_done_callback(lambda f: done_callback(key, f))

    @property
    def ncores(self):
        """Total number of cores available in elfi.client."""
        return sum(elfi_client().ncores().values())