orca-sdk 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

orca_sdk/job.py

@@ -7,7 +7,7 @@ from typing import Callable, Generic, TypedDict, TypeVar, cast
 
 from tqdm.auto import tqdm
 
-from .client import orca_api
+from .client import OrcaClient
 
 
 class JobConfig(TypedDict):
@@ -17,7 +17,7 @@ class JobConfig(TypedDict):
 
 
 class Status(Enum):
-    """Status of a cloud job in the task queue"""
+    """Status of a cloud job in the job queue"""
 
     # the INITIALIZED state should never be returned by the API
     INITIALIZED = "INITIALIZED"
@@ -140,8 +140,9 @@ class Job(Generic[TResult]):
         Returns:
             List of jobs matching the given filters
         """
-        paginated_tasks = orca_api.GET(
-            "/task",
+        client = OrcaClient._resolve_client()
+        paginated_jobs = client.GET(
+            "/job",
             params={
                 "status": (
                     [s.value for s in status]
@@ -174,7 +175,7 @@ class Job(Generic[TResult]):
                     obj,
                 )[-1]
             )(t)
-            for t in paginated_tasks["items"]
+            for t in paginated_jobs["items"]
         ]
 
     def __init__(self, id: str, get_value: Callable[[], TResult | None] | None = None):
@@ -186,27 +187,30 @@ class Job(Generic[TResult]):
             get_value: Optional function to customize how the value is resolved, if not provided the result will be a dict
         """
         self.id = id
-        task = orca_api.GET("/task/{task_id}", params={"task_id": id})
-
-        self._get_value = get_value or (
-            lambda: cast(TResult | None, orca_api.GET("/task/{task_id}", params={"task_id": id})["result"])
-        )
-        self.type = task["type"]
-        self.status = Status(task["status"])
-        self.steps_total = task["steps_total"]
-        self.steps_completed = task["steps_completed"]
-        self.exception = task["exception"]
+        client = OrcaClient._resolve_client()
+        job = client.GET("/job/{job_id}", params={"job_id": id})
+
+        def default_get_value():
+            client = OrcaClient._resolve_client()
+            return cast(TResult | None, client.GET("/job/{job_id}", params={"job_id": id})["result"])
+
+        self._get_value = get_value or default_get_value
+        self.type = job["type"]
+        self.status = Status(job["status"])
+        self.steps_total = job["steps_total"]
+        self.steps_completed = job["steps_completed"]
+        self.exception = job["exception"]
         self.value = (
             None
-            if task["status"] != "COMPLETED"
+            if job["status"] != "COMPLETED"
             else (
                 get_value()
                 if get_value is not None
-                else cast(TResult, task["result"]) if task["result"] is not None else None
+                else cast(TResult, job["result"]) if job["result"] is not None else None
             )
         )
-        self.updated_at = datetime.fromisoformat(task["updated_at"])
-        self.created_at = datetime.fromisoformat(task["created_at"])
+        self.updated_at = datetime.fromisoformat(job["updated_at"])
+        self.created_at = datetime.fromisoformat(job["created_at"])
         self.refreshed_at = datetime.now()
 
     def refresh(self, throttle: float = 0):
@@ -222,7 +226,8 @@ class Job(Generic[TResult]):
             return
         self.refreshed_at = current_time
 
-        status_info = orca_api.GET("/task/{task_id}/status", params={"task_id": self.id})
+        client = OrcaClient._resolve_client()
+        status_info = client.GET("/job/{job_id}/status", params={"job_id": self.id})
         self.status = Status(status_info["status"])
         if status_info["steps_total"] is not None:
             self.steps_total = status_info["steps_total"]
@@ -333,5 +338,6 @@ def abort(self, show_progress: bool = False, refresh_interval: int = 1, max_wait
         refresh_interval: Polling interval in seconds while waiting for the job to abort
         max_wait: Maximum time to wait for the job to abort in seconds
     """
-    orca_api.DELETE("/task/{task_id}/abort", params={"task_id": self.id})
+    client = OrcaClient._resolve_client()
+    client.DELETE("/job/{job_id}/abort", params={"job_id": self.id})
     self.wait(show_progress, refresh_interval, max_wait)