Merge branch 'master' into DBAAS-5102

Author: Myles Novick

splicemachine/features/feature_store.py

@@ -293,7 +293,8 @@ def get_training_set_from_view(self, training_view: str, features: Union[List[Fe
         r = make_request(self._FS_URL, Endpoints.TRAINING_SET_FROM_VIEW, RequestType.POST, self._basic_auth, { "view": training_view }, 
                         { "features": features, "start_time": start_time, "end_time": end_time })
         sql = r["sql"]
-        tvw = r["training_view"]
+        tvw = TrainingView(**r["training_view"])
+        features = [Feature(**f) for f in r["features"]]
 
         # Link this to mlflow for model deployment
         if self.mlflow_ctx and not return_sql:
@@ -528,16 +529,19 @@ def get_training_set_from_deployment(self, schema_name: str, table_name: str):
 
         r = make_request(self._FS_URL, Endpoints.TRAINING_SET_FROM_DEPLOYMENT, RequestType.GET, self._basic_auth, 
             { "schema": schema_name, "table": table_name })
-        metadata = r['metadata']
 
+        metadata = r['metadata']
         sql = r['sql']
-        features = r['features']
+
         tv_name = metadata['name']
         start_time = metadata['training_set_start_ts']
         end_time = metadata['training_set_end_ts']
 
+        tv = TrainingView(**r['training_view']) if 'training_view' in r else None
+        features = [Feature(**f) for f in r['features']]
+
         if self.mlflow_ctx:
-            self.link_training_set_to_mlflow(features, start_time, end_time, tv_name)
+            self.link_training_set_to_mlflow(features, start_time, end_time, tv)
         return self.splice_ctx.df(sql)
 
     def remove_feature(self, name: str):
@@ -551,14 +555,23 @@ def remove_feature(self, name: str):
         """
         make_request(self._FS_URL, Endpoints.FEATURES, RequestType.DELETE, self._basic_auth, { "name": name })
 
-    def get_training_set_features(self, training_set: str = None):
+    def get_deployments(self, schema_name: str = None, table_name: str = None, training_set: str = None):
         """
-        Returns a list of all features from an available Training Set, as well as details about that Training Set
+        Returns a list of all (or specified) available deployments
         :param schema_name: model schema name
         :param table_name: model table name
         :param training_set: training set name
         :return: List[Deployment] the list of Deployments as dicts
         """
+        return make_request(self._FS_URL, Endpoints.DEPLOYMENTS, RequestType.GET, self._basic_auth, 
+            { 'schema': schema_name, 'table': table_name, 'name': training_set })
+      
+    def get_training_set_features(self, training_set: str = None):
+        """
+        Returns a list of all features from an available Training Set, as well as details about that Training Set
+        :param training_set: training set name
+        :return: TrainingSet as dict
+        """
         r = make_request(self._FS_URL, Endpoints.TRAINING_SET_FEATURES, RequestType.GET, self._basic_auth, 
             { 'name': training_set })
         r['features'] = [Feature(**f) for f in r['features']]
@@ -798,6 +811,7 @@ def link_training_set_to_mlflow(self, features: Union[List[Feature], List[str]],
 
         self.mlflow_ctx._active_training_set: TrainingSet = ts
         ts._register_metadata(self.mlflow_ctx)
+
 
     def set_feature_store_url(self, url: str):
         self._FS_URL = url

splicemachine/features/training_set.py

@@ -48,7 +48,8 @@ def _register_metadata(self, mlflow_ctx):
                                              "Training Set was logged to the current active run. If you call "
                                              "fs.get_training_set or fs.get_training_set_from_view before starting an "
                                              "mlflow run, all following runs will assume that Training Set to be the "
-                                             "active Training Set, and will log the Training Set as metadata. For more "
-                                             "information, refer to the documentation. If you'd like to use a new "
-                                             "Training Set, end the current run, call one of the mentioned functions, "
-                                             "and start your new run.") from None
+                                             "active Training Set (until the next call to either of those functions), "
+                                             "and will log the Training Set as metadata. For more information, "
+                                             "refer to the documentation. If you'd like to use a new Training Set, "
+                                             "end the current run, call one of the mentioned functions, and start "
+                                             "your new run. Or, call mlflow.remove_active_training_set()") from None

splicemachine/features/utils/http_utils.py

@@ -28,6 +28,7 @@ class Endpoints:
     """
     Enum for Feature Store Endpoints
     """
+    DEPLOYMENTS: str = "deployments"
     FEATURES: str = "features"
     FEATURE_SETS: str = "feature-sets"
     FEATURE_SET_DESCRIPTIONS: str = "feature-set-descriptions"

splicemachine/mlflow_support/mlflow_support.py

@@ -438,6 +438,16 @@ def _start_run(run_id=None, tags=None, experiment_id=None, run_name=None, nested
 
     return SpliceActiveRun(active_run)
 
+@_mlflow_patch('remove_active_training_set')
+def _remove_active_training_set():
+    """
+    Removes the active training set from mlflow. This function deletes mlflows active training set (retrieved from
+    the feature store), which will in turn stop the automated logging of features to the active mlflow run. To recreate
+    an active training set, call fs.get_training_set or fs.get_training_set_from_view in the Feature Store.
+    """
+    if hasattr(mlflow,'_active_training_set'):
+        del mlflow._active_training_set
+
 
 @_mlflow_patch('log_pipeline_stages')
 def _log_pipeline_stages(pipeline):
@@ -1002,7 +1012,8 @@ def apply_patches():
     targets = [_register_feature_store, _register_splice_context, _lp, _lm, _timer, _log_artifact, _log_feature_transformations,
                _log_model_params, _log_pipeline_stages, _log_model, _load_model, _download_artifact,
                _start_run, _current_run_id, _current_exp_id, _deploy_aws, _deploy_azure, _deploy_db, _login_director,
-               _get_run_ids_by_name, _get_deployed_models, _deploy_kubernetes, _fetch_logs, _watch_job, _end_run, _set_mlflow_uri]
+               _get_run_ids_by_name, _get_deployed_models, _deploy_kubernetes, _fetch_logs, _watch_job, _end_run,
+               _set_mlflow_uri, _remove_active_training_set]
 
     for target in targets:
         gorilla.apply(gorilla.Patch(mlflow, target.__name__.lstrip('_'), target, settings=_GORILLA_SETTINGS))