Azure ML Studio SDK2中导入mlflow遇Dataset导入错误求助

问题概述 #

在Azure ML Studio（SDK2）中创建包含组件的管道时，运行阶段触发ImportError：组件脚本执行import mlflow时，无法从mlflow.entities导入Dataset，报错栈显示实际运行环境为Python 3.8，但自定义conda环境配置指定的是Python 3.10，且MLflow版本为1.26.1，azureml-mlflow版本为1.42.0。

核心原因 #

1. 版本兼容性冲突：MLflow 1.26.1的mlflow.entities模块中尚未包含Dataset类，而azureml-mlflow 1.42.0依赖的MLflow版本范围与1.26.1不匹配，导致内部导入逻辑出错。
2. 环境未正确加载：报错显示运行环境使用Python 3.8，与配置的Python 3.10不符，说明管道组件未正确引用自定义conda环境，而是使用了默认环境。

解决思路 #

1. 对齐MLflow与AzureML-MLflow版本 #

azureml-mlflow与mlflow存在严格的版本绑定关系，需确保两者版本兼容：

• 将mlflow版本调整为与azureml-mlflow 1.42.0匹配的版本，例如mlflow==1.30.0（该版本与azureml-mlflow 1.42.0兼容，且mlflow.entities已包含Dataset类）。
• 修改后的conda环境配置示例：

channels:
  - conda-forge
dependencies:
  - python=3.10
  - numpy=1.21.2
  - pip=21.2.4
  - scikit-learn=0.24.2
  - scipy=1.7.1
  - 'pandas>=1.1,<1.2'
  - pip:
      - 'inference-schema[numpy-support]==1.3.0'
      - xlrd==2.0.1
      - mlflow==1.30.0
      - azureml-mlflow==1.42.0
      - 'psutil>=5.8,<5.9'
      - 'tqdm>=4.59,<4.60'
      - ipykernel~=6.0
      - matplotlib
name: model-env

2. 确保组件正确引用自定义环境 #

在定义管道组件时，明确指定使用自定义conda环境：

• 在组件定义代码中，通过environment参数引用已创建的自定义环境对象，例如：

from azure.ai.ml import command, Environment

custom_env = Environment(
    conda_file="path/to/your/conda.yaml",
    image="mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04:latest"
)

component = command(
    code="./component_code",
    command="python com_row_to_col.py",
    environment=custom_env,
    # 其他组件参数
)

• 避免使用默认的AzureML环境，确保组件运行时加载指定的Python 3.10环境。

3. 清理环境缓存并重新构建 #

AzureML可能会缓存旧的环境镜像，需强制重新构建：

• 删除已创建的自定义环境，重新基于修改后的yaml创建新环境。
• 创建环境时设置build_local=False（远程构建），确保生成全新的环境镜像。
• 重启计算实例，确保新环境被加载。

4. 验证运行环境一致性 #

在组件脚本开头添加版本验证代码，确认实际运行环境与配置一致：

import sys
import mlflow
print(f"Python version: {sys.version}")
print(f"MLflow version: {mlflow.__version__}")

• 若输出的Python版本为3.8或MLflow版本不是指定的1.30.0，说明环境未正确加载，需检查组件的环境引用配置。

内容的提问来源于stack exchange，提问作者Paweł