diff --git a/docs/source/models/user_define.md b/docs/source/models/user_define.md index 9abda46fe..930b52b48 100644 --- a/docs/source/models/user_define.md +++ b/docs/source/models/user_define.md @@ -1,77 +1,86 @@ # 自定义模型 -## 编写模型proto文件 - -TorchEasyRec使用 [Protocol Buffer](https://developers.google.com/protocol-buffers/docs/pythontutorial) 定义配置文件格式。 - -在 `tzrec/protos/models/rank_model.proto` 中增加一个 `CustomRankModel` Message来定义模型配置 - -```protobuf -message CustomRankModel { - required MLP mlp = 1; - ... -}; +TorchEasyRec 支持从独立 Python 包加载自定义模型。用户无需修改 +`tzrec/protos/model.proto` 或 `tzrec/protos/models/` 下的公共配置，模型代码和 +配置 proto 可以单独维护，减少升级 TorchEasyRec 时的代码冲突。 + +## 目录结构 + +默认自定义包名为 `tzrec_custom`，推荐使用以下结构： + +```text +tzrec_custom/ +├── __init__.py +├── models/ +│ ├── __init__.py +│ └── custom_rank_model.py +└── protos/ + ├── __init__.py + └── custom_rank_model.proto ``` -在 `tzrec/protos/model.proto的在` 的 `oneof model`里面增加 `CustomRankModel` - -```protobuf -message ModelConfig { - ... - - oneof model { - ... - CustomRankModel custom_rank_model = 1001; - ... - } - ... -} -``` +一个自定义包可以包含多个模型。TorchEasyRec 启动时会递归导入 +`tzrec_custom.models` 下的非测试模块，具体使用哪个模型由 pipeline config 中的 +`class_path` 决定。 -生成proto python `*_pb2.py` 文件 +如需使用其他包名，在第一次导入 `tzrec` 前设置环境变量： ```bash -bash scripts/gen_proto.sh +export TZREC_CUSTOM_PACKAGE=my_project.tzrec_custom ``` -## 编写模型文件 +环境变量填写完整 Python 包名。运行命令时还需通过 `PYTHONPATH` 或安装 wheel +确保该包可以被 Python 导入。显式配置的包不存在时，TorchEasyRec 会直接报错； +未配置环境变量且默认 `tzrec_custom` 不存在时则保持原有行为。 -### 继承 +## 编写模型配置 proto -继承 `tzrec.models.model.BaseModel` 来实现自定义模型，需重载以下函数 +自定义 proto 可以直接复用 TorchEasyRec 的公共 message。以下模型配置使用了 +`tzrec.protos.MLP`： + +```protobuf +syntax = "proto2"; +package tzrec_custom.protos; -### 初始化: \_\_init\_\_ +import "tzrec/protos/module.proto"; -- 根据模型配置`model_config`和特征配置`features`构建子模块 +message CustomRankModelConfig { + required tzrec.protos.MLP mlp = 1; +} +``` -### 前向: predict +生成 Python binding： -- 根据输入的`batch`数据，进行前向推理，得到`predictions` - - `batch`为`tzrec.datasets.utils.Batch`的数据结构，包含`dense_features`（稠密特征）、`sparse_features`（稀疏特征）、`sequence_dense_features` (序列稠密特征) - - 一般可以将`batch` 传给`EmbeddingGroup`模块`tzrec.modules.embedding.EmbeddingGroup`得到分组的Embedding结果后，再进行进一步前向推理 +```bash +PYTHONPATH=. bash scripts/gen_proto.sh +``` -### 损失: init_loss & loss +`scripts/gen_proto.sh` 会在生成公共 proto 后，检查 +`tzrec_custom/protos/*.proto` 并生成对应的 `*_pb2.py` 和 `*_pb2.pyi`。自定义包名 +不是 `tzrec_custom` 时，生成命令需要使用相同的环境变量： -- `init_loss`函数用于根据模型损失函数配置初始化loss模块，写入到`self._loss_modules`中 -- `loss`函数用于根据输入的`predictions`和`batch`中的label，实际计算每个step的loss，返回一个`loss_dict` +```bash +TZREC_CUSTOM_PACKAGE=my_project.tzrec_custom \ +PYTHONPATH=. bash scripts/gen_proto.sh +``` -### 评估: init_metric & update_metric +如果自定义包作为独立 wheel 发布，也可以在该包的构建流程中自行生成 binding。 -- `init_metric`函数用于根据模型初始化metric模块，写入到`self._metric_modules`中 -- `update_metric`函数用于根据输入的`predictions`和`batch`中的label，更新metric模块的状态 +## 编写模型 -### 常用继承 +自定义模型需要继承 `tzrec.models.model.BaseModel`。排序、多目标排序和召回场景 +通常可以分别继承： -在排序、多目标排序、召回的场景下，可以直接继承以下子模型，可以只用重置前向推理函数 +- `tzrec.models.rank_model.RankModel` +- `tzrec.models.multi_task_rank.MultiTaskRank` +- `tzrec.models.match_model.MatchModel` -- 排序模型可直接继承 `tzrec.models.rank_model.RankModel` -- 多目标模型可直接继承 `tzrec.models.multi_task_rank.MultiTaskRank` -- 召回模型可直接继承 `tzrec.models.match_model.MatchModel` +模型模块必须在顶层导入对应的 `*_pb2.py`。这样自动导入模型时会同时注册 +protobuf descriptor，pipeline config 中的 `Any` 配置才能被解析。 -以排序模型为例 +以下代码展示了排序模型的主要结构： ```python -# tzrec/models/custom_rank_model.py from typing import Any, Dict, List, Optional import torch @@ -84,16 +93,11 @@ from tzrec.modules.embedding import EmbeddingGroup from tzrec.modules.mlp import MLP from tzrec.protos.model_pb2 import ModelConfig from tzrec.utils.config_util import config_to_kwargs +from tzrec_custom.protos.custom_rank_model_pb2 import CustomRankModelConfig class CustomRankModel(RankModel): - """CustomRankModel. - - Args: - model_config (ModelConfig): an instance of ModelConfig. - features (list): list of features. - labels (list): list of label names. - """ + """Example custom ranking model.""" def __init__( self, @@ -101,98 +105,74 @@ class CustomRankModel(RankModel): features: List[BaseFeature], labels: List[str], sample_weights: Optional[List[str]] = None, + custom_config: Optional[CustomRankModelConfig] = None, **kwargs: Any, ) -> None: - super().__init__(model_config, features, labels, sample_weights, **kwargs) - # 构建EmbeddingGroup - self.embedding_group = EmbeddingGroup( - features, list(model_config.feature_groups) + super().__init__( + model_config, + features, + labels, + sample_weights, + custom_config=custom_config, + **kwargs, + ) + self.embedding_group = EmbeddingGroup(features, self.feature_groups) + input_dim = sum( + self.embedding_group.group_total_dim(name) + for name in self.embedding_group.group_names() ) - # 构建MLP - total_in_dim = sum(self.embedding_group.group_total_dim(n) for n in self.embedding_group.group_names()) self.mlp = MLP( - in_features=total_in_dim, + in_features=input_dim, **config_to_kwargs(self._model_config.mlp), ) - final_dim = self.mlp.output_dim() - self.output_mlp = nn.Linear(final_dim, self._num_class) - # 初始化其他模块 - ... - + self.output = nn.Linear(self.mlp.output_dim(), self._num_class) def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Forward the model. - - Args: - batch (Batch): input batch data. - - Return: - predictions (dict): a dict of predicted result. - """ - grouped_features = self.embedding_group( - batch + """Run model prediction.""" + grouped_features = self.embedding_group(batch) + features = torch.cat( + [ + grouped_features[name] + for name in self.embedding_group.group_names() + ], + dim=-1, ) - features = torch.cat([grouped_features[name] for name in self.embedding_group.group_names()], dim=-1) - tower_output = self.mlp(features) - y = self.output_mlp(tower_output) - # 其他前向推理 - ... - - return self._output_to_prediction(y) + return self._output_to_prediction(self.output(self.mlp(features))) ``` -## 测试 - -编写 custom_rank_model.config - -``` +框架会将 `custom_model.config` 解包成强类型 message，并通过 `custom_config` 参数 +传入模型。`BaseModel` 同时会将 `self._model_config` 设置为该 message，因此继承 +现有基础模型时可以继续使用相同的配置访问方式。 -# 数据相关参数配置 -data_config { - ... -} +自定义模型仍需按所继承基础模型的要求实现或复用 `predict`、`init_loss`、 +`loss`、`init_metric` 和 `update_metric` 等接口。 -# 特征相关参数配置 -feature_configs { - ... -} -feature_configs { - ... -} +## 配置 pipeline -# 训练相关的参数配置 -train_config { - ... -} +`custom_model.config` 使用 `google.protobuf.Any` 保存用户定义的强类型配置： -# 评估相关参数配置 -eval_config { - ... -} - -# 模型相关参数配置 +```protobuf model_config { feature_groups { - group_name: 'group1' - feature_names: 'f1' - feature_names: 'f2' - ... - wide_deep: DEEP + group_name: "group1" + feature_names: "f1" + feature_names: "f2" + group_type: DEEP } - feature_groups { - group_name: 'group2' - feature_names: 'f3' - feature_names: 'f4' - ... - wide_deep: DEEP - } - ... - custom_rank_model { - mlp { - hidden_units: [64] + + custom_model { + class_path: "tzrec_custom.models.custom_rank_model.CustomRankModel" + config { + [type.googleapis.com/tzrec_custom.protos.CustomRankModelConfig] { + mlp { + hidden_units: 128 + hidden_units: 64 + dropout_ratio: 0.1 + } + } } - ... } + metrics { auc {} } @@ -202,7 +182,13 @@ model_config { } ``` -运行 +方括号中的名称是 proto 文件声明的 package 与 message 名，不是 Python 文件 +路径。模型模块会在 pipeline config 解析前自动导入，因此对应 descriptor 已经 +注册。 + +## 运行 + +默认使用 `tzrec_custom` 包时： ```bash PYTHONPATH=. torchrun --master_addr=localhost --master_port=32555 \ @@ -214,6 +200,10 @@ PYTHONPATH=. torchrun --master_addr=localhost --master_port=32555 \ --model_dir ${MODEL_DIR} ``` +使用其他自定义包时，训练、评估、预测和导出命令都需要设置同一个 +`TZREC_CUSTOM_PACKAGE`。`torchrun` 启动的各个 worker 会继承该环境变量。 + ### 打包发布 -参考[开发指南](../develop.md) +参考[开发指南](../develop.md)。自定义包需要包含模型代码和生成的 protobuf +binding，并保证运行环境可以导入该包。 diff --git a/scripts/gen_proto.sh b/scripts/gen_proto.sh index 92f7b79a8..3faab22ee 100644 --- a/scripts/gen_proto.sh +++ b/scripts/gen_proto.sh @@ -21,3 +21,10 @@ if [ ! -d protoc ]; then fi python -m grpc_tools.protoc -I . tzrec/protos/*.proto tzrec/protos/models/*.proto --python_out=. --pyi_out=. --doc_out=html,proto.html:docs/source --plugin=protoc-gen-doc=./${PROTO_DIR}/protoc-gen-doc + +CUSTOM_PACKAGE="${TZREC_CUSTOM_PACKAGE:-tzrec_custom}" +CUSTOM_PROTO_DIR="${CUSTOM_PACKAGE//.//}/protos" +CUSTOM_PROTO_FILES=("${CUSTOM_PROTO_DIR}"/*.proto) +if [ -e "${CUSTOM_PROTO_FILES[0]}" ]; then + python -m grpc_tools.protoc -I . "${CUSTOM_PROTO_FILES[@]}" --python_out=. --pyi_out=. +fi diff --git a/tzrec/main.py b/tzrec/main.py index 6da197bfa..699c145d2 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -21,6 +21,8 @@ import pyarrow as pa import torch +from google.protobuf import symbol_database +from google.protobuf.message import Message from torch import distributed as dist from torch import nn, optim from torch.amp import GradScaler @@ -92,6 +94,7 @@ export_model, ) from tzrec.utils.filesystem_util import url_to_fs +from tzrec.utils.load_class import load_by_path from tzrec.utils.logging_util import ProgressLogger, logger from tzrec.utils.online_dense_export_util import OnlineDenseExportManager from tzrec.utils.plan_util import create_planner, get_default_sharders @@ -149,17 +152,39 @@ def _create_model( Return: model: a EasyRec Model. """ - model_cls_name = config_util.which_msg(model_config, "model") - # pyre-ignore [16] - model_cls = BaseModel.create_class(model_cls_name) - - model: BaseModel = model_cls( - model_config, - features, - labels, - sample_weights=sample_weights, - sampler_type=sampler_type, - ) + custom_config: Optional[Message] = None + if model_config.WhichOneof("model") == "custom_model": + custom_model_config = model_config.custom_model + model_cls = load_by_path(custom_model_config.class_path) + if not isinstance(model_cls, type) or not issubclass(model_cls, BaseModel): + raise ValueError( + f"Custom model class {custom_model_config.class_path} must inherit " + "BaseModel." + ) + any_config = custom_model_config.config + try: + config_cls = symbol_database.Default().GetSymbol(any_config.TypeName()) + except KeyError as e: + raise ValueError( + f"Custom model config type {any_config.TypeName()} is not registered." + ) from e + custom_config = config_cls() + if not any_config.Unpack(custom_config): + raise ValueError( + f"Failed to unpack custom model config {any_config.TypeName()}." + ) + else: + model_cls_name = config_util.which_msg(model_config, "model") + # pyre-ignore [16] + model_cls = BaseModel.create_class(model_cls_name) + + model_kwargs = { + "sample_weights": sample_weights, + "sampler_type": sampler_type, + } + if custom_config is not None: + model_kwargs["custom_config"] = custom_config + model: BaseModel = model_cls(model_config, features, labels, **model_kwargs) kernel = Kernel[KernelProto.Name(model_config.kernel)] model.set_kernel(kernel) diff --git a/tzrec/main_test.py b/tzrec/main_test.py index 2afbbb133..df6600fc4 100644 --- a/tzrec/main_test.py +++ b/tzrec/main_test.py @@ -23,11 +23,14 @@ from parameterized import parameterized from tzrec.datasets.utils import RecordBatchTensor -from tzrec.main import _train_and_evaluate, predict, predict_checkpoint +from tzrec.main import _create_model, _train_and_evaluate, predict, predict_checkpoint +from tzrec.models.model import BaseModel from tzrec.optim.ema import DenseEMA from tzrec.protos.data_pb2 import DataConfig from tzrec.protos.eval_pb2 import EvalConfig from tzrec.protos.export_pb2 import ExportConfig +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.module_pb2 import MLP from tzrec.protos.optimizer_pb2 import DenseOptimizer, EMAConfig from tzrec.protos.pipeline_pb2 import EasyRecConfig from tzrec.utils import predict_util @@ -37,6 +40,30 @@ class MainTest(unittest.TestCase): """Tests for tzrec.main orchestration.""" + def test_create_custom_model(self) -> None: + """A custom model receives its unpacked protobuf configuration.""" + model_config = ModelConfig() + model_config.custom_model.class_path = "custom.models.CustomModel" + model_config.custom_model.config.Pack(MLP(hidden_units=[32, 16])) + + with mock.patch("tzrec.main.load_by_path", return_value=BaseModel): + model = _create_model(model_config, [], []) + + self.assertIsInstance(model, BaseModel) + self.assertEqual(list(model._model_config.hidden_units), [32, 16]) + + def test_create_custom_model_requires_base_model(self) -> None: + """Reject custom classes that do not implement the model contract.""" + model_config = ModelConfig() + model_config.custom_model.class_path = "custom.models.InvalidModel" + model_config.custom_model.config.Pack(MLP(hidden_units=[32])) + + with ( + mock.patch("tzrec.main.load_by_path", return_value=torch.nn.Linear), + self.assertRaisesRegex(ValueError, "must inherit BaseModel"), + ): + _create_model(model_config, [], []) + def test_train_and_evaluate_closes_exporter_and_ckpt_on_exception(self) -> None: """A training exception must still drain the exporter and ckpt manager. diff --git a/tzrec/models/model.py b/tzrec/models/model.py index 26ec63dbc..9b48f7758 100644 --- a/tzrec/models/model.py +++ b/tzrec/models/model.py @@ -17,6 +17,7 @@ import torch import torchmetrics +from google.protobuf.message import Message from torch import nn from torchrec.modules.embedding_modules import ( EmbeddingBagCollectionInterface, @@ -46,6 +47,7 @@ class BaseModel(BaseModule, metaclass=_meta_cls): features (list): list of features. labels (list): list of label names. sample_weights (list): sample weight names. + custom_config (Message, optional): unpacked custom model config. """ def __init__( @@ -54,6 +56,7 @@ def __init__( features: List[BaseFeature], labels: List[str], sample_weights: Optional[List[str]] = None, + custom_config: Optional[Message] = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -62,9 +65,11 @@ def __init__( self._features = features self._feature_groups = list(model_config.feature_groups) self._labels = labels - self._model_config = ( - getattr(model_config, self._model_type) if self._model_type else None - ) + self._model_config = custom_config + if self._model_config is None: + self._model_config = ( + getattr(model_config, self._model_type) if self._model_type else None + ) self._metric_modules = nn.ModuleDict() self._loss_modules = nn.ModuleDict() diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index d2c34ae0f..2d5223d5b 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -1,6 +1,7 @@ syntax = "proto2"; package tzrec.protos; +import "google/protobuf/any.proto"; import "tzrec/protos/models/rank_model.proto"; import "tzrec/protos/models/multi_task_rank.proto"; import "tzrec/protos/models/match_model.proto"; @@ -43,6 +44,11 @@ enum Kernel { CUTLASS = 2; } +message CustomModel { + required string class_path = 1; + required google.protobuf.Any config = 2; +} + message ModelConfig { repeated FeatureGroupConfig feature_groups = 1; @@ -81,6 +87,8 @@ message ModelConfig { // SID generation models SidRqvae sid_rqvae = 600; SidRqkmeans sid_rqkmeans = 601; + + CustomModel custom_model = 1000; } optional uint32 num_class = 2 [default = 1]; diff --git a/tzrec/utils/config_util_test.py b/tzrec/utils/config_util_test.py index a107ecd6c..88c0fb2d9 100644 --- a/tzrec/utils/config_util_test.py +++ b/tzrec/utils/config_util_test.py @@ -9,13 +9,42 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import unittest +from tzrec.protos.module_pb2 import MLP from tzrec.protos.pipeline_pb2 import EasyRecConfig from tzrec.utils import config_util +from tzrec.utils.test_util import make_test_dir class ConfigUtilTest(unittest.TestCase): + def test_load_custom_model_any_config(self): + config_path = os.path.join(make_test_dir(), "custom_model.config") + with open(config_path, "w") as f: + f.write( + """ + model_config { + custom_model { + class_path: "tzrec_custom.models.CustomRankModel" + config { + [type.googleapis.com/tzrec.protos.MLP] { + hidden_units: 32 + hidden_units: 16 + } + } + } + } + """ + ) + + pipeline_config = config_util.load_pipeline_config(config_path) + mlp_config = MLP() + self.assertTrue( + pipeline_config.model_config.custom_model.config.Unpack(mlp_config) + ) + self.assertEqual(list(mlp_config.hidden_units), [32, 16]) + def test_get_inference_batch_size(self): pipeline_config = EasyRecConfig() pipeline_config.data_config.batch_size = 16 diff --git a/tzrec/utils/load_class.py b/tzrec/utils/load_class.py index fe488e6d0..5d2029189 100644 --- a/tzrec/utils/load_class.py +++ b/tzrec/utils/load_class.py @@ -9,12 +9,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib import os import pkgutil import pydoc import traceback from abc import ABCMeta +_CUSTOM_PACKAGE_ENV = "TZREC_CUSTOM_PACKAGE" +_DEFAULT_CUSTOM_PACKAGE = "tzrec_custom" + def import_pkg(pkg_info, prefix_to_remove=None): """Import package. @@ -50,6 +54,34 @@ def import_pkg(pkg_info, prefix_to_remove=None): ) from e +def auto_import_package(package_name): + """Import every non-test module in a Python package recursively. + + Args: + package_name: fully qualified Python package name. + """ + package = importlib.import_module(package_name) + if not hasattr(package, "__path__"): + raise ValueError(f"{package_name} is not a Python package") + prefix = package.__name__ + "." + for pkg_info in pkgutil.walk_packages(package.__path__, prefix): + if not pkg_info.name.endswith("_test"): + importlib.import_module(pkg_info.name) + + +def _auto_import_custom_models(): + """Import models from the configured optional custom package.""" + configured_package = os.getenv(_CUSTOM_PACKAGE_ENV) + custom_package = configured_package or _DEFAULT_CUSTOM_PACKAGE + models_package = f"{custom_package}.models" + try: + auto_import_package(models_package) + except ModuleNotFoundError as e: + missing_optional_package = e.name in {custom_package, models_package} + if configured_package is not None or not missing_optional_package: + raise + + def auto_import(user_path=None): """Auto import python files. @@ -98,6 +130,8 @@ def auto_import(user_path=None): for pkg_info in pkgutil.iter_modules([dirname]): import_pkg(pkg_info, prefix_to_remove) + _auto_import_custom_models() + def register_class(class_map, class_name, cls): """Register a class into class_map. diff --git a/tzrec/utils/load_class_test.py b/tzrec/utils/load_class_test.py index 617a0d7b0..84afad035 100644 --- a/tzrec/utils/load_class_test.py +++ b/tzrec/utils/load_class_test.py @@ -9,14 +9,76 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import pkgutil import unittest +from types import SimpleNamespace +from unittest import mock import torch -from tzrec.utils.load_class import load_by_path +from tzrec.utils import load_class +from tzrec.utils.load_class import ( + _auto_import_custom_models, + auto_import_package, + load_by_path, +) class LoadClassTest(unittest.TestCase): + def test_auto_import_package(self): + package = SimpleNamespace( + __name__="tzrec_custom.models", __path__=["custom/models"] + ) + modules = [ + pkgutil.ModuleInfo(None, "tzrec_custom.models.rank", False), + pkgutil.ModuleInfo(None, "tzrec_custom.models.rank_test", False), + pkgutil.ModuleInfo(None, "tzrec_custom.models.match", False), + ] + with ( + mock.patch.object( + load_class.importlib, + "import_module", + side_effect=[package, mock.Mock(), mock.Mock()], + ) as import_module, + mock.patch.object( + load_class.pkgutil, + "walk_packages", + return_value=modules, + ), + ): + auto_import_package("tzrec_custom.models") + + self.assertEqual( + [call.args[0] for call in import_module.call_args_list], + [ + "tzrec_custom.models", + "tzrec_custom.models.rank", + "tzrec_custom.models.match", + ], + ) + + def test_default_custom_package_is_optional(self): + error = ModuleNotFoundError( + "No module named 'tzrec_custom'", name="tzrec_custom" + ) + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("tzrec.utils.load_class.auto_import_package", side_effect=error), + ): + _auto_import_custom_models() + + def test_configured_custom_package_is_required(self): + error = ModuleNotFoundError("No module named 'my_models'", name="my_models") + with ( + mock.patch.dict( + os.environ, {"TZREC_CUSTOM_PACKAGE": "my_models"}, clear=True + ), + mock.patch("tzrec.utils.load_class.auto_import_package", side_effect=error), + self.assertRaises(ModuleNotFoundError), + ): + _auto_import_custom_models() + def test_load_by_path(self): loaded_cls = load_by_path("nn.ReLU") self.assertEqual(loaded_cls, torch.nn.ReLU)