Trial sharing simulations across tests

[willGraham01]

Related to #1144 | Leaving as a draft for now since this is mainly a proof of concept.

This is a proof of concept for how we might decide to share simulation objects across tests, with a view to speeding up the execution of the test suite.

This PR creates the _BaseSharedSim class, whose purpose is to be used as a base class for the module-level tests. We also create the session-scoped shared_small_sim fixture, which is a simulation object that is created and run once right at the start of the test session. _BaseSharedSim then defines an autouse fixture _attach_to_shared_sim that allows any tests within the class (and any tests in derived classes) to access the shared simulation object.

As a proof of concept, this idea has been implemented in two modules (Diarrhoea and Stunting) for the "0-day-simulations" that appear semi-frequently across the test suite: simulations that have the same start date as end date, and a particularly low initial population (100). Two new test classes, TestStunting_SharedSim and TestDiarrhoea_SharedSim are defined, which inherit from the base class.

Complications due to simulation sharing

There are a number of complications regarding sharing simulation objects.
The most obvious is that the shared simulation object has to register all the modules (and their dependencies) that then wish to hook into this shared simulation to run tests. This isn't particularly problematic beyond remembering to adjust the dependencies if a new module is created and wishes to hook into the shared object.

Multiple Persistent Simulation Objects

If we want several different persistent simulation objects, we will need a _BaseSharedSim class for each one (this is due to how pytest handles setup and teardown fixtures). The current implementation has been designed with this possible generalisation in mind; everything except the _attach_to_shared_sim method just needs to be moved into a further abstract class - say __AbstractSharedSim which is then derived from:

@pytest.fixture(scope="session")
def shared_sim_1():
    # foo

@pytest.fixture(scope="session")
def shared_sim_2():
    # bar

class __AbstractSharedSim:
    # Everything in the current _BaseSharedSim except _attach_to_shared_sim

class _BaseShareSim1(__AbstractSharedSim):
    # Inherit from this to have access to the shared_sim_1 object
    @pytest.fixture(scope="class")
    def _attached_to_shared_sim(self, shared_sim_1):
        self.sim = shared_sim_1


class _BaseShareSim2(__AbstractSharedSim):
    # Inherit from this to have access to the shared_sim_2 object
    @pytest.fixture(scope="class")
    def _attached_to_shared_sim(self, shared_sim_2):
        self.sim = shared_sim_2

Why is there a need for the _BaseSharedSim class at all given the small_shared_sim fixture is session-scoped?

Technically there isn't, but:

• It's convenient for bundling all the functionality concerning sharing simulations into a single place.
• It will reduce code bloat if we do decide to create multiple persistent simulation objects. In particular because we will not have to redefine the "setup/teardown" fixtures that are discussed in reference to "changing simulation properties" below.
• It's easier to inherit a class (and a bunch of features that come with said class) than to manually type out the name of a bunch of global fixtures each time we want to use them.
• pytest internals are a mystery and this was the only way I could get this to work.

Changing Simulation Properties / State

A lot of tests need a simulation object for the sole reason that they want to check if a particular event is scheduled under certain conditions.
Setting up these conditions can involve a number of changes to the simulation object, including:

• Adjusting module parameters.
• Adjusting (even clearing) the HSI queue or simulation event queue.
• Adjusting a row in the population dataframe to artificially introduce symptoms.

As such, it is necessary for us to have "setup/teardown" fixtures that allow us to restore any potential edits to the shared object state. This is cheaper than recreating the simulation object from scratch, but still introduces some overhead. It does have the effect of forcing each test to highlight which aspects of the simulation object it expects to change though, which is quite useful from a debugging perspective.

These safe "setup/teardown" fixtures are defined in the _BaseSharedSim class so that they can be inherited by module-level test classes.

[willGraham01]

Performance Summary

• Pre shared simulation figures are quoted from 5c5bb25a2, run locally.
• Post shared simulation figures are quoted from 3501a70 (lost in the rebase) 9b4673f7, run locally.

