separate mea to mal import class

Author: Matt Sokoloff

labelbox/schema/annotation_import.py

@@ -318,6 +318,88 @@ def _create_mea_import_from_bytes(
         return cls(client, res["createModelErrorAnalysisPredictionImport"])
 
 
+class MEAToMALPredictionImport(AnnotationImport):
+    project = Relationship.ToOne("Project", cache=True)
+
+    @property
+    def parent_id(self) -> str:
+        """
+        Identifier for this import. Used to refresh the status
+        """
+        return self.project().uid
+
+    @classmethod
+    def create_for_model_run_data_rows(cls, client: "labelbox.Client",
+                                       model_run_id: str,
+                                       data_row_ids: List[str], project_id: str,
+                                       name: str) -> "MEAToMALPredictionImport":
+        """
+        Create an MEA to MAL prediction import job from a list of data row ids of a specific model run
+
+        Args:
+            client: Labelbox Client for executing queries
+            data_row_ids: A list of data row ids
+            model_run_id: model run id
+        Returns:
+            MEAToMALPredictionImport
+        """
+        query_str = cls._get_model_run_data_rows_mutation()
+        return cls(
+            client,
+            client.execute(query_str,
+                           params={
+                               "dataRowIds": data_row_ids,
+                               "modelRunId": model_run_id,
+                               "projectId": project_id,
+                               "name": name
+                           })["createMalPredictionImportForModelRunDataRows"])
+
+    @classmethod
+    def from_name(cls,
+                  client: "labelbox.Client",
+                  project_id: str,
+                  name: str,
+                  as_json: bool = False) -> "MEAToMALPredictionImport":
+        """
+        Retrieves an MEA to MAL import job.
+
+        Args:
+            client: Labelbox Client for executing queries
+            project_id:  ID used for querying import jobs
+            name: Name of the import job.
+        Returns:
+            MALPredictionImport
+        """
+        query_str = """query getMEAToMALPredictionImportPyApi($projectId : ID!, $name: String!) {
+            meaToMalPredictionImport(
+                where: {projectId: $projectId, name: $name}){
+                    %s
+                }}""" % query.results_query_part(cls)
+        params = {
+            "projectId": project_id,
+            "name": name,
+        }
+        response = client.execute(query_str, params)
+        if response is None:
+            raise labelbox.exceptions.ResourceNotFoundError(
+                MALPredictionImport, params)
+        response = response["meaToMalPredictionImport"]
+        if as_json:
+            return response
+        return cls(client, response)
+
+    @classmethod
+    def _get_model_run_data_rows_mutation(cls) -> str:
+        return """mutation createMalPredictionImportForModelRunDataRowsPyApi($dataRowIds: [ID!]!, $name: String!, $modelRunId: ID!, $projectId:ID!) {
+            createMalPredictionImportForModelRunDataRows(data: {
+                name: $name
+                modelRunId: $modelRunId
+                dataRowIds: $dataRowIds
+                projectId: $projectId
+            }) {%s}
+        }""" % query.results_query_part(cls)
+
+
 class MALPredictionImport(AnnotationImport):
     project = Relationship.ToOne("Project", cache=True)
 
@@ -401,32 +483,6 @@ def create_from_url(cls, client: "labelbox.Client", project_id: str,
         else:
             raise ValueError(f"Url {url} is not reachable")
 
-    @classmethod
-    def create_for_model_run_data_rows(cls, client: "labelbox.Client",
-                                       model_run_id: str,
-                                       data_row_ids: List[str], project_id: str,
-                                       name: str) -> "MALPredictionImport":
-        """
-        Create an MAL prediction import job from a list of data row ids of a specific model run
-
-        Args:
-            client: Labelbox Client for executing queries
-            data_row_ids: A list of data row ids
-            model_run_id: model run id
-        Returns:
-            MALPredictionImport
-        """
-        query_str = cls._get_model_run_data_rows_mutation()
-        return cls(
-            client,
-            client.execute(query_str,
-                           params={
-                               "dataRowIds": data_row_ids,
-                               "modelRunId": model_run_id,
-                               "projectId": project_id,
-                               "name": name
-                           })["createMalPredictionImportForModelRunDataRows"])
-
     @classmethod
     def from_name(cls,
                   client: "labelbox.Client",
@@ -471,17 +527,6 @@ def _get_url_mutation(cls) -> str:
             }) {%s}
         }""" % query.results_query_part(cls)
 
-    @classmethod
-    def _get_model_run_data_rows_mutation(cls) -> str:
-        return """mutation createMalPredictionImportForModelRunDataRowsPyApi($dataRowIds: [ID!]!, $name: String!, $modelRunId: ID!, $projectId:ID!) {
-            createMalPredictionImportForModelRunDataRows(data: {
-                name: $name
-                modelRunId: $modelRunId
-                dataRowIds: $dataRowIds
-                projectId: $projectId
-            }) {id importType inputFileUrl errorFileUrl project { id name } name statusFileUrl state progress}
-        }"""
-
     @classmethod
     def _get_file_mutation(cls) -> str:
         return """mutation createMALPredictionImportByFilePyApi($projectId : ID!, $name: String!, $file: Upload!, $contentLength: Int!) {

labelbox/schema/model_run.py

@@ -140,7 +140,7 @@ def upsert_predictions_and_send_to_project(
             project_id (str): id of the project to import into
             priority (int): priority of the job
         Returns:
-            (AnnotationImport, Project)
+            (MEAPredictionImport, Batch, MEAToMALPredictionImport)
         """
         kwargs = dict(client=self.client, model_run_id=self.uid, name=name)
         project = self.client.get_project(project_id)
@@ -165,7 +165,7 @@ def upsert_predictions_and_send_to_project(
         try:
             batch = project.create_batch(name, mea_to_mal_data_rows, priority)
             try:
-                mal_prediction_import = Entity.MALPredictionImport.create_for_model_run_data_rows(
+                mal_prediction_import = Entity.MEAToMALPredictionImport.create_for_model_run_data_rows(
                     data_row_ids=mea_to_mal_data_rows,
                     project_id=project_id,
                     **kwargs)
@@ -316,11 +316,11 @@ def update_status(self,
 
     @experimental
     def update_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
-        """ 
+        """
          Updates the Model Run's training metadata config
-         Args: 
+         Args:
              config (dict): A dictionary of keys and values
-         Returns: 
+         Returns:
              Model Run id and updated training metadata
          """
         data: Dict[str, Any] = {'config': config}
@@ -337,9 +337,9 @@ def update_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
 
     @experimental
     def reset_config(self) -> Dict[str, Any]:
-        """ 
+        """
          Resets Model Run's training metadata config
-         Returns: 
+         Returns:
              Model Run id and reset training metadata
          """
         res = self.client.execute(
@@ -352,10 +352,10 @@ def reset_config(self) -> Dict[str, Any]:
 
     @experimental
     def get_config(self) -> Dict[str, Any]:
-        """ 
-         Gets Model Run's training metadata 
-         Returns: 
-             training metadata as a dictionary 
+        """
+         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}

tests/integration/annotation_import/conftest.py

@@ -158,17 +158,15 @@ def configured_project_pdf(client, ontology, rand_gen, pdf_url):
 
 
 @pytest.fixture
-def configured_project_without_data_rows(client, ontology, rand_gen):
+def configured_project_without_data_rows(client, configured_project, rand_gen):
     project = client.create_project(name=rand_gen(str))
-    dataset = client.create_dataset(name=rand_gen(str))
     editor = list(
         client.get_labeling_frontends(
             where=LabelingFrontend.name == "editor"))[0]
-    project.setup(editor, ontology)
+    project.setup_editor(configured_project.ontology())
     project.update(queue_mode=project.QueueMode.Batch)
     yield project
     project.delete()
-    dataset.delete()
 
 
 @pytest.fixture
@@ -436,6 +434,7 @@ def model_run_with_model_run_data_rows(client, configured_project,
     model_run.upsert_labels(label_ids)
     time.sleep(3)
     yield model_run
+    model_run.delete()
     # TODO: Delete resources when that is possible ..
 
 

tests/integration/annotation_import/test_upsert_prediction_import.py

@@ -17,8 +17,19 @@ def test_create_from_url(client, tmp_path, object_predictions,
     name = str(uuid.uuid4())
     file_name = f"{name}.json"
     file_path = tmp_path / file_name
+
+    model_run_data_rows = [
+        mrdr.data_row().uid
+        for mrdr in model_run_with_model_run_data_rows.model_run_data_rows()
+    ]
+    predictions = [
+        p for p in object_predictions
+        if p['dataRow']['id'] in model_run_data_rows
+    ]
     with file_path.open("w") as f:
-        ndjson.dump(object_predictions, f)
+        ndjson.dump(predictions, f)
+
+    # Needs to have data row ids
 
     with open(file_path, "r") as f:
         url = client.upload_data(content=f.read(),
@@ -33,55 +44,74 @@ def test_create_from_url(client, tmp_path, object_predictions,
         priority=5)
 
     assert annotation_import.model_run_id == model_run_with_model_run_data_rows.uid
-    annotation_import_test_helpers.check_running_state(annotation_import, name,
-                                                       url)
     annotation_import.wait_until_done()
+    assert not annotation_import.errors
+    assert annotation_import.statuses
+
+    assert batch
+    assert batch.project().uid == configured_project_without_data_rows.uid
+
+    assert mal_prediction_import
+    mal_prediction_import.wait_until_done()
 
-    if batch:
-        assert batch.project().uid == configured_project_without_data_rows.uid
-    if mal_prediction_import:
-        mal_prediction_import.wait_until_done()
+    assert not mal_prediction_import.errors
+    assert mal_prediction_import.statuses
 
 
 def test_create_from_objects(model_run_with_model_run_data_rows,
                              configured_project_without_data_rows,
                              object_predictions,
                              annotation_import_test_helpers):
     name = str(uuid.uuid4())
-
+    model_run_data_rows = [
+        mrdr.data_row().uid
+        for mrdr in model_run_with_model_run_data_rows.model_run_data_rows()
+    ]
+    predictions = [
+        p for p in object_predictions
+        if p['dataRow']['id'] in model_run_data_rows
+    ]
     annotation_import, batch, mal_prediction_import = model_run_with_model_run_data_rows.upsert_predictions_and_send_to_project(
         name=name,
-        predictions=object_predictions,
+        predictions=predictions,
         project_id=configured_project_without_data_rows.uid,
         priority=5)
 
     assert annotation_import.model_run_id == model_run_with_model_run_data_rows.uid
-    annotation_import_test_helpers.check_running_state(annotation_import, name)
-    annotation_import_test_helpers.assert_file_content(
-        annotation_import.input_file_url, object_predictions)
     annotation_import.wait_until_done()
+    assert not annotation_import.errors
+    assert annotation_import.statuses
 
-    if batch:
-        assert batch.project().uid == configured_project_without_data_rows.uid
+    assert batch
+    assert batch.project().uid == configured_project_without_data_rows.uid
 
-    if mal_prediction_import:
-        annotation_import_test_helpers.check_running_state(
-            mal_prediction_import, name)
-        mal_prediction_import.wait_until_done()
+    assert mal_prediction_import
+    mal_prediction_import.wait_until_done()
+
+    assert not mal_prediction_import.errors
+    assert mal_prediction_import.statuses
 
 
 def test_create_from_local_file(tmp_path, model_run_with_model_run_data_rows,
                                 configured_project_without_data_rows,
                                 object_predictions,
                                 annotation_import_test_helpers):
+
     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_with_model_run_data_rows.add_predictions(
-        name=name, predictions=str(file_path))
+    model_run_data_rows = [
+        mrdr.data_row().uid
+        for mrdr in model_run_with_model_run_data_rows.model_run_data_rows()
+    ]
+    predictions = [
+        p for p in object_predictions
+        if p['dataRow']['id'] in model_run_data_rows
+    ]
+
+    with file_path.open("w") as f:
+        ndjson.dump(predictions, f)
 
     annotation_import, batch, mal_prediction_import = model_run_with_model_run_data_rows.upsert_predictions_and_send_to_project(
         name=name,
@@ -90,42 +120,15 @@ def test_create_from_local_file(tmp_path, model_run_with_model_run_data_rows,
         priority=5)
 
     assert annotation_import.model_run_id == model_run_with_model_run_data_rows.uid
-    annotation_import_test_helpers.check_running_state(annotation_import, name)
-    annotation_import_test_helpers.assert_file_content(
-        annotation_import.input_file_url, object_predictions)
     annotation_import.wait_until_done()
+    assert not annotation_import.errors
+    assert annotation_import.statuses
 
-    if batch:
-        assert batch.project().uid == configured_project_without_data_rows.uid
-
-    if mal_prediction_import:
-        annotation_import_test_helpers.check_running_state(
-            mal_prediction_import, name)
-        mal_prediction_import.wait_until_done()
+    assert batch
+    assert batch.project().uid == configured_project_without_data_rows.uid
 
+    assert mal_prediction_import
+    mal_prediction_import.wait_until_done()
 
-@pytest.mark.slow
-def test_wait_till_done(model_run_predictions,
-                        model_run_with_model_run_data_rows):
-    name = str(uuid.uuid4())
-    annotation_import = model_run_with_model_run_data_rows.add_predictions(
-        name=name, predictions=model_run_predictions)
-
-    assert len(annotation_import.inputs) == len(model_run_predictions)
-    annotation_import.wait_until_done()
-    assert annotation_import.state == AnnotationImportState.FINISHED
-    # Check that the status files are being returned as expected
-    assert len(annotation_import.errors) == 0
-    assert len(annotation_import.inputs) == len(model_run_predictions)
-    input_uuids = [
-        input_annot['uuid'] for input_annot in annotation_import.inputs
-    ]
-    inference_uuids = [pred['uuid'] for pred in model_run_predictions]
-    assert set(input_uuids) == set(inference_uuids)
-    assert len(annotation_import.statuses) == len(model_run_predictions)
-    for status in annotation_import.statuses:
-        assert status['status'] == 'SUCCESS'
-    status_uuids = [
-        input_annot['uuid'] for input_annot in annotation_import.statuses
-    ]
-    assert set(input_uuids) == set(status_uuids)
+    assert not mal_prediction_import.errors
+    assert mal_prediction_import.statuses