[AL-4866] Upsert data rows to model run using global keys

Author: whistler

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+# Version 3.40.0 (YYYY-MM-DD)
+
+## Added 
+* Upsert data rows to model runs using global keys
+
 # Version 3.39.0 (2023-02-28)
 ## Added
 * New method `Project.task_queues()` to obtain the task queues for a project.

labelbox/schema/model_run.py

@@ -80,27 +80,30 @@ def upsert_labels(self, label_ids, timeout_seconds=3600):
             }})['MEALabelRegistrationTaskStatus'],
                                      timeout_seconds=timeout_seconds)
 
-    def upsert_data_rows(self, data_row_ids, timeout_seconds=3600):
+    def upsert_data_rows(self,
+                         data_row_ids=None,
+                         global_keys=None,
+                         timeout_seconds=3600):
         """ Adds data rows to a Model Run without any associated labels
         Args:
-            data_row_ids (list): data row ids to add to mea
+            data_row_ids (list): data row ids to add to model run
+            global_keys (list): global keys for data rows to add to model run
             timeout_seconds (float): Max waiting time, in seconds.
         Returns:
             ID of newly generated async task
         """
 
-        if len(data_row_ids) < 1:
-            raise ValueError("Must provide at least one data row id")
-
         mutation_name = 'createMEAModelRunDataRowRegistrationTask'
-        create_task_query_str = """mutation createMEAModelRunDataRowRegistrationTaskPyApi($modelRunId: ID!, $dataRowIds : [ID!]!) {
-          %s(where : { id : $modelRunId}, data : {dataRowIds: $dataRowIds})}
+        create_task_query_str = """mutation createMEAModelRunDataRowRegistrationTaskPyApi($modelRunId: ID!, $dataRowIds: [ID!], $globalKeys: [ID!]) {
+          %s(where : { id : $modelRunId}, data : {dataRowIds: $dataRowIds, globalKeys: $globalKeys})}
           """ % (mutation_name)
 
-        res = self.client.execute(create_task_query_str, {
-            'modelRunId': self.uid,
-            'dataRowIds': data_row_ids
-        })
+        res = self.client.execute(
+            create_task_query_str, {
+                'modelRunId': self.uid,
+                'dataRowIds': data_row_ids,
+                'globalKeys': global_keys
+            })
         task_id = res[mutation_name]
 
         status_query_str = """query MEADataRowRegistrationTaskStatusPyApi($where: WhereUniqueIdInput!){

tests/integration/annotation_import/test_model_run.py

@@ -1,11 +1,9 @@
-import json
 import time
 import os
 import pytest
 
 from collections import Counter
 
-import requests
 from labelbox import DataSplit, ModelRun
 
 
@@ -99,6 +97,14 @@ def test_model_run_upsert_data_rows(dataset, model_run):
     assert n_model_run_data_rows == 1
 
 
+def test_model_run_upsert_data_rows_using_global_keys(model_run, data_rows):
+    global_keys = [dr.global_key for dr in data_rows]
+    assert model_run.upsert_data_rows(global_keys=global_keys)
+    model_run_data_rows = list(model_run.model_run_data_rows())
+    added_data_rows = [mdr.data_row() for mdr in model_run_data_rows]
+    assert set(added_data_rows) == set(data_rows)
+
+
 def test_model_run_upsert_data_rows_with_existing_labels(
         model_run_with_model_run_data_rows):
     model_run_data_rows = list(

tests/integration/conftest.py

@@ -212,6 +212,17 @@ def datarow(dataset, image_url):
     dr.delete()
 
 
+@pytest.fixture()
+def data_rows(dataset, image_url):
+    dr1 = dataset.create_data_row(row_data=image_url,
+                                  global_key=f"global-key-{uuid.uuid4()}")
+    dr2 = dataset.create_data_row(row_data=image_url,
+                                  global_key=f"global-key-{uuid.uuid4()}")
+    yield [dr1, dr2]
+    dr1.delete()
+    dr2.delete()
+
+
 @pytest.fixture
 def iframe_url(environ) -> str:
     if environ in [Environ.PROD, Environ.LOCAL]: