MLflow 1.20.2从Azure Blob存储下载大体积模型工件失败问题咨询

系统信息 #

• 操作系统平台及版本：Windows 10
• MLflow安装方式：pip安装
• MLflow版本：1.20.2
• Python版本：Python 3.9.7

问题描述 #

我保存了一个.h5格式的keras模型，执行mlflow.keras.load_model("run:/id_run/model")加载模型时，等待近一小时仍未完成，终止执行后返回如下报错：

ERROR:root:Internal Python error in the inspect module.
Below is the traceback from this internal error.

ERROR:root:Internal Python error in the inspect module.
Below is the traceback from this internal error.

Traceback (most recent call last):
  File "~\anaconda3\envs\python_38\lib\site-packages\IPython\core\interactiveshell.py", line 3441, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-3-277d37cc6084>", line 1, in <module>
    keras_model = mlflow.keras.load_model("runs:/483745e28a864eceb738c852cf062774/model")
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\keras.py", line 585, in load_model
    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri)
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\tracking\artifact_utils.py", line 83, in _download_artifact_from_uri
    return get_artifact_repository(artifact_uri=root_uri).download_artifacts(
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\store\artifact\runs_artifact_repo.py", line 125, in download_artifacts
    return self.repo.download_artifacts(artifact_path, dst_path)
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\store\artifact\artifact_repo.py", line 180, in download_artifacts
    return download_artifact_dir(
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\store\artifact\artifact_repo.py", line 147, in download_artifact_dir
    download_artifact_dir(
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\store\artifact\artifact_repo.py", line 152, in download_artifact_dir
    download_artifact(
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\store\artifact\artifact_repo.py", line 129, in download_artifact
    self._download_file(
  File "~\AppData\Roaming\Python\Python38\site-packages\mlflow\store\artifact\azure_blob_artifact_repo.py", line 136, in _download_file
    container_client.download_blob(remote_full_path).readinto(file)
  File "~\anaconda3\envs\python_38\lib\site-packages\azure\storage\blob\_download.py", line 617, in readinto
    downloader.process_chunk(chunk)
  File "~\anaconda3\envs\python_38\lib\site-packages\azure\storage\blob\_download.py", line 129, in process_chunk
    chunk_data = self._download_chunk(chunk_start, chunk_end - 1)
  File "~\anaconda3\envs\python_38\lib\site-packages\azure\storage\blob\_download.py", line 211, in _download_chunk
    chunk_data = process_content(response, offset[0], offset[1], self.encryption_options)
  File "~\anaconda3\envs\python_38\lib\site-packages\azure\storage\blob\_download.py", line 52, in process_content
    content = b"".join(list(data))
  File "~\anaconda3\envs\python_38\lib\site-packages\azure\core\pipeline\transport\_requests_basic.py", line 158, in __next__
    chunk = next(self.iter_content_func)
  File "~\anaconda3\envs\python_38\lib\site-packages\requests\models.py", line 758, in generate
    for chunk in self.raw.stream(chunk_size, decode_content=True):
  File "~\anaconda3\envs\python_38\lib\site-packages\urllib3\response.py", line 576, in stream
    data = self.read(amt=amt, decode_content=decode_content)
  File "~\anaconda3\envs\python_38\lib\site-packages\urllib3\response.py", line 519, in read
    data = self._fp.read(amt) if not fp_closed else b""
  File "~\anaconda3\envs\python_38\lib\http\client.py", line 459, in read
    n = self.readinto(b)
  File "~\anaconda3\envs\python_38\lib\http\client.py", line 503, in readinto
    n = self.fp.readinto(b)
  File "~\anaconda3\envs\python_38\lib\socket.py", line 669, in readinto
    return self._sock.recv_into(b)
  File "~\anaconda3\envs\python_38\lib\ssl.py", line 1241, in recv_into
    return self.read(nbytes, buffer)
  File "~\anaconda3\envs\python_38\lib\ssl.py", line 1099, in read
    return self._sslobj.read(len, buffer)
KeyboardInterrupt

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "~\anaconda3\envs\python_38\lib\site-packages\IPython\core\interactiveshell.py", line 2061, in showtraceback
    stb = value._render_traceback_()
AttributeError: 'KeyboardInterrupt' object has no attribute '_render_traceback_'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "~\anaconda3\envs\python_38\lib\site-packages\IPython\core\ultratb.py", line 1101, in get_records
    return _fixed_getinnerframes(etb, number_of_lines_of_context, tb_offset)
  File "~\anaconda3\envs\python_38\lib\site-packages\IPython\core\ultratb.py", line 248, in wrapped
    return f(*args, **kwargs)
  File "~\anaconda3\envs\python_38\lib\site-packages\IPython\core\ultratb.py", line 281, in _fixed_getinnerframes
    records = fix_frame_records_filenames(inspect.getinnerframes(etb, context))
  File "~\anaconda3\envs\python_38\lib\inspect.py", line 1515, in getinnerframes
    frameinfo = (tb.tb_frame,) + getframeinfo(tb, context)
  File "~\anaconda3\envs\python_38\lib\inspect.py", line 1473, in getframeinfo
    filename = getsourcefile(frame) or getfile(frame)
  File "~\anaconda3\envs\python_38\lib\inspect.py", line 708, in getsourcefile
    if getattr(getmodule(object, filename), '__loader__', None) is not None:
  File "~\anaconda3\envs\python_38\lib\inspect.py", line 754, in getmodule
    os.path.realpath(f)] = module.__name__
  File "~\anaconda3\envs\python_38\lib\ntpath.py", line 647, in realpath
    path = _getfinalpathname(path)
KeyboardInterrupt

我的工件存储使用Azure Blob存储，MLflow服务部署在外部服务器上。
我在MLflow UI中确认模型存在，尝试通过UI的下载按钮下载该模型时，下载进度多次中断后自动重启，该模型大小为355MB。我后续将模型目录压缩为相近大小的文件作为工件上传后，下载时也出现了相同问题。

问题复现步骤 #

上传大小接近355MB的文件作为MLflow工件后，尝试下载该工件即可复现该问题。

解决方案 #

从报错栈可以定位到问题出在Azure Blob存储的大文件下载环节，MLflow 1.20.2版本内置的Azure Artifact仓库适配器默认没有配置分块下载和重试策略，网络波动就会导致下载中断反复重传，可通过以下方式解决：

1. 升级MLflow版本：直接升级到1.24.0及以上版本，该版本官方优化了Azure Blob的下载逻辑，默认开启分块下载和失败重试，大幅提升大文件传输稳定性。
2. 手动配置Azure存储客户端参数：如果暂时无法升级MLflow，可在执行加载代码前添加如下配置：

   import os
   # 配置Azure存储超时时间，单位为秒
   os.environ['AZURE_STORAGE_CONNECTION_TIMEOUT'] = '300'
   os.environ['AZURE_STORAGE_READ_TIMEOUT'] = '300'
   
   from azure.storage.blob import BlobClient
   # 设置单块下载最大大小为4MB，最大重试次数为10次
   BlobClient._default_configuration['max_single_get_size'] = 4 * 1024 * 1024
   BlobClient._default_configuration['retry_total'] = 10
   

3. 临时替代方案：如果上述配置仍不生效，可以先用Azure Blob官方客户端直接下载模型文件到本地，再调用keras.models.load_model()直接加载本地文件，跳过MLflow内置的下载逻辑。

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