Merge pull request #88 from scaleapi/list_jobs

Add list_jobs to nucleus-python-client api

Author: ardila

README.md

@@ -50,7 +50,10 @@ dataset = client.create_dataset("My Dataset")
 ```python
 datasets = client.list_datasets()
 ```
-
+### List Jobs
+```python
+jobs = client.list_jobs()
+```
 ### Delete a Dataset
 
 By specifying target dataset id.

nucleus/__init__.py

@@ -86,11 +86,16 @@
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
+    JOB_ID_KEY,
+    JOB_LAST_KNOWN_STATUS_KEY,
+    JOB_TYPE_KEY,
+    JOB_CREATION_TIME_KEY,
     IMAGE_KEY,
     IMAGE_URL_KEY,
     ITEM_METADATA_SCHEMA_KEY,
     ITEMS_KEY,
     KEEP_HISTORY_KEY,
+    MESSAGE_KEY,
     MODEL_RUN_ID_KEY,
     NAME_KEY,
     NUCLEUS_ENDPOINT,
@@ -110,6 +115,7 @@
     NotFoundError,
     NucleusAPIError,
 )
+from .job import AsyncJob
 from .model import Model
 from .model_run import ModelRun
 from .payload_constructor import (
@@ -199,6 +205,26 @@ 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[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 [
+            AsyncJob(
+                job_id=job[JOB_ID_KEY],
+                job_last_known_status=job[JOB_LAST_KNOWN_STATUS_KEY],
+                job_type=job[JOB_TYPE_KEY],
+                job_creation_time=job[JOB_CREATION_TIME_KEY],
+                client=self,
+            )
+            for job in job_objects
+        ]
+
     def get_dataset_items(self, dataset_id) -> List[DatasetItem]:
         """
         Gets all the dataset items inside your repo as a json blob.

nucleus/constants.py

@@ -42,6 +42,10 @@
 ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
 JOB_ID_KEY = "job_id"
 KEEP_HISTORY_KEY = "keep_history"
+JOB_STATUS_KEY = "job_status"
+JOB_LAST_KNOWN_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

@@ -23,7 +23,6 @@
     DATASET_SLICES_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     EXPORTED_ROWS,
-    JOB_ID_KEY,
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
@@ -181,8 +180,7 @@ def annotate(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/annotate?async=1",
             )
-
-            return AsyncJob(response[JOB_ID_KEY], self._client)
+            return AsyncJob.from_json(response, self._client)
 
         return self._client.annotate_dataset(
             self.id, annotations, update=update, batch_size=batch_size
@@ -241,7 +239,7 @@ 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.from_json(response, self._client)
 
         return self._client.populate_dataset(
             self.id,
@@ -368,4 +366,4 @@ def delete_annotations(
         response = self._client.delete_annotations(
             self.id, reference_ids, keep_history
         )
-        return AsyncJob(response[JOB_ID_KEY], self._client)
+        return AsyncJob.from_json(response, self._client)

nucleus/job.py

@@ -1,28 +1,39 @@
 from dataclasses import dataclass
 import time
 from typing import Dict, List
-
 import requests
+from nucleus.constants import (
+    JOB_CREATION_TIME_KEY,
+    JOB_ID_KEY,
+    JOB_LAST_KNOWN_STATUS_KEY,
+    JOB_TYPE_KEY,
+    STATUS_KEY,
+)
 
 JOB_POLLING_INTERVAL = 5
 
 
 @dataclass
 class AsyncJob:
-    id: str
+    job_id: str
+    job_last_known_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(
+        response = self.client.make_request(
             payload={},
-            route=f"job/{self.id}",
+            route=f"job/{self.job_id}",
             requests_command=requests.get,
         )
+        self.job_last_known_status = response[STATUS_KEY]
+        return response
 
     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,
         )
 
@@ -42,6 +53,16 @@ def sleep_until_complete(self, verbose_std_out=True):
         if final_status["status"] == "Errored":
             raise JobError(final_status, self)
 
+    @classmethod
+    def from_json(cls, payload: dict, client):
+        return cls(
+            job_id=payload[JOB_ID_KEY],
+            job_last_known_status=payload[JOB_LAST_KNOWN_STATUS_KEY],
+            job_type=payload[JOB_TYPE_KEY],
+            job_creation_time=payload[JOB_CREATION_TIME_KEY],
+            client=client,
+        )
+
 
 class JobError(Exception):
     def __init__(self, job_status: Dict[str, str], job: AsyncJob):

nucleus/model_run.py

@@ -8,7 +8,6 @@
     ANNOTATIONS_KEY,
     BOX_TYPE,
     DEFAULT_ANNOTATION_UPDATE_MODE,
-    JOB_ID_KEY,
     POLYGON_TYPE,
     REQUEST_ID_KEY,
     SEGMENTATION_TYPE,
@@ -115,8 +114,7 @@ def predict(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"modelRun/{self.model_run_id}/predict?async=1",
             )
-
-            return AsyncJob(response[JOB_ID_KEY], self._client)
+            return AsyncJob.from_json(response, self._client)
         else:
             return self._client.predict(self.model_run_id, annotations, update)
 

nucleus/slice.py

@@ -6,7 +6,9 @@
 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,
+)
 
 
 class Slice:
@@ -122,7 +124,7 @@ 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.from_json(response, self._client)
 
 
 def check_annotations_are_in_slice(

tests/test_dataset.py

@@ -169,12 +169,12 @@ def test_dataset_append_async(dataset: Dataset):
     status = job.status()
     status["message"]["PayloadUrl"] = ""
     assert status == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "PayloadUrl": "",
             "image_upload_step": {"errored": 0, "pending": 0, "completed": 5},
-            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.id}",
+            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.job_id}",
             "ingest_to_reupload_queue": {
                 "epoch": 1,
                 "total": 5,
@@ -204,7 +204,7 @@ def test_dataset_append_async_with_1_bad_url(dataset: Dataset):
     status = job.status()
     status["message"]["PayloadUrl"] = ""
     assert status == {
-        "job_id": f"{job.id}",
+        "job_id": f"{job.job_id}",
         "status": "Errored",
         "message": {
             "PayloadUrl": "",
@@ -220,7 +220,7 @@ def test_dataset_append_async_with_1_bad_url(dataset: Dataset):
                 "datasetId": f"{dataset.id}",
                 "processed": 5,
             },
-            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.id}",
+            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.job_id}",
         },
     }
     # The error is fairly detailed and subject to change. What's important is we surface which URLs failed.
@@ -286,7 +286,7 @@ def test_annotate_async(dataset: Dataset):
     )
     job.sleep_until_complete()
     assert job.status() == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "annotation_upload": {
@@ -321,7 +321,7 @@ def test_annotate_async_with_error(dataset: Dataset):
     job.sleep_until_complete()
 
     assert job.status() == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "annotation_upload": {

tests/test_indexing.py

@@ -1,3 +1,4 @@
+from nucleus.job import AsyncJob
 import pytest
 
 from .helpers import (
@@ -42,7 +43,11 @@ def test_index_integration(dataset):
     create_response = dataset.create_custom_index(
         [signed_embeddings_url], embedding_dim=3
     )
-    assert JOB_ID_KEY in create_response
+    job = AsyncJob.from_json(create_response, client="Nucleus Client")
+    assert job.job_id
+    assert job.job_last_known_status
+    assert job.job_type
+    assert job.job_creation_time
     assert MESSAGE_KEY in create_response
     job_id = create_response[JOB_ID_KEY]
 

tests/test_jobs.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+import time
+import pytest
+from nucleus import (
+    AsyncJob,
+    NucleusClient,
+)
+
+
+def test_reprs():
+    # Have to define here in order to have access to all relevant objects
+    def test_repr(test_object: any):
+        assert eval(str(test_object)) == test_object
+
+    client = NucleusClient(api_key="fake_key")
+    test_repr(
+        AsyncJob(
+            client=client,
+            job_id="fake_job_id",
+            job_last_known_status="fake_job_status",
+            job_type="fake_job_type",
+            job_creation_time="fake_job_creation_time",
+        )
+    )
+
+
+def test_job_creation_and_listing(CLIENT):
+    jobs = CLIENT.list_jobs()
+
+    for job in jobs:
+        assert eval(str(job)) == job