simplify interface further

Author: mfeurer

openml/extensions/extension_interface.py

@@ -159,7 +159,7 @@ def _run_model_on_fold(
         add_local_measures: bool,
         X_test: Optional[Union[np.ndarray, scipy.sparse.spmatrix, pd.DataFrame]] = None,
         n_classes: Optional[int] = None,
-    ) -> Tuple[List[List], List[List], 'OrderedDict[str, float]', Any]:
+    ) -> Tuple[List[List], List[List], 'OrderedDict[str, float]', Optional['OpenMLRunTrace']]:
         """Run a model on a repeat,fold,subsample triplet of the task and return prediction information.
 
         Returns the data that is necessary to construct the OpenML Run object. Is used by
@@ -230,21 +230,6 @@ def obtain_parameter_values(
     ################################################################################################
     # Abstract methods for hyperparameter optimization
 
-    def is_hpo_class(self, model: Any) -> bool:
-        """Check whether the model performs hyperparameter optimization.
-
-        Used to check whether an optimization trace can be extracted from the model after running
-        it.
-
-        Parameters
-        ----------
-        model : Any
-
-        Returns
-        -------
-        bool
-        """
-
     @abstractmethod
     def instantiate_model_from_hpo_class(
         self,
@@ -266,25 +251,3 @@ def instantiate_model_from_hpo_class(
         Any
         """
         # TODO a trace belongs to a run and therefore a flow -> simplify this part of the interface!
-
-    @abstractmethod
-    def obtain_arff_trace(
-        self,
-        model: Any,
-        trace_content: List[List],
-    ) -> 'OpenMLRunTrace':
-        """Create arff trace object from a fitted model and the trace content obtained by
-        repeatedly calling ``run_model_on_task``.
-
-        Parameters
-        ----------
-        model : Any
-            A fitted hyperparameter optimization model.
-
-        trace_content : List[List]
-            Trace content obtained by ``openml.runs.run_flow_on_task``.
-
-        Returns
-        -------
-        OpenMLRunTrace
-        """

openml/extensions/sklearn/extension.py

@@ -937,7 +937,7 @@ def _prevent_optimize_n_jobs(self, model):
         model:
             The model that will be fitted
         """
-        if self.is_hpo_class(model):
+        if self._is_hpo_class(model):
             if isinstance(model, sklearn.model_selection.GridSearchCV):
                 param_distributions = model.param_grid
             elif isinstance(model, sklearn.model_selection.RandomizedSearchCV):
@@ -975,7 +975,7 @@ def _can_measure_cputime(self, model: Any) -> bool:
             True if all n_jobs parameters will be either set to None or 1, False otherwise
         """
         if not (
-                isinstance(model, sklearn.base.BaseEstimator) or self.is_hpo_class(model)
+                isinstance(model, sklearn.base.BaseEstimator) or self._is_hpo_class(model)
         ):
             raise ValueError('model should be BaseEstimator or BaseSearchCV')
 
@@ -1002,7 +1002,7 @@ def _can_measure_wallclocktime(self, model: Any) -> bool:
             True if no n_jobs parameters is set to -1, False otherwise
         """
         if not (
-                isinstance(model, sklearn.base.BaseEstimator) or self.is_hpo_class(model)
+                isinstance(model, sklearn.base.BaseEstimator) or self._is_hpo_class(model)
         ):
             raise ValueError('model should be BaseEstimator or BaseSearchCV')
 
@@ -1231,7 +1231,7 @@ def _prediction_to_probabilities(
             else:
                 used_estimator = model_copy
 
-            if self.is_hpo_class(used_estimator):
+            if self._is_hpo_class(used_estimator):
                 model_classes = used_estimator.best_estimator_.classes_
             else:
                 model_classes = used_estimator.classes_
@@ -1283,28 +1283,13 @@ def _prediction_to_probabilities(
         else:
             raise TypeError(type(task))
 
-        return pred_y, proba_y, user_defined_measures, model_copy
+        if self._is_hpo_class(model_copy):
+            trace_data = self._extract_trace_data(model_copy, rep_no, fold_no)
+            trace = self._obtain_arff_trace(model_copy, trace_data)
+        else:
+            trace = None
 
-    def _extract_trace_data(self, model, rep_no, fold_no):
-        arff_tracecontent = []
-        for itt_no in range(0, len(model.cv_results_['mean_test_score'])):
-            # we use the string values for True and False, as it is defined in
-            # this way by the OpenML server
-            selected = 'false'
-            if itt_no == model.best_index_:
-                selected = 'true'
-            test_score = model.cv_results_['mean_test_score'][itt_no]
-            arff_line = [rep_no, fold_no, itt_no, test_score, selected]
-            for key in model.cv_results_:
-                if key.startswith('param_'):
-                    value = model.cv_results_[key][itt_no]
-                    if value is not np.ma.masked:
-                        serialized_value = json.dumps(value)
-                    else:
-                        serialized_value = np.nan
-                    arff_line.append(serialized_value)
-            arff_tracecontent.append(arff_line)
-        return arff_tracecontent
+        return pred_y, proba_y, user_defined_measures, trace
 
     def obtain_parameter_values(
         self,
@@ -1483,7 +1468,7 @@ def _openml_param_name_to_sklearn(
     ################################################################################################
     # Methods for hyperparameter optimization
 
-    def is_hpo_class(self, model: Any) -> bool:
+    def _is_hpo_class(self, model: Any) -> bool:
         """Check whether the model performs hyperparameter optimization.
 
         Used to check whether an optimization trace can be extracted from the model after
