wedata-feature-engineering 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

wedata/feature_store/feature_table_client/feature_table_client.py

@@ -8,15 +8,18 @@ from pyspark.sql.streaming import StreamingQuery
 from pyspark.sql.types import StructType
 import os
 
-from feature_store.constants.constants import APPEND, DEFAULT_WRITE_STREAM_TRIGGER
+from wedata.feature_store.constants.constants import APPEND, DEFAULT_WRITE_STREAM_TRIGGER
+from wedata.feature_store.entities.feature_table import FeatureTable
+from wedata.feature_store.spark_client.spark_client import SparkClient
+from wedata.feature_store.utils import common_utils
 
 
 class FeatureTableClient:
     """特征表操作类"""
 
     def __init__(
-            self,
-            spark: SparkSession
+        self,
+        spark: SparkSession
     ):
         self._spark = spark
 
@@ -46,12 +49,6 @@ class FeatureTableClient:
                     f"DataFrame与schema不匹配。差异字段: {diff_fields if diff_fields else '字段类型不一致'}"
                 )
 
-    @staticmethod
-    def _validate_table_name(name: str):
-        """验证特征表命名规范"""
-        if name.count('.') < 2:
-            raise ValueError("特征表名称需符合<catalog>.<schema>.<table>格式")
-
     @staticmethod
     def _validate_key_conflicts(primary_keys: List[str], timestamp_keys: List[str]):
         """校验主键与时间戳键是否冲突"""
@@ -75,7 +72,8 @@ class FeatureTableClient:
             schema: Optional[StructType] = None,
             description: Optional[str] = None,
             tags: Optional[Dict[str, str]] = None
-    ):
+    ) -> FeatureTable:
+
         """
         创建特征表（支持批流数据写入）
 
@@ -85,6 +83,7 @@ class FeatureTableClient:
             df: 初始数据（可选，用于推断schema）
             timestamp_keys: 时间戳键（用于时态特征）
             partition_columns: 分区列（优化存储查询）
+            schema: 表结构定义（可选，当不提供df时必需）
             description: 业务描述
             tags: 业务标签
 
@@ -94,6 +93,7 @@ class FeatureTableClient:
         Raises:
             ValueError: 当schema与数据不匹配时
         """
+
         # 参数标准化
         primary_keys = self._normalize_params(primary_keys)
         timestamp_keys = self._normalize_params(timestamp_keys)
@@ -101,23 +101,25 @@ class FeatureTableClient:
 
         # 元数据校验
         self._validate_schema(df, schema)
-        #self._validate_table_name(name)
         self._validate_key_conflicts(primary_keys, timestamp_keys)
 
-        # 表名 格式：<catalog>.<schema>.<table>  catalog默认值：DataLakeCatalog，schema默认值：feature_store
-        table_name = f'DataLakeCatalog.feature_store.{name}'
+        # 表名校验
+        common_utils.validate_table_name(name)
+
+        # 构建完整表名
+        table_name = common_utils.build_full_table_name(name)
 
         # 检查表是否存在
         try:
             if self._spark.catalog.tableExists(table_name):
                 raise ValueError(
-                    f"表 '{table_name}' 已存在\n"
-                    "解决方案：\n"
-                    "1. 使用不同的表名\n"
-                    "2. 删除现有表: spark.sql(f'DROP TABLE {name}')\n"
+                    f"Table '{table_name}' already exists\n"
+                    "Solutions:\n"
+                    "1. Use a different table name\n"
+                    "2. Drop the existing table: spark.sql(f'DROP TABLE {name}')\n"
                 )
         except Exception as e:
-            raise ValueError(f"检查表存在性时出错: {str(e)}") from e
+            raise ValueError(f"Error checking table existence: {str(e)}") from e
 
         # 推断表schema
         table_schema = schema or df.schema
@@ -126,7 +128,7 @@ class FeatureTableClient:
         timestamp_keys_ddl = []
         for timestamp_key in timestamp_keys:
             if timestamp_key not in primary_keys:
-                raise ValueError(f"时间戳键 '{timestamp_key}' 必须是主键")
+                raise ValueError(f"Timestamp key '{timestamp_key}' must be a primary key")
             timestamp_keys_ddl.append(f"`{timestamp_key}` TIMESTAMP")
 
         #从环境变量获取额外标签
@@ -185,7 +187,19 @@ class FeatureTableClient:
             if df is not None:
                 df.write.insertInto(table_name)
         except Exception as e:
