Merge pull request #77 from fuzzylabs/SD-163-Easy-Run

OCarrollM · web-flow · commit f06a73205023 · 2025-03-20T16:57:41.000Z
[SD-163] Run without MLflow / DAGsHub download
diff --git a/ecoml/pyproject.toml b/ecoml/pyproject.toml
@@ -44,3 +44,8 @@ torch = [
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
+
+[tool.poetry]
+include = [
+    "src/ecoml/ecoml_models/**"
+]
diff --git a/ecoml/src/ecoml/.DS_Store b/ecoml/src/ecoml/.DS_Store
diff --git a/ecoml/src/ecoml/model_builder/model_inference.py b/ecoml/src/ecoml/model_builder/model_inference.py
@@ -1,13 +1,11 @@
-"""Load inference model from Mlflow registry."""
+"""Load inference model from MLflow registry."""
 
 import os
 from pathlib import Path
-from typing import Any, Literal, Optional
-
+from typing import Any, Literal
 import dagshub
 import mlflow
 import pandas as pd
-from rich import print
 from loguru import logger
 
 from ecoml.data_preparation.features import (
@@ -16,63 +14,78 @@
     get_pooling_features,
 )
 from ecoml.data_preparation.pytorch_utils import PytorchLayer
+import importlib.resources as pkg_resources
+import ecoml
 
 ALLOWED_LAYER_TYPES = Literal["convolutional", "pooling", "dense"]
 
 
 class InferenceModel:
     """Inference Model.
 
-    It downloads model from MLFlow Registry on DagsHub, if not present on the first run.
+    Downloads model from MLflow Registry on DagsHub if not already present.
     """
 
     def __init__(
         self,
         layer_type: ALLOWED_LAYER_TYPES,
         model_version: int,
         verbose: bool = False,
-        dagshub_repo_owner: Optional[str] = "fuzzylabs",
-        dagshub_repo_name: Optional[str] = "edge-vision-power-estimation",
+        dagshub_repo_owner: str = "fuzzylabs",
+        dagshub_repo_name: str = "edge-vision-power-estimation",
+        use_packaged_models: bool = False,
     ):
         self.layer_type = layer_type
         self.model_version = model_version
         self.verbose = verbose
         self.repo_name = dagshub_repo_name
         self.repo_owner = dagshub_repo_owner
-        self.base_model_dir = Path.cwd() / "ecoml_models" / self.layer_type
-        # Download model from MLFlow Registry if not present on first run
+        self.use_packaged_models = use_packaged_models
+
+        if not self.use_packaged_models:
+            self.base_model_dir = Path.cwd() / "ecoml_models" / self.layer_type
+        else:
+            self.base_model_dir = pkg_resources.files(ecoml).joinpath("ecoml_models", self.layer_type)
+
         self.runtime_model = self.load_model(model_type="runtime")
-        
 
-    def _download_model(self, model_uri: str, dst_path: str) -> None:
+    def _download_model(self, model_uri: str, dst_path: Path) -> None:
         """Download model from MLflow registry to local filesystem.
 
         Args:
             model_uri: URI pointing to model artifact.
-            dst_path: Path of the local filesystem destination directory
-                to which to download the specified artifacts.
+            dst_path: Destination directory path.
         """
         if self.verbose:
-            logger.info(f"Downloading model to {dst_path} folder")
+            logger.info(f"Downloading model to {dst_path}")
         dagshub.init(repo_name=self.repo_name, repo_owner=self.repo_owner, mlflow=True)
         mlflow.artifacts.download_artifacts(artifact_uri=model_uri, dst_path=str(dst_path))
 
     def load_model(self, model_type: str) -> Any:
-        """Download and load power or runtime model from MLflow Registry.
-
-        Download is skipped if model exists in the local filesystem.
+        """Download and load runtime or power model from MLflow Registry.
 
         Returns:
-            Power or runtime model from MLflow Registry.
+            Loaded runtime or power model.
         """
+        if self.use_packaged_models:
+            model_path = self.base_model_dir / model_type / "model.pkl"
+            if not model_path.exists():
+                raise FileNotFoundError(
+                    f"Packaged model file not found: {model_path}. "
+                    "Ensure ecoml_models are packaged correctly."
+                )
+            model_dir = model_path.parent
+            if self.verbose:
+                logger.info(f"Loading packaged {model_type} model from {model_dir}")
+            return mlflow.pyfunc.load_model(str(model_dir))
+
+        # Using DagsHub / MLFlow remote model
         model_name = f"{self.layer_type}_{model_type}_model"
         model_uri = f"models:/{model_name}/{self.model_version}"
         dst_path = self.base_model_dir / model_type
-        # TODO: Tighter check to see if current model version is present
-        # instead of checking only if directory exists
         version_file = dst_path / "version.txt"
-        need_download = True
 
+        need_download = True
         if version_file.exists():
             local_version = version_file.read_text().strip()
             if local_version == str(self.model_version):
@@ -83,23 +96,21 @@ def load_model(self, model_type: str) -> Any:
             version_file.write_text(str(self.model_version))
         else:
             if self.verbose:
-                logger.info("Model already downloaded")
+                logger.info(f"{model_type.capitalize()} model version {self.model_version} already downloaded.")
 
         if self.verbose:
-            logger.info(f"Loading the {model_type} trained model from {dst_path} folder")
-        logger.info(f"Loading {model_type} model from {dst_path}")
+            logger.info(f"Loading {model_type} model from {dst_path}")
 
-    def get_features(self, layer_info: PytorchLayer) -> pd.DataFrame:
-        """Get features for the model to run prediction.
+        return mlflow.pyfunc.load_model(str(dst_path))
 
-        Each layer type creates input features required by
-        power and runtime models using Pytorch model summary file.
+    def get_features(self, layer_info: PytorchLayer) -> pd.DataFrame:
+        """Get features for prediction based on layer type.
 
         Args:
-            layer_info: Pydantic class containing all layer information.
+            layer_info: Layer details.
 
         Returns:
-            Pandas dataframe containing input features.
+            DataFrame containing features.
         """
         if self.layer_type == "convolutional":
             features = get_convolutional_features(layer_info)
@@ -108,6 +119,4 @@ def get_features(self, layer_info: PytorchLayer) -> pd.DataFrame:
         elif self.layer_type == "dense":
             features = get_dense_features(layer_info)
         else:
-            raise ValueError(f"Unsupported layer type: {self.layer_type}")
-        
-        return pd.DataFrame([features])
+   
diff --git a/ecoml/src/ecoml/model_summary/test.py b/ecoml/src/ecoml/model_summary/test.py
@@ -0,0 +1,18 @@
+import torch
+import json
+
+# Load PyTorch quantized model
+pt_model = torch.load("sample_data/resnet18_quantised.pt", map_location="cpu")
+
+# Load JSON model summary
+with open("sample_data/resnet18_quantized.json", "r") as f:
+    json_model = json.load(f)
+
+# Print PyTorch model state dictionary keys
+print("\n--- PyTorch Model State Dict Keys ---")
+print(pt_model.state_dict().keys())  # This should list all layer names
+
+# Print first layer in JSON model for comparison
+print("\n--- JSON Model First Layer ---")
+first_layer_name = list(json_model.keys())[0]
+print(first_layer_name, "->", json_model[first_layer_name])
