Merge pull request #5 from dagardner-nv/david-mdd_dfp-example

Cache list of registered models, avoiding lookups for users which don't have models

examples/dfp_workflow/mlflow/Dockerfile

@@ -3,13 +3,13 @@ FROM python:3.8-slim-buster
 # Install curl for health check
 RUN apt update && \
  apt install -y --no-install-recommends \
- curl && \
+ curl libyaml-cpp-dev libyaml-dev && \
  apt autoremove -y && \
  apt clean all && \
  rm -rf /var/cache/apt/* /var/lib/apt/lists/*
 
 # Install python packages
-RUN pip install mlflow boto3 pymysql
+RUN pip install mlflow boto3 pymysql pyyaml
 
 # We run on port 5000
 EXPOSE 5000

examples/dfp_workflow/morpheus/dfp/stages/dfp_inference_stage.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import os
 import threading
 import time
 import typing
@@ -34,6 +35,15 @@
 logger = logging.getLogger("morpheus.{}".format(__name__))
 
 
+def get_registered_models():
+ client = MlflowClient(os.environ.get('DFP_TRACKING_URI'))
+ models = client.list_registered_models()
+ return set(model.name for model in models)
+
+
+REGISTERED_MODELS = get_registered_models()
+
+
 class ModelCache:
 
  def __init__(self, reg_model_name: str, model_uri: str) -> None:
@@ -115,6 +125,8 @@ def load_model(self, client):
  # Our model does not exist, use fallback
  self._child_user_model_cache = self._manager.load_user_model_cache(
  self._fallback_user_ids[0], fallback_user_ids=self._fallback_user_ids[1:])
+ else:
+ return model_cache
 
  # See if we have a child cache and use that
  if (self._child_user_model_cache is not None):
@@ -173,6 +185,9 @@ def load_model_cache(self, client, reg_model_name: str) -> ModelCache:
 
  # Cache miss. Try to check for a model
  try:
+ if reg_model_name not in REGISTERED_MODELS:
+ raise MlflowException("")
+
  latest_versions = client.get_latest_versions(reg_model_name)
 
  # Default to the first returned one