-            raise ValueError(f"建表失败: {str(e)}") from e
+            raise ValueError(f"Failed to create table: {str(e)}") from e
+
+        # 构建并返回FeatureTable对象
+        return FeatureTable(
+            name=name,
+            table_id=table_name,
+            description=description or "",
+            primary_keys=primary_keys,
+            partition_columns=partition_columns or [],
+            features=[field.name for field in table_schema.fields],
+            timestamp_keys=timestamp_keys or [],
+            tags=dict(**tags or {}, **env_tags)
+        )
 
     def write_table(
             self,
@@ -195,6 +209,7 @@ class FeatureTableClient:
             checkpoint_location: Optional[str] = None,
             trigger: Optional[Dict[str, Any]] = DEFAULT_WRITE_STREAM_TRIGGER
     ) -> Optional[StreamingQuery]:
+
         """
         写入特征表数据（支持批处理和流式写入）
 
@@ -215,10 +230,13 @@ class FeatureTableClient:
         # 验证写入模式
         valid_modes = ["append", "overwrite"]
         if mode not in valid_modes:
-            raise ValueError(f"无效的写入模式 '{mode}'，可选值: {valid_modes}")
+            raise ValueError(f"Invalid write mode '{mode}', valid options: {valid_modes}")
+
+        # 表名校验
+        common_utils.validate_table_name(name)
 
-        # 完整表名格式：<catalog>.<schema>.<table>
-        table_name = f'DataLakeCatalog.feature_store.{name}'
+        # 构建完整表名
+        table_name = common_utils.build_full_table_name(name)
 
         # 判断是否是流式DataFrame
         is_streaming = df.isStreaming
@@ -227,7 +245,7 @@ class FeatureTableClient:
             if is_streaming:
                 # 流式写入
                 if not checkpoint_location:
-                    raise ValueError("流式写入必须提供checkpoint_location参数")
+                    raise ValueError("Streaming write requires checkpoint_location parameter")
 
                 writer = df.writeStream \
                     .format("parquet") \
@@ -252,6 +270,7 @@ class FeatureTableClient:
             self,
             name: str
     ) -> DataFrame:
+
         """
         从特征表中读取数据
 
@@ -264,8 +283,12 @@ class FeatureTableClient:
         Raises:
             ValueError: 当表不存在或读取失败时抛出
         """
+
+        # 表名校验
+        common_utils.validate_table_name(name)
+
         # 构建完整表名
-        table_name = f'DataLakeCatalog.feature_store.{name}'
+        table_name = common_utils.build_full_table_name(name)
 
         try:
             # 检查表是否存在
@@ -278,10 +301,8 @@ class FeatureTableClient:
         except Exception as e:
             raise ValueError(f"读取表 '{table_name}' 失败: {str(e)}") from e
 
-    def drop_table(
-            self,
-            name: str
-    ) -> None:
+    def drop_table(self, name: str):
+
         """
         删除特征表（表不存在时抛出异常）
 
