Abstract single state runner #182

arik-shurygin · 2024-06-20T22:16:09Z

Introducing the abstract_azure_runner a light weight framework to abstract away a lot of the annoying parts of running an azure job.

The AbstractAzureRunner is a class which contains 4 functions, but only one of which needs to be overiden by the user:

__init__(azure_output_dir): sets up azure dual logging so users are able to see stdout and stderr on their azure node, but also saves a txt version of the logs to the output directory.
save_inference_posteriors(inferer, save_filename) saves inference chains from inferer in a standard format for future visualization.
save_inference_timelines(inferer, timeline_filename, particles_saved) picks some number (particles_saved) of particles from the inference process, runs those particles, and saves key metrics like vaccination, predicted/observed hospitalization, and many others.
save_static_run_timelines(parameters, sol, timeline_filename) saves the same metrics as save_inference_timelines but for a static run of the model.
process_state(self, state): filled in by the user, this is where you would slot in your existing azure code. The user is encouraged to call the other functions above from process_state to easily save model outputs so visualizer scripts can pull azure output down after the run finishes!

On a tiny sidenote, I added a useful type hint called SEIC_compartments to describe the tuple structure that commonly represents our 4 compartment model. This is why this PR impacts so many files, but many changes are quite minor

…titute

…a-scenarios-model into abstract-single-state-runner

…into abstract-single-state-runner

kokbent

See specific comments

kokbent · 2024-06-21T22:35:13Z

docker_template_configs/config_inferer_covid.json

-    "HOSP_PATH": "/input/data/hospital_220213_220108.csv",
-    "VAX_MODEL_DATA": "/input/data/spline_fits.csv",
-    "VAX_MODEL_NUM_KNOTS": 18,
+    "HOSPITALIZATION_DATA_PATH": "/input/data/hospital_220213_220108.csv",


We haven't actually used this before... wonder if we should just remove it

Will keep for now because we will eventually want to implement this rather than hardcoding in the runner like we did during SMH

I don't know. It's up to the person designing the inferer (likelihood and whatnot) to decide what's being pass into it... Just thought it being there is more confusing than useful

kokbent · 2024-06-28T00:10:26Z

mechanistic_model/abstract_azure_runner.py

+        )
+        df["seasonality_coef"] = seasonality_timeline
+        # save external introductions timeline of shape (num_days_predicted, num_strains)
+        # sum across age groups since we do % of each age bin anyways


We did do age specific introduction though... but it probably is not a big issue that we want age specific visualization for now...

Generally we only introduce our working age population, if we start to introduce multiple age groups that plot would be useful. For now this works though.

kokbent · 2024-06-28T00:34:42Z

mechanistic_model/abstract_azure_runner.py

