fix(client): add "jobType" to Job.abort, Job.retry and delete_job

Graviti-AI · Apr 6, 2022 · b734ebd · b734ebd
1 parent 534ff59
commit b734ebd
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/tensorbay/client/job.py b/tensorbay/client/job.py
@@ -165,11 +165,15 @@ def update(self, until_complete: bool = False) -> None:
 
     def abort(self) -> None:
         """Abort a :class:`Job`."""
-        self._client.open_api_do("POST", f"jobs/{self.job_id}/abort", self._dataset_id)
+        self._client.open_api_do(
+            "POST", f"jobs/{self.job_id}/abort?jobType={self._job_type}", self._dataset_id
+        )
 
     def retry(self) -> None:
         """Retry a :class:`Job`."""
-        self._client.open_api_do("POST", f"jobs/{self.job_id}/retry", self._dataset_id)
+        self._client.open_api_do(
+            "POST", f"jobs/{self.job_id}/retry?jobType={self._job_type}", self._dataset_id
+        )
 
 
 class SquashAndMergeJob(Job):

diff --git a/tensorbay/client/version.py b/tensorbay/client/version.py
@@ -523,6 +523,7 @@ class JobMixin:
     _dataset_id: str
     _client: Client
     _status: Status
+    _JOB_TYPE: str
 
     def _create_job(
         self,
@@ -616,7 +617,9 @@ def delete_job(self, job_id: str) -> None:
             job_id: The Job id.
 
         """
-        self._client.open_api_do("DELETE", f"jobs/{job_id}", self._dataset_id)
+        self._client.open_api_do(
+            "DELETE", f"jobs/{job_id}?jobType={self._JOB_TYPE}", self._dataset_id
+        )
 
 
 class SquashAndMerge(JobMixin):