mfeurer · Oct 1, 2018 · Jun 19, 2018 · Jun 25, 2018 · Jun 25, 2018 · Jun 25, 2018
diff --git a/openml/runs/__init__.py b/openml/runs/__init__.py
@@ -1,8 +1,26 @@
 from .run import OpenMLRun
 from .trace import OpenMLRunTrace, OpenMLTraceIteration
-from .functions import (run_model_on_task, run_flow_on_task, get_run, list_runs,
-                        get_runs, get_run_trace, initialize_model_from_run,
-                        initialize_model_from_trace)
+from .functions import (
+    run_model_on_task,
+    run_flow_on_task,
+    get_run,
+    list_runs,
+    get_runs,
+    get_run_trace,
+    initialize_model_from_run,
+    initialize_model_from_trace,
+)
 
-__all__ = ['OpenMLRun', 'run_model_on_task', 'run_flow_on_task', 'get_run',
-           'list_runs', 'get_runs']
+__all__ = [
+    'OpenMLRun',
+    'OpenMLRunTrace',
+    'OpenMLTraceIteration',
+    'run_model_on_task',
+    'run_flow_on_task',
+    'get_run',
+    'list_runs',
+    'get_runs',
+    'get_run_trace',
+    'initialize_model_from_run',
+    'initialize_model_from_trace'
+]
diff --git a/openml/runs/functions.py b/openml/runs/functions.py
@@ -23,8 +23,7 @@
 from ..exceptions import OpenMLCacheException, OpenMLServerException
 from ..tasks import OpenMLTask
 from .run import OpenMLRun, _get_version_information
-from .trace import OpenMLRunTrace, OpenMLTraceIteration
-
+from .trace import OpenMLRunTrace
 
 # _get_version_info, _get_dict and _create_setup_string are in run.py to avoid
 # circular imports
@@ -128,6 +127,7 @@ def run_flow_on_task(flow, task, avoid_duplicate_runs=True, flow_tags=None,
                              'exist on the server according to flow_exists')
         _publish_flow_if_necessary(flow)
 
+    data_content, trace, fold_evaluations, sample_evaluations = res
     if not isinstance(flow.flow_id, int):
         # This is the usual behaviour, where the flow object was initiated off
         # line and requires some additional information (flow_id, input_id for
@@ -141,19 +141,23 @@ def run_flow_on_task(flow, task, avoid_duplicate_runs=True, flow_tags=None,
         # through "run_model_on_task"
         if flow.flow_id != flow_id:
             # This should never happen, unless user made a flow-creation fault
-            raise ValueError('Result flow_exists and flow.flow_id are not same. ')
+            raise ValueError(
+                "Result from API call flow_exists and flow.flow_id are not "
+                "same: '%s' vs '%s'" % (str(flow.flow_id), str(flow_id))
+            )
 
     run = OpenMLRun(
         task_id=task.task_id,
         flow_id=flow.flow_id,
         dataset_id=dataset.dataset_id,
         model=flow.model,
-        tags=tags,
         flow_name=flow.name,
+        tags=tags,
+        trace=trace,
+        data_content=data_content,
     )
     run.parameter_settings = OpenMLRun._parse_parameters(flow)
 
-    run.data_content, run.trace_content, run.trace_attributes, fold_evaluations, sample_evaluations = res
     # now we need to attach the detailed evaluations
     if task.task_type_id == 3:
         run.sample_evaluations = sample_evaluations
@@ -199,7 +203,7 @@ def get_run_trace(run_id):
     openml.runs.OpenMLTrace
     """
     trace_xml = openml._api_calls._perform_api_call('run/trace/%d' % run_id)
-    run_trace = _create_trace_from_description(trace_xml)
+    run_trace = OpenMLRunTrace.trace_from_xml(trace_xml)
     return run_trace
 
 
@@ -231,7 +235,7 @@ def initialize_model_from_trace(run_id, repeat, fold, iteration=None):
     Parameters
     ----------
     run_id : int
-        The Openml run_id. Should contain a trace file, 
+        The Openml run_id. Should contain a trace file,
         otherwise a OpenMLServerException is raised
 
     repeat: int
@@ -242,7 +246,7 @@ def initialize_model_from_trace(run_id, repeat, fold, iteration=None):
 
     iteration: int
         The iteration nr (column in trace file). If None, the
-        best (selected) iteration will be searched (slow), 
+        best (selected) iteration will be searched (slow),
         according to the selection criteria implemented in
         OpenMLRunTrace.get_selected_iteration
 
@@ -479,15 +483,19 @@ def _prediction_to_probabilities(y, model_classes):
     if isinstance(model_fold, sklearn.model_selection._search.BaseSearchCV):
         # arff_tracecontent is already set
         arff_trace_attributes = _extract_arfftrace_attributes(model_fold)
+        trace = OpenMLRunTrace.generate(
+            arff_trace_attributes,
+            arff_tracecontent,
+        )
     else:
-        arff_tracecontent = None
-        arff_trace_attributes = None
+        trace = None
 
-    return arff_datacontent, \
-           arff_tracecontent, \
-           arff_trace_attributes, \
-           user_defined_measures_per_fold, \
-           user_defined_measures_per_sample
+    return (
+        arff_datacontent,
+        trace,
+        user_defined_measures_per_fold,
+        user_defined_measures_per_sample,
+    )
 
 
 def _run_model_on_fold(model, task, rep_no, fold_no, sample_no, can_measure_runtime, add_local_measures):
@@ -679,8 +687,9 @@ def _extract_arfftrace_attributes(model):
                     raise TypeError('Unsupported param type in param grid: %s' %key)
 
             # we renamed the attribute param to parameter, as this is a required
-            # OpenML convention
-            attribute = ("parameter_" + key[6:], type)
+            # OpenML convention - this also guards against name collisions
+            # with the required trace attributes
+            attribute = (openml.runs.trace.PREFIX + key[6:], type)
             trace_attributes.append(attribute)
     return trace_attributes
 
@@ -748,7 +757,7 @@ def _create_run_from_xml(xml, from_server=True):
     run : OpenMLRun
         New run object representing run_xml.
     """
-    
+
     def obtain_field(xml_obj, fieldname, from_server, cast=None):
         # this function can be used to check whether a field is present in an object.
         # if it is not present, either returns None or throws an error (this is
@@ -769,7 +778,7 @@ def obtain_field(xml_obj, fieldname, from_server, cast=None):
     task_id = int(run['oml:task_id'])
     task_type = obtain_field(run, 'oml:task_type', from_server)
 
-    # even with the server requirement this field may be empty. 
+    # even with the server requirement this field may be empty.
     if 'oml:task_evaluation_measure' in run:
         task_evaluation_measure = run['oml:task_evaluation_measure']
     else:
@@ -877,85 +886,7 @@ def obtain_field(xml_obj, fieldname, from_server, cast=None):
                      tags=tags)
 
 
