sutro 0.1.16__tar.gz → 0.1.18__tar.gz

{sutro-0.1.16 → sutro-0.1.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sutro
-Version: 0.1.16
+Version: 0.1.18
 Summary: Sutro Python SDK
 Project-URL: Homepage, https://sutro.sh
 Project-URL: Documentation, https://docs.sutro.sh

{sutro-0.1.16 → sutro-0.1.18}/pyproject.toml

@@ -9,7 +9,7 @@ installer = "uv"
 
 [project]
 name = "sutro"
-version = "0.1.16"
+version = "0.1.18"
 description = "Sutro Python SDK"
 readme = "README.md"
 requires-python = ">=3.10"

{sutro-0.1.16 → sutro-0.1.18}/sutro/sdk.py

@@ -1,6 +1,7 @@
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
+from enum import Enum
 
 import requests
 import pandas as pd
@@ -17,6 +18,31 @@ import time
 from pydantic import BaseModel
 import json
 
+
+class JobStatus(str, Enum):
+    """Job statuses that will be returned by the API & SDK"""
+
+    UNKNOWN = "UNKNOWN"
+    QUEUED = "QUEUED"  # Job is waiting to start
+    STARTING = "STARTING"  # Job is in the process of starting up
+    RUNNING = "RUNNING"  # Job is actively running
+    SUCCEEDED = "SUCCEEDED"  # Job completed successfully
+    CANCELLING = "CANCELLING"  # Job is in the process of being canceled
+    CANCELLED = "CANCELLED"  # Job was canceled by the user
+    FAILED = "FAILED"  # Job failed
+
+    @classmethod
+    def terminal_statuses(cls) -> list["JobStatus"]:
+        return [
+            cls.SUCCEEDED,
+            cls.FAILED,
+            cls.CANCELLING,
+            cls.CANCELLED,
+        ]
+
+    def is_terminal(self) -> bool:
+        return self in self.terminal_statuses()
+
 # Initialize colorama (required for Windows)
 init()
 
@@ -181,7 +207,7 @@ class Sutro:
         sampling_params: dict = None,
         system_prompt: str = None,
         dry_run: bool = False,
-        stay_attached: bool = False,
+        stay_attached: Optional[bool] = None,
         random_seed_per_input: bool = False,
         truncate_rows: bool = False
     ):
