Benchmark refactoring (#101)

* cleanup: remove duplication of absolute path check

* convenience functionality to validate results

* set upper limit on pandas version since pandas 3.0.0 causes some DrHook tests to fail.

* addressing some PR comments.

* make reference_path Path not str

* fix doc string

* fix tests that were failing with pandas 3.0.0

* Fix frame close validation checks and add more log output. Update tests to match.

* remove some excessively verbose output.

* refactoring of benchmark class: setup_rundir is called as part of run to make it less errorprone if forgotten. Add job as field directly to benchmark, can still be overridden if using same benchmark with different job. This will make it easier to run ensembles of benchmarks with different combinations of setups and jobs

* benchmark.run tests to also test that the TechSetup application is used if it exists

Author: uhager

ifsbench/benchmark.py

@@ -25,7 +25,7 @@
 from ifsbench.launch import Launcher
 
 
-__all__ = ['ScienceSetup', 'TechSetup', 'Benchmark']
+__all__ = ['ScienceSetup', 'TechSetup', 'Benchmark', 'BenchmarkSetup']
 
 
 class ScienceSetup(SerialisationMixin):
@@ -82,6 +82,18 @@ class TechSetup(SerialisationMixin):
     env_handlers: List[EnvHandler] = Field(default_factory=list)
 
 
+class BenchmarkSetup(SerialisationMixin):
+    #: The main (scientific) benchmark setup that describes _what_ gets benchmarked
+    # and _which data_ is used.
+    science: ScienceSetup
+
+    #: The parallel setup for the benchmark.
+    job: Job
+
+    #: Additional technical details that don't alter the results.
+    tech: Optional[TechSetup] = None
+
+
 class BenchmarkSummary(SerialisationMixin):
     """
     Summary of a benchmark run.
@@ -107,12 +119,8 @@ class Benchmark(SerialisationMixin):
       3. Create a result object, using the data in the run directory.
     """
 
-    #: The main (scientific) benchmark setup that describes _what_ gets benchmarked
-    # and _which data_ is used.
-    science: ScienceSetup
-
-    #: Additional technical details that don't alter the results.
-    tech: Optional[TechSetup] = None
+    #: The science and tech setup to run with the given job. The job can be overridden when running.
+    setup: BenchmarkSetup
 
     def setup_rundir(self, run_dir: Path, force: bool = False):
         """
@@ -137,17 +145,17 @@ def setup_rundir(self, run_dir: Path, force: bool = False):
         if exists and not force:
             return
 
-        handlers = self.science.data_handlers_init
-        if self.tech:
-            handlers += self.tech.data_handlers_init
+        handlers = self.setup.science.data_handlers_init
+        if self.setup.tech:
+            handlers += self.setup.tech.data_handlers_init
 
         for handler in handlers:
             handler.execute(run_dir)
 
     def run(
         self,
         run_dir: Path,
-        job: Job,
+        job: Optional[Job] = None,
         arch: Optional[Arch] = None,
         launcher: Optional[Launcher] = None,
         launcher_flags: Optional[List[str]] = None,
@@ -160,7 +168,7 @@ def run(
         run_dir: pathlib.Path
             The path to the run directory.
         job: Job
-            The parallel setup for the benchmark.
+            The parallel setup for the benchmark. If None, the Job from the BenchmarkSetup is used.
         arch: Arch
             A specific architecture that is used.
         launcher: Launcher
@@ -175,11 +183,17 @@ def run(
             of the benchmark.
         """
 
+        # Setup run directory without replacing the current contents if it already exists.
+        self.setup_rundir(run_dir, force=False)
+
         env_pipeline = DefaultEnvPipeline(
-            handlers=self.science.env_handlers, env_initial=os.environ
+            handlers=self.setup.science.env_handlers, env_initial=os.environ
         )
-        if self.tech:
-            env_pipeline.add(self.tech.env_handlers)
+        if self.setup.tech:
+            env_pipeline.add(self.setup.tech.env_handlers)
+
+        if not job:
+            job = self.setup.job
 
         if arch:
             arch_result = arch.process_job(job)
@@ -196,17 +210,17 @@ def run(
         if launcher is None:
             raise ValueError('No launcher was specified!')
 
-        application = self.science.application
-        if self.tech is not None and self.tech.application is not None:
-            application = self.tech.application
+        application = self.setup.science.application
+        if self.setup.tech is not None and self.setup.tech.application is not None:
+            application = self.setup.tech.application
 
         cmd = application.get_command(run_dir, job)
 
         library_paths = application.get_library_paths(run_dir, job)
 
-        data_handlers = list(self.science.data_handlers_runtime)
-        if self.tech:
-            data_handlers += self.tech.data_handlers_runtime
+        data_handlers = list(self.setup.science.data_handlers_runtime)
+        if self.setup.tech:
+            data_handlers += self.setup.tech.data_handlers_runtime
         data_handlers += application.get_data_handlers(run_dir, job)
 
         for handler in data_handlers:

ifsbench/tests/test_benchmark.py

@@ -10,14 +10,17 @@
 """
 
 from pathlib import Path
+from typing import List
 from time import sleep
 import sys
+import yaml
 
 import pytest
 from typing_extensions import Literal
 
 from ifsbench import (
     Benchmark,
+    BenchmarkSetup,
     ScienceSetup,
     TechSetup,
     DefaultApplication,
@@ -74,11 +77,13 @@ def test_defaultbenchmark_setup_rundir(tmp_path, test_setup_files, force, use_te
     """
 
     science, tech, science_list, tech_list = test_setup_files
+    job = Job(tasks=2)
 
     if use_tech:
-        benchmark = Benchmark(science=science, tech=tech)
+        setup = BenchmarkSetup(science=science, job=job, tech=tech)
     else:
-        benchmark = Benchmark(science=science)
+        setup = BenchmarkSetup(science=science, job=job)
+    benchmark = Benchmark(setup=setup)
 
     # Create the run directory and check that all files in file list
     # actually exist.
@@ -111,28 +116,49 @@ def test_defaultbenchmark_setup_rundir(tmp_path, test_setup_files, force, use_te
             assert stat.st_atime == stats[file].st_atime
 
 
-@pytest.fixture(name='test_run_setup')
-def fixture_test_run_setup():
-    env_handlers = [
-        EnvHandler(mode=EnvOperation.CLEAR),
-    ]
+class DummyApplication(DefaultApplication):
+    def get_command(self, run_dir: Path, job: Job) -> List[str]:
+        job_config = job.dump_config()
 
-    application = DefaultApplication(
-        command=[
+        command = [
             sys.executable,
             '-c',
-            "from pathlib import Path; Path('test.txt').touch()",
+            'from pathlib import Path; import yaml; '
+            f'f = open("{self.command[0]}", "w"); '
+            f'yaml.dump({job_config}, f); '
+            'f.close()',
         ]
+        return command
+
+
+@pytest.fixture(name='tech_setup_application')
+def fixture_tech_setup_application(request):
+    tech_application = None
+    if request.param:
+        tech_application = DummyApplication(command=['tech.txt'])
+
+    tech = TechSetup(
+        application=tech_application,
+        env_handlers=[EnvHandler(mode=EnvOperation.SET, key='KEY', value='VALUE')],
     )
 
+    return tech
+
+
+@pytest.fixture(name='test_run_science_setup')
+def fixture_test_run_science_setup():
+    env_handlers = [
+        EnvHandler(mode=EnvOperation.CLEAR),
+    ]
+
+    science_application = DummyApplication(command=['science.txt'])
+
     science = ScienceSetup(
-        application=application,
+        application=science_application,
         env_handlers=env_handlers,
     )
 
-    tech = TechSetup(env_handlers=[EnvHandler(mode=EnvOperation.SET, key='KEY', value='VALUE')])
-
-    return science, tech
+    return science
 
 
 class _DummyLauncher(Launcher):
@@ -155,15 +181,26 @@ def prepare(
         return LaunchData(run_dir=run_dir, cmd=cmd)
 
 
-@pytest.mark.parametrize('job', [Job(tasks=2)])
+@pytest.mark.parametrize('job_override', [None, Job(tasks=2, account='override')])
 @pytest.mark.parametrize('use_arch', [False, True])
 @pytest.mark.parametrize(
     'use_launcher, launcher_flags',
     [(False, None), (True, None), (True, ['something'])],
 )
-@pytest.mark.parametrize('use_tech', [True, False])
+@pytest.mark.parametrize(
+    'use_tech, tech_setup_application',
+    [(True, True), (True, False), (False, False)],
+    indirect=['tech_setup_application'],
+)
 def test_defaultbenchmark_run(
-    tmp_path, test_run_setup, job, use_arch, use_launcher, launcher_flags, use_tech
+    tmp_path,
+    test_run_science_setup,
+    job_override,
+    use_arch,
+    use_launcher,
+    launcher_flags,
+    use_tech,
+    tech_setup_application,
 ):
     """
     Test the Benchmark.run function.
@@ -181,24 +218,22 @@ def test_defaultbenchmark_run(
         else None
     )
 
-    science, tech = test_run_setup
+    science = test_run_science_setup
+    tech = tech_setup_application
+    job = Job(tasks=5, account='default')
 
     if use_tech:
-        benchmark = Benchmark(science=science, tech=tech)
-
+        setup = BenchmarkSetup(science=science, job=job, tech=tech)
     else:
-        benchmark = Benchmark(science=science)
-
-    # Create the run directory and check that all files in file list
-    # actually exist.
-    benchmark.setup_rundir(tmp_path)
+        setup = BenchmarkSetup(science=science, job=job)
+    benchmark = Benchmark(setup=setup)
 
     if arch is None and launcher is None:
         with pytest.raises(ValueError):
-            benchmark.run(tmp_path, job, arch, launcher, launcher_flags)
+            benchmark.run(tmp_path, job_override, arch, launcher, launcher_flags)
         return
 
-    benchmark.run(tmp_path, job, arch, launcher, launcher_flags)
+    benchmark.run(tmp_path, job_override, arch, launcher, launcher_flags)
 
     if launcher is not None:
         assert launcher._prepare_called is True
@@ -207,4 +242,16 @@ def test_defaultbenchmark_run(
     elif arch is not None:
         assert arch.get_default_launcher()._prepare_called is True
 
-    assert (tmp_path / 'test.txt').exists()
+    # Confirm the correct application was run, ScienceSetup vs TechSetup
+    out_file = 'tech.txt' if tech.application else 'science.txt'
+    output_path = tmp_path / out_file
+    assert (output_path).exists()
+
+    # Confirm the correct job was run, the benchmark member or the override
+    with output_path.open('r') as f:
+        config = yaml.safe_load(f)
+    executed_job = Job.from_config(config)
+    if job_override:
+        assert executed_job == job_override
+    else:
+        assert executed_job == job