instrumented Python code with Elastic APM

Author: Ahil PonArul

src/penn_chime/model/sir.py

@@ -195,9 +195,9 @@ def get_argmin_doubling_time(self, p: Parameters, dts):
             intrinsic_growth_rate = get_growth_rate(i_dt)
             self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
             self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
-
+    
             raw = self.run_projection(p, self.gen_policy(p))
-
+    
             # Skip values the would put the fit past peak
             peak_admits_day = raw["admits_hospitalized"].argmax()
             if peak_admits_day < 0:
@@ -228,7 +228,7 @@ def gen_policy(self, p: Parameters) -> Sequence[Tuple[float, int]]:
             (self.beta,   pre_mitigation_days),
             (self.beta_t, post_mitigation_days),
         ]
-
+    
     def run_projection(self, p: Parameters, policy: Sequence[Tuple[float, int]]):
         raw = sim_sir(
             self.susceptible,
@@ -245,20 +245,17 @@ def run_projection(self, p: Parameters, policy: Sequence[Tuple[float, int]]):
 
         return raw
 
-
 def get_loss(current_hospitalized, predicted) -> float:
     """Squared error: predicted vs. actual current hospitalized."""
     return (current_hospitalized - predicted) ** 2.0
 
-
 def get_argmin_ds(census, current_hospitalized: float) -> float:
     # By design, this forbids choosing a day after the peak
     # If that's a problem, see #381
     peak_day = census.argmax()
     losses = (census[:peak_day] - current_hospitalized) ** 2.0
     return losses.argmin()
 
-
 def get_beta(
     intrinsic_growth_rate: float,
     gamma: float,
@@ -271,14 +268,12 @@ def get_beta(
         * (1.0 - relative_contact_rate)
     )
 
-
 def get_growth_rate(doubling_time: Optional[float]) -> float:
     """Calculates average daily growth rate from doubling time."""
     if doubling_time is None or doubling_time == 0.0:
         return 0.0
     return (2.0 ** (1.0 / doubling_time) - 1.0)
 
-
 def sir(
     s: float, i: float, r: float, beta: float, gamma: float, n: float
 ) -> Tuple[float, float, float]:
@@ -289,7 +284,6 @@ def sir(
     scale = n / (s_n + i_n + r_n)
     return s_n * scale, i_n * scale, r_n * scale
 
-
 def sim_sir(
     s: float, i: float, r: float, gamma: float, i_day: int, policies: Sequence[Tuple[float, int]]
 ):
@@ -334,7 +328,6 @@ def sim_sir(
         "ever_infected": i_a + r_a
     }
 
-
 def build_sim_sir_w_date_df(
     raw_df: pd.DataFrame,
     current_date: datetime,
@@ -350,7 +343,6 @@ def build_sim_sir_w_date_df(
         }
     })
 
-
 def build_floor_df(df, keys, prefix):
     """Build floor sim sir w date."""
     return pd.DataFrame({
@@ -362,7 +354,6 @@ def build_floor_df(df, keys, prefix):
         }
     })
 
-
 def calculate_dispositions(
     raw: Dict,
     rates: Dict[str, float],
@@ -373,7 +364,6 @@ def calculate_dispositions(
         raw["ever_" + key] = raw["ever_infected"] * rate * market_share
         raw[key] = raw["ever_infected"] * rate * market_share
 
-
 def calculate_admits(raw: Dict, rates):
     """Build admits dataframe from dispositions."""
     for key in rates.keys():
@@ -384,7 +374,6 @@ def calculate_admits(raw: Dict, rates):
         raw["admits_"+key] = admit
         raw[key] = admit
 
-
 def calculate_census(
     raw: Dict,
     lengths_of_stay: Dict[str, int],

src/penn_chime/view/charts.py

@@ -4,10 +4,11 @@
 import pandas as pd
 import i18n
 import numpy as np
+import elasticapm
 
 from ..constants import DATE_FORMAT
 
-
+@elasticapm.capture_span()
 def build_admits_chart(
     *, alt, admits_floor_df: pd.DataFrame, max_y_axis: Optional[int] = None
 ) -> Chart:
@@ -51,7 +52,7 @@ def build_admits_chart(
         .interactive()
     )
 
-
+@elasticapm.capture_span()
 def build_census_chart(
     *, alt, census_floor_df: pd.DataFrame, max_y_axis: Optional[int] = None
 ) -> Chart:
@@ -95,7 +96,7 @@ def build_census_chart(
         .interactive()
     )
 
-
+@elasticapm.capture_span()
 def build_sim_sir_w_date_chart(
     *, alt, sim_sir_w_date_floor_df: pd.DataFrame, max_y_axis: Optional[int] = None
 ) -> Chart:
@@ -139,6 +140,7 @@ def build_sim_sir_w_date_chart(
         .interactive()
     )
 
+@elasticapm.capture_span()
 def build_table(
     *, df: pd.DataFrame, labels: Dict[str, str], modulo: int = 1
 ) -> pd.DataFrame:

src/penn_chime/view/st_app.py

@@ -6,6 +6,8 @@
 import streamlit as st  # type: ignore
 import i18n  # type: ignore
 
+from elasticapm import Client
+
 i18n.set('filename_format', '{locale}.{format}')
 i18n.set('locale', 'en')
 i18n.set('fallback', 'en')
@@ -28,6 +30,10 @@
 
 
 def main():
+    #client = Client(server_url='http://apm-server:8200',service_name='chime_local')
+    #client = Client(server_url='${ELASTIC_APM_SERVER_URL}',service_name='${ELASTIC_APM_SERVICE_NAME}')
+    client = Client()
+    client.begin_transaction('main_page')
     # This is somewhat dangerous:
     # Hide the main menu with "Rerun", "run on Save", "clear cache", and "record a screencast"
     # This should not be hidden in prod, but removed
@@ -37,8 +43,9 @@ def main():
     d = Parameters.create(os.environ, [])
     p = display_sidebar(st, d)
     m = Sir(p)
-
+    
     display_header(st, m, p)
+   
 
     st.subheader(i18n.t("app-new-admissions-title"))
     st.markdown(i18n.t("app-new-admissions-text"))
@@ -73,3 +80,4 @@ def main():
         df=m.sim_sir_w_date_df,
     )
     display_footer(st)
+    client.end_transaction('main_page')