Add model run config on create [AL-3060]

Author: mmarinwilcke

examples/model_diagnostics/model_diagnostics_guide.ipynb

@@ -350,7 +350,8 @@
    "source": [
     "lb_model = client.create_model(name=f\"{project.name}-model\",\n",
     "                               ontology_id=project.ontology().uid)\n",
-    "lb_model_run = lb_model.create_model_run(\"0.0.0\")\n",
+    "lb_model_run_hyperparameters = {\"batch_size\": 1000}\n",
+    "lb_model_run = lb_model.create_model_run(\"0.0.0\", lb_model_run_hyperparameters)\n",
     "lb_model_run.upsert_labels([label.uid for label in labels])"
    ]
   },

labelbox/schema/model.py

@@ -18,23 +18,26 @@ class Model(DbObject):
     name = Field.String("name")
     model_runs = Relationship.ToMany("ModelRun", False)
 
-    def create_model_run(self, name) -> "ModelRun":
+    def create_model_run(self, name, config) -> "ModelRun":
         """ Creates a model run belonging to this model.
 
         Args:
             name (string): The name for the model run.
+            config (json): Model run's training metadata config
         Returns:
             ModelRun, the created model run.
         """
         name_param = "name"
+        config_param = "config"
         model_id_param = "modelId"
         ModelRun = Entity.ModelRun
         query_str = """mutation CreateModelRunPyApi($%s: String!, $%s: ID!) {
-            createModelRun(data: {name: $%s, modelId: $%s}) {%s}}""" % (
-            name_param, model_id_param, name_param, model_id_param,
-            query.results_query_part(ModelRun))
+            createModelRun(data: {name: $%s, trainingMetadata: $%s, modelId: $%s}) {%s}}""" % (
+            name_param, model_id_param, name_param, config_param,
+            model_id_param, query.results_query_part(ModelRun))
         res = self.client.execute(query_str, {
             name_param: name,
+            config_param: config,
             model_id_param: self.uid
         })
         return ModelRun(self.client, res["createModelRun"])

labelbox/schema/model_run.py

@@ -32,6 +32,7 @@ class ModelRun(DbObject):
     created_at = Field.DateTime("created_at")
     created_by_id = Field.String("created_by_id", "createdBy")
     model_id = Field.String("model_id")
+    training_metadata = Field.Json("training_metadata")
 
     class Status(Enum):
         EXPORTING_DATA = "EXPORTING_DATA"

tests/integration/annotation_import/test_model_run.py

@@ -14,8 +14,10 @@ def test_model_run(client, configured_project_with_label, rand_gen):
     model = client.create_model(data["name"], data["ontology_id"])
 
     name = rand_gen(str)
-    model_run = model.create_model_run(name)
+    config = {"batch_size": 100}
+    model_run = model.create_model_run(name, config)
     assert model_run.name == name
+    assert model_run.training_metadata["batch_size"] == config["batch_size"]
     assert model_run.model_id == model.uid
     assert model_run.created_by_id == client.get_user().uid