Merge pull request #669 from Labelbox/mmw/AL-3061

Update config [AL-3061]

Author: mmarinwilcke

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+# Version 0.0.0 (YYYY-MM-DD)
+### Added
+* `ModelRun.update_config()`
+    * Updates model run training metadata
+* `ModelRun.reset_config()`
+    * Resets model run training metadata
+* `ModelRun.get_config()`
+    * Fetches model run training metadata
+    
+### Changed
+* `Model.create_model_run()`
+    * Add training metadata config as a model run creation param
+
 # Version 3.26.0 (2022-08-15)
 ## Added
 * `Batch.delete()` which will delete an existing `Batch`
@@ -663,7 +676,3 @@ a `Label`. Default value is 0.0.
 
 ## Version 2.2 (2019-10-18)
 Changelog not maintained before version 2.2.
-
-### Changed
-* `Model.create_model_run()`
-    * Add training metadata config as a model run creation param

labelbox/schema/model_run.py

@@ -42,7 +42,7 @@ class Status(Enum):
         FAILED = "FAILED"
 
     def upsert_labels(self, label_ids, timeout_seconds=60):
-        """ Adds data rows and labels to a model run
+        """ Adds data rows and labels to a Model Run
         Args:
             label_ids (list): label ids to insert
             timeout_seconds (float): Max waiting time, in seconds.
@@ -75,7 +75,7 @@ def upsert_labels(self, label_ids, timeout_seconds=60):
                                      timeout_seconds=timeout_seconds)
 
     def upsert_data_rows(self, data_row_ids, timeout_seconds=60):
-        """ Adds data rows to a model run without any associated labels
+        """ Adds data rows to a Model Run without any associated labels
         Args:
             data_row_ids (list): data row ids to add to mea
             timeout_seconds (float): Max waiting time, in seconds.
@@ -167,7 +167,7 @@ def model_run_data_rows(self):
             ['annotationGroups', 'pageInfo', 'endCursor'])
 
     def delete(self):
-        """ Deletes specified model run.
+        """ Deletes specified Model Run.
 
         Returns:
             Query execution success.
@@ -178,10 +178,10 @@ def delete(self):
         self.client.execute(query_str, {ids_param: str(self.uid)})
 
     def delete_model_run_data_rows(self, data_row_ids: List[str]):
-        """ Deletes data rows from model runs.
+        """ Deletes data rows from Model Runs.
 
         Args:
-            data_row_ids (list): List of data row ids to delete from the model run.
+            data_row_ids (list): List of data row ids to delete from the Model Run.
         Returns:
             Query execution success.
         """
@@ -262,6 +262,56 @@ def update_status(self,
             },
             experimental=True)
 
+    @experimental
+    def update_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
+        """ 
+         Updates the Model Run's training metadata config
+         Args: 
+             config (dict): A dictionary of keys and values
+         Returns: 
+             Model Run id and updated training metadata
+         """
+        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) -> Dict[str, Any]:
+        """ 
+         Resets Model Run's training metadata config
+         Returns: 
+             Model Run id and reset training metadata
+         """
+        res = self.client.execute(
+            """mutation resetModelRunConfigPyApi($modelRunId: ID!){
+                resetModelRunConfig(modelRun: {id : $modelRunId}){trainingMetadata}
+            }
+        """, {'modelRunId': self.uid},
+            experimental=True)
+        return res["resetModelRunConfig"]
+
+    @experimental
+    def get_config(self) -> Dict[str, Any]:
+        """ 
+         Gets Model Run's training metadata 
+         Returns: 
+             training metadata as a dictionary 
+         """
+        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

@@ -360,6 +360,19 @@ def model_run(rand_gen, model):
         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,
                                        model_run_predictions, model_run):

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"]["batch_size"] == 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"] is None
+
+
+def test_model_run_get_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.get_config()
+    assert res["trainingMetadata"]["batch_size"] == 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]