Models From Code
Models from Code 在 MLflow 2.12.2 及更高版本可用。对于较早的版本,请使用在 Custom Python Model 文档中概述的遗留序列化方法。
Models from Code 专为没有优化权重的模型设计(GenAI 智能体、应用程序、自定义逻辑)。对于具有训练权重的传统 ML/DL 模型,请使用内置的 log_model() APIs 或使用自定义的 PythonModel 并调用 mlflow.pyfunc.log_model()。
Models from Code 改变了你定义、存储和加载自定义模型和应用程序的方式。它不是依赖复杂的序列化,而是将模型保存为可读的 Python 脚本,从而使开发更透明,调试也更加容易。
为什么要从代码中创建模型?
关键区别在于模型在序列化过程中如何表示:
遗留方法 - 使用 cloudpickle 或自定义序列化器对模型对象进行序列化,创建二进制文件,这些文件可能难以调试并且存在兼容性限制。
从代码生成模型 - 保存包含模型定义的简单 Python 脚本,使其可读、可调试,并可跨环境移植。

关键优势
透明性和可读性 - 你的模型代码以纯 Python 脚本的形式存储,便于在 MLflow UI 中直接理解和调试。
减少调试复杂性 - 不再需要针对序列化问题进行反复试错。你所编写的正是被执行的内容。
更好的兼容性 - 消除 pickle/cloudpickle 的限制,例如 Python 版本依赖、复杂对象序列化问题和性能瓶颈。
增强的安全性 - 可读性更高的代码使在部署前更容易审计和验证模型行为。
核心要求
理解这些关键概念将帮助你更有效地使用 Models from Code:
脚本执行
你的模型脚本会在记录期间执行以验证其正确性。请确保在日志记录环境中正确配置任何外部依赖或认证。
导入管理
只包含你实际使用的导入。MLflow 会从所有顶级导入推断依赖,因此未使用的导入会不必要地使模型的依赖变得臃肿。
外部依赖
无法通过 pip 安装的包必须通过 code_paths 指定。系统不会自动捕获标准包导入以外的外部引用。
在开发时使用 linter 来识别未使用的导入。这可以保持模型的依赖清晰并使部署更轻量。
模型代码以纯文本形式存储。切勿在脚本中包含敏感信息,例如 API 密钥或密码。请改用环境变量或安全的配置管理。
在 Jupyter 笔记本中进行开发
Jupyter 笔记本非常适合 AI 开发,但 Models from Code 需要 Python 脚本(.py 文件)。幸运的是,IPython 的 %%writefile magic 命令完美弥补了这一差距。
使用 %%writefile
%%writefile 魔法命令会捕获单元格内容并将其写入文件:
# %%writefile "./hello.py" # Uncomment to create the file locally
print("hello!")
这将创建一个包含以下内容的hello.py文件:
print("hello!")
Jupyter 的最佳实践
覆盖而非追加 - 使用默认的 %%writefile 行为而不是 -a 追加选项,以避免重复代码和调试混淆。
逐单元格开发 - 每个 %%writefile 单元格会创建一个脚本文件。这使你的模型定义保持简洁并更聚焦。
即时测试 - 您可以在编写生成的脚本后立即运行它,以验证其是否正常工作。
示例与模式
- 简单自定义模型
- 多文件依赖
- LangChain 集成
此示例通过一个简单的数学模型演示了来自代码的模型的基础知识。
创建模型脚本
# 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

这个示例演示了 MLflow 对 LangChain Models from Code 的原生支持,用于构建人工智能应用。
创建一个 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)

Troubleshooting Common Issues
- 依赖管理
- 安全与日志记录
- 性能优化
加载模型时出现名称错误
Problem: Getting NameError when loading your saved model.
解决方案: 确保在你的模型脚本中定义了所有必需的导入:
# ❌ 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 with External Dependencies
Problem: ImportError when loading models with external dependencies.
解决方案:使用 code_paths 来处理非 PyPI 依赖:
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 密钥或机密信息。
立即操作:
- 删除 MLflow 运行:使用 UI 或
delete_run()API - 清理工件:运行
mlflow gcCLI 命令 - 轮换被泄露的凭证
- 使用适当的安全实践重新记录模型
预防:
# ❌ 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进行快速原型开发 - 在记录日志之前,独立测试你的脚本
- 使用 linters 来捕获未使用的导入和其他问题
附加资源
有关相关主题的更多信息: