fix-auto-log

Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED

Author: Vperiodt

sdk/python/feast/feature_store.py

@@ -311,13 +311,10 @@ def _resolve_feature_service_name(self, feature_refs: List[str]) -> Optional[str
             _logger.debug("Failed to resolve feature service name: %s", e)
             return None
 
-    def _auto_log_entity_df_info(self, entity_df, start_date=None, end_date=None):
-        """Log entity_df info to MLflow for reproducibility.
+    def _log_entity_df_metadata(self, entity_df, start_date=None, end_date=None):
+        """Log lightweight entity_df metadata to MLflow (type, row count, columns, query, dates).
 
-        Handles three entity_df types:
-        - pd.DataFrame: saves metadata + full parquet artifact (within configured limit)
-        - str (SQL query): logs the query as a param
-        - None (range-based): logs start_date/end_date
+        Always called during historical retrieval regardless of auto_log_entity_df.
         """
         try:
             import mlflow
@@ -350,15 +347,6 @@ def _auto_log_entity_df_info(self, entity_df, start_date=None, end_date=None):
                     cols = cols[:MLFLOW_PARAM_TRUNCATION_SLICE] + "..."
                 client.log_param(run_id, "feast.entity_df_columns", cols)
 
-                max_rows = mlflow_cfg.entity_df_max_rows
-                if len(entity_df) <= max_rows:
-                    import tempfile
-
-                    with tempfile.TemporaryDirectory() as tmp_dir:
-                        path = os.path.join(tmp_dir, "entity_df.parquet")
-                        entity_df.to_parquet(path, index=False)
-                        client.log_artifact(run_id, path)
-
             elif entity_df is None and (start_date or end_date):
                 client.set_tag(run_id, "feast.entity_df_type", "range")
                 if start_date:
@@ -367,7 +355,38 @@ def _auto_log_entity_df_info(self, entity_df, start_date=None, end_date=None):
                     client.log_param(run_id, "feast.end_date", str(end_date))
 
         except Exception as e:
-            _logger.debug("Failed to log entity_df info to MLflow: %s", e)
+            _logger.debug("Failed to log entity_df metadata to MLflow: %s", e)
+
+    def _log_entity_df_artifact(self, entity_df):
+        """Upload entity DataFrame as a parquet artifact to MLflow.
+
+        Only called when auto_log_entity_df is True and entity_df is a DataFrame
+        within the configured row limit.
+        """
+        try:
+            import mlflow
+
+            if mlflow.active_run() is None:
+                return
+            if not isinstance(entity_df, pd.DataFrame):
+                return
+
+            mlflow_cfg = self.config.mlflow
+            tracking_uri = mlflow_cfg.get_tracking_uri()
+            client = mlflow.MlflowClient(tracking_uri=tracking_uri)
+            run_id = mlflow.active_run().info.run_id
+
+            max_rows = mlflow_cfg.entity_df_max_rows
+            if len(entity_df) <= max_rows:
+                import tempfile
+
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    path = os.path.join(tmp_dir, "entity_df.parquet")
+                    entity_df.to_parquet(path, index=False)
+                    client.log_artifact(run_id, path)
+
+        except Exception as e:
+            _logger.debug("Failed to log entity_df artifact to MLflow: %s", e)
 
     def _init_openlineage_emitter(self) -> Optional[Any]:
         """Initialize OpenLineage emitter if configured and enabled."""
@@ -1750,10 +1769,12 @@ def get_historical_features(
                         tracking_uri=self.config.mlflow.get_tracking_uri(),
                     )
 
+                    self._log_entity_df_metadata(
+                        entity_df, start_date=start_date, end_date=end_date
+                    )
+
                     if self.config.mlflow.auto_log_entity_df:
-                        self._auto_log_entity_df_info(
-                            entity_df, start_date=start_date, end_date=end_date
-                        )
+                        self._log_entity_df_artifact(entity_df)
         except Exception as e:
             _logger.debug("MLflow auto-log failed for historical retrieval: %s", e)
 
@@ -3036,7 +3057,9 @@ async def get_online_features_async(
         """
         provider = self._get_provider()
 
-        return await provider.get_online_features_async(
+        _retrieval_start = time.monotonic()
+
+        response = await provider.get_online_features_async(
             config=self.config,
             features=features,
             entity_rows=entity_rows,
@@ -3046,6 +3069,52 @@ async def get_online_features_async(
             include_feature_view_version_metadata=include_feature_view_version_metadata,
         )
 
+        try:
+            if (
+                self.config.mlflow is not None
+                and self.config.mlflow.enabled
+                and self.config.mlflow.auto_log
+            ):
+                _log_fn = _get_mlflow_log_fn()
+                if _log_fn is not None:
+                    _duration = time.monotonic() - _retrieval_start
+                    _feature_refs = utils._get_features(
+                        self.registry, self.project, features, allow_cache=True
+                    )
+                    if isinstance(entity_rows, list):
+                        _entity_count = len(entity_rows)
+                    elif isinstance(entity_rows, Mapping):
+                        try:
+                            _first_col = next(iter(entity_rows.values()))
+                            if isinstance(_first_col, RepeatedValue):
+                                _entity_count = len(_first_col.val)
+                            else:
+                                _entity_count = len(_first_col)
+                        except Exception:
+                            _entity_count = 0
+                    else:
+                        _entity_count = 0
+                    _fs = features if isinstance(features, FeatureService) else None
+                    _fs_name = (
+                        features.name
+                        if isinstance(features, FeatureService)
+                        else self._resolve_feature_service_name(_feature_refs)
+                    )
+                    _log_fn(
+                        feature_refs=_feature_refs,
+                        entity_count=_entity_count,
+                        duration_seconds=_duration,
+                        retrieval_type="online",
+                        feature_service=_fs,
+                        feature_service_name=_fs_name,
+                        project=self.project,
+                        tracking_uri=self.config.mlflow.get_tracking_uri(),
+                    )
+        except Exception as e:
+            _logger.debug("MLflow auto-log failed for online retrieval: %s", e)
+
+        return response
+
     def retrieve_online_documents(
         self,
         query: Union[str, List[float]],

sdk/python/feast/mlflow_integration/client.py

@@ -144,9 +144,7 @@ def _log_required_features(self):
                 path = os.path.join(tmp_dir, "required_features.json")
                 with open(path, "w") as f:
                     json.dump(features, f)
-                self._client.log_artifact(
-                    run.info.run_id, path, artifact_path=""
-                )
+                self._client.log_artifact(run.info.run_id, path, artifact_path="")
         except Exception as e:
             _logger.debug("Failed to log required_features.json: %s", e)
 

sdk/python/feast/mlflow_integration/logger.py

@@ -207,9 +207,7 @@ def log_apply_to_mlflow(
             elif isinstance(obj, Entity) and obj.name != "__dummy":
                 entity_names.append(obj.name)
 
-        run = client.create_run(
-            experiment_id, run_name=f"apply_{project}"
-        )
+        run = client.create_run(experiment_id, run_name=f"apply_{project}")
         run_id = run.info.run_id
         try:
             client.set_tag(run_id, "feast.operation", "apply")
@@ -236,9 +234,7 @@ def log_apply_to_mlflow(
             client.log_metric(
                 run_id, "feast.apply.feature_services_count", len(fs_names)
             )
-            client.log_metric(
-                run_id, "feast.apply.entities_count", len(entity_names)
-            )
+            client.log_metric(run_id, "feast.apply.entities_count", len(entity_names))
         finally:
             client.set_terminated(run_id)
 
@@ -276,9 +272,7 @@ def log_materialize_to_mlflow(
         experiment_id = _get_or_create_experiment(client, experiment_name)
 
         op_type = "materialize_incremental" if incremental else "materialize"
-        run = client.create_run(
-            experiment_id, run_name=f"{op_type}_{project}"
-        )
+        run = client.create_run(experiment_id, run_name=f"{op_type}_{project}")
         run_id = run.info.run_id
         try:
             client.set_tag(run_id, "feast.operation", op_type)

sdk/python/tests/integration/test_mlflow_integration.py

@@ -548,11 +548,13 @@ def test_auto_log_entity_df_false_skips_artifact(
                 entity_df=entity_df,
             ).to_df()
 
+        run_data = client.get_run(run.info.run_id).data
         artifacts = [a.path for a in client.list_artifacts(run.info.run_id)]
         assert "entity_df.parquet" not in artifacts
-        assert "feast.entity_df_rows" not in client.get_run(run.info.run_id).data.params
-        tags = client.get_run(run.info.run_id).data.tags
-        assert tags["feast.feature_service"] == "driver_activity_v1"
+
+        assert "feast.entity_df_rows" in run_data.params
+        assert run_data.tags["feast.entity_df_type"] == "dataframe"
+        assert run_data.tags["feast.feature_service"] == "driver_activity_v1"
 
 
 class TestEntityDfBuilder: