add partition and model run status

Author: Matt Sokoloff

labelbox/client.py

@@ -379,7 +379,8 @@ def upload_data(self,
 
         if not file_data or not file_data.get("uploadFile", None):
             raise labelbox.exceptions.LabelboxError(
-                "Failed to upload, message: %s" % file_data.get("error", None))
+                "Failed to upload, message: %s" % file_data or
+                file_data.get("error"))
 
         return file_data["uploadFile"]["url"]
 
@@ -918,4 +919,4 @@ def get_model_run(self, model_run_id: str) -> ModelRun:
         Returns:
             A ModelRun object.
         """
-        return self._get_single(Entity.ModelRun, model_run_id)
+        return self._get_single(Entity.ModelRun, model_run_id)

labelbox/data/serialization/labelbox_v1/label.py

@@ -141,6 +141,7 @@ class LBV1Label(BaseModel):
     has_open_issues: Optional[float] = Extra('Has Open Issues')
     skipped: Optional[bool] = Extra('Skipped')
     media_type: Optional[str] = Extra('media_type')
+    data_split: Optional[str] = Extra('Data Split')
 
     def to_common(self) -> Label:
         if isinstance(self.label, list):

labelbox/schema/model_run.py

@@ -180,6 +180,38 @@ def delete_model_run_data_rows(self, data_row_ids):
             data_row_ids_param: data_row_ids
         })
 
+    @experimental
+    def update_status(self,
+                      status: str,
+                      metadata: Optional[Dict[str, str]] = None,
+                      error_message: Optional[Dict[str, str]] = None):
+
+        valid_statuses = [
+            "EXPORTING_DATA", "PREPARING_DATA", "TRAINING_MODEL", "COMPLETE",
+            "FAILED"
+        ]
+        if status not in valid_statuses:
+            raise ValueError(
+                f"Status must be one of : `{valid_statuses}`. Found : `{status}`"
+            )
+
+        data = {'status': status}
+        if error_message:
+            data['errorMessage'] = error_message
+
+        if metadata:
+            data['metadata'] = metadata
+
+        self.client.execute(
+            """mutation setPipelineStatusPyApi($modelRunId: ID!, $data: UpdateTrainingPipelineInput!){
+                updateTrainingPipeline(modelRun: {id : $modelRunId}, data: $data){status}
+            }
+        """, {
+                'modelRunId': self.uid,
+                'data': data
+            },
+            experimental=True)
+
     @experimental
     def export_labels(
         self,
@@ -196,7 +228,7 @@ def export_labels(
         Returns:
             URL of the data file with this ModelRun's labels.
             If download=True, this instead returns the contents as NDJSON format.
-            If the server didn't generate during the `timeout_seconds` period, 
+            If the server didn't generate during the `timeout_seconds` period,
             None is returned.
         """
         sleep_time = 2

tests/integration/annotation_import/test_model_run.py

@@ -1,3 +1,5 @@
+from labelbox.orm.db_object import experimental
+from labelbox.schema.model_run import ModelRun
 import time
 
 
@@ -82,3 +84,36 @@ def test_model_run_upsert_data_rows_with_existing_labels(
 def test_model_run_export_labels(model_run_with_model_run_data_rows):
     labels = model_run_with_model_run_data_rows.export_labels(download=True)
     assert len(labels) == 3
+
+
+def test_model_run_status(model_run_with_model_run_data_rows: ModelRun):
+
+    def get_model_run_status():
+        return model_run_with_model_run_data_rows.client.execute(
+            """query trainingPipelinePyApi($modelRunId: ID!) {
+            trainingPipeline(where: {id : $modelRunId}) {status, errorMessage, metadata}}
+        """, {'modelRunId': model_run_with_model_run_data_rows.uid},
+            experimental=True)['trainingPipeline']
+
+    model_run_status = get_model_run_status()
+    assert model_run_status['status'] is None
+    assert model_run_status['metadata'] is None
+    assert model_run_status['errorMessage'] is None
+
+    status = "COMPLETE"
+    metadata = {'key1': 'value1'}
+    errorMessage = "an error"
+    model_run_with_model_run_data_rows.update_status(status, metadata,
+                                                     errorMessage)
+
+    model_run_status = get_model_run_status()
+    assert model_run_status['status'] == status
+    assert model_run_status['metadata'] == metadata
+    assert model_run_status['errorMessage'] == errorMessage
+
+    extra_metadata = {'key2': 'value2'}
+    model_run_with_model_run_data_rows.update_status(status, extra_metadata)
+    model_run_status = get_model_run_status()
+    assert model_run_status['status'] == status
+    assert model_run_status['metadata'] == {**metadata, **extra_metadata}
+    assert model_run_status['errorMessage'] == errorMessage