[Nucleus] Get job status (#379)

Author: jean-lucas

.flake8

@@ -3,6 +3,9 @@ ignore = E203, E266, E501, W503, F403, F401
 max-line-length = 79
 max-complexity = 18
 select = B,C,E,F,W,T4,B9
+classmethod-decorators =
+    classmethod
+    validator
 exclude =
     # All of these excludes should mirror something in .gitignore
     .git,

CHANGELOG.md

@@ -5,6 +5,41 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.15.1](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.1) - 2023-01-16
+
+### Changed
+- Better filter tuning of `client.list_jobs(args)` method
+
+### Added
+- Dataset method to filter jobs, and statistics on running jobs
+Example:
+```python
+>>> client = nucleus.NucleusClient(API_KEY)
+>>> ds = client.get_dataset(ds_id)
+>>> ds.jobs(show_completed=True, stats_only=True)
+{'autotagInference': {'Cancelled': 1, 'Completed': 11},
+ 'modelRunCommit': {'Completed': 7, 'Errored_Server': 1, 'Running': 1},
+ 'sliceQuery': {'Completed': 40, 'Running': 2}}
+```
+
+Detailed Example
+```python
+>>> from nucleus.job import CustomerJobTypes
+>>> client = nucleus.NucleusClient(API_KEY)
+>>> ds = client.get_dataset(ds_id)
+>>> from_date = "2022-12-20"; to_date = "2023-01-15"
+>>> job_types = [CustomerJobTypes.MODEL_INFERENCE_RUN, CustomerJobTypes.UPLOAD_DATASET_ITEMS]
+>>> ds.jobs(
+  from_date=from_date,
+  to_date=to_date,
+  show_completed=True,
+  job_types=job_types,
+  limit=150
+)
+# ... returns list of AsyncJob objects
+```
+
+
 ## [0.15.0](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.0) - 2022-12-19
 
 ### Changed

nucleus/__init__.py

@@ -40,6 +40,7 @@
     "VideoScene",
 ]
 
+import datetime
 import os
 import warnings
 from typing import Any, Dict, List, Optional, Sequence, Union
@@ -105,6 +106,7 @@
 )
 from .data_transfer_object.dataset_details import DatasetDetails
 from .data_transfer_object.dataset_info import DatasetInfo
+from .data_transfer_object.job_status import JobInfoRequestPayload
 from .dataset import Dataset
 from .dataset_item import DatasetItem
 from .deprecation_warning import deprecated
@@ -116,6 +118,7 @@
     NotFoundError,
     NucleusAPIError,
 )
+from .job import CustomerJobTypes
 from .logger import logger
 from .model import Model
 from .model_run import ModelRun
@@ -251,23 +254,47 @@ 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
+        self,
+        show_completed: bool = False,
+        from_date: Optional[Union[str, datetime.datetime]] = None,
+        to_date: Optional[Union[str, datetime.datetime]] = None,
+        job_types: Optional[List[CustomerJobTypes]] = None,
+        limit: Optional[int] = None,
+        dataset_id: Optional[str] = None,
+        date_limit: Optional[str] = None,
     ) -> List[AsyncJob]:
         """Fetches all of your running jobs in Nucleus.
 
         Parameters:
-            show_completed: Whether to fetch completed and errored jobs or just
-              running jobs. Default behavior is False.
-            date_limit: Only fetch jobs that were started after this date. Default
-              behavior is 2 weeks prior to the current date.
-
-        Returns:
-            List[:class:`AsyncJob`]: List of running asynchronous jobs
-            associated with the client API key.
+            job_types: Filter on set of job types, if None, fetch all types
+            from_date: beginning of date range filter
+            to_date: end of date range filter
+            limit: number of results to fetch, max 50_000
+            show_completed: dont fetch jobs with Completed status
+            stats_only: return overview of jobs, instead of a list of job objects
+            dataset_id: filter on a particular dataset
+            date_limit: Deprecated, do not use
+
+         Returns:
+             List[:class:`AsyncJob`]: List of running asynchronous jobs
+             associated with the client API key.
         """
-        # TODO: What type is date_limit? Use pydantic ...
-        payload = {show_completed: show_completed, date_limit: date_limit}
-        job_objects = self.make_request(payload, "jobs/", requests.get)
+
+        if date_limit is not None:
+            warnings.warn(
+                "Argument `date_limit` is no longer supported. Consider using the `from_date` and `to_date` args."
+            )
+
+        payload = JobInfoRequestPayload(
+            dataset_id=dataset_id,
+            show_completed=show_completed,
+            from_date=from_date,
+            to_date=to_date,
+            limit=limit,
+            job_types=job_types,
+        ).dict()
+
+        job_objects = self.make_request(payload, "jobs/", requests.post)
         return [
             AsyncJob(
                 job_id=job[JOB_ID_KEY],
@@ -1032,7 +1059,7 @@ def make_request(
         route: str,
         requests_command=requests.post,
         return_raw_response: bool = False,
-    ) -> dict:
+    ) -> Union[dict, Any]:
         """Makes a request to a Nucleus API endpoint.
 
         Logs a warning if not successful.

nucleus/constants.py

@@ -77,6 +77,7 @@
 K4_KEY = "k4"
 KEEP_HISTORY_KEY = "keep_history"
 LENGTH_KEY = "length"
+JOB_REQ_LIMIT = 50_000
 JOB_STATUS_KEY = "job_status"
 JOB_LAST_KNOWN_STATUS_KEY = "job_last_known_status"
 JOB_TYPE_KEY = "job_type"

nucleus/data_transfer_object/job_status.py

@@ -0,0 +1,55 @@
+# pylint: disable=E0213
+
+from datetime import datetime
+from typing import List, Optional, Union
+
+from dateutil.parser import ParserError, parse
+from pydantic import validator
+
+from nucleus.constants import JOB_REQ_LIMIT
+from nucleus.job import CustomerJobTypes
+from nucleus.pydantic_base import ImmutableModel
+
+
+class JobInfoRequestPayload(ImmutableModel):
+    dataset_id: Optional[str]
+    job_types: Optional[List[CustomerJobTypes]]
+    from_date: Optional[Union[str, datetime]]
+    to_date: Optional[Union[str, datetime]]
+    limit: Optional[int]
+    show_completed: bool
+
+    @validator("from_date", "to_date")
+    def ensure_date_format(cls, date):
+        if date is None:
+            return None
+        if isinstance(date, datetime):
+            return str(date)
+        try:
+            parse(date)
+        except ParserError as err:
+            raise ValueError(
+                f"Date {date} not a valid date. Try using YYYY-MM-DD format."
+            ) from err
+        return date
+
+    @validator("limit")
+    def ensure_limit(cls, limit):
+        if limit is None:
+            return JOB_REQ_LIMIT
+        if limit > JOB_REQ_LIMIT:
+            raise ValueError(f"Max request limit is 50,000, but got: {limit}.")
+        return limit
+
+    @validator("job_types")
+    def ensure_job_type(cls, job_types):
+        if job_types is None:
+            return []
+        try:
+            assert all(t in CustomerJobTypes for t in job_types)
+        except AssertionError as badType:
+            raise ValueError(
+                f"Job types must be one of: {CustomerJobTypes.options()}"
+            ) from badType
+
+        return [t.value for t in job_types]

nucleus/dataset.py

@@ -1,5 +1,16 @@
+import datetime
 import os
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import requests
 
@@ -32,6 +43,7 @@
     EXPORTED_ROWS,
     FRAME_RATE_KEY,
     ITEMS_KEY,
+    JOB_REQ_LIMIT,
     KEEP_HISTORY_KEY,
     MESSAGE_KEY,
     NAME_KEY,
@@ -54,6 +66,7 @@
 from .dataset_item_uploader import DatasetItemUploader
 from .deprecation_warning import deprecated
 from .errors import NotFoundError, NucleusAPIError
+from .job import CustomerJobTypes, jobs_status_overview
 from .metadata_manager import ExportMetadataType, MetadataManager
 from .payload_constructor import (
     construct_append_scenes_payload,
@@ -70,6 +83,9 @@
 )
 from .upload_response import UploadResponse
 
+if TYPE_CHECKING:
+    from . import NucleusClient
+
 # TODO: refactor to reduce this file to under 1000 lines.
 # pylint: disable=C0302
 
@@ -107,7 +123,7 @@ class Dataset:
         existing_dataset = client.get_dataset("YOUR_DATASET_ID")
     """
 
-    def __init__(self, dataset_id, client, name=None):
+    def __init__(self, dataset_id, client: "NucleusClient", name=None):
         self.id = dataset_id
         self._client = client
         # NOTE: Optionally set name on creation such that the property access doesn't need to hit the server
@@ -144,7 +160,7 @@ def is_scene(self) -> bool:
             {}, f"dataset/{self.id}/is_scene", requests.get
         )[DATASET_IS_SCENE_KEY]
         self._is_scene = response
-        return self._is_scene
+        return self._is_scene  # type: ignore
 
     @property
     def model_runs(self) -> List[str]:
@@ -153,7 +169,7 @@ def model_runs(self) -> List[str]:
         response = self._client.make_request(
             {}, f"dataset/{self.id}/model_runs", requests.get
         )
-        return response
+        return response  # type: ignore
 
     @property
     def slices(self) -> List[Slice]:
@@ -885,7 +901,7 @@ def build_slice(
         sample_size: int,
         sample_method: Union[str, SliceBuilderMethods],
         filters: Optional[SliceBuilderFilters] = None,
-    ) -> Union[str, Tuple[AsyncJob, str]]:
+    ) -> Union[str, Tuple[AsyncJob, str], dict]:
         """Build a slice using Nucleus' Smart Sample tool. Allowing slices to be built
         based on certain criteria, and filters.
 
@@ -1926,3 +1942,36 @@ def delete_tracks(self, track_reference_ids: List[str]) -> None:
             route=f"dataset/{self.id}/tracks",
             requests_command=requests.delete,
         )
+
+    def jobs(
+        self,
+        job_types: Optional[List[CustomerJobTypes]] = None,
+        from_date: Optional[Union[str, datetime.datetime]] = None,
+        to_date: Optional[Union[str, datetime.datetime]] = None,
+        limit: int = JOB_REQ_LIMIT,
+        show_completed: bool = False,
+        stats_only: bool = False,
+    ):
+        """
+        Fetch jobs pertaining to this particular dataset.
+
+        Parameters:
+            job_types: Filter on set of job types, if None, fetch all types, ie: ['uploadDatasetItems']
+            from_date: beginning of date range, as a string 'YYYY-MM-DD' or datetime object.
+                For example: '2021-11-05', parser.parse('Nov 5 2021'), or datetime(2021,11,5)
+            to_date: end of date range
+            limit: number of results to fetch, max 50_000
+            show_completed: dont fetch jobs with Completed status
+            stats_only: return overview of jobs, instead of a list of job objects
+        """
+        job_objects = self._client.list_jobs(
+            dataset_id=self.id,
+            show_completed=show_completed,
+            from_date=from_date,
+            to_date=to_date,
+            limit=limit,
+            job_types=job_types,
+        )
+        if stats_only:
+            return jobs_status_overview(job_objects)
+        return job_objects

nucleus/job.py

@@ -0,0 +1,55 @@
+from collections import defaultdict
+from enum import Enum
+from typing import Any, Dict, List
+
+from .async_job import AsyncJob
+
+
+class CustomerJobTypes(str, Enum):
+    UPLOAD_DATASET_ITEMS = "uploadDatasetItems"
+    UPLOAD_PREDICTIONS = "uploadPredictions"
+    UPLOAD_ANNOTATIONS = "uploadAnnotations"
+    UPLOAD_LIDAR_SCENE = "uploadLidarScene"
+    UPLOAD_VIDEO_SCENE = "uploadVideoScene"
+    MODEL_INFERENCE_RUN = "modelInferenceRun"
+    INDEXING_IMAGE = "indexingImage"
+    INDEXING_OBJECT = "indexingObject"
+    ANNOTATION_DELETION = "annotationDeletion"
+    SEND_TO_LABELING = "sendToLabeling"
+    INGEST_TASKS = "ingestTasks"
+    CUSTOM_INDEXING = "customIndexing"
+    EMBEDDING_DELETE = "embeddingDelete"
+    TEST_EVALUATION = "testEvaluation"
+    VALIDATE_METRICS = "modelMetrics"
+    MODEL_RUN_COMMIT = "modelRunCommit"
+    AUTOTAG_INFERENCE = "autotagInference"
+    SLICE_QUERY = "sliceQuery"
+    CLONE_DATASET = "cloneDataset"
+    METADATA_UPDATE = "metadataUpdate"
+    TRIGGER_EVALUATE = "triggerEvaluate"
+
+    def __contains__(self, item):
+        try:
+            self(item)
+        except ValueError:
+            return False
+        return True
+
+    @staticmethod
+    def options():
+        return list(map(lambda c: c.value, CustomerJobTypes))
+
+
+def jobs_status_overview(jobs: List[AsyncJob]) -> Dict[str, Any]:
+    jobs_by_type = defaultdict(list)
+    for job in jobs:
+        jobs_by_type[job.job_type].append(job)
+
+    jobs_status = {}
+    for job_type, job_collection in jobs_by_type.items():
+        overview = defaultdict(int)  # type: Dict[str, int]
+        for job in job_collection:
+            overview[job.job_last_known_status] += 1
+        jobs_status[job_type] = dict(overview)
+
+    return jobs_status