For the implementation on this branch, durations <0.01s are not listed, but are broken down into setup, call, and teardown sections. This is what the "worst case maximum" column is accounting for in relation to the "post shared sim" column.

So for example, a "post" time of 1.0 (call) + 0.02 (teardown) could take as long as 1.23s to execute, since the setup step might have taken 0.01 seconds or less. Similarly, a figure of 0.02 (setup) could correspond to a total execution time of 0.04s (0.02 from setup plus just under 0.01 for each of the unlisted call and teardown).

Test names have been truncated so that the table formats properly, see the list below for full names (in row order).

Name of test	Module	Pre shared sim	Post shared sim	Worst case max
stunted_and_correctly_diagnosed	Stunting	1.34	1.60 (setup)	1.62
stunted_but_no_checking	Stunting	1.29	-	0.03
not_stunted	Stunting	1.33	-	0.03
diarrhoea_severe_dehydration	Diarrhoea	1.93	2.18 (setup) + 0.01 (teardown)	2.20
diarrhoea_severe_dehydration_dxtest_notfunctional	Diarrhoea	2.08	0.01 (teardown)	0.03
diarrhoea_non_severe_dehydration	Diarrhoea	1.90	0.01 (teardown)	0.03
test_run_each_of_the_HSI	Diarrhoea	2.00	0.01 (call) + 0.01 (teardown)	0.03
test_effect_of_vaccine	Diarrhoea	2.86	-	0.03
perfect_treatment_leads_to_zero_risk_of_death	Diarrhoea	2.82	0.06 (call)	0.08
treatment_for_those_that_will_not_die	Diarrhoea	3.02	0.02 (call) + 0.01 (teardown)	0.04

Conclusions

• The new implementation is slower for the "first" test in each of the Test classes, almost surely due to the overheads in setting up the test class.
• After this initially slower test, the remaining tests that use the same simulation are much faster. Almost surely because the tests which are refactored to use the same simulation do not have to re-make the simulation object each time!

As such, there appears to be a net time save to be obtained from sharing simulations.

Full test names

1. test_routine_assessment_for_chronic_undernutrition_if_stunted_and_correctly_diagnosed
2. test_routine_assessment_for_chronic_undernutrition_if_stunted_but_no_checking
3. test_routine_assessment_for_chronic_undernutrition_if_not_stunted
4. test_do_when_presentation_with_diarrhoea_severe_dehydration
5. test_do_when_presentation_with_diarrhoea_severe_dehydration_dxtest_notfunctional
6. test_do_when_presentation_with_diarrhoea_non_severe_dehydration
7. test_run_each_of_the_HSI
8. test_effect_of_vaccine
9. test_check_perfect_treatment_leads_to_zero_risk_of_death
10. test_do_treatment_for_those_that_will_not_die

[matt-graham]

Thanks @willGraham01 for your work on this, there are some interesting ideas here!

I've added some initial thoughts and comments. Overall I think the idea of providing standardized infrastructure for allowing sharing simulations across tests safely when the test functions involves manipulating the simulation object is nice. If we could avoid requiring refactoring test functions into test classes I think that would be good, as we generally have not used test classes, so this would be both quite a lot code refactoring to apply to existing tests and also make it less familiar for existing users when adding new tests.

Also as our biggest test computational burden are the GitHub Actions runs which run a matrix with one test module per runner, a session scoped shared simulation fixture will in reality have the same computational burden as a module scoped one. This to some degree limits how much computational gain we could get from this sort of approach, but also means we don't necessarily need to find a 'lowest common denominator' shared simulation across multiple modules, but can instead just look at sharing simulations within a single module, which presumably should be easier to manage in terms of specifying the Module subclasses to register, population size etc. appropriately. In this case having standard fixtures to allow safely updating properties of a shared_sim fixture within a test function without persisting those changes would still be useful.

[matt-graham]

Possibly naming start_date or similar would make a bit more descriptive (and also make more sense if we later changed default start date).

[matt-graham]