+        df["total_infection_incidence"] = np.insert(
+            infection_incidence, [0], [0]
+        )
+        strain_proportions, _ = utils.get_timeline_from_solution_with_command(


If I'm not wrong the method for strain prevalence in this function is quite slow, be great to change it to a faster one (see https://github.com/cdcent/cfa-scenarios-model/blob/dbb3f615e2e6f7fbeebc9a67b5693b8cd8a9c6a9/exp/fifty_state_2202_2307_3strain/postaz_process.py#L359)

Referenced in a separate issue to be addressed later.

kokbent · 2024-06-28T00:40:00Z

mechanistic_model/abstract_azure_runner.py

+        None
+        """
+        # if inference complete, convert jnp/np arrays to list, then json dump
+        if inferer.infer_complete:


not sure if we should raise a warning or error here if it's not complete. Most likely something gone haywire already at this point 🤔

we dont currently raise a warning, we just do nothing if they are trying to save inference posteriors before fitting. I will modify to warn though

mechanistic_model/abstract_azure_runner.py

kokbent · 2024-06-28T00:46:23Z

mechanistic_model/abstract_azure_runner.py

+            timeline_filename,
+        )
+        all_particles_df = pd.DataFrame()
+        for particle in range(particles_saved):


This reads like you can only run particle 0, 1, 2, 3, up to particles_saved, instead of random sets of particles?

You are right, fixing this to pick random particles

kokbent · 2024-06-28T15:57:26Z

utils.py

@@ -971,7 +971,9 @@ def drop_sample_chains(samples: dict, dropped_chain_vals: list):
    return filtered_dict


-def flatten_list_parameters(samples):
+def flatten_list_parameters(


I don't know the sorcery here, rubber stamping this for now lol

it looks nasty but it basically flattens things that use numpyro.plate into 2 dimensions (chain, particle)

kokbent · 2024-06-28T19:47:49Z

mechanistic_model/abstract_azure_runner.py

+        ):
+            # select this random particle for each of our chains
+            chain_particle_pairs = [
+                (particle, chain)


(chain, particle) I think

You are right sorry, fixed. Thanks for the catch

kokbent

lglglgtm

arik-shurygin and others added 19 commits June 12, 2024 12:12

modifying flatten_list_parameters to flatten any dim sampled parameters

3b3bfa8

first run at an abstract azure runner class

8104f1d

checkpoint, load_posterior_particle half way

78c9443

redesigning the load_posterior_particles to use numpyro.handlers.subs…

b4d1ac3

…titute

Merge branch 'run-posterior-particle' of https://github.com/cdcent/cf…

4d7bb2a

…a-scenarios-model into abstract-single-state-runner

adding type hinting

9e03f4b

Merge branch 'run-posterior-particle' of https://github.com/cdcent/cf…

0af3ab5

…a-scenarios-model into abstract-single-state-runner

checkpoint, merging in another branch

42973d3

installing pre-commit on this wsl

79e5ae6

implementing PR feedback

2522e74

tests caught a bug!

9ddfc85

merging main into this branch

09ae7e9

Merge branch 'run-posterior-particle' of https://github.com/cdcent/cf…

dd26d21

…a-scenarios-model into abstract-single-state-runner

checkpoint, working on example

7b87be1

forgot to add example

760d253

turning on pre-commit

a2425e2

Merge branch 'main' of https://github.com/cdcent/cfa-scenarios-model …

7c7d93b

…into abstract-single-state-runner

better type aliasing, completed setup of abstrac_azure_runner class

c1df54b

Merge branch 'main' of https://github.com/cdcent/cfa-scenarios-model …

28d4284

…into abstract-single-state-runner

arik-shurygin requested review from kokbent and tjhladish June 20, 2024 22:19

fixing type hint

139ac8e

arik-shurygin mentioned this pull request Jun 26, 2024

Reduce image uploads #184

Merged

arik-shurygin added 2 commits June 26, 2024 17:47

adding save_config function to easily save config classes

ead9038

fixing bug with save_config

c2336b1

arik-shurygin mentioned this pull request Jun 27, 2024

Shiny Azure Visualizers #185

Merged

kokbent reviewed Jun 28, 2024

View reviewed changes

arik-shurygin mentioned this pull request Jun 28, 2024

revamp utils.get_timeline_from_solution_with_command #181

Open

implementing changes requested

36e9805

arik-shurygin requested a review from kokbent June 28, 2024 19:34

kokbent reviewed Jun 28, 2024

View reviewed changes

small bugfix

410187b

kokbent approved these changes Jun 28, 2024

View reviewed changes

arik-shurygin merged commit 3b2fcab into main Jun 28, 2024
2 checks passed

arik-shurygin deleted the abstract-single-state-runner branch June 28, 2024 20:10

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Abstract single state runner #182

Abstract single state runner #182

arik-shurygin commented Jun 20, 2024

kokbent left a comment

kokbent Jun 21, 2024

arik-shurygin Jun 28, 2024

kokbent Jun 28, 2024

kokbent Jun 28, 2024

arik-shurygin Jun 28, 2024

kokbent Jun 28, 2024

arik-shurygin Jun 28, 2024

arik-shurygin Jun 28, 2024

kokbent Jun 28, 2024

arik-shurygin Jun 28, 2024 •

edited

Loading

kokbent Jun 28, 2024

arik-shurygin Jun 28, 2024

kokbent Jun 28, 2024

arik-shurygin Jun 28, 2024

kokbent Jun 28, 2024

arik-shurygin Jun 28, 2024

kokbent left a comment

Abstract single state runner #182

Abstract single state runner #182

Conversation

arik-shurygin commented Jun 20, 2024

kokbent left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

arik-shurygin Jun 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kokbent left a comment

Choose a reason for hiding this comment

arik-shurygin Jun 28, 2024 •

edited

Loading