diff --git a/api/models.py b/api/models.py index dd5f6cc..38c71a9 100644 --- a/api/models.py +++ b/api/models.py @@ -91,11 +91,16 @@ class Dataset(BaseModel): example={'postive_tests': 'infected'}, ) +class InterventionObject(BaseModel): + timestep: float + name: str + value: float class SimulatePostRequest(BaseModel): engine: Engine = Field(..., example="ciemss") model_config_id: str = Field(..., example="ba8da8d4-047d-11ee-be56") timespan: Timespan + interventions: List[InterventionObject] = Field(default_factory=list, example=[{"timestep":1,"name":"beta","value":.4}]) extra: SimulateExtra = Field( None, description="optional extra system specific arguments for advanced use cases", diff --git a/api/server.py b/api/server.py index 928d49a..8812a24 100644 --- a/api/server.py +++ b/api/server.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from fastapi import FastAPI, Response, status from fastapi.middleware.cors import CORSMiddleware @@ -13,6 +14,9 @@ ) +logging.basicConfig() +logging.getLogger().setLevel(logging.DEBUG) + def build_api(*args) -> FastAPI: api = FastAPI( @@ -53,7 +57,7 @@ def get_status(simulation_id: str) -> StatusSimulationIdGetResponse: from utils import fetch_job_status status = fetch_job_status(simulation_id) - print(status) + logging.info(status) if not isinstance(status, str): return status @@ -71,6 +75,10 @@ def simulate_model(body: SimulatePostRequest) -> JobResponse: model_config_id = body.model_config_id start = body.timespan.start end = body.timespan.end + interventions = [ + (intervention.timestep, intervention.name, intervention.value) for intervention in body.interventions + ] + operation_name = "operations.simulate" options = { @@ -79,11 +87,14 @@ def simulate_model(body: SimulatePostRequest) -> JobResponse: "start": start, "end": end, "extra": body.extra.dict(), - "visual_options": True + "visual_options": True, + "interventions": interventions } resp = create_job(operation_name=operation_name, options=options) + if len(interventions) > 0: + logging.info(f"{resp['id']} used interventions: {interventions}") response = {"simulation_id": resp["id"]} return response @@ -97,7 +108,7 @@ def calibrate_model(body: CalibratePostRequest) -> JobResponse: from utils import create_job # Parse request body - print(body) + logging.info(body) engine = str(body.engine).lower() model_config_id = body.model_config_id dataset = body.dataset @@ -105,7 +116,6 @@ def calibrate_model(body: CalibratePostRequest) -> JobResponse: end = body.timespan.end extra = body.extra.dict() - operation_name = "operations.calibrate_then_simulate" options = { "engine": engine, @@ -114,7 +124,7 @@ def calibrate_model(body: CalibratePostRequest) -> JobResponse: "end": end, "dataset": dataset.dict(), "extra": extra, - "visual_options": True + "visual_options": True, } resp = create_job(operation_name=operation_name, options=options) @@ -135,3 +145,4 @@ def create_ensemble(body: EnsemblePostRequest) -> JobResponse: ) + diff --git a/api/utils.py b/api/utils.py index 63cda45..f33131f 100644 --- a/api/utils.py +++ b/api/utils.py @@ -53,7 +53,7 @@ def create_job(operation_name: str, options: Optional[Dict[Any, Any]] = None): job = q.fetch_job(job_id) if STANDALONE: - print(f"OPTIONS: {options}") + logging.info(f"OPTIONS: {options}") ex_payload = { "engine": "ciemss", "model_config_id": options.get("model_config_id"), @@ -73,9 +73,9 @@ def create_job(operation_name: str, options: Optional[Dict[Any, Any]] = None): "engine": "ciemss", "workflow_id": job_id, } - print(payload) + logging.info(payload) sys.stdout.flush() - print(requests.put(post_url, json=json.loads(json.dumps(payload))).content) + logging.info(requests.put(post_url, json=json.loads(json.dumps(payload))).content) if job and force_restart: job.cleanup(ttl=0) # Cleanup/remove data immediately