From f6a0556d03ebe6140f7b3c7c979b9ee2f6abcab0 Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Fri, 15 Aug 2025 10:58:30 +0100 Subject: [PATCH 01/12] Pass through command-line arguments to scenario --- src/tlo/cli.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/tlo/cli.py b/src/tlo/cli.py index 9436d060ff..6f83c82076 100644 --- a/src/tlo/cli.py +++ b/src/tlo/cli.py @@ -107,15 +107,16 @@ def parse_log(log_directory): with open(path / f"{key}.pickle", "wb") as f: pickle.dump(output, f) -@cli.command() +@cli.command(context_settings=dict(ignore_unknown_options=True)) @click.argument("scenario_file", type=click.Path(exists=True)) @click.option("--asserts-on", type=bool, default=False, is_flag=True, help="Enable assertions in simulation run.") @click.option("--more-memory", type=bool, default=False, is_flag=True, help="Request machine wth more memory (for larger population sizes).") @click.option("--image-tag", type=str, help="Tag of the Docker image to use.") @click.option("--keep-pool-alive", type=bool, default=False, is_flag=True, hidden=True) +@click.argument('scenario_args', nargs=-1, type=click.UNPROCESSED) @click.pass_context -def batch_submit(ctx, scenario_file, asserts_on, more_memory, keep_pool_alive, image_tag=None): +def batch_submit(ctx, scenario_file, asserts_on, more_memory, keep_pool_alive, image_tag=None, scenario_args=None): """Submit a scenario to the batch system. SCENARIO_FILE is path to file containing scenario class. @@ -132,6 +133,10 @@ def batch_submit(ctx, scenario_file, asserts_on, more_memory, keep_pool_alive, i scenario = load_scenario(scenario_file) + # if we have other scenario arguments, parse them + if scenario_args is not None: + scenario.parse_arguments(scenario_args) + # get the commit we're going to submit to run on batch, and save the run config for that commit # it's the most recent commit on current branch repo = Repo(".") From 70d721b3605d9d28c402cb959a97469fc50470f5 Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Fri, 15 Aug 2025 11:00:14 +0100 Subject: [PATCH 02/12] Expand environment variables in specified path if provided (quick fix to work with AZ_* env vars. --- src/tlo/scenario.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/tlo/scenario.py b/src/tlo/scenario.py index cd83d0e63f..1735722d11 100644 --- a/src/tlo/scenario.py +++ b/src/tlo/scenario.py @@ -407,6 +407,9 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): hasattr(self.scenario, "resume_simulation") and self.scenario.resume_simulation is not None ): + if "$" in self.scenario.resume_simulation: + self.scenario.resume_simulation = os.path.expandvars(self.scenario.resume_simulation) + suspended_simulation_path = ( Path(self.scenario.resume_simulation) / str(draw_number) From 81c36d97a10a8644b5a6dd72e0010e41295d5013 Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Fri, 15 Aug 2025 11:01:11 +0100 Subject: [PATCH 03/12] Close the log file handle before pickling the simulation --- src/tlo/scenario.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/tlo/scenario.py b/src/tlo/scenario.py index 1735722d11..38c2ea2e67 100644 --- a/src/tlo/scenario.py +++ b/src/tlo/scenario.py @@ -63,6 +63,7 @@ def draw_parameters(self, draw_number, rng): import argparse import datetime import json +import os from collections.abc import Iterable from itertools import product from pathlib import Path, PurePosixPath @@ -445,12 +446,13 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): ): sim.run_simulation_to(to_date=self.scenario.suspend_date) suspended_simulation_path = Path(log_config["directory"]) / "suspended_simulation.pickle" - sim.save_to_pickle(pickle_path=suspended_simulation_path) - sim.close_output_file() logger.info( key="message", - data=f"Simulation suspended at {self.scenario.suspend_date} and saved to {suspended_simulation_path}", + data=f"Suspending simulation at {self.scenario.suspend_date} and saving to {suspended_simulation_path}." + f"Note, output file handle will be closed first & no more output logged", ) + sim.close_output_file() + sim.save_to_pickle(pickle_path=suspended_simulation_path) else: sim.run_simulation_to(to_date=self.scenario.end_date) sim.finalise() From 4723a4e0586d59eb6c1c67b6f13b9b6d2ead945e Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Fri, 15 Aug 2025 11:02:14 +0100 Subject: [PATCH 04/12] Scenario for testing suspend/resume --- .../dev/scenarios/suspend-resume-test.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 src/scripts/dev/scenarios/suspend-resume-test.py diff --git a/src/scripts/dev/scenarios/suspend-resume-test.py b/src/scripts/dev/scenarios/suspend-resume-test.py new file mode 100644 index 0000000000..6bec04e417 --- /dev/null +++ b/src/scripts/dev/scenarios/suspend-resume-test.py @@ -0,0 +1,54 @@ +""" +This file defines a batch run of a large population for a long time with all disease modules and full use of HSIs +It's used for calibrations (demographic patterns, health burdens and healthsystem usage) + +Run on the batch system using: +```tlo batch-submit src/scripts/calibration_analyses/scenarios/long_run_all_diseases.py``` + +or locally using: + ```tlo scenario-run src/scripts/calibration_analyses/scenarios/long_run_all_diseases.py``` + +""" + +from tlo import Date, logging +from tlo.analysis.utils import get_parameters_for_status_quo +from tlo.methods.fullmodel import fullmodel +from tlo.scenario import BaseScenario + + +class SuspendResumeTest(BaseScenario): + def __init__(self): + super().__init__() + self.seed = 0 + self.start_date = Date(2010, 1, 1) + self.end_date = Date(2012, 1, 1) # The simulation will stop before reaching this date. + self.pop_size = 1_000 + self.number_of_draws = 2 + self.runs_per_draw = 2 + + def log_configuration(self): + return { + 'filename': 'suspend_resume_test', # <- (specified only for local running) + 'directory': './outputs', # <- (specified only for local running) + 'custom_levels': { + '*': logging.WARNING, + 'tlo.methods.demography': logging.INFO, + 'tlo.methods.demography.detail': logging.WARNING, + 'tlo.methods.healthburden': logging.INFO, + 'tlo.methods.healthsystem': logging.INFO, + 'tlo.methods.healthsystem.summary': logging.INFO, + "tlo.methods.contraception": logging.INFO, + } + } + + def modules(self): + return fullmodel() + + def draw_parameters(self, draw_number, rng): + return get_parameters_for_status_quo() + + +if __name__ == '__main__': + from tlo.cli import scenario_run + + scenario_run([__file__]) From fb3c80dd8ebfe18c0f7dd80c82df8ca32fdefeb6 Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Mon, 25 Aug 2025 14:30:20 +0100 Subject: [PATCH 05/12] Rewrite the argument for resume-simulation argument - takes the supplied job id and makes it the full path to the saved job - not perfect --- src/tlo/cli.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/tlo/cli.py b/src/tlo/cli.py index 6f83c82076..3dc19ee680 100644 --- a/src/tlo/cli.py +++ b/src/tlo/cli.py @@ -133,8 +133,22 @@ def batch_submit(ctx, scenario_file, asserts_on, more_memory, keep_pool_alive, i scenario = load_scenario(scenario_file) + config = load_config(ctx.obj['config_file']) + + # Directory where the file share will be mounted, relative to + # ${AZ_BATCH_NODE_MOUNTS_DIR}. + file_share_mount_point = "mnt" + # if we have other scenario arguments, parse them if scenario_args is not None: + # we rewrite the path to the simulation to resume + if '--resume-simulation' in scenario_args: + i = scenario_args.index('--resume-simulation') + path_to_job = (f"${{AZ_BATCH_NODE_MOUNTS_DIR}}/" + f"{file_share_mount_point}/" + f"{config['DEFAULT']['USERNAME']}/" + f"{scenario_args[i+1]}") + scenario_args = scenario_args[:i + 1] + (path_to_job, ) + scenario_args[i + 2:] scenario.parse_arguments(scenario_args) # get the commit we're going to submit to run on batch, and save the run config for that commit @@ -145,8 +159,6 @@ def batch_submit(ctx, scenario_file, asserts_on, more_memory, keep_pool_alive, i print(">Setting up batch\r", end="") - config = load_config(ctx.obj['config_file']) - # ID of the Batch job. timestamp = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ") job_id = scenario.get_log_config()["filename"] + "-" + timestamp @@ -226,10 +238,6 @@ def batch_submit(ctx, scenario_file, asserts_on, more_memory, keep_pool_alive, i # Options for running the Docker container container_run_options = "--rm --workdir /TLOmodel" - # Directory where the file share will be mounted, relative to - # ${AZ_BATCH_NODE_MOUNTS_DIR}. - file_share_mount_point = "mnt" - azure_file_share_configuration = batch_models.AzureFileShareConfiguration( account_name=config["STORAGE"]["NAME"], azure_file_url=azure_file_url, From e0408946cd005148038c247019e221fdc8da3c33 Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Mon, 25 Aug 2025 15:17:40 +0100 Subject: [PATCH 06/12] Use printing because logger only works when configured via simulation - Okay after simulation is setup --- src/tlo/scenario.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/tlo/scenario.py b/src/tlo/scenario.py index 38c2ea2e67..ad4268f456 100644 --- a/src/tlo/scenario.py +++ b/src/tlo/scenario.py @@ -161,7 +161,7 @@ def parse_arguments(self, extra_arguments: List[str]) -> None: for key, value in vars(arguments).items(): if value is not None: if hasattr(self, key): - logger.info(key="message", data=f"Overriding attribute: {key}: {getattr(self, key)} -> {value}") + print(f"Overriding attribute with argument value: {key}: {getattr(self, key)} -> {value}") setattr(self, key, value) def add_arguments(self, parser: argparse.ArgumentParser) -> None: @@ -360,8 +360,8 @@ def __init__(self, run_configuration_path): self.scenario = ScenarioLoader(self.run_config["scenario_script_path"]).get_scenario() if self.run_config["arguments"] is not None: self.scenario.parse_arguments(self.run_config["arguments"]) - logger.info(key="message", data=f"Loaded scenario using {run_configuration_path}") - logger.info(key="message", data=f"Found {self.number_of_draws} draws; {self.runs_per_draw} runs/draw") + print(f"Loaded scenario from config at {run_configuration_path}") + print(f"Found {self.number_of_draws} draws with {self.runs_per_draw} runs-per-draw.") @property def number_of_draws(self): @@ -398,10 +398,7 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): sample = self.get_sample(draw, sample_number) log_config = self.scenario.get_log_config(output_directory) - logger.info( - key="message", - data=f"Running draw {sample['draw_number']}, sample {sample['sample_number']}", - ) + print(f"Running draw {sample['draw_number']}, run {sample['sample_number']}.") # if user has specified a restore simulation, we load it from a pickle file if ( @@ -417,11 +414,13 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): / str(sample_number) / "suspended_simulation.pickle" ) + + sim = Simulation.load_from_pickle(pickle_path=suspended_simulation_path, log_config=log_config) + logger.info( key="message", data=f"Loading pickled suspended simulation from {suspended_simulation_path}", ) - sim = Simulation.load_from_pickle(pickle_path=suspended_simulation_path, log_config=log_config) else: sim = Simulation( start_date=self.scenario.start_date, From 16b2618908f3982ba8851bea17f2112dbd140ac7 Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Mon, 25 Aug 2025 15:21:11 +0100 Subject: [PATCH 07/12] Improve message --- src/tlo/scenario.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tlo/scenario.py b/src/tlo/scenario.py index ad4268f456..7cda663e61 100644 --- a/src/tlo/scenario.py +++ b/src/tlo/scenario.py @@ -419,7 +419,7 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): logger.info( key="message", - data=f"Loading pickled suspended simulation from {suspended_simulation_path}", + data=f"Loading suspended simulation from {suspended_simulation_path}", ) else: sim = Simulation( From 91fa8010e3ca8a4c94427621890c277387192a15 Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Mon, 25 Aug 2025 15:39:34 +0100 Subject: [PATCH 08/12] Allow user to specific a specific draw to resume simulation from - useful when all draws can be resumed from a a single specified draw --- src/tlo/scenario.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/tlo/scenario.py b/src/tlo/scenario.py index 7cda663e61..97cdccaeba 100644 --- a/src/tlo/scenario.py +++ b/src/tlo/scenario.py @@ -408,12 +408,13 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): if "$" in self.scenario.resume_simulation: self.scenario.resume_simulation = os.path.expandvars(self.scenario.resume_simulation) - suspended_simulation_path = ( - Path(self.scenario.resume_simulation) - / str(draw_number) - / str(sample_number) - / "suspended_simulation.pickle" - ) + suspended_simulation_path = Path(self.scenario.resume_simulation) + + # if the resume_simulation doesn't end with a draw number, we are resuming all draws + if not self.scenario.resume_simulation.rstrip("/").split("/")[-1].isdigit(): + suspended_simulation_path = suspended_simulation_path / str(draw_number) + + suspended_simulation_path = suspended_simulation_path / str(sample_number) / "suspended_simulation.pickle" sim = Simulation.load_from_pickle(pickle_path=suspended_simulation_path, log_config=log_config) From 509fc7661920568b027cd9020869a3a7913a4b7d Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Mon, 25 Aug 2025 15:52:14 +0100 Subject: [PATCH 09/12] Override parameters when restoring simulation - we may be restoring simulation from a baseline without any parameters - the draw itself may override parameters --- src/tlo/scenario.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/tlo/scenario.py b/src/tlo/scenario.py index 97cdccaeba..2ff60a13ec 100644 --- a/src/tlo/scenario.py +++ b/src/tlo/scenario.py @@ -422,6 +422,10 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): key="message", data=f"Loading suspended simulation from {suspended_simulation_path}", ) + + # if parameters are specified, we override them + if sample["parameters"] is not None: + self.override_parameters(sim, sample["parameters"]) else: sim = Simulation( start_date=self.scenario.start_date, From 6a2ac5d32cfafe4afbffd18aae06fa791e13e8ea Mon Sep 17 00:00:00 2001 From: Asif Tamuri Date: Tue, 26 Aug 2025 16:04:16 +0100 Subject: [PATCH 10/12] Handle resuming from multi-digit draws --- src/tlo/scenario.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/tlo/scenario.py b/src/tlo/scenario.py index 2ff60a13ec..f623bcdeb3 100644 --- a/src/tlo/scenario.py +++ b/src/tlo/scenario.py @@ -405,13 +405,17 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): hasattr(self.scenario, "resume_simulation") and self.scenario.resume_simulation is not None ): + # expand any environment variables in the path if "$" in self.scenario.resume_simulation: self.scenario.resume_simulation = os.path.expandvars(self.scenario.resume_simulation) suspended_simulation_path = Path(self.scenario.resume_simulation) # if the resume_simulation doesn't end with a draw number, we are resuming all draws - if not self.scenario.resume_simulation.rstrip("/").split("/")[-1].isdigit(): + last_component = self.scenario.resume_simulation.rstrip("/").split("/")[-1] + try: + int(last_component) + except ValueError: suspended_simulation_path = suspended_simulation_path / str(draw_number) suspended_simulation_path = suspended_simulation_path / str(sample_number) / "suspended_simulation.pickle" @@ -453,7 +457,7 @@ def run_sample_by_number(self, output_directory, draw_number, sample_number): logger.info( key="message", data=f"Suspending simulation at {self.scenario.suspend_date} and saving to {suspended_simulation_path}." - f"Note, output file handle will be closed first & no more output logged", + f" Note, output file handle will be closed first and no more output logged", ) sim.close_output_file() sim.save_to_pickle(pickle_path=suspended_simulation_path) From 44cdc96ec46dd386d764f9b6f45fc542c8de52fe Mon Sep 17 00:00:00 2001 From: Margherita Molaro <48129834+marghe-molaro@users.noreply.github.com> Date: Wed, 7 Jan 2026 11:20:59 +0000 Subject: [PATCH 11/12] Create scenario to test suspend/resume in HealthSystem --- .../scenario_test_resume_with_change.py | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 src/scripts/test_suspend_resume/scenario_test_resume_with_change.py diff --git a/src/scripts/test_suspend_resume/scenario_test_resume_with_change.py b/src/scripts/test_suspend_resume/scenario_test_resume_with_change.py new file mode 100644 index 0000000000..28c657cec4 --- /dev/null +++ b/src/scripts/test_suspend_resume/scenario_test_resume_with_change.py @@ -0,0 +1,148 @@ +"""This Scenario file run the model under different assumptions for HR capabilities expansion in order to estimate the +impact that is achieved under each. + +Run on the batch system using: +``` +tlo batch-submit src/scripts/healthsystem/impact_of_policy/scenario_impact_of_const_capabilities_expansion.py +``` + +or locally using: +``` +tlo scenario-run src/scripts/healthsystem/impact_of_policy/scenario_impact_of_const_capabilities_expansion.py +``` + +""" +from pathlib import Path +from typing import Dict + +import pandas as pd + +from tlo import Date, logging +from tlo.analysis.utils import get_parameters_for_status_quo, mix_scenarios +from tlo.methods.fullmodel import fullmodel +from tlo.scenario import BaseScenario + + +class ImpactOfHealthSystemMode(BaseScenario): + def __init__(self): + super().__init__() + self.seed = 0 + self.start_date = Date(2010, 1, 1) + self.end_date = self.start_date + pd.DateOffset(years=5) + self.pop_size = 20_000 + self._scenarios = self._get_scenarios() + self.number_of_draws = len(self._scenarios) + self.runs_per_draw = 3 + + def log_configuration(self): + return { + 'filename': 'effect_of_capabilities_scaling', + 'directory': Path('./outputs'), # <- (specified only for local running) + 'custom_levels': { + '*': logging.WARNING, + 'tlo.methods.demography': logging.INFO, + 'tlo.methods.demography.detail': logging.WARNING, + 'tlo.methods.healthburden': logging.INFO, + 'tlo.methods.healthsystem.summary': logging.INFO, + } + } + + def modules(self): + return ( + fullmodel() + ) + + def draw_parameters(self, draw_number, rng): + if draw_number < self.number_of_draws: + return list(self._scenarios.values())[draw_number] + else: + return + + def _get_scenarios(self) -> Dict[str, Dict]: + """Return the Dict with values for the parameters that are changed, keyed by a name for the scenario. + """ + + self.YEAR_OF_CHANGE = 2012 + + return { + + # =========== STATUS QUO ============ + "Baseline": + mix_scenarios( + self._baseline(), + ), + } + """ + "Mode 2 no rescaling": + mix_scenarios( + self._baseline(), + { + "HealthSystem": { + "mode_appt_constraints_postSwitch": 2, # <-- Mode 2 post-change to show effects of HRH + }, + } + ), + + "Mode 2 with rescaling": + mix_scenarios( + self._baseline(), + { + "HealthSystem": { + "mode_appt_constraints_postSwitch": 2, # <-- Mode 2 post-change to show effects of HRH + "scale_to_effective_capabilities": True, + }, + } + ), + + "Mode 2 with rescaling and funded plus": + mix_scenarios( + self._baseline(), + { + "HealthSystem": { + "mode_appt_constraints_postSwitch": 2, # <-- Mode 2 post-change to show effects of HRH + "scale_to_effective_capabilities": True, + "use_funded_or_actual_staffing_postSwitch": "funded_plus", + }, + } + ), + + "Mode 1 perfect consumables": + mix_scenarios( + self._baseline(), + { + "HealthSystem": { + "cons_availability_postSwitch":"perfect", + }, + } + ), + + } + """ + def _baseline(self) -> Dict: + """Return the Dict with values for the parameter changes that define the baseline scenario. """ + return mix_scenarios( + get_parameters_for_status_quo(), + { + "HealthSystem": { + "year_mode_switch":self.YEAR_OF_CHANGE, + "year_cons_availability_switch":self.YEAR_OF_CHANGE, + "year_use_funded_or_actual_staffing_switch":self.YEAR_OF_CHANGE, + "mode_appt_constraints": 1, + "cons_availability": "default", + "use_funded_or_actual_staffing": "actual", + "mode_appt_constraints_postSwitch": 1, + "use_funded_or_actual_staffing_postSwitch":"actual", + "cons_availability_postSwitch":"default", + "policy_name": "Naive", + "tclose_overwrite": 1, + "tclose_days_offset_overwrite": 7, + + } + }, + ) + + +if __name__ == '__main__': + from tlo.cli import scenario_run + + scenario_run([__file__]) From 827b3622da7f695ac12f71b22d70537beb1711a1 Mon Sep 17 00:00:00 2001 From: Margherita Molaro <48129834+marghe-molaro@users.noreply.github.com> Date: Wed, 7 Jan 2026 13:25:45 +0000 Subject: [PATCH 12/12] Modify scenario file to run the complete analysis upon resuming --- .../test_suspend_resume/scenario_test_resume_with_change.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/scripts/test_suspend_resume/scenario_test_resume_with_change.py b/src/scripts/test_suspend_resume/scenario_test_resume_with_change.py index 28c657cec4..5f8f26a804 100644 --- a/src/scripts/test_suspend_resume/scenario_test_resume_with_change.py +++ b/src/scripts/test_suspend_resume/scenario_test_resume_with_change.py @@ -71,8 +71,7 @@ def _get_scenarios(self) -> Dict[str, Dict]: mix_scenarios( self._baseline(), ), - } - """ + "Mode 2 no rescaling": mix_scenarios( self._baseline(), @@ -117,7 +116,7 @@ def _get_scenarios(self) -> Dict[str, Dict]: ), } - """ + def _baseline(self) -> Dict: """Return the Dict with values for the parameter changes that define the baseline scenario. """ return mix_scenarios(