@@ -211,7 +237,7 @@ class Sutro:
 
         """
         input_data = self.handle_data_helper(data, column)
-        stay_attached = stay_attached or job_priority == 0
+        stay_attached = stay_attached if stay_attached is not None else job_priority == 0
 
         # Convert BaseModel to dict if needed
         if output_schema is not None:
@@ -240,11 +266,6 @@ class Sutro:
             "random_seed_per_input": random_seed_per_input,
             "truncate_rows": truncate_rows
         }
-        if dry_run:
-            spinner_text = to_colored_text("Retrieving cost estimates...")
-        else:
-            t = f"Creating priority {job_priority} job"
-            spinner_text = to_colored_text(t)
 
         # There are two gotchas with yaspin:
         # 1. Can't use print while in spinner is running
@@ -253,6 +274,8 @@ class Sutro:
         # Terminal size {self._terminal_width} is too small to display spinner with the given settings.
         # https://github.com/pavdmyt/yaspin/blob/9c7430b499ab4611888ece39783a870e4a05fa45/yaspin/core.py#L568-L571
         job_id = None
+        t = f"Creating {'[dry run] ' if dry_run else ''}priority {job_priority} job"
+        spinner_text = to_colored_text(t)
         with yaspin(SPINNER, text=spinner_text, color=YASPIN_COLOR) as spinner:
             response = requests.post(
                 endpoint, data=json.dumps(payload), headers=headers
@@ -264,15 +287,21 @@ class Sutro:
                 )
                 spinner.stop()
                 print(to_colored_text(response.json(), state="fail"))
-                return
+                return None
             else:
+                job_id = response_data["results"]
                 if dry_run:
                     spinner.write(
-                        to_colored_text("✔ Cost estimates retrieved", state="success")
+                        to_colored_text(f"Awaiting cost estimates with job ID: {job_id}. You can safely detach and retrieve the cost estimates later.", state="info")
                     )
-                    return response_data["results"]
+                    spinner.stop()
+                    self.await_job_completion(job_id, obtain_results=False)
+                    cost_estimate = self._get_job_cost_estimate(job_id)
+                    spinner.write(
+                        to_colored_text(f"✔ Cost estimates retrieved for job {job_id}: ${cost_estimate}", state="success")
+                    )
+                    return job_id
                 else:
-                    job_id = response_data["results"]
                     spinner.write(
                         to_colored_text(
                             f"🛠️  Priority {job_priority} Job created with ID: {job_id}",
@@ -289,6 +318,13 @@ class Sutro:
 
         success = False
         if stay_attached and job_id is not None:
+            spinner.write(to_colored_text("Awaiting job start...", "info"))
+            started = self._await_job_start(job_id)
+            if not started:
+                failure_reason = self._get_failure_reason(job_id)
+                spinner.write(to_colored_text(f"Failure reason: {failure_reason['message']}", "fail"))
+                return None
+            
             s = requests.Session()
             payload = {
                 "job_id": job_id,
@@ -375,7 +411,7 @@ class Sutro:
                         )
                     )
                     spinner.stop()
-                    return
+                    return None
 
                 results = job_results_response.json()["results"]
 
@@ -401,6 +437,8 @@ class Sutro:
                     return data
 
                 return results
+            return None
+        return None
 
     def register_stream_listener(self, job_id: str) -> str:
         """Register a new stream listener and get a session token."""
@@ -691,40 +729,94 @@ class Sutro:
                 return
         return response.json()["jobs"]
 
-    def get_job_status(self, job_id: str):
+    def _list_jobs_helper(self):
         """
-        Get the status of a job by its ID.
+        Helper function to list jobs.
+        """
+        endpoint = f"{self.base_url}/list-jobs"
+        headers = {
+            "Authorization": f"Key {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        response = requests.get(endpoint, headers=headers)
+        if response.status_code != 200:
+            return None
+        return response.json()["jobs"]
 
-        This method retrieves the status of a job using its unique identifier.
+    def _get_job_cost_estimate(self, job_id: str):
+        """
+        Get the cost estimate for a job.
+        """
+        all_jobs = self._list_jobs_helper()
+        for job in all_jobs:
+            if job["job_id"] == job_id:
+                return job["cost_estimate"]
+        return None
+    
+    def _get_failure_reason(self, job_id: str):
+        """
+        Get the failure reason for a job.
+        """
+        all_jobs = self._list_jobs_helper()
+        for job in all_jobs:
+            if job["job_id"] == job_id:
+                return job["failure_reason"]
+        return None
+
+    def _fetch_job_status(self, job_id: str):
+        """
+        Core logic to fetch job status from the API.
 
         Args:
             job_id (str): The ID of the job to retrieve the status for.
 
         Returns:
