diff --git a/tests/ert/performance_tests/test_obs_and_responses_performance.py b/tests/ert/performance_tests/test_obs_and_responses_performance.py index 096239dc7cf..8c70728b415 100644 --- a/tests/ert/performance_tests/test_obs_and_responses_performance.py +++ b/tests/ert/performance_tests/test_obs_and_responses_performance.py @@ -1,8 +1,6 @@ import datetime import random -import time from dataclasses import dataclass -from textwrap import dedent import memray import numpy as np @@ -228,9 +226,7 @@ def __str__(self) -> str: @dataclass class _ExpectedPerformance: memory_limit_mb: float - time_limit_s: float last_measured_memory_mb: float | None = None # For bookkeeping - last_measured_time: float | None = None # For bookkeeping @dataclass @@ -262,15 +258,11 @@ class _Benchmark: ), expected_join_performance=_ExpectedPerformance( memory_limit_mb=100, - time_limit_s=2, - last_measured_time=0.03, last_measured_memory_mb=17, ), expected_update_performance=_ExpectedPerformance( - last_measured_time=0.09, last_measured_memory_mb=7.13, memory_limit_mb=100, - time_limit_s=2, ), ), _Benchmark( @@ -288,14 +280,10 @@ class _Benchmark: ), expected_join_performance=_ExpectedPerformance( memory_limit_mb=1500, - time_limit_s=30, - last_measured_time=2.78, last_measured_memory_mb=1027, ), expected_update_performance=_ExpectedPerformance( memory_limit_mb=3100, - time_limit_s=30, - last_measured_time=5.34, last_measured_memory_mb=2230, ), ), @@ -314,15 +302,11 @@ class _Benchmark: ), expected_join_performance=_ExpectedPerformance( memory_limit_mb=4500, - time_limit_s=45, last_measured_memory_mb=1710, - last_measured_time=3.59, ), expected_update_performance=_ExpectedPerformance( memory_limit_mb=4000, - time_limit_s=40, last_measured_memory_mb=3088, - last_measured_time=6.23, ), ), _Benchmark( @@ -340,14 +324,10 @@ class _Benchmark: ), expected_join_performance=_ExpectedPerformance( memory_limit_mb=3300, - time_limit_s=55, - last_measured_time=3.96, last_measured_memory_mb=1715, ), expected_update_performance=_ExpectedPerformance( memory_limit_mb=4500, - time_limit_s=50, - last_measured_time=6.98, last_measured_memory_mb=3115, ), ), @@ -367,39 +347,28 @@ class _Benchmark: # ), # expected_join_performance=_ExpectedPerformance( # memory_limit_mb=55000, - # time_limit_s=60, - # last_measured_time=16.03, # last_measured_memory_mb=48504, # ), # expected_update_performance=_ExpectedPerformance( # memory_limit_mb=55000, - # time_limit_s=60, - # last_measured_time=13, # last_measured_memory_mb=45315, # ), # ), ] -@pytest.mark.parametrize( - "config, expected_performance, alias", - [ - pytest.param( +@pytest.fixture( + params=[ + ( + b.alias, b.config, b.expected_join_performance, - f"{b.alias}:{b.config!s}", - id=f"{b.alias}:{b.config!s}", ) for b in _BenchMarks ], ) -@pytest.mark.integration_test -def test_performance_of_joining_observations_and_responses( - config: _UpdatePerfTestConfig, - expected_performance: _ExpectedPerformance, - alias, - tmp_path, -): +def setup_benchmark(tmp_path, request): + alias, config, expected_performance = request.param info = create_experiment_args( config.num_parameters, config.num_gen_data_keys, @@ -428,39 +397,54 @@ def test_performance_of_joining_observations_and_responses( ens.save_response("summary", info.summary_responses.clone(), real) ens.save_response("gen_data", info.gen_data_responses.clone(), real) - with memray.Tracker(tmp_path / "memray.bin"): - t0 = time.time() - ens.get_observations_and_responses( - experiment.observation_keys, np.array(range(config.num_realizations)) - ) - t1 = time.time() - time_elapsed_s = t1 - t0 + yield ( + alias, + ens, + experiment.observation_keys, + np.array(range(config.num_realizations)), + expected_performance, + ) + + +def test_memory_performance_of_joining_observations_and_responses( + setup_benchmark, tmp_path +): + _, ens, observation_keys, mask, expected_performance = setup_benchmark + + with memray.Tracker(tmp_path / "memray.bin"): + ens.get_observations_and_responses(observation_keys, mask) stats = memray._memray.compute_statistics(str(tmp_path / "memray.bin")) mem_usage_mb = stats.total_memory_allocated / (1024**2) - print( - dedent(f"join,'{alias}',{round(time_elapsed_s,2)}s,{round(mem_usage_mb,2)}mb") - ) assert mem_usage_mb < expected_performance.memory_limit_mb - assert time_elapsed_s < expected_performance.time_limit_s -@pytest.mark.parametrize( - "config, expected_performance, alias", - [ - pytest.param( +def test_time_performance_of_joining_observations_and_responses( + setup_benchmark, benchmark +): + alias, ens, observation_keys, mask, _ = setup_benchmark + + if alias != "small": + pytest.skip() + + def run(): + ens.get_observations_and_responses(observation_keys, mask) + + benchmark(run) + + +@pytest.fixture( + params=[ + ( + b.alias, b.config, b.expected_update_performance, - f"{b.alias}:{b.config!s}", - id=f"{b.alias}:{b.config!s}", ) for b in _BenchMarks ], ) -@pytest.mark.integration_test -def test_performance_of_doing_es_update( - config: _UpdatePerfTestConfig, expected_performance, alias, tmp_path -): +def setup_es_benchmark(tmp_path, request): + alias, config, expected_performance = request.param info = create_experiment_args( config.num_parameters, config.num_gen_data_keys, @@ -504,21 +488,43 @@ def test_performance_of_doing_es_update( prior_ensemble=prior, iteration=1, ) - with memray.Tracker(tmp_path / "memray.bin"): - t0 = time.time() - smoother_update( - prior, - posterior, - prior.experiment.observation_keys, - [info.gen_kw_config.name], - ) - t1 = time.time() - time_elapsed_s = t1 - t0 + + yield ( + alias, + prior, + posterior, + info.gen_kw_config.name, + expected_performance, + ) + + +def test_memory_performance_of_doing_es_update(setup_es_benchmark, tmp_path): + _, prior, posterior, gen_kw_name, expected_performance = setup_es_benchmark + with memray.Tracker(tmp_path / "memray.bin"): + smoother_update( + prior, + posterior, + prior.experiment.observation_keys, + [gen_kw_name], + ) stats = memray._memray.compute_statistics(str(tmp_path / "memray.bin")) mem_usage_mb = stats.total_memory_allocated / (1024**2) - print( - dedent(f"update,'{alias}',{round(time_elapsed_s,2)}s,{round(mem_usage_mb,2)}mb") - ) assert mem_usage_mb < expected_performance.memory_limit_mb - assert time_elapsed_s < expected_performance.time_limit_s + + +def test_speed_performance_of_doing_es_update(setup_es_benchmark, benchmark): + alias, prior, posterior, gen_kw_name, _ = setup_es_benchmark + + if alias != "small": + pytest.skip() + + def run(): + smoother_update( + prior, + posterior, + prior.experiment.observation_keys, + [gen_kw_name], + ) + + benchmark(run)