-def _create_trace_from_description(xml):
-    result_dict = xmltodict.parse(xml, force_list=('oml:trace_iteration',))['oml:trace']
-
-    run_id = result_dict['oml:run_id']
-    trace = collections.OrderedDict()
-
-    if 'oml:trace_iteration' not in result_dict:
-        raise ValueError('Run does not contain valid trace. ')
-
-    assert type(result_dict['oml:trace_iteration']) == list, \
-        type(result_dict['oml:trace_iteration'])
-
-    for itt in result_dict['oml:trace_iteration']:
-        repeat = int(itt['oml:repeat'])
-        fold = int(itt['oml:fold'])
-        iteration = int(itt['oml:iteration'])
-        setup_string = json.loads(itt['oml:setup_string'])
-        evaluation = float(itt['oml:evaluation'])
-
-        selectedValue = itt['oml:selected']
-        if selectedValue == 'true':
-            selected = True
-        elif selectedValue == 'false':
-            selected = False
-        else:
-            raise ValueError('expected {"true", "false"} value for '\
-                             'selected field, received: %s' %selectedValue)
-
-        current = OpenMLTraceIteration(repeat, fold, iteration,
-                                        setup_string, evaluation,
-                                        selected)
-        trace[(repeat, fold, iteration)] = current
-
-    return OpenMLRunTrace(run_id, trace)
-
-
-def _create_trace_from_arff(arff_obj):
-    """
-    Creates a trace file from arff obj (for example, generated by a local run)
-
-    Parameters
-    ----------
-    arff_obj : dict
-        LIAC arff obj, dict containing attributes, relation, data and description
-
-    Returns
-    -------
-    run : OpenMLRunTrace
-        Object containing None for run id and a dict containing the trace iterations
-    """
-    trace = collections.OrderedDict()
-    attribute_idx = {att[0]: idx for idx, att in enumerate(arff_obj['attributes'])}
-    for required_attribute in ['repeat', 'fold', 'iteration', 'evaluation', 'selected']:
-        if required_attribute not in attribute_idx:
-            raise ValueError('arff misses required attribute: %s' %required_attribute)
-
-    for itt in arff_obj['data']:
-        repeat = int(itt[attribute_idx['repeat']])
-        fold = int(itt[attribute_idx['fold']])
-        iteration = int(itt[attribute_idx['iteration']])
-        evaluation = float(itt[attribute_idx['evaluation']])
-        selectedValue = itt[attribute_idx['selected']]
-        if selectedValue == 'true':
-            selected = True
-        elif selectedValue == 'false':
-            selected = False
-        else:
-            raise ValueError('expected {"true", "false"} value for selected field, received: %s' % selectedValue)
-
-        # TODO: if someone needs it, he can use the parameter
-        # fields to revive the setup_string as well
-        # However, this is usually done by the OpenML server
-        # and if we are going to duplicate this functionality
-        # it needs proper testing
-
-        current = OpenMLTraceIteration(repeat, fold, iteration, None, evaluation, selected)
-        trace[(repeat, fold, iteration)] = current
 
