Update config [AL-3061]

Author: mmarinwilcke

labelbox/schema/model_run.py

@@ -262,6 +262,45 @@ def update_status(self,
             },
             experimental=True)
 
+    @experimental
+    def update_config(self,
+                      config: Dict[str, Any]):
+        data: Dict[str, Any] = {'config': config}
+        res = self.client.execute(
+            """mutation updateModelRunConfigPyApi($modelRunId: ID!, $data: UpdateModelRunConfigInput!){
+                updateModelRunConfig(modelRun: {id : $modelRunId}, data: $data){trainingMetadata}
+            }
+        """, {
+                'modelRunId': self.uid,
+                'data': data
+            },
+            experimental=True)
+        return  res["updateModelRunConfig"]
+
+    @experimental
+    def reset_config(self):
+        res = self.client.execute(
+            """mutation resetModelRunConfigPyApi($modelRunId: ID!){
+                resetModelRunConfig(modelRun: {id : $modelRunId}){trainingMetadata}
+            }
+        """, {
+                'modelRunId': self.uid
+            },
+            experimental=True)
+        return  res["resetModelRunConfig"]
+    
+    @experimental
+    def config(self):
+        res = self.client.execute(
+            """query ModelRunPyApi($modelRunId: ID!){
+                modelRun(where: {id : $modelRunId}){trainingMetadata}
+            }
+        """, {
+                'modelRunId': self.uid
+            },
+            experimental=True)
+        return  res["modelRun"]
+
     @experimental
     def export_labels(
         self,

tests/integration/annotation_import/conftest.py

@@ -359,6 +359,18 @@ def model_run(rand_gen, model):
         # Already was deleted by the test
         pass
 
+@pytest.fixture
+def model_run_with_training_metadata(rand_gen, model):
+    name = rand_gen(str)
+    training_metadata = {"batch_size": 1000}
+    model_run = model.create_model_run(name, training_metadata)
+    yield model_run
+    try:
+        model_run.delete()
+    except:
+        # Already was deleted by the test
+        pass
+
 
 @pytest.fixture
 def model_run_with_model_run_data_rows(client, configured_project,

tests/integration/annotation_import/test_model_run.py

@@ -56,6 +56,24 @@ def test_model_run_delete(client, model_run):
     assert len(before) == len(after) + 1
 
 
+def test_model_run_update_config(model_run_with_training_metadata):
+    new_config = {"batch_size": 2000}
+    res = model_run_with_training_metadata.update_config(new_config)
+    assert res["trainingMetadata"]["batchSize"] == new_config["batch_size"]
+
+
+def test_model_run_reset_config(model_run_with_training_metadata):
+    res = model_run_with_training_metadata.reset_config()
+    assert res["trainingMetadata"] == None
+
+
+def test_model_run_fetch_config(model_run_with_training_metadata):
+    new_config = {"batch_size": 2000}
+    model_run_with_training_metadata.update_config(new_config)
+    res = model_run_with_training_metadata.config()
+    assert res["trainingMetadata"]["batchSize"] == new_config["batch_size"]
+
+
 def test_model_run_data_rows_delete(client, model_run_with_model_run_data_rows):
     models = list(client.get_models())
     model = models[0]