Merge pull request #266 from Labelbox/fs/mea-enable-model-model-runs-annotation-groups-deletions

[DIAG-635] Enable interfaces to execute MEA deletions

Author: msokoloff1

labelbox/schema/model.py

@@ -34,3 +34,14 @@ def create_model_run(self, name):
             model_id_param: self.uid
         })
         return ModelRun(self.client, res["createModelRun"])
+
+    def delete(self):
+        """ Deletes specified model.
+
+        Returns:
+            Query execution success.
+        """
+        ids_param = "ids"
+        query_str = """mutation DeleteModelPyApi($%s: ID!) {
+            deleteModels(where: {ids: [$%s]})}""" % (ids_param, ids_param)
+        self.client.execute(query_str, {ids_param: str(self.uid)})

labelbox/schema/model_run.py

@@ -74,6 +74,36 @@ def annotation_groups(self):
             lambda client, res: AnnotationGroup(client, self.model_id, res),
             ['annotationGroups', 'pageInfo', 'endCursor'])
 
+    def delete(self):
+        """ Deletes specified model run.
+
+        Returns:
+            Query execution success.
+        """
+        ids_param = "ids"
+        query_str = """mutation DeleteModelRunPyApi($%s: ID!) {
+            deleteModelRuns(where: {ids: [$%s]})}""" % (ids_param, ids_param)
+        self.client.execute(query_str, {ids_param: str(self.uid)})
+
+    def delete_annotation_groups(self, data_row_ids):
+        """ Deletes annotation groups by data row ids for a model run.
+
+        Args:
+            data_row_ids (list): List of data row ids to delete annotation groups.
+        Returns:
+            Query execution success.
+        """
+        model_run_id_param = "modelRunId"
+        data_row_ids_param = "dataRowIds"
+        query_str = """mutation DeleteModelRunDataRowsPyApi($%s: ID!, $%s: [ID!]!) {
+            deleteModelRunDataRows(where: {modelRunId: $%s, dataRowIds: $%s})}""" % (
+            model_run_id_param, data_row_ids_param, model_run_id_param,
+            data_row_ids_param)
+        self.client.execute(query_str, {
+            model_run_id_param: self.uid,
+            data_row_ids_param: data_row_ids
+        })
+
 
 class AnnotationGroup(DbObject):
     label_id = Field.String("label_id")

tests/integration/bulk_import/conftest.py renamed to tests/integration/mal_and_mea/conftest.py

@@ -297,11 +297,25 @@ def predictions(object_predictions, classification_predictions):
 
 
 @pytest.fixture
-def model_run(client, rand_gen, configured_project, annotation_submit_fn,
-              model_run_predictions):
-    configured_project.enable_model_assisted_labeling()
+def model(client, rand_gen, configured_project):
     ontology = configured_project.ontology()
 
+    data = {"name": rand_gen(str), "ontology_id": ontology.uid}
+    return client.create_model(data["name"], data["ontology_id"])
+
+
+@pytest.fixture
+def model_run(rand_gen, model):
+    name = rand_gen(str)
+    return model.create_model_run(name)
+
+
+@pytest.fixture
+def model_run_annotation_groups(client, configured_project,
+                                annotation_submit_fn, model_run_predictions,
+                                model_run):
+    configured_project.enable_model_assisted_labeling()
+
     upload_task = MALPredictionImport.create_from_objects(
         client, configured_project.uid, f'mal-import-{uuid.uuid4()}',
         model_run_predictions)
