来自代码的模型

可用性

从代码构建模型功能适用于 MLflow 2.12.2 及以上版本。对于早期版本，请参考自定义 Python 模型文档中概述的旧版序列化方法。

目标用例

从代码构建模型功能适用于没有优化权重（GenAI Agents、应用程序、自定义逻辑）的模型。对于具有训练权重的传统 ML/DL 模型，请使用内置的 log_model() API 或自定义 PythonModel 配合 mlflow.pyfunc.log_model()。

从代码构建模型功能改变了您定义、存储和加载自定义模型及应用程序的方式。它不再依赖复杂的序列化，而是将您的模型保存为可读的 Python 脚本，从而使开发更加透明，调试也更加容易。

为什么选择从代码构建模型？

关键区别在于序列化期间模型表示方式的不同

旧版方法 - 使用 cloudpickle 或自定义序列化器对模型对象进行序列化，生成难以调试且存在兼容性限制的二进制文件。

从代码构建模型 - 保存带有模型定义的简单 Python 脚本，使其可读、可调试且可在不同环境间移植。

主要优势

透明度和可读性 - 您的模型代码以纯 Python 脚本的形式存储，方便在 MLflow UI 中直接理解和调试。

降低调试复杂度 - 告别因序列化问题导致的反复试错。您编写的代码将精确地执行。

更好的兼容性 - 解决了 pickle/cloudpickle 的限制，例如 Python 版本依赖、复杂对象序列化问题以及性能瓶颈。

增强的安全性 - 人类可读的代码使得部署前更容易审计和验证模型行为。

核心要求

理解这些关键概念将帮助您有效使用从代码构建模型功能

脚本执行

您的模型脚本将在日志记录期间执行以验证正确性。请确保在日志记录环境中正确配置了任何外部依赖项或身份验证。

导入管理

仅包含您实际使用的导入。MLflow 会从所有顶级导入中推断出需求，因此未使用的导入会不必要地增加模型依赖项的负担。

外部依赖

无法通过 pip 安装的包必须通过 code_paths 指定。系统不会自动捕获标准包导入之外的外部引用。

开发工作流程

在开发过程中，请使用 linter 来识别未使用的导入。这将使您的模型需求保持清洁，并使部署轻量化。

安全注意事项

模型代码以纯文本形式存储。切勿在脚本中包含 API 密钥或密码等敏感信息。请改用环境变量或安全的配置管理。

在 Jupyter Notebook 中开发

Jupyter Notebook 非常适合 AI 开发，但从代码构建模型功能需要 Python 脚本（.py 文件）。幸运的是，IPython 的 %%writefile 魔法命令可以完美地弥合这一差距。

使用 %%writefile

%%writefile 魔法命令会捕获单元格内容并将其写入文件

python

# %%writefile "./hello.py"  # Uncomment to create the file locally

print("hello!")

这将创建一个包含以下内容的 hello.py 文件

python

print("hello!")

Jupyter 的最佳实践

覆盖，而不是追加 - 使用默认的 %%writefile 行为，而不是 -a 追加选项，以避免代码重复和调试混乱。

逐单元格开发 - 每个 %%writefile 单元格创建一个脚本文件。这有助于保持模型定义清晰和专注。

即时测试 - 写入后可以立即运行生成的脚本以验证其是否正常工作。

示例和模式

简单的自定义模型

此示例通过一个简单的数学模型演示了从代码构建模型功能的基础知识。

创建模型脚本

python

# If running in a Jupyter notebook, uncomment the next line:
# %%writefile "./basic.py"

import pandas as pd
from typing import List, Dict
from mlflow.pyfunc import PythonModel
from mlflow.models import set_model


class BasicModel(PythonModel):
    def exponential(self, numbers):
        return {f"{x}": 2**x for x in numbers}

    def predict(self, context, model_input) -> Dict[str, float]:
        if isinstance(model_input, pd.DataFrame):
            model_input = list(model_input.iloc[0].values())
        return self.exponential(model_input)


# This tells MLflow which object to use for inference
set_model(BasicModel())

记录模型

python

import mlflow

mlflow.set_experiment("Basic Model From Code")

model_info = mlflow.pyfunc.log_model(
    python_model="basic.py",  # Path to your script
    name="arithmetic_model",
    input_example=[42.0, 24.0],
)

使用模型

python

# Load and use the model
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)

# Make predictions
result = loaded_model.predict([2.2, 3.1, 4.7])
print(result)  # {'2.2': 4.59, '3.1': 8.57, '4.7': 25.99}

多文件依赖

此示例展示了如何使用 code_paths 功能处理多个 Python 文件。

创建辅助函数

首先，创建一个包含共享函数的实用程序文件

