orca-sdk 0.0.97__py3-none-any.whl → 0.0.98__py3-none-any.whl

orca_sdk/job.py

@@ -3,12 +3,11 @@ from __future__ import annotations
 import time
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Any, Callable, Generic, TypedDict, TypeVar, cast
+from typing import Callable, Generic, TypedDict, TypeVar, cast
 
 from tqdm.auto import tqdm
 
-from ._generated_api_client.api import abort_task, get_task, get_task_status, list_tasks
-from ._generated_api_client.models import TaskStatus
+from .client import TaskStatus, orca_api
 
 
 class JobConfig(TypedDict):
@@ -21,6 +20,8 @@ class Status(Enum):
     """Status of a cloud job in the task queue"""
 
     # the INITIALIZED state should never be returned by the API
+    INITIALIZED = "INITIALIZED"
+    """The job has been initialized"""
 
     DISPATCHED = "DISPATCHED"
     """The job has been queued and is waiting to be processed"""
@@ -136,17 +137,20 @@ class Job(Generic[TResult]):
         Returns:
             List of jobs matching the given filters
         """
-        paginated_tasks = list_tasks(
-            status=(
-                [TaskStatus(s.value) for s in status]
-                if isinstance(status, list)
-                else TaskStatus(status.value) if isinstance(status, Status) else None
-            ),
-            type=type,
-            limit=limit,
-            offset=offset,
-            start_timestamp=start,
-            end_timestamp=end,
+        paginated_tasks = orca_api.GET(
+            "/task",
+            params={
+                "status": (
+                    [s.value for s in status]
+                    if isinstance(status, list)
+                    else status.value if isinstance(status, Status) else None
+                ),
+                "type": type,
+                "limit": limit,
+                "offset": offset,
+                "start_timestamp": start.isoformat() if start is not None else None,
+                "end_timestamp": end.isoformat() if end is not None else None,
+            },
         )
 
         # can't use constructor because it makes an API call, so we construct the objects manually
@@ -154,20 +158,20 @@ class Job(Generic[TResult]):
             (
                 lambda t: (
                     obj := cls.__new__(cls),
-                    setattr(obj, "id", t.id),
-                    setattr(obj, "type", t.type),
-                    setattr(obj, "status", Status(t.status.value)),
-                    setattr(obj, "steps_total", t.steps_total),
-                    setattr(obj, "steps_completed", t.steps_completed),
-                    setattr(obj, "exception", t.exception),
-                    setattr(obj, "value", cast(TResult, t.result.to_dict()) if t.result is not None else None),
-                    setattr(obj, "updated_at", t.updated_at),
-                    setattr(obj, "created_at", t.created_at),
+                    setattr(obj, "id", t["id"]),
+                    setattr(obj, "type", t["type"]),
+                    setattr(obj, "status", Status(t["status"])),
+                    setattr(obj, "steps_total", t["steps_total"]),
+                    setattr(obj, "steps_completed", t["steps_completed"]),
+                    setattr(obj, "exception", t["exception"]),
+                    setattr(obj, "value", cast(TResult, t["result"]) if t["result"] is not None else None),
+                    setattr(obj, "updated_at", datetime.fromisoformat(t["updated_at"])),
+                    setattr(obj, "created_at", datetime.fromisoformat(t["created_at"])),
                     setattr(obj, "refreshed_at", datetime.now()),
                     obj,
                 )[-1]
             )(t)
-            for t in paginated_tasks.items
+            for t in paginated_tasks["items"]
         ]
 
     def __init__(self, id: str, get_value: Callable[[], TResult | None] | None = None):
@@ -179,27 +183,27 @@ 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 = get_task(self.id)
+        task = orca_api.GET("/task/{task_id}", params={"task_id": id})
 
         self._get_value = get_value or (
-            lambda: (r := get_task(id).result) and (cast(TResult, r.to_dict()) if r else None)
+            lambda: cast(TResult | None, orca_api.GET("/task/{task_id}", params={"task_id": id})["result"])
         )
-        self.type = task.type
-        self.status = Status(task.status.value)
-        self.steps_total = task.steps_total
-        self.steps_completed = task.steps_completed
-        self.exception = task.exception
+        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"]
         self.value = (
             None
-            if task.status != TaskStatus.COMPLETED
+            if task["status"] != "COMPLETED"
             else (
                 get_value()
                 if get_value is not None
-                else cast(TResult, task.result.to_dict()) if task.result is not None else None
+                else cast(TResult, task["result"]) if task["result"] is not None else None
             )
         )
-        self.updated_at = task.updated_at
-        self.created_at = task.created_at
+        self.updated_at = datetime.fromisoformat(task["updated_at"])
+        self.created_at = datetime.fromisoformat(task["created_at"])
         self.refreshed_at = datetime.now()
 
     def refresh(self, throttle: float = 0):
@@ -215,17 +219,17 @@ class Job(Generic[TResult]):
             return
         self.refreshed_at = current_time
 
-        status_info = get_task_status(self.id)
-        self.status = Status(status_info.status.value)
-        if status_info.steps_total is not None:
-            self.steps_total = status_info.steps_total
-        if status_info.steps_completed is not None:
-            self.steps_completed = status_info.steps_completed
+        status_info = orca_api.GET("/task/{task_id}/status", params={"task_id": self.id})
+        self.status = Status(status_info["status"])
+        if status_info["steps_total"] is not None:
+            self.steps_total = status_info["steps_total"]
+        if status_info["steps_completed"] is not None:
+            self.steps_completed = status_info["steps_completed"]
 
-        self.exception = status_info.exception
-        self.updated_at = status_info.updated_at
+        self.exception = status_info["exception"]
+        self.updated_at = datetime.fromisoformat(status_info["updated_at"])
 
-        if status_info.status == TaskStatus.COMPLETED:
+        if status_info["status"] == "COMPLETED":
             self.value = self._get_value()
 
     def __getattribute__(self, name: str):
@@ -326,5 +330,5 @@ 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
     """
-    abort_task(self.id)
+    orca_api.DELETE("/task/{task_id}/abort", params={"task_id": self.id})
     self.wait(show_progress, refresh_interval, max_wait)

orca_sdk/job_test.py

@@ -5,6 +5,19 @@ from .datasource import Datasource
 from .job import Job, Status
 
 
+def wait_for_jobs_status(job_ids, expected_statuses, timeout=10, poll_interval=0.2):
+    """
+    Wait until all jobs reach one of the expected statuses or timeout is reached.
+    """
+    start = time.time()
+    while time.time() - start < timeout:
+        jobs = [Job(job_id) for job_id in job_ids]
+        if all(job.status in expected_statuses for job in jobs):
+            return
+        time.sleep(poll_interval)
+    raise TimeoutError(f"Jobs did not reach statuses {expected_statuses} within {timeout} seconds")
+
+
 def test_job_creation(classification_model: ClassificationModel, datasource: Datasource):
     job = classification_model.evaluate(datasource, background=True)
     assert job.id is not None
@@ -56,6 +69,9 @@ def test_job_query_pagination(classification_model: ClassificationModel, datasou
         job = classification_model.evaluate(datasource, background=True)
         jobs_created.append(job.id)
 
+    # Wait for jobs to be at least PROCESSING or COMPLETED
+    wait_for_jobs_status(jobs_created, expected_statuses=[Status.PROCESSING, Status.COMPLETED])
+
     # Test basic pagination with limit
     jobs_page1 = Job.query(type="EVALUATE_MODEL", limit=2)
     assert len(jobs_page1) == 2