refactor: update model logging approach and clean up utility functions

Author: NickyJhames

deep-learning/classification-with-keras/notebooks/register-model.ipynb

@@ -124,7 +124,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "006f7912",
    "metadata": {},
    "outputs": [
@@ -172,7 +172,10 @@
     "# ------------------------ Utils Import ------------------------\n",
     "sys.path.append(\"../src\")\n",
     "from onnx_utils import ModelExportConfig\n",
-    "from utils import load_config\n"
+    "from utils import load_config\n",
+    "\n",
+    "# ------------------------ MLflow Logger Import ------------------------\n",
+    "from src.mlflow import Logger"
    ]
   },
   {
@@ -243,7 +246,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "id": "9c6d53c6",
    "metadata": {},
    "outputs": [
@@ -257,116 +260,39 @@
     }
    ],
    "source": [
-    "class MNISTModel(mlflow.pyfunc.PythonModel):\n",
-    "    def load_context(self, context):\n",
-    "        \"\"\"\n",
-    "        Load model and configuration from artifacts.\n",
-    "        \"\"\"\n",
-    "        try:\n",
-    "            import tensorflow as tf\n",
-    "            import yaml\n",
-    "            \n",
-    "            # Load the model from artifacts\n",
-    "            self.model = tf.keras.models.load_model(context.artifacts[\"mnist_model\"])\n",
-    "            \n",
-    "            # Load configuration from artifacts (like vanilla-rag does)\n",
-    "            config_path = context.artifacts[\"config\"]\n",
-    "            with open(config_path, 'r') as f:\n",
-    "                self.config = yaml.safe_load(f)\n",
-    "            \n",
-    "            logger.info(\"✅ Model and configuration loaded successfully\")\n",
-    "\n",
-    "        except Exception as e:\n",
-    "            logger.error(f\"❌ Error loading context: {str(e)}\")\n",
-    "            raise\n",
-    "\n",
-    "    def predict(self, context, model_input, params=None):\n",
-    "        \"\"\"\n",
-    "        Computes the predicted digit, by converting the base64 to a numpy array.\n",
-    "        \"\"\"\n",
-    "        try:\n",
-    "            if isinstance(model_input, pd.DataFrame):\n",
-    "                image_input = model_input.iloc[0, 0]\n",
-    "            elif isinstance(model_input, list):\n",
-    "                image_input = model_input[0]\n",
-    "            else:\n",
-    "                image_input = str(model_input)\n",
-    "                \n",
-    "            base64_array = base64_to_numpy(image_input)\n",
-    "\n",
-    "            predictions = self.model.predict(base64_array)\n",
-    "            predicted_classes = np.argmax(predictions, axis=1)\n",
-    "            \n",
-    "            return predicted_classes.tolist()\n",
+    "def log_mnist_model_to_mlflow(model, artifact_path, config_path, demo_folder):\n",
+    "    \"\"\"\n",
+    "    Log MNIST model to MLflow using the new Logger approach.\n",
+    "    \"\"\"\n",
+    "    try:\n",
+    "        # Define input and output schema\n",
+    "        input_schema = Schema([\n",
+    "            ColSpec(\"string\", name=\"digit\"),\n",
+    "        ])\n",
+    "        output_schema = Schema([\n",
+    "            ColSpec(\"long\", name=\"prediction\"),\n",
+    "        ])\n",
     "        \n",
-    "        except Exception as e:\n",
-    "            logger.error(f\"❌ Error performing prediction: {str(e)}\")\n",
-    "            raise\n",
-    "    \n",
-    "    @classmethod\n",
-    "    def log_model(cls,artifact_path, config_path, demo_folder):\n",
-    "        \"\"\"\n",
-    "        Logs the model to MLflow with appropriate artifacts and schema.\n",
-    "        Now uses in-memory model loading for ONNX export efficiency.\n",
-    "        \"\"\"\n",
+    "        # Define model signature\n",
+    "        signature = ModelSignature(inputs=input_schema, outputs=output_schema)\n",
     "        \n",
-    "        try:\n",
-    "            sys.path.append(\"../src\")\n",
-    "            from onnx_utils import ModelExportConfig,log_model\n",
-    "            \n",
-    "            # Define input and output schema\n",
-    "            input_schema = Schema([\n",
-    "                ColSpec(\"string\", name=\"digit\"),\n",
-    "            ])\n",
-    "            output_schema = Schema([\n",
-    "                ColSpec(\"long\", name=\"prediction\"),\n",
-    "            ])\n",
-    "            \n",
-    "            # Define model signature\n",
-    "            signature = ModelSignature(inputs=input_schema, outputs=output_schema)\n",
-    "            \n",
-    "            # Save the model to disk for artifacts (still needed for MLflow artifacts)\n",
-    "            model.save(MODEL_PATH)\n",
-    "          \n",
-    "            artifacts = {\n",
-    "                \"mnist_model\": MODEL_PATH,\n",
-    "                \"config\": config_path\n",
-    "            }\n",
-    "            \n",
-    "            if demo_folder and os.path.exists(demo_folder):\n",
-    "                artifacts[\"demo\"] = demo_folder\n",
-    "                logger.info(f\"✅ Demo folder added to artifacts: {demo_folder}\")\n",
-    "      \n",
-    "            \n",
-    "            model_configs = [\n",
-    "                ModelExportConfig(\n",
-    "                    model=model,                            # 🚀 Pre-loaded Keras model object!\n",
-    "                    model_name=\"keras_mnist_onnx\",          # ONNX file naming\n",
-    "                    input_sample = np.random.random((1, 28, 28, 1)).astype(np.float32)\n",
-    "                  \n",
-    "                )    \n",
-    "            ]\n",
-    "\n",
-    "            log_model(\n",
-    "                    artifact_path=artifact_path,\n",
-    "                    python_model=cls(),\n",
-    "                    artifacts=artifacts,\n",
-    "                    signature=signature,\n",
-    "                    models_to_convert_onnx=model_configs,\n",
-    "                    pip_requirements=[\n",
-    "                    \"tensorflow>=2.0.0\",\n",
-    "                    \"numpy\",\n",
-    "                    \"pillow\",\n",
-    "                    \"streamlit>=1.28.0\",\n",
-    "                    \"pyyaml\"\n",
-    "                    ]\n",
-    "                )\n",
-    "            \n",
-    "            logger.info(\"✅ Model and artifacts successfully registered in MLflow\")\n",
+    "        # Save the model to disk for artifacts\n",
+    "        model.save(MODEL_PATH)\n",
+    "        \n",
+    "        # Use the new Logger to log the model\n",
+    "        Logger.log_model(\n",
+    "            artifact_path=artifact_path,\n",
+    "            config_path=config_path,\n",
+    "            demo_folder=demo_folder,\n",
+    "            signature=signature,\n",
+    "            model_path=MODEL_PATH\n",
+    "        )\n",
+    "        \n",
+    "        logger.info(\"✅ Model and artifacts successfully registered in MLflow using new Logger\")\n",
     "\n",
-    "        except Exception as e:\n",
-    "            logger.error(f\"❌ Error logging model: {str(e)}\")\n",
-    "            raise"
+    "    except Exception as e:\n",
+    "        logger.error(f\"❌ Error logging model: {str(e)}\")\n",
+    "        raise"
    ]
   },
   {
@@ -505,7 +431,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "id": "326467bc",
    "metadata": {},
    "outputs": [
@@ -630,8 +556,9 @@
     "    \n",
     "    logger.info(f\"📊 Test accuracy: {test_accuracy:.4f}\")\n",
     "\n",
-    "    # Log the model to MLflow using vanilla-rag pattern\n",
-    "    MNISTModel.log_model(\n",
+    "    # Log the model to MLflow using new Logger approach\n",
+    "    log_mnist_model_to_mlflow(\n",
+    "        model=model,\n",
     "        artifact_path=MODEL_NAME,\n",
     "        config_path=CONFIG_PATH,\n",
     "        demo_folder=DEMO_FOLDER\n",

deep-learning/classification-with-keras/requirements.txt

@@ -2,4 +2,4 @@ onnx>=1.14.0
 onnxruntime>=1.15.0
 numpy>=1.21.0
 tf2onnx>=1.15.0
-mlflow==2.21.2
+mlflow==3.1.0 #TODO: remove me when you finish

deep-learning/classification-with-keras/src/mlflow/loader.py

@@ -26,28 +26,22 @@ def _load_pyfunc(data_path: str):
         Model: Initialized MNIST model instance ready for prediction
     """
     from src.mlflow.model import Model
-    from src.utils import load_config, get_model_path
+    from src.utils import load_config
 
     logger.info(f"Loading MNIST Model from artifacts at: {data_path}")
 
-    # Set MODEL_ARTIFACTS_PATH for model loading
-    os.environ["MODEL_ARTIFACTS_PATH"] = data_path
-    
     config_path = os.path.join(data_path, "config.yaml")
     if not os.path.exists(config_path):
         raise FileNotFoundError(f"Configuration file not found at: {config_path}")
 
     config = load_config(config_path)
     logger.info("Configuration loaded successfully")
 
-    # Get model path from artifacts
-    model_path = get_model_path("model_keras_mnist.keras")
+    # The model file is always named "model_keras_mnist.keras" in artifacts
+    model_path = os.path.join(data_path, "model_keras_mnist.keras")
 
     if not os.path.exists(model_path):
-        # Fallback to direct path in artifacts
-        model_path = os.path.join(data_path, "model_keras_mnist.keras")
-        if not os.path.exists(model_path):
-            raise FileNotFoundError(f"Model file not found at: {model_path}")
+        raise FileNotFoundError(f"Model file not found at: {model_path}")
 
     logger.info(f"Model path resolved to: {model_path}")
 

deep-learning/classification-with-keras/src/utils.py

@@ -83,22 +83,3 @@ def get_ports_config(config: Dict[str, Any]) -> Dict[str, Any]:
         Ports configuration dictionary.
     """
     return config.get("ports", {})
-
-
-def get_model_path(model_name: str) -> str:
-    """
-    Get the full path to the model file using the artifacts path and model name.
-
-    Args:
-        model_name: Name of the model file or full path (will extract filename)
-
-    Returns:
-        Full path to the model file
-    """
-    # Extract just the filename if model_name contains a path
-    filename = os.path.basename(model_name)
-
-    artifacts_path = os.environ.get("MODEL_ARTIFACTS_PATH", "")
-    model_path = os.path.join(artifacts_path, filename)
-
-    return model_path