-    return OpenMLRunTrace(None, trace)
 
 
 def _get_cached_run(run_id):

diff --git a/openml/runs/run.py b/openml/runs/run.py
@@ -27,7 +27,7 @@ class OpenMLRun(object):
     def __init__(self, task_id, flow_id, dataset_id, setup_string=None,
                  output_files=None, setup_id=None, tags=None, uploader=None, uploader_name=None,
                  evaluations=None, fold_evaluations=None, sample_evaluations=None,
-                 data_content=None, trace_attributes=None, trace_content=None,
+                 data_content=None, trace=None,
                  model=None, task_type=None, task_evaluation_measure=None, flow_name=None,
                  parameter_settings=None, predictions_url=None, task=None,
                  flow=None, run_id=None):
@@ -47,8 +47,7 @@ def __init__(self, task_id, flow_id, dataset_id, setup_string=None,
         self.sample_evaluations = sample_evaluations
         self.data_content = data_content
         self.output_files = output_files
-        self.trace_attributes = trace_attributes
-        self.trace_content = trace_content
+        self.trace = trace
         self.error_message = None
         self.task = task
         self.flow = flow
@@ -119,10 +118,7 @@ def from_filesystem(cls, folder, expect_model=True):
                 run.model = pickle.load(fp)
 
         if os.path.isfile(trace_path):
-            trace_arff = openml.runs.OpenMLRunTrace._from_filesystem(trace_path)
-
-            run.trace_attributes = trace_arff['attributes']
-            run.trace_content = trace_arff['data']
+            run.trace = openml.runs.OpenMLRunTrace._from_filesystem(trace_path)
 
         return run
 
@@ -167,10 +163,8 @@ def to_filesystem(self, output_directory, store_model=True):
             with open(os.path.join(output_directory, 'model.pkl'), 'wb') as f:
                 pickle.dump(self.model, f)
 
-        if self.trace_content is not None:
-            trace_arff = arff.dumps(self._generate_trace_arff_dict())
-            with open(os.path.join(output_directory, 'trace.arff'), 'w') as f:
-                f.write(trace_arff)
+        if self.trace is not None:
+            self.trace._to_filesystem(output_directory)
 
     def _generate_arff_dict(self):
         """Generates the arff dictionary for uploading predictions to the server.
@@ -204,29 +198,6 @@ def _generate_arff_dict(self):
         arff_dict['relation'] = 'openml_task_' + str(task.task_id) + '_predictions'
         return arff_dict
 
-    def _generate_trace_arff_dict(self):
-        """Generates the arff dictionary for uploading predictions to the server.
-
-        Assumes that the run has been executed.
-
-        Returns
-        -------
-        arf_dict : dict
-            Dictionary representation of the ARFF file that will be uploaded.
-            Contains information about the optimization trace.
-        """
-        if self.trace_content is None or len(self.trace_content) == 0:
-            raise ValueError('No trace content available.')
-        if len(self.trace_attributes) != len(self.trace_content[0]):
-            raise ValueError('Trace_attributes and trace_content not compatible')
-
-        arff_dict = OrderedDict()
-        arff_dict['attributes'] = self.trace_attributes
-        arff_dict['data'] = self.trace_content
-        arff_dict['relation'] = 'openml_task_' + str(self.task_id) + '_predictions'
-
-        return arff_dict
-
     def get_metric_fn(self, sklearn_fn, kwargs={}):
         """Calculates metric scores based on predicted values. Assumes the
         run has been executed locally (and contains run_data). Furthermore,
@@ -328,9 +299,15 @@ def publish(self):
         self : OpenMLRun
         """
         if self.model is None:
-            raise PyOpenMLError("OpenMLRun obj does not contain a model. (This should never happen.) ");
+            raise PyOpenMLError(
+                "OpenMLRun obj does not contain a model. "
+                "(This should never happen.) "
+            )
         if self.flow_id is None:
-            raise PyOpenMLError("OpenMLRun obj does not contain a flow id. (Should have been added while executing the task.) ");
+            raise PyOpenMLError(
+                "OpenMLRun obj does not contain a flow id. "
+                "(Should have been added while executing the task.) "
+            )
 
         description_xml = self._create_description_xml()
         file_elements = {'description': ("description.xml", description_xml)}
@@ -339,8 +316,8 @@ def publish(self):
             predictions = arff.dumps(self._generate_arff_dict())
             file_elements['predictions'] = ("predictions.arff", predictions)
 
-        if self.trace_content is not None:
-            trace_arff = arff.dumps(self._generate_trace_arff_dict())
+        if self.trace is not None:
+            trace_arff = arff.dumps(self.trace.trace_to_arff())
             file_elements['trace'] = ("trace.arff", trace_arff)
 
         return_value = openml._api_calls._perform_api_call("/run/", file_elements=file_elements)