来自代码的模型

可用性

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

目标用例

代码模型专为没有优化权重（GenAI 代理、应用程序、自定义逻辑）的模型设计。对于具有训练权重的传统 ML/DL 模型，请使用内置的 log_model() API 或带有 mlflow.pyfunc.log_model() 的自定义 PythonModel。

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

为什么选择代码模型？

主要区别在于模型在序列化时的表示方式

传统方法 - 使用 cloudpickle 或自定义序列化程序序列化模型对象，创建难以调试且具有兼容性限制的二进制文件。

代码模型 - 保存包含模型定义的简单 Python 脚本，使其可读、可调试且可在不同环境中移植。

主要优势

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

降低调试复杂性 - 不再需要反复尝试解决序列化问题。您编写的内容就是执行的内容。

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

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

核心要求

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

脚本执行

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

导入管理

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

外部依赖项

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

开发流程

在开发时使用 linter 来识别未使用的导入。这可以保持模型的依赖项干净，并使部署轻量化。

安全考虑

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

在 Jupyter Notebooks 中进行开发

Jupyter Notebooks 对于 AI 开发非常出色，但代码模型需要 Python 脚本（.py 文件）。幸运的是，IPython 的 %%writefile 魔术命令完美地弥补了这一差距。

使用 %%writefile

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

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

print("hello!")

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

print("hello!")

Jupyter 最佳实践

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

逐单元格开发 - 每个 %%writefile 单元格创建一个脚本文件。这使您的模型定义保持简洁和专注。

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

示例和模式

简单自定义模型

此示例演示了代码模型的基本知识和一个简单的数学模型。

创建模型脚本

# 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())

记录模型

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],
)

使用模型

# 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 文件。

创建辅助函数

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

# 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

创建主模型

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

# 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())

使用依赖项进行日志记录

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},
    )

测试多功能模型

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

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

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

# 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。

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

# ❌ 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

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 包含不必要的包。

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

# ❌ 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. 使用正确的安全实践重新记录模型

预防:

# ❌ 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 通过执行您的代码来验证其正确性。请确保您的日志记录环境已正确配置身份验证。

最佳实践:

# 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

性能优化

优化模型加载

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

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)

最小化依赖

条件导入：只导入特定操作所需的模块

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)

内存管理

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

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

从传统序列化迁移

如果您当前正在使用传统模型序列化，以下是迁移方法

之前 (传统)

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")

之后 (代码模型)

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