[auto_trainer] fix get_offline deprecated function usage (#787)

mlrun · Feb 6, 2024 · 754c7cb · 754c7cb
1 parent 438128d
commit 754c7cb
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 9 deletions.
diff --git a/auto_trainer/auto_trainer.py b/auto_trainer/auto_trainer.py
@@ -72,10 +72,9 @@ def _get_dataframe(
     # Getting the dataset:
     if mlrun.utils.StorePrefix.FeatureVector == store_uri_prefix:
         # FeatureVector case:
+        fv = mlrun.datastore.get_store_resource(dataset.artifact_url)
         label_columns = label_columns or dataset.meta.status.label_column
-        dataset = fs.get_offline_features(
-            dataset.meta.uri, drop_columns=drop_columns
-        ).to_dataframe()
+        dataset = fv.get_offline_features(drop_columns=drop_columns).to_dataframe()
 
         context.logger.info(f"label columns: {label_columns}")