@@ -296,8 +317,12 @@ class FeatureTableClient:
             # 基本删除
             drop_table("user_features")
         """
+
+        # 表名校验
+        common_utils.validate_table_name(name)
+
         # 构建完整表名
-        table_name = f'DataLakeCatalog.feature_store.{name}'
+        table_name = common_utils.build_full_table_name(name)
 
         try:
             # 检查表是否存在
@@ -311,3 +336,33 @@ class FeatureTableClient:
             raise  # 直接抛出已知的ValueError
         except Exception as e:
             raise RuntimeError(f"删除表 '{table_name}' 失败: {str(e)}") from e
+
+    def get_table(
+            self,
+            name: str,
+            spark_client: SparkClient
+    ) -> FeatureTable:
+
+        """获取特征表元数据信息
+
+        参数:
+            name: 特征表名称
+            spark_client: Spark客户端
+
+        返回:
+            FeatureTable对象
+
+        异常:
+            ValueError: 当表不存在或获取失败时抛出
+        """
+
+        # 表名校验
+        common_utils.validate_table_name(name)
+
+        # 构建完整表名
+        table_name = common_utils.build_full_table_name(name)
+
+        try:
+            return spark_client.get_feature_table(table_name)
+        except Exception as e:
+            raise ValueError(f"获取表'{name}'元数据失败: {str(e)}") from e

wedata/feature_store/spark_client/spark_client.py

@@ -6,73 +6,52 @@ from pyspark.sql.catalog import Column
 from pyspark.sql.functions import when, isnull
 from pyspark.sql.types import StructType, StringType, StructField
 
-from feature_store.entities.feature import Feature
-from feature_store.entities.feature_table import FeatureTable
-from feature_store.entities.function_info import FunctionParameterInfo, FunctionInfo
-from feature_store.utils.common_utils import unsanitize_identifier
-from feature_store.utils.utils import sanitize_multi_level_name
+from wedata.feature_store.entities.feature import Feature
+from wedata.feature_store.entities.feature_table import FeatureTable
+from wedata.feature_store.entities.function_info import FunctionParameterInfo, FunctionInfo
+from wedata.feature_store.utils.common_utils import unsanitize_identifier, sanitize_multi_level_name
 
 
 class SparkClient:
     def __init__(self, spark: SparkSession):
         self._spark = spark
 
-    def createDataFrame(self, data, schema) -> DataFrame:
-        return self._spark.createDataFrame(data, schema)
-
-    def read_table(
-            self, qualified_table_name, as_of_delta_timestamp=None, streaming=False
-    ):
-        """
-        Reads a Delta table, optionally as of some timestamp.
-        """
-        if streaming and as_of_delta_timestamp:
-            raise ValueError(
-                "Internal error: as_of_delta_timestamp cannot be specified when"
-                " streaming=True."
-            )
-
-        base_reader = (
-            # By default, Structured Streaming only handles append operations. Because
-            # we have a notion of primary keys, most offline feature store operations
-            # are not appends. For example, FeatureStoreClient.write_table(mode=MERGE)
-            # will issue a MERGE operation.
-            # In order to propagate the non-append operations to the
-            # readStream, we set ignoreChanges to "true".
-            # For more information,
-            # see https://docs.databricks.com/delta/delta-streaming.html#ignore-updates-and-deletes
-            self._spark.readStream.format("delta").option("ignoreChanges", "true")
-            if streaming
-            else self._spark.read.format("delta")
-        )
-
-        if as_of_delta_timestamp:
-            return base_reader.option("timestampAsOf", as_of_delta_timestamp).table(
-                sanitize_multi_level_name(qualified_table_name)
-            )
-        else:
-            return base_reader.table(sanitize_multi_level_name(qualified_table_name))
 
     def get_current_catalog(self):
         """
-        Get current set catalog in the spark context.
+        获取当前Spark会话的catalog名称（使用spark.catalog.currentCatalog属性）
+
+        返回:
+            str: 当前catalog名称，如果未设置则返回None
         """
         try:
-            df = self._spark.sql("SELECT CURRENT_CATALOG()").collect()
-            return unsanitize_identifier(df[0][0])
-        except Exception as e:
+            return unsanitize_identifier(self._spark.catalog.currentCatalog())
+        except Exception:
             return None
 
     def get_current_database(self):
         """
-        Get current set database in the spark context.
+        获取Spark上下文中当前设置的数据库名称
+
+        返回:
+            str: 当前数据库名称，如果获取失败则返回None
         """
         try:
-            df = self._spark.sql("SELECT CURRENT_DATABASE()").collect()
-            return unsanitize_identifier(df[0][0])
-        except Exception as e:
+            # 使用Spark SQL查询当前数据库
+            df = self._spark.sql("SELECT CURRENT_DATABASE()")
+            # 获取第一行第一列的值并去除特殊字符
+            return unsanitize_identifier(df.first()[0])
+        except Exception:
+            # 捕获所有异常并返回None
             return None
 
+
+
+
+    def createDataFrame(self, data, schema) -> DataFrame:
+        return self._spark.createDataFrame(data, schema)
+
+
     def read_table(self, table_name):
         """读取Spark表数据
 
@@ -134,11 +113,6 @@ class SparkClient:
             ) for row in columns
         ]
 
-    def get_online_stores(self, table_name):
-        return None
-
-
-
     def get_feature_table(self, table_name):
 
         # 获取表元数据
@@ -170,19 +144,19 @@ class SparkClient:
         return FeatureTable(
             name=table_name,
             table_id=table_properties.get("table_id", table_name),
-            description=table.description or table_properties.get("description", table_name),
+            description=table.description or table_properties.get("comment", table_name),
             primary_keys=table_properties.get("primaryKeys", "").split(",") if table_properties.get("primaryKeys") else [],
             partition_columns=table.partitionColumnNames if hasattr(table, 'partitionColumnNames') else [],
             features=features,
             creation_timestamp=None,  # Spark表元数据不包含创建时间戳
-            online_stores=self.get_online_stores(table_name),
+            online_stores=None,
             notebook_producers=None,
             job_producers=None,
             table_data_sources=None,
             path_data_sources=None,
             custom_data_sources=None,
             timestamp_keys=table_properties.get("timestamp_keys"),
-            tags=table_properties.get("tags")
+            tags=table_properties
         )
 
     def _get_routines_with_parameters(self, full_routine_names: List[str]) -> DataFrame:

wedata/feature_store/training_set_client/training_set_client.py

@@ -1,35 +1,28 @@
-import json
 import logging
 import os
-from collections import defaultdict
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Set, Union
+from typing import Any, List, Optional, Set, Union
 
 import mlflow
-import yaml
-from mlflow.models import Model, ModelSignature
-from mlflow.utils.file_utils import TempDir, YamlSafeDumper, read_yaml
+from mlflow.models import Model
+from mlflow.utils.file_utils import TempDir, read_yaml
 from pyspark.sql import DataFrame
-from pyspark.sql.functions import struct
 
-from feature_store.entities.feature_column_info import FeatureColumnInfo
-from feature_store.entities.feature_function import FeatureFunction
-from feature_store.entities.feature_lookup import FeatureLookup
-from feature_store.entities.feature_spec import FeatureSpec
-from feature_store.entities.training_set import TrainingSet
-from feature_store.spark_client.spark_client import SparkClient
+from wedata.feature_store.constants import constants
+from wedata.feature_store.entities.feature_function import FeatureFunction
+from wedata.feature_store.entities.feature_lookup import FeatureLookup
+from wedata.feature_store.entities.feature_spec import FeatureSpec
+from wedata.feature_store.entities.training_set import TrainingSet
+from wedata.feature_store.spark_client.spark_client import SparkClient
 
