Merge pull request #63 from BoostV/exclude-pickled-model

Add option to exclude pickled result

optimizerapi/optimizer.py

@@ -179,6 +179,8 @@ def process_result(result, optimizer, dimensions, cfg, extras, data, space):
 
     objective_pars = extras.get("objectivePars", "result")
 
+    pickle_model = json.loads(extras.get("includeModel", "true").lower())
+
     # In the following section details that should be reported to
     # clients should go into the "resultDetails" dictionary and plots
     # go into the "plots" list (this is handled by calling the "addPlot" function)
@@ -220,7 +222,8 @@ def process_result(result, optimizer, dimensions, cfg, extras, data, space):
                 plot_Pareto(optimizer)
                 add_plot(plots, "pareto")
 
-    result_details["pickled"] = pickleToString(result, get_crypto())
+    if pickle_model:
+        result_details["pickled"] = pickleToString(result, get_crypto())
 
     add_version_info(result_details["extras"])
 

tests/test_optimizer.py

@@ -105,3 +105,29 @@ def test_can_accept_multi_objective_data():
     validateResult(result)
     assert len(result["result"]["models"]) > 1
     assert len(result["plots"]) == 5
+
+
+def test_deselecting_pickled_model():
+    # If includeModel is false, pickled data should not be included in result
+    result = optimizer.run(
+        body={
+            "data": sampleData,
+            "optimizerConfig": sampleConfig,
+            "extras": {"includeModel": "false"},
+        }
+    )
+    assert "pickled" in result["result"]
+    assert len(result["result"]["pickled"]) == 0
+
+
+def test_selecting_pickled_model():
+    # If includeModel is true, pickled data should be included in result
+    result = optimizer.run(
+        body={
+            "data": sampleData,
+            "optimizerConfig": sampleConfig,
+            "extras": {"includeModel": "true"},
+        }
+    )
+    assert "pickled" in result["result"]
+    assert len(result["result"]["pickled"]) > 0