Is dict being used here to force a copy? If so using the copy method would perhaps be more explicit. Also is there any need for calling dict with old_param_values as presumably we can just restore the original (copy)?

[matt-graham]

Using this module with stunting.Stunting may be problematic as both define some of the same properties ( diarrhoea.DiarrhoeaPropertiesOfOtherModules is intended as a stand-in for test purposes for a number of modules, including Alri, Epi, Hiv, Stunting and Wasting).

[matt-graham]

Suggested change

	@pytest.fixture(scope="function")
	def changes_module_properties(self):
	"""
	The Diarrhoea.models property takes a copy of the
	module's parameters on initialisation, so is not
	affected by our temporary parameter change from the
	base class.
	Overwriting the base fixture to remedy this.
	"""
	old_param_values = dict(self.this_module.parameters)
	yield
	self.this_module.parameters = dict(old_param_values)
	self.this_module.models.p = self.this_module.parameters

Don't think this is necessary, though maybe I'm wrong if this was specifically added because you found things weren't working as expected. tlo.diarrhoea.Models has an attribute p which points to the same dict instance as the parameters attribute of the tlo.diarrhoea.Diarrhoea instance from which it is initialised - as far as I can see there is no copy involved.

[matt-graham]

The idea of using fixtures to safely mediate updates to the shared simulation fixture is cool, I hadn't come across the use of yield in fixtures to allow setup / teardown before.

What I'm less sure of is if we need the _BaseSharedSim class and corresponding subclasses? We could do something similar with just test functions, for example

@pytest.fixture(scope="session")
def start_date() -> Date:
    return Date(2010, 1, 1)

@pytest.fixture(scope="session")
def resource_filepath() -> Path:
    return (Path(os.path.dirname(__file__)) / "../resources").resolve()

@pytest.fixture(scope="session")
def shared_sim(seed, start_date, resource_filepath):
    ...

@pytest.fixture
def clear_shared_sim_hsi_event_queue(shared_sim):
    cached_queue = shared_sim_healthsystem.HSI_EVENT_QUEUE.copy()
    shared_sim.modules["HealthSystem"].reset_queue()
    yield
    shared_sim.modules["HealthSystem"].HSI_EVENT_QUEUE = cached_queue

@pytest.fixture
def changes_shared_sim_event_queue(shared_sim):
    cached_event_queue = shared_sim.event_queue.copy()
    yield
    shared_sim.event_queue = cached_event_queue

@pytest.fixture
def changes_shared_sim_module_parameters(shared_sim, changed_module_name):
    cached_parameter_values = shared_sim.modules[changed_module_name].parameters.copy()
    yield
    shared_sim.modules[changed_module_name] = cached_parameter_values 

@pytest.fixture
def changes_shared_sim_person_properties(shared_sim, changed_person_id):
    cached_values = shared_sim_df.loc[changed_person_id].copy()
    yield
    shared_sim_df.loc[changed_person_id] = cached_values

This assumes a changed_module_name / changed_person_id fixture would be defined in the scope of the test function using the changes_shared_sim_module_parameters / changes_shared_sim_person_properties fixture, which would work when we want to use the same module / person ID values across a test module (or test class). If we wanted to also be able to set these at a per test function level we could use the ability to use marks to pass data to fixtures and have defaults for the fixtures like

@pytest.fixture
def changed_module_name(request):
    return request.node.get_closest_marker("changed_module_name").args[0]

@pytest.fixture
def changed_person_id(request):
    return request.node.get_closest_marker("changed_person_id").args[0]

with an example test function usage then something like

@pytest.mark.changed_module_name("Diarrhoea")
@pytest.mark.changed_person_id(0)
def test_do_when_presentation_with_diarrhoea_severe_dehydration(
    shared_sim,
    changed_person_id,
    changes_shared_sim_module_parameters,
    changes_shared_sim_person_properties,
    clear_shared_sim_hsi_event_queue,
):
    ...

[matt-graham]

Also if we do keep the base class approach, I would say we don't want to use an underscrore prefixed name as it is intended for use outside this module?