-from feature_store.constants.constants import (
+from wedata.feature_store.constants.constants import (
     _NO_RESULT_TYPE_PASSED,
-    _PREBUILT_ENV_URI,
-    _USE_SPARK_NATIVE_JOIN,
-    _WARN,
-    MODEL_DATA_PATH_ROOT,
-    PREDICTION_COLUMN_NAME,
+    _USE_SPARK_NATIVE_JOIN
 )
 
-from feature_store.utils import common_utils, training_set_utils
-from feature_store.utils.feature_spec_utils import convert_to_yaml_string
+from wedata.feature_store.utils import common_utils, training_set_utils, uc_utils
+from wedata.feature_store.utils.signature_utils import get_mlflow_signature_from_feature_spec, \
+    drop_signature_inputs_and_invalid_params
 
 _logger = logging.getLogger(__name__)
 
@@ -46,7 +39,6 @@ class TrainingSetClient:
     def create_training_set(
         self,
         feature_spec: FeatureSpec,
-        feature_column_infos: List[FeatureColumnInfo],
         label_names: List[str],
         df: DataFrame,
         ft_metadata: training_set_utils._FeatureTableMetadata,
@@ -57,18 +49,6 @@ class TrainingSetClient:
             {odci.udf_name for odci in feature_spec.on_demand_column_infos},
         )
 
-        # TODO(divyagupta-db): Move validation from _validate_join_feature_data in feature_lookup_utils.py
-        #  to a helper function called here and in score_batch.
-
-        # Add consumer of each feature and instrument as final step
-        consumer_feature_table_map = defaultdict(list)
-        for feature in feature_column_infos:
-            consumer_feature_table_map[feature.table_name].append(feature.feature_name)
-        consumed_udf_names = [f.udf_name for f in feature_spec.function_infos]
-
-        # Spark query planning is known to cause spark driver to crash if there are many feature tables to PiT join.
-        # See https://docs.google.com/document/d/1EyA4vvlWikTJMeinsLkxmRAVNlXoF1eqoZElOdqlWyY/edit
-        # So we disable native join by default.
         training_set_utils.warn_if_non_photon_for_native_spark(
             kwargs.get(_USE_SPARK_NATIVE_JOIN, False), self._spark_client
         )
@@ -96,6 +76,12 @@ class TrainingSetClient:
         feature_lookups = [f for f in features if isinstance(f, FeatureLookup)]
         feature_functions = [f for f in features if isinstance(f, FeatureFunction)]
 
+        # 最多支持100个FeatureFunctions
+        if len(feature_functions) > training_set_utils.MAX_FEATURE_FUNCTIONS:
+            raise ValueError(
+                f"A maximum of {training_set_utils.MAX_FEATURE_FUNCTIONS} FeatureFunctions are supported."
+            )
+
         # 如果未提供标签，则用空列表初始化label_names
         label_names = common_utils.as_list(label, [])
         del label
@@ -137,7 +123,6 @@ class TrainingSetClient:
 
         return self.create_training_set(
             feature_spec,
-            column_infos.feature_column_infos,
             label_names,
             df,
             ft_metadata,
@@ -145,9 +130,6 @@ class TrainingSetClient:
         )
 
 
-
-
-
     def create_feature_spec(
         self,
         name: str,
@@ -194,3 +176,192 @@ class TrainingSetClient:
         )
 
         return feature_spec
+
+
+    def log_model(
+            self,
+            model: Any,
+            artifact_path: str,
+            *,
+            flavor: ModuleType,
+            training_set: Optional[TrainingSet],
+            registered_model_name: Optional[str],
+            await_registration_for: int,
+            infer_input_example: bool,
+            **kwargs,
+    ):
+        # Validate only one of the training_set, feature_spec_path arguments is provided.
+        # Retrieve the FeatureSpec, then remove training_set, feature_spec_path from local scope.
+        feature_spec_path = kwargs.pop("feature_spec_path", None)
+        if (training_set is None) == (feature_spec_path is None):
+            raise ValueError(
+                "Either 'training_set' or 'feature_spec_path' must be provided, but not both."
+            )
+        # Retrieve the FeatureSpec and then reformat tables in local metastore to 2L before serialization.
+        # This will make sure the format of the feature spec with local metastore tables is always consistent.
+        if training_set:
+            all_uc_tables = all(
+                [
+                    uc_utils.is_uc_entity(table_info.table_name)
+                    for table_info in training_set.feature_spec.table_infos
+                ]
+            )
+            # training_set.feature_spec is guaranteed to be 3L from FeatureStoreClient.create_training_set.
+            feature_spec = uc_utils.get_feature_spec_with_reformat_full_table_names(
+                training_set.feature_spec
+            )
+            label_type_map = training_set._label_data_types
+
+            labels = training_set._labels
+            df_head = training_set._df.drop(*labels).head()
+        else:
+            # FeatureSpec.load expects the root directory of feature_spec.yaml
+            root_dir, file_name = os.path.split(feature_spec_path)
+            if file_name != FeatureSpec.FEATURE_ARTIFACT_FILE:
+                raise ValueError(
+                    f"'feature_spec_path' must be a path to {FeatureSpec.FEATURE_ARTIFACT_FILE}."
+                )
+            feature_spec = FeatureSpec.load(root_dir)
+
+            # The loaded FeatureSpec is not guaranteed to be 3L.
+            # First call get_feature_spec_with_full_table_names to append the default metastore to 2L names,
+            # as get_feature_spec_with_reformat_full_table_names expects full 3L table names and throws otherwise.
+            # TODO (ML-26593): Consolidate this into a single function that allows either 2L/3L names.
+            feature_spec_with_full_table_names = (
+                uc_utils.get_feature_spec_with_full_table_names(feature_spec)
+            )
+            all_uc_tables = all(
+                [
+                    uc_utils.is_uc_entity(table_info.table_name)
+                    for table_info in feature_spec_with_full_table_names.table_infos
+                ]
+            )
+            feature_spec = uc_utils.get_feature_spec_with_reformat_full_table_names(
+                feature_spec_with_full_table_names
+            )
+            label_type_map = None
+            df_head = None
+        del training_set, feature_spec_path
+
+        override_output_schema = kwargs.pop("output_schema", None)
+        params = kwargs.pop("params", {})
+        params["result_type"] = params.get("result_type", _NO_RESULT_TYPE_PASSED)
+        # Signatures will ony be supported for UC-table-only models to
+        # mitigate new online scoring behavior from being a breaking regression for older
+        # models.
+        # See https://docs.google.com/document/d/1L5tLY-kRreRefDfuAM3crXvYlirkcPuUUppU8uIMVM0/edit#
+        try:
+            if all_uc_tables:
+                signature = get_mlflow_signature_from_feature_spec(
+                    feature_spec, label_type_map, override_output_schema, params
+                )
+            else:
+                _logger.warning(
+                    "Model could not be logged with a signature because the training set uses feature tables in "
+                    "Hive Metastore. Migrate the feature tables to Unity Catalog for model to be logged "
+                    "with a signature. "
+                    "See https://docs.databricks.com/en/machine-learning/feature-store/uc/upgrade-feature-table-to-uc.html for more details."
+                )
+                signature = None
+        except Exception as e:
+            _logger.warning(f"Model could not be logged with a signature: {e}")
+            signature = None
+
+        with TempDir() as tmp_location:
+            data_path = os.path.join(tmp_location.path(), "feature_store")
+            raw_mlflow_model = Model(
+                signature=drop_signature_inputs_and_invalid_params(signature)
+            )
+            raw_model_path = os.path.join(
+                data_path, constants.RAW_MODEL_FOLDER
+            )
+            if flavor.FLAVOR_NAME != mlflow.pyfunc.FLAVOR_NAME:
+                flavor.save_model(
+                    model, raw_model_path, mlflow_model=raw_mlflow_model, **kwargs
+                )
+            else:
+                flavor.save_model(
+                    raw_model_path,
+                    mlflow_model=raw_mlflow_model,
+                    python_model=model,
+                    **kwargs,
+                )
+            if not "python_function" in raw_mlflow_model.flavors:
+                raise ValueError(
+                    f"FeatureStoreClient.log_model does not support '{flavor.__name__}' "
+                    f"since it does not have a python_function model flavor."
+                )
+
+            # Re-use the conda environment from the raw model for the packaged model. Later, we may
+            # add an additional requirement for the Feature Store library. At the moment, however,
+            # the databricks-feature-store package is not available via conda or pip.
+            model_env = raw_mlflow_model.flavors["python_function"][mlflow.pyfunc.ENV]
+            if isinstance(model_env, dict):
+                # mlflow 2.0 has multiple supported environments
+                conda_file = model_env[mlflow.pyfunc.EnvType.CONDA]
+            else:
+                conda_file = model_env
+
+            conda_env = read_yaml(raw_model_path, conda_file)
+
+            # Check if databricks-feature-lookup version is specified in conda_env
+            lookup_client_version_specified = False
+            for dependency in conda_env.get("dependencies", []):
+                if isinstance(dependency, dict):
+                    for pip_dep in dependency.get("pip", []):
+                        if pip_dep.startswith(
+                                constants.FEATURE_LOOKUP_CLIENT_PIP_PACKAGE
+                        ):
+                            lookup_client_version_specified = True
+                            break
+
+            # If databricks-feature-lookup version is not specified, add default version
+            if not lookup_client_version_specified:
+                # Get the pip package string for the databricks-feature-lookup client
+                default_databricks_feature_lookup_pip_package = common_utils.pip_depependency_pinned_major_version(
+                    pip_package_name=constants.FEATURE_LOOKUP_CLIENT_PIP_PACKAGE,
+                    major_version=constants.FEATURE_LOOKUP_CLIENT_MAJOR_VERSION,
+                )
+                common_utils.add_mlflow_pip_depependency(
+                    conda_env, default_databricks_feature_lookup_pip_package
+                )
+
+            try:
+                if df_head is not None and infer_input_example:
+                    input_example = df_head.asDict()
+                else:
+                    input_example = None
+            except Exception:
+                input_example = None
+
+            # todo:
+            #feature_spec.save(data_path)
+
+            # Log the packaged model. If no run is active, this call will create an active run.
+            mlflow.pyfunc.log_model(
+                artifact_path=artifact_path,
+                loader_module=constants.MLFLOW_MODEL_NAME,
+                data_path=data_path,
+                code_path=None,
+                conda_env=conda_env,
+                signature=signature,
+                input_example=input_example,
+            )
+        if registered_model_name is not None:
+            # The call to mlflow.pyfunc.log_model will create an active run, so it is safe to
+            # obtain the run_id for the active run.
+            run_id = mlflow.tracking.fluent.active_run().info.run_id
+
+            # If the user provided an explicit model_registry_uri when constructing the FeatureStoreClient,
+            # we respect this by setting the registry URI prior to reading the model from Model
+            # Registry.
+            # todo:
+            # if self._model_registry_uri:
+            #     # This command will override any previously set registry_uri.
+            #     mlflow.set_registry_uri(self._model_registry_uri)
+
+            mlflow.register_model(
+                "runs:/%s/%s" % (run_id, artifact_path),
+                registered_model_name,
+                await_registration_for=await_registration_for,
+                )