add status, type, and creation time to async job

Author: Bihan Jiang

nucleus/__init__.py

@@ -86,6 +86,10 @@
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
+    JOB_ID_KEY,
+    JOB_STATUS_KEY,
+    JOB_TYPE_KEY,
+    JOB_CREATION_TIME_KEY,
     IMAGE_KEY,
     IMAGE_URL_KEY,
     ITEM_METADATA_SCHEMA_KEY,
@@ -110,7 +114,7 @@
     NotFoundError,
     NucleusAPIError,
 )
-from .job import Job
+from .job import AsyncJob
 from .model import Model
 from .model_run import ModelRun
 from .payload_constructor import (
@@ -200,19 +204,21 @@ def list_datasets(self) -> Dict[str, Union[str, List[str]]]:
         """
         return self.make_request({}, "dataset/", requests.get)
 
-    def list_jobs(self, show_completed=None, date_limit=None) -> List[Job]:
+    def list_jobs(
+        self, show_completed=None, date_limit=None
+    ) -> List[AsyncJob]:
         """
         Lists jobs for user.
         :return: jobs
         """
         payload = {show_completed: show_completed, date_limit: date_limit}
         job_objects = self.make_request(payload, "jobs/", requests.get)
         return [
-            Job(
-                job_id=job["job_id"],
-                job_status=["job_status"],
-                job_type=["job_type"],
-                job_creation_time=["job_creation_time"],
+            AsyncJob(
+                job_id=job[JOB_ID_KEY],
+                job_status=job[JOB_STATUS_KEY],
+                job_type=job[JOB_TYPE_KEY],
+                job_creation_time=job[JOB_CREATION_TIME_KEY],
                 client=self,
             )
             for job in job_objects

nucleus/constants.py

@@ -42,6 +42,9 @@
 ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
 JOB_ID_KEY = "job_id"
 KEEP_HISTORY_KEY = "keep_history"
+JOB_STATUS_KEY = "job_last_known_status"
+JOB_TYPE_KEY = "job_type"
+JOB_CREATION_TIME_KEY = "job_creation_time"
 LABEL_KEY = "label"
 MASK_URL_KEY = "mask_url"
 MESSAGE_KEY = "message"

nucleus/dataset.py

@@ -24,6 +24,9 @@
     DEFAULT_ANNOTATION_UPDATE_MODE,
     EXPORTED_ROWS,
     JOB_ID_KEY,
+    JOB_STATUS_KEY,
+    JOB_TYPE_KEY,
+    JOB_CREATION_TIME_KEY,
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
@@ -182,7 +185,13 @@ def annotate(
                 route=f"dataset/{self.id}/annotate?async=1",
             )
 
-            return AsyncJob(response[JOB_ID_KEY], self._client)
+            return AsyncJob(
+                response[JOB_ID_KEY],
+                response[JOB_STATUS_KEY],
+                response[JOB_TYPE_KEY],
+                response[JOB_CREATION_TIME_KEY],
+                self._client,
+            )
 
         return self._client.annotate_dataset(
             self.id, annotations, update=update, batch_size=batch_size
@@ -241,7 +250,13 @@ def append(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/append?async=1",
             )
-            return AsyncJob(response[JOB_ID_KEY], self._client)
+            return AsyncJob(
+                response[JOB_ID_KEY],
+                response[JOB_STATUS_KEY],
+                response[JOB_TYPE_KEY],
+                response[JOB_CREATION_TIME_KEY],
+                self._client,
+            )
 
         return self._client.populate_dataset(
             self.id,

nucleus/job.py

@@ -9,20 +9,23 @@
 
 @dataclass
 class AsyncJob:
-    id: str
+    job_id: str
+    job_status: str
+    job_type: str
+    job_creation_time: str
     client: "NucleusClient"  # type: ignore # noqa: F821
 
     def status(self) -> Dict[str, str]:
         return self.client.make_request(
             payload={},
-            route=f"job/{self.id}",
+            route=f"job/{self.job_id}",
             requests_command=requests.get,
         )
 
     def errors(self) -> List[str]:
         return self.client.make_request(
             payload={},
-            route=f"job/{self.id}/errors",
+            route=f"job/{self.job_id}/errors",
             requests_command=requests.get,
         )
 
@@ -43,22 +46,6 @@ def sleep_until_complete(self, verbose_std_out=True):
             raise JobError(final_status, self)
 
 
-class Job:
-    def __init__(
-        self,
-        job_id,
-        job_status,
-        job_type,
-        job_creation_time,
-        client,
-    ) -> None:
-        self.job_id = job_id
-        self.job_status = job_status
-        self.job_type = job_type
-        self.job_creation_type = job_creation_time
-        self._client = client
-
-
 class JobError(Exception):
     def __init__(self, job_status: Dict[str, str], job: AsyncJob):
         final_status_message = job_status["message"]

nucleus/model_run.py

@@ -9,6 +9,9 @@
     BOX_TYPE,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     JOB_ID_KEY,
+    JOB_STATUS_KEY,
+    JOB_TYPE_KEY,
+    JOB_CREATION_TIME_KEY,
     POLYGON_TYPE,
     REQUEST_ID_KEY,
     SEGMENTATION_TYPE,
@@ -116,7 +119,13 @@ def predict(
                 route=f"modelRun/{self.model_run_id}/predict?async=1",
             )
 
-            return AsyncJob(response[JOB_ID_KEY], self._client)
+            return AsyncJob(
+                response[JOB_ID_KEY],
+                response[JOB_STATUS_KEY],
+                response[JOB_TYPE_KEY],
+                response[JOB_CREATION_TIME_KEY],
+                self._client,
+            )
         else:
             return self._client.predict(self.model_run_id, annotations, update)
 

nucleus/slice.py

@@ -6,7 +6,13 @@
 from nucleus.dataset_item import DatasetItem
 from nucleus.job import AsyncJob
 from nucleus.utils import convert_export_payload, format_dataset_item_response
-from nucleus.constants import EXPORTED_ROWS
+from nucleus.constants import (
+    EXPORTED_ROWS,
+    JOB_ID_KEY,
+    JOB_STATUS_KEY,
+    JOB_TYPE_KEY,
+    JOB_CREATION_TIME_KEY,
+)
 
 
 class Slice:
@@ -122,7 +128,13 @@ def send_to_labeling(self, project_id: str):
         response = self._client.make_request(
             {}, f"slice/{self.slice_id}/{project_id}/send_to_labeling"
         )
-        return AsyncJob(response["job_id"], self._client)
+        return AsyncJob(
+            response[JOB_ID_KEY],
+            response[JOB_STATUS_KEY],
+            response[JOB_TYPE_KEY],
+            response[JOB_CREATION_TIME_KEY],
+            self._client,
+        )
 
 
 def check_annotations_are_in_slice(