def test_sig_slice(lt_ctx, backend, udf_class, tileshape, success):
if backend == 'cupy':
d = detect()
cudas = detect()['cudas']
if not d['cudas'] or not d['has_cupy']:
pytest.skip("No CUDA device or no CuPy, skipping CuPy test")
data = _mk_random(size=(30, 3, 7), dtype="float32")
ref_res = data.sum(axis=(-1, -2))
dataset = MemoryDataSet(data=data,
tileshape=tileshape,
num_partitions=2,
sig_dims=2)
try:
if backend == 'cupy':
set_use_cuda(cudas[0])
udf = udf_class()
if success:
res = lt_ctx.run_udf(udf=udf, dataset=dataset)
assert np.allclose(res['checksum'].raw_data, ref_res)
else:
with pytest.raises(Exception):
lt_ctx.run_udf(udf=udf, dataset=dataset)
finally:
set_use_cpu(0)
def test_multi_mask_stack_dense(lt_ctx, TYPE, backend):
if backend == 'cupy':
d = detect()
cudas = detect()['cudas']
if not d['cudas'] or not d['has_cupy']:
pytest.skip("No CUDA device or no CuPy, skipping CuPy test")
try:
if backend == 'cupy':
set_use_cuda(cudas[0])
data = _mk_random(size=(16, 16, 16, 16), dtype="<u2")
masks = _mk_random(size=(2, 16, 16))
expected = _naive_mask_apply(masks, data)
dataset = MemoryDataSet(data=data,
tileshape=(4 * 4, 4, 4),
num_partitions=2)
analysis = lt_ctx.create_mask_analysis(
dataset=dataset,
factories=lambda: masks,
mask_count=2,
)
analysis.TYPE = TYPE
results = lt_ctx.run(analysis)
assert np.allclose(
results.mask_0.raw_data,
expected[0],
)
assert np.allclose(
results.mask_1.raw_data,
expected[1],
)
finally:
set_use_cpu(0)
def test_noncontiguous_tiles(lt_ctx, backend):
if backend == 'cupy':
d = detect()
cudas = detect()['cudas']
if not d['cudas'] or not d['has_cupy']:
pytest.skip("No CUDA device or no CuPy, skipping CuPy test")
data = _mk_random(size=(30, 3, 7), dtype="float32")
dataset = MemoryDataSet(data=data,
tileshape=(3, 2, 2),
num_partitions=2,
sig_dims=2)
try:
if backend == 'cupy':
set_use_cuda(cudas[0])
udf = ReshapedViewUDF()
res = lt_ctx.run_udf(udf=udf, dataset=dataset)
partition = next(dataset.get_partitions())
p_udf = udf.copy_for_partition(partition=partition, roi=None)
# Enabling debug=True checks for disjoint cache keys
UDFRunner([p_udf], debug=True).run_for_partition(
partition=partition,
roi=None,
corrections=None,
env=Environment(threads_per_worker=1),
)
finally:
set_use_cpu(0)
assert np.all(res["sigbuf"].data == 1)
def test_holo_reconstruction(lt_ctx, backend):
if backend == 'cupy':
d = detect()
cudas = detect()['cudas']
if not d['cudas'] or not d['has_cupy']:
pytest.skip("No CUDA device or no CuPy, skipping CuPy test")
# Prepare image parameters and mesh
nx, ny = (5, 7)
sx, sy = (64, 64)
slice_crop = (slice(None),
slice(None),
slice(sx // 4, sx // 4 * 3),
slice(sy // 4, sy // 4 * 3))
lnx = np.arange(nx)
lny = np.arange(ny)
lsx = np.arange(sx)
lsy = np.arange(sy)
mnx, mny, msx, msy = np.meshgrid(lnx, lny, lsx, lsy)
# Prepare phase image
phase_ref = np.pi * msx * (mnx.max() - mnx) * mny / sx**2 \
+ np.pi * msy * mnx * (mny.max() - mny) / sy**2
# Generate holograms
holo = np.zeros_like(phase_ref)
ref = np.zeros_like(phase_ref)
for i in range(nx):
for j in range(ny):
holo[j, i, :, :] = hologram_frame(np.ones((sx, sy)), phase_ref[j, i, :, :])
ref[j, i, :, :] = hologram_frame(np.ones((sx, sy)), np.zeros((sx, sy)))
# Prepare LT datasets and do reconstruction
dataset_holo = MemoryDataSet(data=holo, tileshape=(ny, sx, sy),
num_partitions=2, sig_dims=2)
dataset_ref = MemoryDataSet(data=ref, tileshape=(ny, sx, sy),
num_partitions=1, sig_dims=2)
sb_position = [11, 6]
sb_size = 6.26498204
holo_job = HoloReconstructUDF(out_shape=(sx, sy),
sb_position=sb_position,
sb_size=sb_size)
try:
if backend == 'cupy':
set_use_cuda(cudas[0])
w_holo = lt_ctx.run_udf(dataset=dataset_holo, udf=holo_job)['wave'].data
w_ref = lt_ctx.run_udf(dataset=dataset_ref, udf=holo_job)['wave'].data
finally:
set_use_cpu(0)
w = w_holo / w_ref
phase = np.angle(w)
assert np.allclose(phase_ref[slice_crop], phase[slice_crop], rtol=0.12)
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 local_cluster_url():
"""
Shared dask cluster, can be used repeatedly by different executors.
This allows numba caching across tests, without sharing the executor,
for example
"""
cluster_port = find_unused_port()
devices = detect()
spec = cluster_spec(
# Only use at most 2 CPUs and 1 GPU
cpus=devices['cpus'][:2],
cudas=devices['cudas'][:1],
has_cupy=devices['has_cupy'])
cluster_kwargs = {
'silence_logs': logging.WARN,
'scheduler': {
'cls': Scheduler,
'options': {
'port': cluster_port
},
},
}
cluster = dd.SpecCluster(workers=spec, **(cluster_kwargs or {}))
yield 'tcp://localhost:%d' % cluster_port
cluster.close()
def make_local(cls, spec=None, cluster_kwargs=None, client_kwargs=None):
"""
Spin up a local dask cluster
interesting cluster_kwargs:
threads_per_worker
n_workers
Returns
-------
DaskJobExecutor
the connected JobExecutor
"""
if spec is None:
spec = cluster_spec(**detect())
if client_kwargs is None:
client_kwargs = {}
if client_kwargs.get('set_as_default') is None:
client_kwargs['set_as_default'] = False
if cluster_kwargs is None:
cluster_kwargs = {}
if cluster_kwargs.get('silence_logs') is None:
cluster_kwargs['silence_logs'] = logging.WARN
cluster = dd.SpecCluster(workers=spec, **(cluster_kwargs or {}))
client = dd.Client(cluster, **(client_kwargs or {}))
return cls(client=client, is_local=True, lt_resources=True)
def test_detect_broken(monkeypatch):
try:
import cupy
except Exception:
pytest.skip("Importable CuPy required for this test.")
# CuPy can throw all kinds of exceptions, depending on what exactly goes wrong
class FunkyException(Exception):
pass
def badfunc(*args, **kwargs):
raise FunkyException()
monkeypatch.setattr(cupy, 'array', badfunc)
monkeypatch.setattr(cupy, 'zeros', badfunc)
# Self test
with pytest.raises(FunkyException):
cupy.array(cupy.zeros(1,))
monkeypatch.delitem(sys.modules, 'libertem.utils.devices', raising=False)
from libertem.utils.devices import detect
with pytest.warns(RuntimeWarning, match='FunkyException'):
result = detect()
assert not result['has_cupy']
def test_start_local_default(hdf5_ds_1, local_cluster_ctx):
mask = _mk_random(size=(16, 16))
d = detect()
cudas = d['cudas']
with hdf5_ds_1.get_reader().get_h5ds() as h5ds:
data = h5ds[:]
expected = _naive_mask_apply([mask], data)
ctx = local_cluster_ctx
analysis = ctx.create_mask_analysis(dataset=hdf5_ds_1,
factories=[lambda: mask])
num_cores_ds = ctx.load('memory', data=np.zeros((2, 3, 4, 5)))
workers = ctx.executor.get_available_workers()
cpu_count = len(workers.has_cpu())
gpu_count = len(workers.has_cuda())
assert num_cores_ds._cores == max(cpu_count, gpu_count)
# Based on ApplyMasksUDF, which is CuPy-enabled
hybrid = ctx.run(analysis)
_ = ctx.run_udf(udf=DebugDeviceUDF(), dataset=hdf5_ds_1)
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'numpy')),
dataset=hdf5_ds_1)
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', 'numpy')),
dataset=hdf5_ds_1)
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'cuda', 'numpy')),
dataset=hdf5_ds_1)
if cudas:
cuda_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', 'numpy')),
dataset=hdf5_ds_1,
backends=('cuda', ))
if d['has_cupy']:
cupy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy',
'numpy')),
dataset=hdf5_ds_1,
backends=('cupy', ))
else:
with pytest.raises(RuntimeError):
cupy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy',
'numpy')),
dataset=hdf5_ds_1,
backends=('cupy', ))
cupy_only = None
numpy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('numpy', )),
dataset=hdf5_ds_1)
assert np.allclose(hybrid.mask_0.raw_data, expected)
if cudas:
assert np.all(cuda_only['device_class'].data == 'cuda')
if cupy_only is not None:
assert np.all(cupy_only['device_class'].data == 'cuda')
assert np.all(numpy_only['device_class'].data == 'cpu')
def test_udf_noncontiguous_tiles(lt_ctx, backend, benchmark):
if backend == 'cupy':
d = detect()
cudas = detect()['cudas']
if not d['cudas'] or not d['has_cupy']:
pytest.skip("No CUDA device or no CuPy, skipping CuPy test")
data = np.zeros(shape=(30, 3, 256), dtype="float32")
dataset = MemoryDataSet(data=data,
tileshape=(3, 2, 2),
num_partitions=2,
sig_dims=2)
try:
if backend == 'cupy':
set_use_cuda(cudas[0])
udf = NoopSigUDF()
res = benchmark(lt_ctx.run_udf, udf=udf, dataset=dataset)
finally:
set_use_cpu(0)
assert np.all(res["sigbuf"].data == 0)
def get_config(self):
detected_devices = devices.detect()
return {
"version": libertem.__version__,
"resultFileFormats": ResultFormatRegistry.get_available_formats(),
"revision": libertem.revision,
"localCores": self.get_local_cores(),
"devices": detected_devices,
"cwd": os.getcwd(),
# '/' works on Windows, too.
"separator": '/'
}
async def put(self):
# TODO: extract json request data stuff into mixin?
request_data = tornado.escape.json_decode(self.request.body)
connection = request_data['connection']
pool = AsyncAdapter.make_pool()
if connection["type"].lower() == "tcp":
try:
sync_executor = await sync_to_async(partial(
DaskJobExecutor.connect,
scheduler_uri=connection['address'],
),
pool=pool)
except Exception as e:
msg = Message(self.state).cluster_conn_error(msg=str(e))
log_message(msg)
self.write(msg)
return None
elif connection["type"].lower() == "local":
devices = detect()
options = {"local_directory": self.state.get_local_directory()}
if "numWorkers" in connection:
devices["cpus"] = range(connection["numWorkers"])
devices["cudas"] = connection.get("cudas", [])
sync_executor = await sync_to_async(partial(
DaskJobExecutor.make_local,
spec=cluster_spec(**devices,
options=options,
preload=self.state.get_preload())),
pool=pool)
else:
raise ValueError("unknown connection type")
executor = AsyncAdapter(wrapped=sync_executor, pool=pool)
await self.state.executor_state.set_executor(executor, request_data)
await self.state.dataset_state.verify()
datasets = await self.state.dataset_state.serialize_all()
msg = Message(self.state).initial_state(
jobs=self.state.job_state.serialize_all(),
datasets=datasets,
analyses=self.state.analysis_state.serialize_all(),
compound_analyses=self.state.compound_analysis_state.serialize_all(
),
)
log_message(msg)
# FIXME: don't broadcast, only send to the websocket that matches this HTTP connection
# (is this even possible?)
self.event_registry.broadcast_event(msg)
await self.engine.send_existing_job_results()
self.write({
"status": "ok",
"connection": connection,
})
async def put(self):
# TODO: extract json request data stuff into mixin?
request_data = tornado.escape.json_decode(self.request.body)
connection = request_data['connection']
if connection["type"].lower() == "tcp":
try:
sync_executor = await sync_to_async(partial(DaskJobExecutor.connect,
scheduler_uri=connection['address'],
))
except Exception as e:
msg = Message(self.state).cluster_conn_error(msg=str(e))
log_message(msg)
self.write(msg)
return None
elif connection["type"].lower() == "local":
devices = detect()
options = {
"local_directory": self.state.get_local_directory()
}
if "numWorkers" in connection:
devices["cpus"] = range(connection["numWorkers"])
# Deactivate GPU support in local cluster until GUI allows deactivation
# to not interfere with other applications using the GPU
# FIXME implement GUI interface https://github.com/LiberTEM/LiberTEM/issues/803
devices["cudas"] = []
sync_executor = await sync_to_async(
partial(
DaskJobExecutor.make_local,
spec=cluster_spec(**devices, options=options),
)
)
else:
raise ValueError("unknown connection type")
executor = AsyncAdapter(wrapped=sync_executor)
await self.state.executor_state.set_executor(executor, request_data)
await self.state.dataset_state.verify()
datasets = await self.state.dataset_state.serialize_all()
msg = Message(self.state).initial_state(
jobs=self.state.job_state.serialize_all(),
datasets=datasets, analyses=self.state.analysis_state.serialize_all(),
compound_analyses=self.state.compound_analysis_state.serialize_all(),
)
log_message(msg)
# FIXME: don't broadcast, only send to the websocket that matches this HTTP connection
# (is this even possible?)
self.event_registry.broadcast_event(msg)
await self.send_existing_job_results()
self.write({
"status": "ok",
"connection": connection,
})
def get_available_workers(self):
resources = {"compute": 1, "CPU": 1}
# We don't know at this time,
# but assume one worker per CPU
devices = detect()
return WorkerSet([
Worker(
name='delayed',
host='localhost',
resources=resources,
nthreads=len(devices['cpus']),
)
])
def test_multi_masks(lt_ctx, backend):
if backend == 'cupy':
d = detect()
cudas = detect()['cudas']
if not d['cudas'] or not d['has_cupy']:
pytest.skip("No CUDA device or no CuPy, skipping CuPy test")
try:
if backend == 'cupy':
set_use_cuda(cudas[0])
data = _mk_random(size=(16, 16, 16, 16), dtype="<u2")
mask0 = _mk_random(size=(16, 16))
mask1 = sp.csr_matrix(_mk_random(size=(16, 16)))
mask2 = sparse.COO.from_numpy(_mk_random(size=(16, 16)))
expected = _naive_mask_apply([mask0, mask1, mask2], data)
dataset = MemoryDataSet(data=data,
tileshape=(4 * 4, 4, 4),
num_partitions=2)
analysis = lt_ctx.create_mask_analysis(
dataset=dataset,
factories=[lambda: mask0, lambda: mask1, lambda: mask2],
)
results = lt_ctx.run(analysis)
assert np.allclose(
results.mask_0.raw_data,
expected[0],
)
assert np.allclose(
results.mask_1.raw_data,
expected[1],
)
assert np.allclose(
results.mask_2.raw_data,
expected[2],
)
finally:
set_use_cpu(0)
def make_local(cls):
"""
Create a local ConcurrentJobExecutor backed by
a :class:`python:concurrent.futures.ThreadPoolExecutor`
Returns
-------
ConcurrentJobExecutor
the connected JobExecutor
"""
devices = detect()
n_threads = len(devices['cpus'])
client = concurrent.futures.ThreadPoolExecutor(max_workers=n_threads)
return cls(client=client, is_local=True)
def main(kind, scheduler, local_directory, n_cpus, cudas, has_cupy, name,
log_level, preload: Tuple[str]):
from libertem.common.threading import set_num_threads_env
with set_num_threads_env(1):
from libertem.utils.devices import detect
from libertem.cli_tweaks import console_tweaks
from libertem.executor.dask import cli_worker
console_tweaks()
if kind != 'dask':
raise NotImplementedError(
f"Currently only worker kind 'dask' is implemented, got {kind}."
)
numeric_level = getattr(logging, log_level.upper(), None)
if not isinstance(numeric_level, int):
raise click.UsageError(
f'Invalid log level: {log_level}.\n{log_values}')
defaults = detect()
if n_cpus is None:
cpus = list(defaults['cpus'])
else:
cpus = list(range(n_cpus))
if cudas == '':
cudas = []
elif cudas is None:
cudas = list(defaults['cudas'])
else:
cudas = list(map(int, cudas.split(',')))
if has_cupy is None:
has_cupy = defaults['has_cupy']
if not name:
name = socket.gethostname()
cli_worker(
scheduler=scheduler,
local_directory=local_directory,
cpus=cpus,
cudas=cudas,
has_cupy=has_cupy,
name=name,
log_level=numeric_level,
preload=preload,
)
def test_detect_notfound(monkeypatch):
monkeypatch.delitem(sys.modules, 'cupy', raising=False)
monkeypatch.delitem(sys.modules, 'cupy.cuda', raising=False)
monkeypatch.setattr(builtins, '__import__', monkey_notfound)
# Self test
with pytest.raises(ModuleNotFoundError):
import cupy # NOQA: F401
monkeypatch.delitem(sys.modules, 'libertem.utils.devices', raising=False)
from libertem.utils.devices import detect
result = detect()
assert not result['has_cupy']
def test_detect_error(monkeypatch):
monkeypatch.delitem(sys.modules, 'cupy', raising=False)
monkeypatch.delitem(sys.modules, 'cupy.cuda', raising=False)
monkeypatch.setattr(builtins, '__import__', monkey_importerror)
# Self test
with pytest.raises(ImportError):
import cupy # NOQA: F401
monkeypatch.delitem(sys.modules, 'libertem.utils.devices', raising=False)
with pytest.warns(RuntimeWarning, match="Mocked import error cupy"):
from libertem.utils.devices import detect
result = detect()
assert not result['has_cupy']
def test_start_local_default(hdf5_ds_1):
mask = _mk_random(size=(16, 16))
d = detect()
cudas = d['cudas']
with hdf5_ds_1.get_reader().get_h5ds() as h5ds:
data = h5ds[:]
expected = _naive_mask_apply([mask], data)
with api.Context() as ctx:
analysis = ctx.create_mask_analysis(dataset=hdf5_ds_1,
factories=[lambda: mask])
# Based on ApplyMasksUDF, which is CuPy-enabled
hybrid = ctx.run(analysis)
_ = ctx.run_udf(udf=DebugDeviceUDF(), dataset=hdf5_ds_1)
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'numpy')),
dataset=hdf5_ds_1)
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', 'numpy')),
dataset=hdf5_ds_1)
_ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'cuda', 'numpy')),
dataset=hdf5_ds_1)
if cudas:
cuda_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda',
'numpy')),
dataset=hdf5_ds_1,
backends=('cuda', ))
if d['has_cupy']:
cupy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy',
'numpy')),
dataset=hdf5_ds_1,
backends=('cupy', ))
else:
with pytest.raises(RuntimeError):
cupy_only = ctx.run_udf(
udf=DebugDeviceUDF(backends=('cupy', 'numpy')),
dataset=hdf5_ds_1,
backends=('cupy', ))
cupy_only = None
numpy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('numpy', )),
dataset=hdf5_ds_1)
assert np.allclose(hybrid.mask_0.raw_data, expected)
if cudas:
assert np.all(cuda_only['device_class'].data == 'cuda')
if cupy_only is not None:
assert np.all(cupy_only['device_class'].data == 'cuda')
assert np.all(numpy_only['device_class'].data == 'cpu')
def test_detect_nocuda(monkeypatch):
try:
import cupy
except Exception:
pytest.skip("Importable CuPy required for this test.")
monkeypatch.delattr(sys.modules['cupy'], 'cuda', raising=False)
# Self test
with pytest.raises(Exception):
cupy.cuda
monkeypatch.delitem(sys.modules, 'libertem.utils.devices', raising=False)
with pytest.warns(RuntimeWarning, match="module 'cupy' has no attribute 'cuda'"):
from libertem.utils.devices import detect
result = detect()
assert not result['has_cupy']
def get_available_workers(self):
resources = {"compute": 1, "CPU": 1}
if get_use_cuda() is not None:
resources["CUDA"] = 1
return WorkerSet([
Worker(name='concurrent',
host='localhost',
resources=resources,
nthreads=1)
])
else:
devices = detect()
return WorkerSet([
Worker(
name='concurrent',
host='localhost',
resources=resources,
nthreads=len(devices['cpus']),
)
])
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 make_local(cls, spec=None, cluster_kwargs=None, client_kwargs=None):
"""
Spin up a local dask cluster
interesting cluster_kwargs:
threads_per_worker
n_workers
Returns
-------
DaskJobExecutor
the connected JobExecutor
"""
# Distributed doesn't adjust the event loop policy when being run
# from within pytest as of version 2.21.0. For that reason we
# adjust the policy ourselves here.
adjust_event_loop_policy()
if spec is None:
from libertem.utils.devices import detect
spec = cluster_spec(**detect())
if client_kwargs is None:
client_kwargs = {}
if client_kwargs.get('set_as_default') is None:
client_kwargs['set_as_default'] = False
if cluster_kwargs is None:
cluster_kwargs = {}
if cluster_kwargs.get('silence_logs') is None:
cluster_kwargs['silence_logs'] = logging.WARN
cluster = dd.SpecCluster(workers=spec, **(cluster_kwargs or {}))
client = dd.Client(cluster, **(client_kwargs or {}))
client.wait_for_workers(len(spec))
return cls(client=client, is_local=True, lt_resources=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)))
@pytest.mark.functional
@pytest.mark.skipif(not detect()['cudas'], reason="No CUDA devices")
@pytest.mark.skipif(not has_cupy(), reason="No functional CuPy")
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,
from libertem import api as lt
from libertem.executor.inline import InlineJobExecutor
from libertem.io.dataset.memory import MemoryDataSet
from libertem.masks import circular
from libertem.corrections.coordinates import identity, rotate_deg
from libertem.common.container import MaskContainer
from ptychography.reconstruction.ssb import SSB_UDF, generate_masks, get_results
from ptychography.reconstruction.ssb.trotters import mask_tile_pair
from ptychography.reconstruction.common import wavelength
try:
from libertem.utils.devices import detect, has_cupy
from libertem.common.backend import set_use_cpu, set_use_cuda
use_cupy = detect()['cudas'] and has_cupy()
except ImportError:
use_cupy = False
DATA_PATH = "/storage/holo/clausen/testdata/ER-C-1/projects/ptycho-4.0/data/RefData/slice_00002_thick_0.312293_nm_blocksz.mat"
PARAM_PATH = "/storage/holo/clausen/testdata/ER-C-1/projects/ptycho-4.0/data/slice_00002_thick_0.312293_nm_blocksz.params.json"
has_real = isfile(DATA_PATH) and isfile(PARAM_PATH)
if has_real:
@pytest.fixture(scope="session")
def real_params():
with open(PARAM_PATH) as f:
params = json.load(f)
params['transformation'] = np.array(params['transformation'])
def make_local(cls,
spec: Optional[dict] = None,
cluster_kwargs: Optional[dict] = None,
client_kwargs: Optional[dict] = None,
preload: Optional[Tuple[str]] = None):
"""
Spin up a local dask cluster
Parameters
----------
spec
Dask cluster spec, see
http://distributed.dask.org/en/stable/api.html#distributed.SpecCluster
for more info.
:func:`libertem.utils.devices.detect` allows to detect devices that can be used
with LiberTEM, and :func:`cluster_spec` can be used to create a :code:`spec`
with customized parameters.
cluster_kwargs
Passed to :class:`distributed.SpecCluster`.
client_kwargs
Passed to :class:`distributed.Client`. Pass
:code:`client_kwargs={'set_as_default': True}` to set the Client as the
default Dask scheduler.
preload: Optional[Tuple[str]]
Passed to :func:`cluster_spec` if :code:`spec` is :code:`None`.
Returns
-------
DaskJobExecutor
the connected JobExecutor
"""
# Distributed doesn't adjust the event loop policy when being run
# from within pytest as of version 2.21.0. For that reason we
# adjust the policy ourselves here.
adjust_event_loop_policy()
if spec is None:
from libertem.utils.devices import detect
spec = cluster_spec(**detect(), preload=preload)
else:
if preload is not None:
raise ValueError(
"Passing both spec and preload is not supported. "
"Instead, include preloading specification in the spec")
if client_kwargs is None:
client_kwargs = {}
if client_kwargs.get('set_as_default') is None:
client_kwargs['set_as_default'] = False
if cluster_kwargs is None:
cluster_kwargs = {}
if cluster_kwargs.get('silence_logs') is None:
cluster_kwargs['silence_logs'] = logging.WARN
with set_num_threads_env(n=1):
cluster = dd.SpecCluster(workers=spec, **(cluster_kwargs or {}))
client = dd.Client(cluster, **(client_kwargs or {}))
client.wait_for_workers(len(spec))
is_local = not client_kwargs['set_as_default']
return cls(client=client, is_local=is_local, lt_resources=True)
