Skip to content

Commit 4e90460

Browse files
anth-volkanth-volk
committed
fix: Enable running API v1 via URL parameter
1 parent cae3184 commit 4e90460

6 files changed

Lines changed: 175 additions & 568 deletions

File tree

‎policyengine_api/jobs/calculate_economy_simulation_job.py‎

Lines changed: 160 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -73,20 +73,18 @@
7373
datasets["us"][dataset] = f"{datasets['us'][dataset]}@{us_dataset_version}"
7474

7575

76-
run_api_v2 = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") is not None
77-
78-
if not run_api_v2:
79-
logger.log_text(
80-
"Didn't find any GOOGLE_APPLICATION_CREDENTIALS, so will not run results with APIv2.",
81-
severity="WARNING",
82-
)
83-
84-
8576
class CalculateEconomySimulationJob(BaseJob):
8677
def __init__(self):
8778
super().__init__()
88-
if run_api_v2:
89-
self.api_v2 = SimulationAPIv2()
79+
if os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") is None:
80+
logger.log_text(
81+
"GOOGLE_APPLICATION_CREDENTIALS not set; unable to run simulation API.",
82+
severity="ERROR",
83+
)
84+
raise ValueError(
85+
"GOOGLE_APPLICATION_CREDENTIALS not set; unable to run simulation API."
86+
)
87+
self.sim_api = SimulationAPI()
9088

