Track progress of calibrate using RabbitMQ (#9)

* merging things

* working service with latest changes

* Pika status

* set pyciemms to offical repo

* working

* working

* removed print

* remove test

* removed print

* Fix pika upstream

* Fix pika usage

* Deploy in separate image

* Fix mock consumer and README

---------

Co-authored-by: Five Grant <5@fivegrant.com>

README.md

@@ -17,11 +17,20 @@ in order to pull the PyCIEMSS repository in as a submodule and set up the enviro
 to start the containers and the API. The API url will be `http://localhost:8010` by default 
 
 ## Notes
+
+### Result Files
 Every operation saves 3 files to S3
 - `result.csv`
 - `eval.csv`
 - `visualization.json`
 
+### RabbitMQ
+Only the `calibrate` operation reports progress to RabbitMQ. This is to 
+the `simulation-status` queue with a payload that looks like `{"job_id": "some string", "progress": "float between 0 and 1"}`.
+
+The Docker Compose starts rabbitmq AND a mock consumer for the messages. The 
+mock consumer is only helpful for testing without the full stack. 
+
 
 ## License
 

docker/Dockerfile.api

@@ -8,6 +8,9 @@ RUN poetry config virtualenvs.create false && \
     poetry install --no-root --no-cache
 
 COPY service service
+COPY README.md README.md
+
+RUN poetry install
 
 ENV REDIS_HOST redis
 ENV REDIS_PORT 6379

docker/Dockerfile.worker

@@ -7,7 +7,7 @@ COPY poetry.lock poetry.lock
 RUN poetry config virtualenvs.create false && \
     poetry install --no-root --no-cache --extras worker
 
-RUN pip install --no-cache-dir git+https://github.com/ciemss/pyciemss.git
+RUN pip install --no-cache-dir git+https://github.com/fivegrant/pyciemss.git@fg/remove-pika
 COPY service service
 COPY README.md README.md
 

docker/docker-compose.yaml

@@ -20,8 +20,9 @@ services:
       - data-api
     depends_on:
       - redis
+      - rabbitmq
     volumes:
-      - $PWD/api:/api
+      - $PWD/service:/service
     extra_hosts:
       - "host.docker.internal:host-gateway"
   redis:
@@ -41,6 +42,32 @@ services:
     depends_on:
       - redis
       - api
+      - rabbitmq
     networks:
       - data-api
       - pyciemss
+  rabbitmq:
+    container_name: rabbitmq
+    hostname: rabbitmq
+    image: 'rabbitmq:3'
+    healthcheck:
+      test: ["CMD", "rabbitmqctl", "status"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+    ports:
+      - "5672:5672"
+    networks:
+      - pyciemss
+  rabbitmq-mock-consumer:
+    container_name: rabbitmq-mock-consumer
+    build:
+      context: ..
+      dockerfile: docker/Dockerfile.api
+    command: poetry run mockrabbitmq 
+    env_file:
+      - ../.env
+    networks:
+      - pyciemss
+    depends_on:
+      - rabbitmq

env.sample

@@ -1,4 +1,5 @@
+TDS_URL=http://data-service-api:8000
 REDIS_HOST=redis
 REDIS_PORT=6379
-TDS_URL=http://data-service-api:8000
-PYCIEMSS_OUTPUT_FILEPATH=result.csv
+RABBITMQ_HOST=rabbitmq.pyciemss
+RABBITMQ_PORT=5672

pyproject.toml

@@ -4,7 +4,7 @@ version = "1.6.0"
 description = "PyCIEMSS simulation service to run CIEMSS simulations"
 authors = ["Powell Fendley", "Five Grant"]
 readme = "README.md"
-packages = [{include = "api"}, {include = "workers"}]
+packages = [{include = "service"}]
 
 [tool.poetry.dependencies]
 python = "^3.9"
@@ -14,13 +14,18 @@ rq = "^1.15.0"
 redis = "^4.5.5"
 boto3 = "^1.26.147"
 uvicorn = "^0.22.0"
+pika = "1.3.2"
 pandas = "^2.0.0"
 torch = { version = "^2.0.1", optional = true }
 numpy = { version = "^1.24.0", optional = true }
+filelock = "^3.12.2"
 
 [tool.poetry.extras]
 worker = ["torch", "numpy"]
 
+[tool.poetry.scripts]
+mockrabbitmq = "service.utils.rabbitmq:mock_rabbitmq_consumer"
+
 
 [build-system]
 requires = ["poetry-core"]

service/operations.py

@@ -9,6 +9,7 @@
 from settings import settings
 from utils.rq_helpers import update_status_on_job_fail
 from utils.tds import update_tds_status, cleanup_job_dir, fetch_dataset, fetch_model, attach_files
+from utils.rabbitmq import gen_rabbitmq_hook
 
 from pyciemss.PetriNetODE.interfaces import (
     load_and_calibrate_and_sample_petri_model,
@@ -80,6 +81,7 @@ def calibrate_then_simulate(request, *, job_id):
         petri_model_or_path=amr_path,
         timepoints=timepoints,
         data_path=dataset_path,
+        progress_hook=gen_rabbitmq_hook(job_id),
         visual_options=True,
         **request.extra.dict()
     )

service/server.py

@@ -13,8 +13,14 @@
     EnsembleCalibratePostRequest,
     StatusSimulationIdGetResponse,
 )
+import os
+import redis
+import sys
+from threading import Thread
+import time
 
 from utils.rq_helpers import create_job, fetch_job_status, kill_job
+from utils.rabbitmq import mock_rabbitmq_consumer
 
 
 logging.basicConfig()
@@ -79,7 +85,6 @@ def cancel_job(simulation_id: str) -> StatusSimulationIdGetResponse:
     return {"status": Status.from_rq(status)}
 
 
-import logging
 @app.post("/simulate", response_model=JobResponse)
 def simulate_model(body: SimulatePostRequest) -> JobResponse:
     """
@@ -118,4 +123,3 @@ def create_calibrate_ensemble(body: EnsembleCalibratePostRequest) -> JobResponse
     resp = create_job("operations.ensemble_calibrate", body, "ensemble-calibrate")
     response = {"simulation_id": resp["id"]}
     return response
-

service/settings.py

@@ -10,9 +10,11 @@ class Settings(BaseSettings):
     """
     pyciemss-service configuration
     """
+    TDS_URL: str = "http://data-service-api:8000"
     REDIS_HOST: str = "redis"
     REDIS_PORT: int = 6379
-    TDS_URL: str = "http://data-service-api:8000"
+    RABBITMQ_HOST: str = "rabbitmq.pyciemss"
+    RABBITMQ_PORT: int = 5672
 
 
 settings = Settings()

service/utils/rabbitmq.py

@@ -0,0 +1,43 @@
+import pika, sys, os
+import json
+import redis
+import time
+import logging
+
+from settings import settings
+
+conn_config = pika.ConnectionParameters(host=settings.RABBITMQ_HOST, port=settings.RABBITMQ_PORT)
+
+
+def mock_rabbitmq_consumer():
+    logging.basicConfig()
+    logging.getLogger().setLevel(logging.DEBUG)
+    # TODO: Conditionally start on status of rabbitmq
+    time.sleep(10)
+    connection = pika.BlockingConnection(conn_config)
+    channel = connection.channel()
+
+    channel.queue_declare(queue='simulation-status')
+
+    def callback(ch, method, properties, body):
+        resp = json.loads(body)
+        logging.info("job_id:%s; progress:%s", resp['job_id'], resp['progress'])
+
+
+    channel.basic_consume(queue='simulation-status', on_message_callback=callback, auto_ack=True)
+
+    logging.info(' [*] Waiting for messages. To exit press CTRL+C')
+    channel.start_consuming()
+
+
+def gen_rabbitmq_hook(job_id):
+    connection = pika.BlockingConnection(conn_config)
+    channel = connection.channel()
+
+    def hook(progress):
+        channel.basic_publish(
+            exchange='',
+            routing_key='simulation-status',
+            body=json.dumps({"job_id":job_id, "progress":progress})
+        )
+    return hook