@@ -1518,7 +1503,7 @@ def instantiate_model_from_hpo_class(
         -------
         Any
         """
-        if not self.is_hpo_class(model):
+        if not self._is_hpo_class(model):
             raise AssertionError(
                 'Flow model %s is not an instance of sklearn.model_selection._search.BaseSearchCV'
                 % model
@@ -1527,7 +1512,28 @@ def instantiate_model_from_hpo_class(
         base_estimator.set_params(**trace_iteration.get_parameters())
         return base_estimator
 
-    def obtain_arff_trace(
+    def _extract_trace_data(self, model, rep_no, fold_no):
+        arff_tracecontent = []
+        for itt_no in range(0, len(model.cv_results_['mean_test_score'])):
+            # we use the string values for True and False, as it is defined in
+            # this way by the OpenML server
+            selected = 'false'
+            if itt_no == model.best_index_:
+                selected = 'true'
+            test_score = model.cv_results_['mean_test_score'][itt_no]
+            arff_line = [rep_no, fold_no, itt_no, test_score, selected]
+            for key in model.cv_results_:
+                if key.startswith('param_'):
+                    value = model.cv_results_[key][itt_no]
+                    if value is not np.ma.masked:
+                        serialized_value = json.dumps(value)
+                    else:
+                        serialized_value = np.nan
+                    arff_line.append(serialized_value)
+            arff_tracecontent.append(arff_line)
+        return arff_tracecontent
+
+    def _obtain_arff_trace(
         self,
         model: Any,
         trace_content: List,
@@ -1547,7 +1553,7 @@ def obtain_arff_trace(
         -------
         OpenMLRunTrace
         """
-        if not self.is_hpo_class(model):
+        if not self._is_hpo_class(model):
             raise AssertionError(
                 'Flow model %s is not an instance of sklearn.model_selection._search.BaseSearchCV'
                 % model

openml/runs/functions.py

@@ -381,6 +381,7 @@ def _run_task_get_arffcontent(
 ]:
     arff_datacontent = []  # type: List[List]
     arff_tracecontent = []  # type: List[List]
+    traces = []  # type: List[OpenMLRunTrace]
     # stores fold-based evaluation measures. In case of a sample based task,
     # this information is multiple times overwritten, but due to the ordering
     # of tne loops, eventually it contains the information based on the full
@@ -396,9 +397,11 @@ def _run_task_get_arffcontent(
     num_reps, num_folds, num_samples = task.get_split_dimensions()
     n_classes = None
 
+    n_fit = 0
     for rep_no in range(num_reps):
         for fold_no in range(num_folds):
             for sample_no in range(num_samples):
+                n_fit += 1
 
                 train_indices, test_indices = task.get_train_test_split_indices(
                     repeat=rep_no, fold=fold_no, sample=sample_no)
@@ -422,7 +425,7 @@ def _run_task_get_arffcontent(
                     pred_y,
                     proba_y,
                     user_defined_measures_fold,
-                    model_fold,
+                    trace,
                 ) = extension._run_model_on_fold(
                     model=model,
                     task=task,
@@ -437,12 +440,8 @@ def _run_task_get_arffcontent(
                 )
 
                 arff_datacontent_fold = []  # type: List[List]
-                # extract trace, if applicable
-                arff_tracecontent_fold = []  # type: List[List]
-                if extension.is_hpo_class(model_fold):
-                    arff_tracecontent_fold.extend(
-                        extension._extract_trace_data(model_fold, rep_no, fold_no)
-                    )
+                if trace is not None:
+                    traces.append(trace)
 
                 # add client-side calculated metrics. These is used on the server as
                 # consistency check, only useful for supervised tasks
@@ -489,7 +488,6 @@ def _calculate_local_measure(sklearn_fn, openml_name):
                     raise TypeError(type(task))
 
                 arff_datacontent.extend(arff_datacontent_fold)
-                arff_tracecontent.extend(arff_tracecontent_fold)
 
                 for measure in user_defined_measures_fold:
 
@@ -511,10 +509,13 @@ def _calculate_local_measure(sklearn_fn, openml_name):
                     user_defined_measures_per_sample[measure][rep_no][fold_no][
                         sample_no] = user_defined_measures_fold[measure]
 
-    # Note that we need to use a fitted model (i.e., model_fold, and not model)
-    # here, to ensure it contains the hyperparameter data (in cv_results_)
-    if extension.is_hpo_class(model):
-        trace = extension.obtain_arff_trace(model_fold, arff_tracecontent)  # type: Optional[OpenMLRunTrace]  # noqa E501
+    if len(traces) > 0:
+        if len(traces) != n_fit:
+            raise ValueError(
+                'Did not find enough traces (expected %d, found %d)' % (n_fit, len(traces))
+            )
+        else:
+            trace = OpenMLRunTrace.merge_traces(traces)
     else:
         trace = None
 

openml/runs/trace.py

@@ -1,8 +1,10 @@
-import arff
+from collections import OrderedDict
 import json
 import os
+from typing import List
+
+import arff
 import xmltodict
-from collections import OrderedDict
 
 PREFIX = 'parameter_'
 REQUIRED_ATTRIBUTES = [
@@ -344,11 +346,26 @@ def trace_from_xml(cls, xml):
             )
             trace[(repeat, fold, iteration)] = current
 
-        return cls(run_id, trace)
+        return cls(None, trace)
+
+    @classmethod
+    def merge_traces(cls, traces: List['OpenMLRunTrace']):
+        for i in range(1, len(traces)):
+            if traces[i] != traces[i - 1]:
+                raise ValueError('Cannot merge traces!')
+
+        merged_trace = OrderedDict()
+
+        for trace in traces:
+            for iteration in trace:
+                merged_trace[(iteration.repeat, iteration.fold, iteration.iteration)] = iteration
+
+        return cls(None, merged_trace)
+
 
     def __str__(self):
         return '[Run id: %d, %d trace iterations]' % (
-            self.run_id,
+            -1 if self.run_id is None else self.run_id,
             len(self.trace_iterations),
         )
 
@@ -448,3 +465,14 @@ def __str__(self):
             self.evaluation,
             self.selected,
         )
+
+    def __eq__(self, other):
+        if not isinstance(other, OpenMLTraceIteration):
+            return False
+        attributes = [
+            'repeat', 'fold', 'iteration', 'setup_string', 'evaluation', 'selected', 'paramaters',
+        ]
+        for attr in attributes:
+            if getattr(self, attr) != getattr(other, attr):
+                return False
+        return True

tests/test_extensions/test_sklearn_extension/test_sklearn_extension.py

@@ -1374,7 +1374,7 @@ def test__extract_trace_data(self):
         self.assertIn(clf.best_estimator_.hidden_layer_sizes, param_grid['hidden_layer_sizes'])
 
         trace_list = self.extension._extract_trace_data(clf, rep_no=0, fold_no=0)
-        trace = self.extension.obtain_arff_trace(clf, trace_list)
+        trace = self.extension._obtain_arff_trace(clf, trace_list)
 
         self.assertIsInstance(trace, OpenMLRunTrace)
         self.assertIsInstance(trace_list, list)