python

# If running in a Jupyter notebook, uncomment the next line:
# %%writefile "./calculator.py"


def add(x, y):
    return x + y


def multiply(x, y):
    return x * y


def calculate_compound_interest(principal, rate, time):
    return principal * (1 + rate) ** time

创建主模型

接下来，创建使用辅助函数的模型

python

# If running in a Jupyter notebook, uncomment the next line:
# %%writefile "./math_model.py"

from mlflow.pyfunc import PythonModel
from mlflow.models import set_model
from calculator import add, multiply, calculate_compound_interest


class MathModel(PythonModel):
    def predict(self, context, model_input, params=None):
        operation = model_input.get("operation", "add")

        if operation == "add":
            return add(model_input["x"], model_input["y"])
        elif operation == "multiply":
            return multiply(model_input["x"], model_input["y"])
        elif operation == "compound_interest":
            return calculate_compound_interest(
                model_input["principal"], model_input["rate"], model_input["time"]
            )
        else:
            raise ValueError(f"Unknown operation: {operation}")


set_model(MathModel())

带依赖项的日志记录

python

import mlflow

mlflow.set_experiment("Math Model From Code")

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        python_model="math_model.py",
        name="math_model",
        code_paths=["calculator.py"],  # Include dependency
        input_example={"operation": "add", "x": 5, "y": 3},
    )

测试多函数模型

python

loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)

# Test different operations
print(loaded_model.predict({"operation": "add", "x": 10, "y": 5}))  # 15
print(loaded_model.predict({"operation": "multiply", "x": 4, "y": 7}))  # 28
print(
    loaded_model.predict(
        {"operation": "compound_interest", "principal": 1000, "rate": 0.05, "time": 10}
    )
)  # 1628.89

LangChain 集成

此示例演示了 MLflow 对 LangChain 从代码构建模型功能的原生支持，用于构建 AI 应用程序。

创建 LangChain 模型

python

# If running in a Jupyter notebook, uncomment the next line:
# %%writefile "./landscape_advisor.py"

import os
from operator import itemgetter
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI
import mlflow


def get_region(input_data):
    """Extract region from input, with fallback default."""
    default = "Virginia, USA"
    if isinstance(input_data[0], dict):
        return input_data[0].get("content", {}).get("region", default)
    return default


def get_area(input_data):
    """Extract area from input, with fallback default."""
    default = "5000 square feet"
    if isinstance(input_data[0], dict):
        return input_data[0].get("content", {}).get("area", default)
    return default


# Define the prompt template
prompt = PromptTemplate(
    template="""You are a highly accomplished landscape designer providing suggestions
    for landscape design decisions in a particular geographic region.

    Your goal is to suggest low-maintenance hardscape and landscape options using
    materials and plants native to the specified region. Include a general cost
    estimate based on the square footage provided.

    Region: {region}
    Square Footage: {area}

    Provide recommendations for a moderately sophisticated suburban housing community.""",
    input_variables=["region", "area"],
)

# Initialize the language model
model = ChatOpenAI(model="gpt-4o", temperature=0.95, max_tokens=4096)

# Create the chain using LangChain Expression Language (LCEL)
chain = (
    {
        "region": itemgetter("messages") | RunnableLambda(get_region),
        "area": itemgetter("messages") | RunnableLambda(get_area),
    }
    | prompt
    | model
    | StrOutputParser()
)

# Set this chain as the model
mlflow.models.set_model(chain)

记录 LangChain 模型

python

import mlflow

mlflow.set_experiment("Landscape Design Advisor")

input_example = {
    "messages": [
        {
            "role": "user",
            "content": {
                "region": "Austin, TX, USA",
                "area": "1750 square feet",
            },
        }
    ]
}

with mlflow.start_run():
    model_info = mlflow.langchain.log_model(
        lc_model="landscape_advisor.py",  # Path to your script
        name="landscape_chain",
        input_example=input_example,
    )

使用 LangChain 模型

python

# Load the model
landscape_advisor = mlflow.langchain.load_model(model_info.model_uri)

# Create a query
query = {
    "messages": [
        {
            "role": "user",
            "content": {
                "region": "Raleigh, North Carolina USA",
                "area": "3850 square feet",
            },
        },
    ]
}

# Get landscape recommendations
response = landscape_advisor.invoke(query)
print(response)

故障排除常见问题

依赖管理

加载模型时出现 NameError

问题：加载保存的模型时出现 NameError。

解决方案：确保所有必需的导入都在模型脚本中定义

python

# ❌ Bad - imports missing in script
def predict(self, context, model_input):
    return pd.DataFrame(model_input)  # NameError: pd not defined


# ✅ Good - imports included
import pandas as pd


def predict(self, context, model_input):
    return pd.DataFrame(model_input)

外部依赖项的 ImportError

问题：加载具有外部依赖项的模型时出现 ImportError。

解决方案：对非 PyPI 依赖项使用 code_paths

python

mlflow.pyfunc.log_model(
    python_model="my_model.py",
    name="model",
    code_paths=["utils.py", "helpers/"],  # Include external files
    extra_pip_requirements=["custom-package==1.0.0"],  # Manual requirements
)

膨胀的 requirements 文件

问题：requirements.txt 包含不必要的包。

解决方案：清理您的导入，仅包含您使用的内容

python

# ❌ Bad - unused imports
import pandas as pd
import numpy as np
import tensorflow as tf
import torch
from sklearn.ensemble import RandomForestClassifier


def predict(self, context, model_input):
    return {"result": model_input * 2}  # Only uses basic operations


# ✅ Good - minimal imports
def predict(self, context, model_input):
    return {"result": model_input * 2}

安全与日志记录

意外包含敏感数据

问题：API 密钥或密码已包含在模型脚本中。

立即采取行动:

1. 删除 MLflow 运行：使用 UI 或 delete_run() API
2. 清理工件：运行 mlflow gc CLI 命令
3. 轮换受损凭证
4. 使用适当的安全实践重新记录模型

预防:

python

# ❌ Bad - hardcoded secrets
api_key = "sk-abc123..."
model = ChatOpenAI(api_key=api_key)

# ✅ Good - environment variables
import os

model = ChatOpenAI(api_key=os.getenv("OPENAI_API_KEY"))

日志记录期间的模型执行

问题：您的模型在日志记录期间会进行外部调用。

解释：MLflow 通过执行您的代码来验证其正确性。请确保您的日志记录环境已配置正确的身份验证。

最佳实践:

python

# Handle authentication gracefully
try:
    # External service calls during initialization
    client = ExternalService(api_key=os.getenv("API_KEY"))
    client.validate_connection()
except Exception as e:
    print(f"Warning: Could not validate external service: {e}")
    # Continue with model definition

性能优化

优化模型加载

延迟加载：仅在需要时初始化昂贵的资源

python

class OptimizedModel(PythonModel):
    def __init__(self):
        self._expensive_resource = None

    @property
    def expensive_resource(self):
        if self._expensive_resource is None:
            self._expensive_resource = load_expensive_model()
        return self._expensive_resource

    def predict(self, context, model_input):
        return self.expensive_resource.predict(model_input)

最小化依赖项

条件导入：仅导入特定操作所需的库

python

def predict(self, context, model_input):
    operation = model_input.get("type", "simple")

    if operation == "complex":
        import tensorflow as tf  # Only import when needed

        return self._tensorflow_predict(model_input)
    else:
        return self._simple_predict(model_input)

内存管理

资源清理：在模型中正确管理资源

python

class ResourceManagedModel(PythonModel):
    def __enter__(self):
        self.connection = create_connection()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if hasattr(self, "connection"):
            self.connection.close()

    def predict(self, context, model_input):
        # Use self.connection for predictions
        pass

从旧版序列化迁移

如果您目前正在使用旧版模型序列化，请按以下方式进行迁移

之前（旧版）

python

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        return model_input * 2


# Log object instance
model_instance = MyModel()
mlflow.pyfunc.log_model(python_model=model_instance, name="model")

之后（从代码构建模型）

python

# Save as script: my_model.py
# %%writefile "./my_model.py"
import mlflow
from mlflow.pyfunc import PythonModel
from mlflow.models import set_model


class MyModel(PythonModel):
    def predict(self, context, model_input):
        return model_input * 2


set_model(MyModel())

# Log script path
mlflow.pyfunc.log_model(python_model="my_model.py", name="model")

最佳实践摘要

代码组织

• 保持模型脚本的专注和精简
• 为模型文件和函数使用描述性名称
• 使用 code_paths 将相关功能组织到单独的模块中

安全

• 切勿在模型脚本中硬编码敏感信息
• 使用环境变量进行配置
• 在记录之前审查代码，确保不包含任何敏感信息

性能

• 仅导入您需要的内容，以最大程度地减少依赖项
• 对昂贵的资源使用延迟加载
• 为长时间运行的模型考虑内存管理

开发工作流

• 在 Jupyter 中使用 %%writefile 进行快速原型开发
• 在记录前独立测试您的脚本
• 使用 linter 来捕获未使用的导入和其他问题

其他资源

有关相关主题的更多信息

• 模型 API 文档
• 在 MLflow 模型中管理依赖项
• 自定义 Python 模型