From d884ebcdd5081148e76f094385f48d0ee587dad7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A2=A8=E8=A6=8B=E4=BA=AE?= <ryokazami72@gmail.com>
Date: Wed, 1 Nov 2023 16:23:21 +0900
Subject: [PATCH] feat get training jobs

---
 README.md             | 84 +++++++++++++++++++++++++++++++++++++++++++
 fastlabel/__init__.py | 31 ++++++++++++++--
 2 files changed, 113 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index f89e6ab..bb56609 100644
--- a/README.md
+++ b/README.md
@@ -2964,6 +2964,90 @@ for image_file_path in glob.iglob(os.path.join(input_dataset_path, "**/**.jpg"),
 
 > Please check const.COLOR_PALLETE for index colors.
 
+## Model
+
+### Get training jobs
+
+Get training jobs.
+
+```python
+def get_training_jobs() -> list[dict]:
+    all_training_jobs = []
+    offset = None
+    while True:
+        time.sleep(1)
+
+        training_jobs = client.get_training_jobs(offset=offset)
+        all_training_jobs.extend(training_jobs)
+
+        if len(training_jobs) > 0:
+            offset = len(all_training_jobs)
+        else:
+            break
+    return all_training_jobs
+
+```
+
+#### Response
+
+Example of two training jobs.
+
+```python
+[
+    {
+        "trainingJobId": "f40c5838-4c3a-482f-96b7-f77e16c96fed",
+        "status": "in_progress",
+        "baseModelName": "FastLabel Object Detection High Accuracy - 汎用",
+        "instanceType": "ml.p3.2xlarge",
+        "epoch": 300,
+        "projects": [
+            "image-bbox"
+        ],
+        "statuses": [],
+        "tags": [],
+        "contentCount": 23,
+        "userName": "Admin",
+        "createdAt": "2023-10-31T07:10:28.306Z",
+        "completedAt": null,
+        "customModel": {
+            "modelId": "",
+            "modelName": "",
+            "modelURL": "",
+            "classes": []
+        }
+    },
+    {
+        "trainingJobId": "1d2bc86a-c7f1-40a5-8e85-48246cc3c8d2",
+        "status": "completed",
+        "baseModelName": "custom-object-detection-image",
+        "instanceType": "ml.p3.2xlarge",
+        "epoch": 300,
+        "projects": [
+            "image-bbox"
+        ],
+        "statuses": [
+            "approved"
+        ],
+        "tags": [
+            "trainval"
+        ],
+        "contentCount": 20,
+        "userName": "Admin",
+        "createdAt": "2023-10-31T06:56:28.112Z",
+        "completedAt": "2023-10-31T07:08:26.000Z",
+        "customModel": {
+            "modelId": "a6728876-2eb7-49b5-9fd8-7dee1b8a81b3",
+            "modelName": "fastlabel_object_detection-2023-10-31-07-08-29",
+            "modelURL": "URL for download model file",
+            "classes": [
+                "person"
+            ]
+        }
+    }
+]
+```
+
+
 ## Execute endpoint
 
 Create the endpoint from the screen at first.
diff --git a/fastlabel/__init__.py b/fastlabel/__init__.py
index 9634f82..a31791c 100644
--- a/fastlabel/__init__.py
+++ b/fastlabel/__init__.py
@@ -1258,7 +1258,8 @@ def create_video_task(
             )
         if not utils.is_video_supported_codec(file_path):
             raise FastLabelInvalidException(
-                "Supported video encoding for registration through the SDK is only H.264.", 422
+                "Supported video encoding for registration through the SDK is only H.264.",
+                422,
             )
 
         file = utils.base64_encode(file_path)
@@ -1326,7 +1327,8 @@ def create_video_classification_task(
             )
         if not utils.is_video_supported_codec(file_path):
             raise FastLabelInvalidException(
-                "Supported video encoding for registration through the SDK is only H.264.", 422
+                "Supported video encoding for registration through the SDK is only H.264.",
+                422,
             )
 
         file = utils.base64_encode(file_path)
@@ -4144,6 +4146,31 @@ def __get_signed_path(
         params = {"project": project, "fileName": file_name, "fileType": file_type}
         return self.api.get_request(endpoint, params)
 
+    def get_training_jobs(
+        self,
+        offset: int = None,
+        limit: int = 100,
+    ) -> list:
+        """
+        Returns a list of training jobs.
+        Returns up to 1000 at a time, to get more, set offset as the starting position
+        to fetch.
+        offset is the starting position number to fetch (Optional).
+        limit is the max number to fetch (Optional).
+        """
+        if limit > 1000:
+            raise FastLabelInvalidException(
+                "Limit must be less than or equal to 1000.", 422
+            )
+        endpoint = "trainings"
+        params = {}
+        if offset:
+            params["offset"] = offset
+        if limit:
+            params["limit"] = limit
+
+        return self.api.get_request(endpoint, params=params)
+
     def execute_endpoint(
         self,
         endpoint_name: str,