Optimize utilize intermediate results

service/models/operations/calibrate.py

@@ -13,7 +13,7 @@
  fetch_and_convert_static_interventions,
  fetch_and_convert_dynamic_interventions,
 )
-from utils.rabbitmq import gen_rabbitmq_hook
+from utils.rabbitmq import gen_calibrate_rabbitmq_hook
 from utils.tds import fetch_dataset, fetch_model
 
 
@@ -77,7 +77,7 @@ def gen_pyciemss_args(self, job_id):
 
  # TODO: Test RabbitMQ
  try:
- hook = gen_rabbitmq_hook(job_id)
+ hook = gen_calibrate_rabbitmq_hook(job_id)
  except (socket.gaierror, AMQPConnectionError):
  logging.warning(
  "%s: Failed to connect to RabbitMQ. Unable to log progress", job_id

service/models/operations/ensemble_calibrate.py

@@ -11,7 +11,7 @@
 
 from models.base import Dataset, OperationRequest, Timespan, ModelConfig
 from models.converters import convert_to_solution_mapping
-from utils.rabbitmq import gen_rabbitmq_hook
+from utils.rabbitmq import gen_calibrate_rabbitmq_hook
 from utils.tds import fetch_dataset, fetch_model
 
 
@@ -52,7 +52,7 @@ def gen_pyciemss_args(self, job_id):
  dataset_path = fetch_dataset(self.dataset.dict(), job_id)
 
  try:
- hook = gen_rabbitmq_hook(job_id)
+ hook = gen_calibrate_rabbitmq_hook(job_id)
  except (socket.gaierror, AMQPConnectionError):
  logging.warning(
  "%s: Failed to connect to RabbitMQ. Unable to log progress", job_id

service/models/operations/optimize.py

@@ -2,6 +2,7 @@
 
 from typing import ClassVar, List, Optional
 from enum import Enum
+from utils.rabbitmq import OptimizeHook
 
 import numpy as np
 import torch
@@ -151,6 +152,11 @@ def gen_pyciemss_args(self, job_id):
  if step_size is not None and solver_method == "euler":
  solver_options["step_size"] = step_size
 
+ total_possible_iterations = (
+ extra_options.get("maxiter") + 1
+ ) * extra_options.get("maxfeval")
+ progress_hook = OptimizeHook(job_id, total_possible_iterations)
+
  return {
  "model_path_or_json": amr_path,
  "logging_step_size": self.logging_step_size,
@@ -173,6 +179,7 @@ def gen_pyciemss_args(self, job_id):
  "n_samples_ouu": n_samples_ouu,
  "solver_method": solver_method,
  "solver_options": solver_options,
+ "progress_hook": progress_hook,
  **extra_options,
  }
 

service/utils/rabbitmq.py

@@ -52,7 +52,7 @@ def callback(ch, method, properties, body):
  channel.start_consuming()
 
 
-def gen_rabbitmq_hook(job_id):
+def gen_calibrate_rabbitmq_hook(job_id):
  connection = pika.BlockingConnection(
  conn_config,
  )
@@ -68,3 +68,30 @@ def hook(progress, loss):
  )
 
  return hook
+
+
+class OptimizeHook:
+ def __init__(self, job_id, total_possible_iterations):
+ connection = pika.BlockingConnection(
+ conn_config,
+ )
+ self.channel = connection.channel()
+ self.job_id = job_id
+ self.result = []
+ self.step = 0
+ self.total_possible_iterations = total_possible_iterations
+
+ def __call__(self, current_results):
+ self.step += 1
+ self.channel.basic_publish(
+ exchange="",
+ routing_key="simulation-status",
+ body=json.dumps(
+ {
+ "job_id": self.job_id,
+ "progress": self.step,
+ "current_results": current_results.tolist(),
+ "total_possible_iterations": self.total_possible_iterations,
+ }
+ ),
+ )