Added Base MLFlow Image in Metadata for Import Model Component

assets/training/model_management/components/convert_model_to_mlflow/spec.yaml

@@ -13,6 +13,7 @@ environment: azureml://registries/azureml/environments/model-management/versions
 
 code: ../../src/
 command: |
+  pip install transformers==4.48.0 torch==2.2.2 numpy==1.23.5 pandas==1.5.3 urllib3==1.26.19 --no-cache-dir --force-reinstall
   # TODO: This has one disadvantage as shell logs wont be pushed to appinsights
   set -ex
   IFS=',' read -ra pip_pkgs <<< "$[[${{inputs.extra_pip_requirements}}]]"
@@ -39,7 +40,7 @@ inputs:
 
   vllm_enabled:
     type: boolean
-    description: Enable vllm in the converted model 
+    description: Enable vllm in the converted model
     default: false
     optional: false
 
@@ -141,7 +142,7 @@ inputs:
       Pip dependencies expressed as below. Do not use quotes for passing.
       eg: pkg1==1.0, pkg2, pkg3==1.0
     optional: true
-  
+
   inference_base_image:
     type: string
     description: |

assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/convertors.py

@@ -114,6 +114,14 @@ def _save(
         :type metadata: Optional[Dict]. Defaults to {}.
         """
         signatures = self._signatures or self.get_model_signature()
+        # Check and update Base MLFlow Image metadata in Import Model Component
+        if not self._vllm_enabled:
+            from azureml.model.mgmt.utils.common_utils import get_mlclient
+            mlclient = get_mlclient("azureml")
+            mlFlow_image = mlclient.environments.get("mlflow-model-inference", label="latest")
+            metadata["azureml.base_image"] = "mcr.microsoft.com/azureml/curated/mlflow-model-inference:" \
+                + str(mlFlow_image.version)
+            logger.info("Metadata: {}".format(metadata))
         # set metadata info
         metadata.update(fetch_mlflow_acft_metadata(
             base_model_name=self._model_id,

assets/training/model_management/src/azureml/model/mgmt/processors/transformers/convertors.py

@@ -183,6 +183,12 @@ def _save(
             metadata["azureml.base_image"] = "mcr.microsoft.com/azureml/curated/foundation-model-inference:" \
                 + str(vllm_image.version)
             logger.info("Metadata: {}".format(metadata))
+        else:
+            mlclient = get_mlclient("azureml")
+            mlFlow_image = mlclient.environments.get("mlflow-model-inference", label="latest")
+            metadata["azureml.base_image"] = "mcr.microsoft.com/azureml/curated/mlflow-model-inference:" \
+                + str(mlFlow_image.version)
+            logger.info("Metadata: {}".format(metadata))
 
         if self._model_flavor == "OSS":
             try: