来自代码的模型

可用性

代码模型在 MLflow 2.12.2 及更高版本中可用。对于早期版本，请使用 自定义 Python 模型 文档中概述的旧版序列化方法。

目标用例

代码模型专为不带优化权重的模型（GenAI 代理、应用程序、自定义逻辑）而设计。对于具有训练权重的传统 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 magic 命令完美地弥合了这一差距。

使用 %%writefile

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

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.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 模型