-            str: The status of the job.
+            dict: The response JSON from the API.
+
+        Raises:
+            requests.HTTPError: If the API returns a non-200 status code.
         """
         endpoint = f"{self.base_url}/job-status/{job_id}"
         headers = {
             "Authorization": f"Key {self.api_key}",
             "Content-Type": "application/json",
         }
+
+        response = requests.get(endpoint, headers=headers)
+        response.raise_for_status()
+
+        return response.json()["job_status"][job_id]
+
+    def get_job_status(self, job_id: str):
+        """
+        Get the status of a job by its ID.
+
+        This method retrieves the status of a job using its unique identifier.
+
+        Args:
+            job_id (str): The ID of the job to retrieve the status for.
+
+        Returns:
+            str: The status of the job.
+        """
         with yaspin(
-            SPINNER,
-            text=to_colored_text(f"Checking job status with ID: {job_id}"),
-            color=YASPIN_COLOR,
+                SPINNER,
+                text=to_colored_text(f"Checking job status with ID: {job_id}"),
+                color=YASPIN_COLOR,
         ) as spinner:
-            response = requests.get(endpoint, headers=headers)
-            if response.status_code != 200:
+            try:
+                response_data = self._fetch_job_status(job_id)
+                spinner.write(to_colored_text("✔ Job status retrieved!", state="success"))
+                return response_data["job_status"][job_id]
+            except requests.HTTPError as e:
                 spinner.write(
                     to_colored_text(
-                        f"Bad status code: {response.status_code}", state="fail"
+                        f"Bad status code: {e.response.status_code}", state="fail"
                     )
                 )
                 spinner.stop()
-                print(to_colored_text(response.json(), state="fail"))
-                return
-            spinner.write(to_colored_text("✔ Job status retrieved!", state="success"))
-        return response.json()["job_status"][job_id]
+                print(to_colored_text(e.response.json(), state="fail"))
+                return None
 
     def get_job_results(
         self,
@@ -1113,3 +1205,96 @@ class Sutro:
                 print(to_colored_text(f"Error: {response.json()}", state="fail"))
                 return
         return response.json()["quotas"]
+
+    def await_job_completion(self, job_id: str, timeout: Optional[int] = 7200, obtain_results: bool = True) -> list | None:
+        """
+        Waits for job completion to occur and then returns the results upon
+        a successful completion.
+
+        Prints out the job's status every 5 seconds.
+
+        Args:
+            job_id (str): The ID of the job to await.
+            timeout (Optional[int]): The max time in seconds the function should wait for job results for. Default is 7200 (2 hours).
+
+        Returns:
+            list: The results of the job.
+        """
+        POLL_INTERVAL = 5
+
+        results = None
+        start_time = time.time()
+        with yaspin(
+            SPINNER, text=to_colored_text("Awaiting job completion"), color=YASPIN_COLOR
+        ) as spinner:
+            while (time.time() - start_time) < timeout:
+                try:
+                    status = self._fetch_job_status(job_id)
+                except requests.HTTPError as e:
+                    spinner.write(
+                        to_colored_text(
+                            f"Bad status code: {e.response.status_code}", state="fail"
+                        )
+                    )
+                    spinner.stop()
+                    print(to_colored_text(e.response.json(), state="fail"))
+                    return None
+
+                spinner.text = to_colored_text(f"Job status is {status} for {job_id}")
+
+                if status == JobStatus.SUCCEEDED:
+                    spinner.write(to_colored_text("Job completed! Retrieving results...", "success"))
+                    spinner.stop() # Stop this spinner as `get_job_results` has its own spinner text
+                    if obtain_results:
+                        results = self.get_job_results(job_id)
+                    break
+                if status == JobStatus.FAILED:
+                    spinner.write(to_colored_text("Job has failed", "fail"))
+                    return None
+                if status == JobStatus.CANCELLED:
+                    spinner.write(to_colored_text("Job has been cancelled"))
+                    return None
+
+
+                time.sleep(POLL_INTERVAL)
+
+        return results
+    
+    def _await_job_start(self, job_id: str, timeout: Optional[int] = 7200):
+        """
+        Waits for job start to occur and then returns the results upon
+        a successful start.
+        
+        """
+        POLL_INTERVAL = 5
+
+        start_time = time.time()
+        with yaspin(
+                SPINNER, text=to_colored_text("Awaiting job completion"), color=YASPIN_COLOR
+        ) as spinner:
+            while (time.time() - start_time) < timeout:
+                try:
+                    status = self._fetch_job_status(job_id)
+                except requests.HTTPError as e:
+                    spinner.write(
+                        to_colored_text(
+                            f"Bad status code: {e.response.status_code}", state="fail"
+                        )
+                    )
+                    spinner.stop()
+                    print(to_colored_text(e.response.json(), state="fail"))
+                    return None
+
+                spinner.text = to_colored_text(f"Job status is {status} for {job_id}")
+
+                if status == JobStatus.RUNNING or status == JobStatus.STARTING:
+                    return True
+                if status == JobStatus.FAILED:
+                    return False
+                if status == JobStatus.CANCELLED:
+                    return False
+
+                time.sleep(POLL_INTERVAL)
+
+        return False
+