9189
def run(
9290
self,
@@ -99,11 +97,16 @@ def run(
9997
options: dict,
10098
baseline_policy: dict,
10199
reform_policy: Annotated[str, "String-formatted JSON"],
100+
run_api_v1: bool = False,
102101
):
103102
job_id = self._set_job_id()
104103
job_setup_options = {
105104
"job_id": job_id,
106-
"job_type": "CALCULATE_ECONOMY_SIMULATION_JOB",
105+
"job_type": (
106+
"CALCULATE_ECONOMY_SIMULATION_JOB"
107+
if not run_api_v1
108+
else "CALCULATE_ECONOMY_SIMULATION_JOB_V1"
109+
),
107110
"baseline_policy_id": baseline_policy_id,
108111
"reform_policy_id": policy_id,
109112
"country_id": country_id,
@@ -113,6 +116,13 @@ def run(
113116
"options": options,
114117
"baseline_policy": baseline_policy,
115118
"reform_policy": reform_policy,
119+
"workflow_id": None,
120+
"model_version": COUNTRY_PACKAGE_VERSIONS[country_id],
121+
"data_version": (
122+
uk_dataset_version
123+
if country_id == "uk"
124+
else us_dataset_version
125+
),
116126
}
117127
logger.log_struct(
118128
{
@@ -122,6 +132,7 @@ def run(
122132
severity="INFO",
123133
)
124134

135+
options_hash = None
125136
try:
126137
# Configure inputs
127138
# Note for anyone modifying options_hash: redis-queue treats ":" as a namespace
@@ -159,17 +170,6 @@ def run(
159170
)
160171
return
161172

162-
# Save identifiers for later commenting on processing status
163-
identifiers = (
164-
country_id,
165-
policy_id,
166-
baseline_policy_id,
167-
region,
168-
dataset,
169-
time_period,
170-
options_hash,
171-
)
172-
173173
logger.log_struct(
174174
{
175175
"message": "No existing completed result found, proceeding with computation",
@@ -224,120 +224,90 @@ def run(
224224
),
225225
)
226226

227+
# TODO: Figure out how to handle this better
228+
# If using API v1, run that job
229+
if run_api_v1:
230+
logger.log_struct(
231+
{
232+
"message": "Running API v1 job",
233+
**job_setup_options,
234+
}
235+
)
236+
impact = self._run_v1_job(
237+
job_id=job_id,
238+
baseline_policy_id=baseline_policy_id,
239+
reform_policy_id=policy_id,
240+
country_id=country_id,
241+
region=region,
242+
dataset=dataset,
243+
time_period=time_period,
244+
options=options,
245+
baseline_policy=baseline_policy,
246+
reform_policy=reform_policy,
247+
)
248+
reform_impacts_service.set_complete_reform_impact(
249+
country_id=country_id,
250+
reform_policy_id=policy_id,
251+
baseline_policy_id=baseline_policy_id,
252+
region=region,
253+
dataset=dataset,
254+
time_period=time_period,
255+
options_hash=options_hash,
256+
reform_impact_json=impact,
257+
)
258+
return
259+
260+
# Set up sim API job
227261
logger.log_struct(
228262
{
229-
"message": "Starting computation of baseline policy",
263+
"message": "Setting up sim API job",
230264
**job_setup_options,
231265
}
232266
)
267+
sim_config: dict[str, Any] = self.sim_api._setup_sim_options(
268+
country_id=country_id,
269+
scope="macro",
270+
reform_policy=reform_policy,
271+
baseline_policy=baseline_policy,
272+
time_period=time_period,
273+
region=region,
274+
dataset=dataset,
275+
model_version=COUNTRY_PACKAGE_VERSIONS[country_id],
276+
data_version=(
277+
uk_dataset_version
278+
if country_id == "uk"
279+
else us_dataset_version
280+
),
281+
)
233282

234-
# If comparing against API v2, start job
235-
if run_api_v2:
236-
237-
try:
238-
# Populate v2/v1 comparison config data; we will pass this
239-
# to GCP logs either on error or success
240-
comparison_data = {
241-
**job_setup_options,
242-
"v1_country_package_version": COUNTRY_PACKAGE_VERSIONS[
243-
country_id
244-
],
245-
"v2_id": None, # Unavailable until job starts
246-
"v2_country_package_version": None, # Unavailable until job completes
247-
}
283+
sim_api_execution = self.sim_api.run(sim_config)
284+
execution_id = self.sim_api.get_execution_id(sim_api_execution)
248285

249-
# Set up APIv2 job
250-
logger.log_struct(
251-
{
252-
"message": "Setting up APIv2 job",
253-
**comparison_data,
254-
}
255-
)
256-
sim_config: dict[str, Any] = (
257-
self.api_v2._setup_sim_options(
258-
country_id=country_id,
259-
scope="macro",
260-
reform_policy=reform_policy,
261-
baseline_policy=baseline_policy,
262-
time_period=time_period,
263-
region=region,
264-
dataset=dataset,
265-
model_version=COUNTRY_PACKAGE_VERSIONS[country_id],
266-
data_version=(
267-
uk_dataset_version
268-
if country_id == "uk"
269-
else us_dataset_version
270-
),
271-
)
272-
)
286+
job_setup_options["workflow_id"] = execution_id
273287

274-
api_v2_execution = self.api_v2.run(sim_config)
275-
execution_id = self.api_v2.get_execution_id(
276-
api_v2_execution
277-
)
288+
progress_log = {
289+
**job_setup_options,
290+
"message": "Sim API job started",
291+
}
292+
logger.log_struct(progress_log, severity="INFO")
278293

279-
comparison_data["v2_id"] = execution_id
280-
281-
# Pass name and status to logs
282-
progress_log: V2V1Comparison = (
283-
V2V1Comparison.model_validate(
284-
{
285-
**comparison_data,
286-
"v1_impact": None,
287-
"v2_impact": None,
288-
"v1_v2_diff": None,
289-
"message": "APIv2 job started",
290-
}
291-
)
292-
)
293-
logger.log_struct(
294-
progress_log.model_dump(mode="json"), severity="INFO"
295-
)
296-
except Exception as e:
297-
try:
298-
trace = traceback.format_exc()
299-
# Send error log to GCP
300-
error_log: V2V1Comparison = (
301-
V2V1Comparison.model_validate(
302-
{
303-
**comparison_data,
304-
"v2_error": trace,
305-
}
306-
)
307-
)
308-
logger.log_struct(
309-
error_log.model_dump(mode="json"), severity="ERROR"
310-
)
311-
except Exception as e2:
312-
print("Something really wrong.", e2)
313-
pass
314294
# Wait for job to complete
315-
if run_api_v2:
316-
317-
try:
318-
execution_id: str = self.api_v2.get_execution_id(
319-
api_v2_execution
320-
)
321-
api_v2_output = self.api_v2.wait_for_completion(
322-
api_v2_execution
323-
)
295+
sim_api_output = None
296+
try:
297+
sim_api_output = self.sim_api.wait_for_completion(
298+
sim_api_execution
299+
)
324300

325-
except Exception as e:
326-
trace = traceback.format_exc()
327-
# If job fails, send error log to GCP
328-
error_log: V2V1Comparison = V2V1Comparison.model_validate(
329-
{
330-
**comparison_data,
331-
"v2_error": trace,
332-
"v1_impact": None,
333-
"v2_impact": None,
334-
"v1_v2_diff": None,
335-
"message": "APIv2 job failed",
336-
}
337-
)
338-
logger.log_struct(
339-
error_log.model_dump(mode="json"), severity="ERROR"
340-
)
301+
except Exception as e:
302+
trace = traceback.format_exc()
303+
# If job fails, send error log to GCP
304+
error_log = {
305+
**job_setup_options,
306+
"message": "Sim API job failed",
307+
"error": str(e),
308+
"traceback": trace,
309+
}
310+
logger.log_struct(error_log, severity="ERROR")
341311

342312
# Finally, update all reform impact rows with the same baseline and reform policy IDs
343313
reform_impacts_service.set_complete_reform_impact(
@@ -348,7 +318,7 @@ def run(
348318
dataset=dataset,
349319
time_period=time_period,
350320
options_hash=options_hash,
351-
reform_impact_json=json.dumps(api_v2_output),
321+
reform_impact_json=json.dumps(sim_api_output),
352322
)
353323

354324
except Exception as e:
@@ -371,6 +341,66 @@ def run(
371341
)
372342
raise e
373343

344+
# Fallback method to be invoked via URL search param; otherwise,
345+
# to be deprecated in favor of new simulation API
346+
def _run_v1_job(
347+
self,
348+
country_id: str,
349+
region: str,
350+
dataset: str,
351+
time_period: str,
352+
options: dict,
353+
baseline_policy: dict,
354+
reform_policy: Annotated[str, "String-formatted JSON"],
355+
job_setup_options: dict[
356+
str, Any
357+
] = {}, # Dictionary of setup options for logging purposes
358+
) -> dict[str, Any]:
359+
# Compute baseline economy
360+
logger.log_struct(
361+
{
362+
"message": "Computing baseline economy using v1...",
363+
**job_setup_options,
364+
}
365+
)
366+
baseline_economy = self._compute_economy_v1(
367+
country_id=country_id,
368+
region=region,
369+
dataset=dataset,
370+
time_period=time_period,
371+
options=options,
372+
policy_json=baseline_policy,
373+
)
374+
375+
# Compute reform economy
376+
logger.log_struct(
377+
{
378+
"message": "Computing reform economy using v1...",
379+
**job_setup_options,
380+
}
381+
)
382+
reform_economy = self._compute_economy_v1(
383+
country_id=country_id,
384+
region=region,
385+
dataset=dataset,
386+
time_period=time_period,
387+
options=options,
388+
policy_json=reform_policy,
389+
)
390+
391+
baseline_economy = baseline_economy["result"]
392+
reform_economy = reform_economy["result"]
393+
logger.log_struct(
394+
{
395+
"message": "Computed baseline and reform economies using v1",
396+
**job_setup_options,
397+
}
398+
)
399+
impact: dict[str, Any] = compare_economic_outputs(
400+
baseline_economy, reform_economy, country_id=country_id
401+
)
402+
return impact
403+
374404
def _set_job_id(self) -> str:
375405
"""
376406
Generate a unique job ID based on the current timestamp and a random number.
@@ -380,7 +410,8 @@ def _set_job_id(self) -> str:
380410
random_number = np.random.randint(1000, 9999)
381411
return f"job_{timestamp}_{random_number}"
382412

383-
def _compute_economy(
413+
# To be deprecated in favor of new simulation API
414+
def _compute_economy_v1(
384415
self, country_id, region, dataset, time_period, options, policy_json
385416
):
386417
try:
@@ -615,7 +646,7 @@ def _compute_cliff_impacts(self, simulation: Microsimulation) -> Dict:
615646
}
616647

617648

618-
class SimulationAPIv2:
649+
class SimulationAPI:
619650
project: str
620651
location: str
621652
workflow: str

‎policyengine_api/routes/economy_routes.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ def get_economic_impact(
3131
region = options.pop("region")
3232
dataset = options.pop("dataset", "default")
3333
time_period = options.pop("time_period")
34+
run_api_v1 = options.pop("run_api_v1", False)
3435
api_version = options.pop(
3536
"version", COUNTRY_PACKAGE_VERSIONS.get(country_id)
3637
)
@@ -44,5 +45,6 @@ def get_economic_impact(
4445
time_period,
4546
options,
4647
api_version,
48+
run_api_v1=run_api_v1,
4749
)
4850
return result

0 commit comments

Comments
 (0)