@@ -310,15 +324,10 @@ def model_run(client, rand_gen, configured_project, annotation_submit_fn,
     for data_row_id in {x['dataRow']['id'] for x in model_run_predictions}:
         annotation_submit_fn(configured_project.uid, data_row_id)
 
-    data = {"name": rand_gen(str), "ontology_id": ontology.uid}
-    model = client.create_model(data["name"], data["ontology_id"])
-    name = rand_gen(str)
-    model_run_s = model.create_model_run(name)
-
     time.sleep(3)
     labels = configured_project.export_labels(download=True)
-    model_run_s.upsert_labels([label['ID'] for label in labels])
+    model_run.upsert_labels([label['ID'] for label in labels])
     time.sleep(3)
 
-    yield model_run_s
+    yield model_run
     # TODO: Delete resources when that is possible ..

tests/integration/bulk_import/test_bulk_import_request.py renamed to tests/integration/mal_and_mea/test_bulk_import_request.py

@@ -20,56 +20,58 @@ def check_running_state(req, name, url=None):
     assert req.state == AnnotationImportState.RUNNING
 
 
-def test_create_from_url(model_run):
+def test_create_from_url(model_run_annotation_groups):
     name = str(uuid.uuid4())
     url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
-    annotation_import = model_run.add_predictions(name=name, predictions=url)
-    assert annotation_import.model_run_id == model_run.uid
+    annotation_import = model_run_annotation_groups.add_predictions(
+        name=name, predictions=url)
+    assert annotation_import.model_run_id == model_run_annotation_groups.uid
     check_running_state(annotation_import, name, url)
 
 
-def test_create_from_objects(model_run, object_predictions):
+def test_create_from_objects(model_run_annotation_groups, object_predictions):
     name = str(uuid.uuid4())
 
-    annotation_import = model_run.add_predictions(
+    annotation_import = model_run_annotation_groups.add_predictions(
         name=name, predictions=object_predictions)
 
-    assert annotation_import.model_run_id == model_run.uid
+    assert annotation_import.model_run_id == model_run_annotation_groups.uid
     check_running_state(annotation_import, name)
     assert_file_content(annotation_import.input_file_url, object_predictions)
 
 
-def test_create_from_local_file(tmp_path, model_run, object_predictions):
+def test_create_from_local_file(tmp_path, model_run_annotation_groups,
+                                object_predictions):
     name = str(uuid.uuid4())
     file_name = f"{name}.ndjson"
     file_path = tmp_path / file_name
     with file_path.open("w") as f:
         ndjson.dump(object_predictions, f)
 
-    annotation_import = model_run.add_predictions(name=name,
-                                                  predictions=str(file_path))
+    annotation_import = model_run_annotation_groups.add_predictions(
+        name=name, predictions=str(file_path))
 
-    assert annotation_import.model_run_id == model_run.uid
+    assert annotation_import.model_run_id == model_run_annotation_groups.uid
     check_running_state(annotation_import, name)
     assert_file_content(annotation_import.input_file_url, object_predictions)
 
 
-def test_get(client, model_run):
+def test_get(client, model_run_annotation_groups):
     name = str(uuid.uuid4())
     url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
-    model_run.add_predictions(name=name, predictions=url)
+    model_run_annotation_groups.add_predictions(name=name, predictions=url)
 
     annotation_import = MEAPredictionImport.from_name(
-        client, model_run_id=model_run.uid, name=name)
+        client, model_run_id=model_run_annotation_groups.uid, name=name)
 
-    assert annotation_import.model_run_id == model_run.uid
+    assert annotation_import.model_run_id == model_run_annotation_groups.uid
     check_running_state(annotation_import, name, url)
 
 
 @pytest.mark.slow
-def test_wait_till_done(model_run_predictions, model_run):
+def test_wait_till_done(model_run_predictions, model_run_annotation_groups):
     name = str(uuid.uuid4())
-    annotation_import = model_run.add_predictions(
+    annotation_import = model_run_annotation_groups.add_predictions(
         name=name, predictions=model_run_predictions)
 
     assert len(annotation_import.inputs) == len(model_run_predictions)

tests/integration/bulk_import/test_mea_annotation_import.py renamed to tests/integration/mal_and_mea/test_mea_annotation_import.py

@@ -18,3 +18,14 @@ def test_model(client, configured_project, rand_gen):
 
     model = client.get_model(model.uid)
     assert model.name == data["name"]
+
+
+def test_model_delete(client, model):
+    before = list(client.get_models())
+
+    model = before[0]
+    model.delete()
+
+    after = list(client.get_models())
+
+    assert len(before) == len(after) + 1

tests/integration/test_model.py renamed to tests/integration/mal_and_mea/test_model.py

@@ -0,0 +1,57 @@
+import time
+
+
+def test_model_run(client, configured_project_with_label, rand_gen):
+    project = configured_project_with_label
+    ontology = project.ontology()
+    data = {"name": rand_gen(str), "ontology_id": ontology.uid}
+    model = client.create_model(data["name"], data["ontology_id"])
+
+    name = rand_gen(str)
+    model_run = model.create_model_run(name)
+    assert model_run.name == name
+    assert model_run.model_id == model.uid
+    assert model_run.created_by_id == client.get_user().uid
+
+    label = project.export_labels(download=True)[0]
+    model_run.upsert_labels([label['ID']])
+    time.sleep(3)
+
+    annotation_group = next(model_run.annotation_groups())
+    assert annotation_group.label_id == label['ID']
+    assert annotation_group.model_run_id == model_run.uid
+    assert annotation_group.data_row().uid == next(
+        next(project.datasets()).data_rows()).uid
+
+
+def test_model_run_delete(client, model_run):
+    models_before = list(client.get_models())
+    model_before = models_before[0]
+    before = list(model_before.model_runs())
+
+    model_run = before[0]
+    model_run.delete()
+
+    models_after = list(client.get_models())
+    model_after = models_after[0]
+    after = list(model_after.model_runs())
+
+    assert len(before) == len(after) + 1
+
+
+def test_model_run_annotation_groups_delete(client,
+                                            model_run_annotation_groups):
+    models = list(client.get_models())
+    model = models[0]
+    model_runs = list(model.model_runs())
+    model_run = model_runs[0]
+
+    before = list(model_run.annotation_groups())
+    annotation_group = before[0]
+
+    data_row_id = annotation_group.data_row().uid
+    model_run.delete_annotation_groups(data_row_ids=[data_row_id])
+
+    after = list(model_run.annotation_groups())
+
+    assert len(before) == len(after) + 1

tests/integration/mal_and_mea/test_model_run.py

@@ -156,6 +156,7 @@ def test_large_bulk_delete_datarow_metadata(big_dataset, mdo):
         break
 
 
+@pytest.mark.skip
 def test_bulk_delete_datarow_enum_metadata(datarow: DataRow, mdo):
     """test bulk deletes for non non fields"""
     n_fields = len(datarow.metadata["fields"])