def test_start_local_cupyonly(hdf5_ds_1):
cudas = detect()['cudas']
# Make sure we have enough partitions
hdf5_ds_1.set_num_cores(len(cudas))
mask = _mk_random(size=(16, 16))
with hdf5_ds_1.get_reader().get_h5ds() as h5ds:
data = h5ds[:]
expected = _naive_mask_apply([mask], data)
spec = cluster_spec(cpus=(), cudas=cudas, has_cupy=True)
with DaskJobExecutor.make_local(spec=spec) as executor:
ctx = api.Context(executor=executor)
# Uses ApplyMasksUDF, which supports CuPy
analysis = ctx.create_mask_analysis(
dataset=hdf5_ds_1, factories=[lambda: mask]
)
results = ctx.run(analysis)
udf_res = ctx.run_udf(udf=DebugDeviceUDF(), dataset=hdf5_ds_1)
# No CPU compute resources
with pytest.raises(RuntimeError):
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('numpy',)), dataset=hdf5_ds_1)
cuda_res = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda',)), dataset=hdf5_ds_1)
assert np.allclose(
results.mask_0.raw_data,
expected
)
found = {}
for val in udf_res['device_id'].data[0].values():
print(val)
# no CPU
assert val["cpu"] is None
# Register which GPUs got work
found[val["cuda"]] = True
for val in cuda_res['device_id'].data[0].values():
print(val)
# no CPU
assert val["cpu"] is None
# Register which GPUs got work
found[val["cuda"]] = True
for val in udf_res['backend'].data[0].values():
# use CuPy
print(val)
assert 'cupy' in val
for val in cuda_res['backend'].data[0].values():
# no CuPy, i.e. NumPy
print(val)
assert 'numpy' in val
# Test if each GPU got work. We have to see if this
# actually works always since this depends on the scheduler behavior
assert set(found.keys()) == set(cudas)
assert np.all(udf_res['device_class'].data == 'cuda')
assert np.allclose(udf_res['on_device'].data, data.sum(axis=(0, 1)))
def _create_local_executor(self):
cores = psutil.cpu_count(logical=False)
if cores is None:
cores = 2
return DaskJobExecutor.make_local(
cluster_kwargs={"threads_per_worker": 1, "n_workers": cores}
)
def get_libertem_executor(self):
cores = psutil.cpu_count(logical=False)
if cores is None:
cores = 2
executor = DaskJobExecutor.make_local(cluster_kwargs={
"threads_per_worker": 1,
"n_workers": cores
})
return AsyncAdapter(wrapped=executor)
def test_start_local_cpuonly(hdf5_ds_1):
# We don't use all since that might be too many
cpus = (0, 1)
hdf5_ds_1.set_num_cores(len(cpus))
mask = _mk_random(size=(16, 16))
with hdf5_ds_1.get_reader().get_h5ds() as h5ds:
data = h5ds[:]
expected = _naive_mask_apply([mask], data)
spec = cluster_spec(cpus=cpus, cudas=(), has_cupy=False)
with DaskJobExecutor.make_local(spec=spec) as executor:
ctx = api.Context(executor=executor)
analysis = ctx.create_mask_analysis(
dataset=hdf5_ds_1, factories=[lambda: mask]
)
results = ctx.run(analysis)
udf_res = ctx.run_udf(udf=DebugDeviceUDF(), dataset=hdf5_ds_1)
# No CuPy resources
with pytest.raises(RuntimeError):
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy',)), dataset=hdf5_ds_1)
assert np.allclose(
results.mask_0.raw_data,
expected
)
found = {}
for val in udf_res['device_id'].data[0].values():
print(val)
# no CUDA
assert val["cuda"] is None
found[val["cpu"]] = True
for val in udf_res['backend'].data[0].values():
print(val)
# no CUDA
assert 'numpy' in val
# Each CPU got work. We have to see if this
# actually works always since this depends on the scheduler behavior
assert set(found.keys()) == set(cpus)
assert np.all(udf_res['device_class'].data == 'cpu')
assert np.allclose(udf_res['on_device'].data, data.sum(axis=(0, 1)))
def test_preload(hdf5_ds_1):
# We don't use all since that might be too many
cpus = (0, 1)
hdf5_ds_1.set_num_cores(len(cpus))
class CheckEnvUDF(NoOpUDF):
def process_tile(self, tile):
assert os.environ['LT_TEST_1'] == 'hello'
assert os.environ['LT_TEST_2'] == 'world'
preloads = (
"import os; os.environ['LT_TEST_1'] = 'hello'",
"import os; os.environ['LT_TEST_2'] = 'world'",
)
spec = cluster_spec(cpus=cpus, cudas=(), has_cupy=False, preload=preloads)
with DaskJobExecutor.make_local(spec=spec) as executor:
ctx = api.Context(executor=executor)
ctx.run_udf(udf=CheckEnvUDF(), dataset=hdf5_ds_1)
def test_start_local_cudaonly(hdf5_ds_1):
cudas = detect()['cudas']
# Make sure we have enough partitions
hdf5_ds_1.set_num_cores(len(cudas))
with hdf5_ds_1.get_reader().get_h5ds() as h5ds:
data = h5ds[:]
spec = cluster_spec(cpus=(), cudas=cudas, has_cupy=False)
with DaskJobExecutor.make_local(spec=spec) as executor:
ctx = api.Context(executor=executor)
udf_res = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', )),
dataset=hdf5_ds_1)
# No CPU compute resources
with pytest.raises(RuntimeError):
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('numpy', )),
dataset=hdf5_ds_1)
# No ndarray (CuPy) resources
with pytest.raises(RuntimeError):
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', )),
dataset=hdf5_ds_1)
found = {}
for val in udf_res['device_id'].data[0].values():
print(val)
# no CPU
assert val["cpu"] is None
# Register which GPUs got work
found[val["cuda"]] = True
for val in udf_res['backend'].data[0].values():
print(val)
# CUDA, but no CuPy, i.e. use NumPy
assert 'numpy' in val
# Test if each GPU got work. We have to see if this
# actually works always since this depends on the scheduler behavior
assert set(found.keys()) == set(cudas)
assert np.all(udf_res['device_class'].data == 'cuda')
assert np.allclose(udf_res['on_device'].data, data.sum(axis=(0, 1)))
def dask_executor():
sync_executor = DaskJobExecutor.make_local()
yield sync_executor
sync_executor.close()
async def aexecutor():
sync_executor = DaskJobExecutor.make_local()
executor = AsyncAdapter(wrapped=sync_executor)
yield executor
await executor.close()
def _create_local_executor(self):
return DaskJobExecutor.make_local()
def make_with(cls, executor_spec: ExecutorSpecType, *args,
**kwargs) -> 'Context':
'''
Create a Context with a specific kind of executor.
.. versionadded:: 0.9.0
This simplifies creating a :class:`Context` for a number of common executor
choices. See :ref:`executors` for general information on executors.
Parameters
----------
executor_spec:
A string identifier for executor variants:
"synchronous", "inline":
Use a single-process, single-threaded
:class:`~libertem.executor.inline.InlineJobExecutor`
"threads":
Use a multi-threaded :class:`~libertem.executor.concurrent.ConcurrentJobExecutor`
"dask-integration":
Use a JobExecutor that is compatible with the currently active Dask scheduler.
See :func:`~libertem.executor.integration.get_dask_integration_executor` for
more information.
"dask-make-default":
Create a local :code:`dask.distributed` cluster and client
using :meth:`~libertem.executor.dask.DaskJobExecutor.make_local`, similar to
the default behaviour of :code:`Context()` called with no arguments.
However, the Client will be set as the default Dask scheduler and will
persist after the LiberTEM Context closes, which is suitable for downstream
computation using :code:`dask.distributed`.
"delayed":
Create a :class:`~libertem.executor.delayed.DelayedJobExecutor` which performs
computation using `dask.delayed <https://docs.dask.org/en/stable/delayed.html>`_.
This functionality is highly experimental at this time, see
:ref:`delayed_udfs` for more information.
*args, **kwargs
Passed to :class:`Context`.
Returns
-------
Instance of :class:`Context` using a new instance of the specified executor.
'''
if executor_spec in ('synchronous', 'inline'):
executor = InlineJobExecutor()
elif executor_spec == 'threads':
executor = ConcurrentJobExecutor.make_local()
elif executor_spec == 'dask-integration':
executor = get_dask_integration_executor()
elif executor_spec == 'dask-make-default':
executor = DaskJobExecutor.make_local(
client_kwargs={"set_as_default": True})
elif executor_spec == 'delayed':
executor = DelayedJobExecutor()
else:
raise ValueError(
f'Argument `executor_spec` is {executor_spec}. Allowed are '
f'synchronous", "inline", "threads", "dask-integration" or "dask-make-default".'
)
return cls(executor=executor, *args, **kwargs)
def get_libertem_executor(self):
executor = DaskJobExecutor.make_local()
return AsyncAdapter(wrapped=executor)