From 328642b495a530fee00c2508f0e1c7d4748e29c2 Mon Sep 17 00:00:00 2001 From: Gerrit Holl Date: Tue, 24 Oct 2023 09:46:50 +0200 Subject: [PATCH] Added option to run CLI with dask profiler For the cli tool recently added in #184, this adds an option to profile the run via dask and visualise the output as a bokeh plot. --- trollflow2/cli.py | 8 ++++---- trollflow2/tests/test_cli.py | 29 ++++++++++++++++++++++++----- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/trollflow2/cli.py b/trollflow2/cli.py index f5dd67ae..724bd48d 100644 --- a/trollflow2/cli.py +++ b/trollflow2/cli.py @@ -56,11 +56,11 @@ def cli(args=None): if args.dask_profiler: profs.append(stack.enter_context(dask.diagnostics.Profiler())) if args.dask_resource_profiler: - profs.append(stack.enter_context(dask.diagnostics.ResourceProfiler(dt=args.dask_resource_profiles))) + profs.append(stack.enter_context(dask.diagnostics.ResourceProfiler(dt=args.dask_resource_profiler))) process_files(args.files, json.loads(args.metadata), args.product_list, produced_files) - if args.dask_profiler: - dask.diagnostics.visualize( - [profs], show=False, save=True, filename=args.dask_profiler) + if args.dask_profiler: + dask.diagnostics.visualize( + profs, show=False, save=True, filename=args.dask_profiler) def _read_log_config(args): diff --git a/trollflow2/tests/test_cli.py b/trollflow2/tests/test_cli.py index a0b760d9..a370153b 100644 --- a/trollflow2/tests/test_cli.py +++ b/trollflow2/tests/test_cli.py @@ -37,6 +37,16 @@ """ +@pytest.fixture +def product_list_filename(tmp_path): + """Filename for a test product list, with contents.""" + product_list = "my_product_list.yaml" + filename = os.fspath(tmp_path / product_list) + with open(filename, "w") as fd: + fd.write(yaml_test_noop) + return filename + + def test_arg_parsing(): """Test parsing args.""" product_list = "my_product_list.yaml" @@ -91,13 +101,9 @@ def test_cli_raises_an_error_when_product_list_is_empty(tmp_path, caplog, empty_ assert "check YAML file" in caplog.text -def test_cli_starts_processing_when_files_are_provided(tmp_path): +def test_cli_starts_processing_when_files_are_provided(tmp_path, product_list_filename): """Test that the cli start processing when files are provided.""" - product_list = "my_product_list.yaml" files = ["file1", "file2"] - product_list_filename = os.fspath(tmp_path / product_list) - with open(product_list_filename, "w") as fd: - fd.write(yaml_test_noop) from trollflow2.launcher import process_files new_process = mock.Mock(wraps=process_files) mda = {"dish": "pizza"} @@ -105,3 +111,16 @@ def test_cli_starts_processing_when_files_are_provided(tmp_path): with mock.patch("trollflow2.cli.Queue") as q_mock: cli(["-p", os.fspath(product_list_filename), "-m", json.dumps(mda), *files]) new_process.assert_called_once_with(files, mda, product_list_filename, q_mock.return_value) + + +def test_cli_dask_profiler(product_list_filename, tmp_path): + """Test that dask profiles are written.""" + from trollflow2.launcher import process_files + new_process = mock.Mock(wraps=process_files) + proffile = tmp_path / "dask-prof.html" + with (mock.patch("trollflow2.cli.process_files", new=new_process), + mock.patch("trollflow2.cli.Queue")): + cli(["-p", os.fspath(product_list_filename), "--dask-profiler", + os.fspath(proffile), "--dask-resource-profiler", "0.1", + "-m", json.dumps({"food": "soy"}), "aquafaba", "tempeh"]) + assert proffile.exists()