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)