sutro 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

sutro/sdk.py

@@ -97,6 +97,16 @@ def to_colored_text(
             # Default to blue for normal/processing states
             return f"{Fore.BLUE}{text}{Style.RESET_ALL}"
 
+# Isn't fully support in all terminals unfortunately. We should switch to Rich
+# at some point, but even Rich links aren't clickable on MacOS Terminal
+def make_clickable_link(url, text=None):
+    """
+    Create a clickable link for terminals that support OSC 8 hyperlinks.
+    Falls back to plain text for terminals that don't support it.
+    """
+    if text is None:
+        text = url
+    return f"\033]8;;{url}\033\\{text}\033]8;;\033\\"
 
 class Sutro:
     def __init__(
@@ -266,11 +276,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
@@ -279,6 +284,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
@@ -292,16 +299,22 @@ class Sutro:
                 print(to_colored_text(response.json(), state="fail"))
                 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")
+                    )
+                    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 response_data["results"]
+                    return job_id
                 else:
-                    job_id = response_data["results"]
                     spinner.write(
                         to_colored_text(
-                            f"🛠️  Priority {job_priority} Job created with ID: {job_id}",
+                            f"🛠 Priority {job_priority} Job created with ID: {job_id}.",
                             state="success",
                         )
                     )
@@ -309,12 +322,20 @@ class Sutro:
                         spinner.write(
                             to_colored_text(
                                 f"Use `so.get_job_status('{job_id}')` to check the status of the job."
+                                )
                             )
-                        )
                         return job_id
 
         success = False
         if stay_attached and job_id is not None:
+            spinner.write(to_colored_text("Awaiting job start...", ))
+            spinner.write(to_colored_text(f'Progress can also be monitored at: {make_clickable_link(f'https://app.sutro.sh/jobs/{job_id}')}'))
+            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,
@@ -589,6 +610,7 @@ class Sutro:
                     text=to_colored_text("Awaiting status updates..."),
                     color=YASPIN_COLOR,
                 )
+                spinner.write(to_colored_text(f'Progress can also be monitored at: {make_clickable_link(f'https://app.sutro.sh/jobs/{job_id}')}'))
                 spinner.start()
                 for line in streaming_response.iter_lines():
                     if line:
@@ -719,6 +741,40 @@ class Sutro:
                 return
         return response.json()["jobs"]
 
+    def _list_jobs_helper(self):
+        """
+        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"]
+
+    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.
@@ -775,10 +831,12 @@ class Sutro:
                 return None
 
     def get_job_results(
-        self,
-        job_id: str,
-        include_inputs: bool = False,
-        include_cumulative_logprobs: bool = False,
+            self,
+            job_id: str,
+            include_inputs: bool = False,
+            include_cumulative_logprobs: bool = False,
+            with_original_df: pl.DataFrame | pd.DataFrame = None,
+            output_column: str = "inference_result",
     ):
         """
         Get the results of a job by its ID.
@@ -789,9 +847,11 @@ class Sutro:
             job_id (str): The ID of the job to retrieve the results for.
             include_inputs (bool, optional): Whether to include the inputs in the results. Defaults to False.
             include_cumulative_logprobs (bool, optional): Whether to include the cumulative logprobs in the results. Defaults to False.
+            with_original_df (pd.DataFrame | pl.DataFrame, optional): Original DataFrame to concatenate with results. Defaults to None.
+            output_column (str, optional): Name of the output column. Defaults to "inference_result".
 
         Returns:
-            list: The results of the job.
+            Union[pl.DataFrame, pd.DataFrame]: The results as a DataFrame. By default, returns polars.DataFrame; when with_original_df is an instance of pandas.DataFrame, returns pandas.DataFrame.
         """
         endpoint = f"{self.base_url}/job-results"
         payload = {
@@ -804,18 +864,14 @@ class Sutro:
             "Content-Type": "application/json",
         }
         with yaspin(
-            SPINNER,
-            text=to_colored_text(f"Gathering results from job: {job_id}"),
-            color=YASPIN_COLOR,
+                SPINNER,
+                text=to_colored_text(f"Gathering results from job: {job_id}"),
+                color=YASPIN_COLOR,
         ) as spinner:
             response = requests.post(
                 endpoint, data=json.dumps(payload), headers=headers
             )
-            if response.status_code == 200:
-                spinner.write(
-                    to_colored_text("✔ Job results retrieved", state="success")
-                )
-            else:
+            if response.status_code != 200:
                 spinner.write(
                     to_colored_text(
                         f"Bad status code: {response.status_code}", state="fail"
@@ -823,8 +879,56 @@ class Sutro:
                 )
                 spinner.stop()
                 print(to_colored_text(response.json(), state="fail"))
-                return
-        return response.json()["results"]
+                return None
+
+            spinner.write(
+                to_colored_text("✔ Job results retrieved", state="success")
+            )
+
+        response_data = response.json()
+        results_df = pl.DataFrame(response_data["results"])
+
+
+        if len(results_df.columns ) == 1:
+            # Default column when API is only returning a list, and we construct the df
+            # from that
+            original_results_column = 'column_0'
+        else:
+            original_results_column = 'outputs'
+
+        results_df = results_df.rename({original_results_column: output_column})
+
+        # Ordering inputs col first seems most logical/useful
+        column_config = [
+            ('inputs', include_inputs),
+            (output_column, True),
+            ('cumulative_logprobs', include_cumulative_logprobs),
+        ]
+
+        columns_to_keep = [col for col, include in column_config
+                           if include and col in results_df.columns]
+
+        results_df = results_df.select(columns_to_keep)
+
+        # Handle concatenation with original DataFrame
+        if with_original_df is not None:
+            if isinstance(with_original_df, pd.DataFrame):
+                # Convert to polars for consistent handling
+                original_pl = pl.from_pandas(with_original_df)
+
+                combined_df = original_pl.with_columns(results_df)
+
+                # Convert back to pandas to match input type
+                return combined_df.to_pandas()
+
+            elif isinstance(with_original_df, pl.DataFrame):
+                return with_original_df.with_columns(results_df)
+
+        # Return pd.DataFrame type when appropriate
+        if with_original_df is None and isinstance(with_original_df, pd.DataFrame):
+            return results_df.to_pandas()
+
+        return results_df
 
     def cancel_job(self, job_id: str):
         """
@@ -1162,7 +1266,7 @@ class Sutro:
                 return
         return response.json()["quotas"]
 
-    def await_job_completion(self, job_id: str, timeout: Optional[int] = 7200) -> list | None:
+    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.
@@ -1181,8 +1285,9 @@ class Sutro:
         results = None
         start_time = time.time()
         with yaspin(
-                SPINNER, text=to_colored_text("Awaiting job completion"), color=YASPIN_COLOR
+            SPINNER, text=to_colored_text("Awaiting job completion"), color=YASPIN_COLOR
         ) as spinner:
+            spinner.write(to_colored_text(f'Progress can also be monitored at: {make_clickable_link(f'https://app.sutro.sh/jobs/{job_id}')}'))
             while (time.time() - start_time) < timeout:
                 try:
                     status = self._fetch_job_status(job_id)
@@ -1201,7 +1306,8 @@ class Sutro:
                 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
-                    results = self.get_job_results(job_id)
+                    if obtain_results:
+                        results = self.get_job_results(job_id)
                     break
                 if status == JobStatus.FAILED:
                     spinner.write(to_colored_text("Job has failed", "fail"))
@@ -1213,4 +1319,43 @@ class Sutro:
 
                 time.sleep(POLL_INTERVAL)
 
-        return results
+        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
+            

{sutro-0.1.17.dist-info → sutro-0.1.19.dist-info}/METADATA

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

sutro-0.1.19.dist-info/RECORD

@@ -0,0 +1,8 @@
+sutro/__init__.py,sha256=yUiVwcZ8QamSqDdRHgzoANyTZ-x3cPzlt2Fs5OllR_w,402
+sutro/cli.py,sha256=6Qy9Vwaaho92HeO8YA_z1De4zp1dEFkSX3bEnLvdbkE,13203
+sutro/sdk.py,sha256=1FLepL3M7afptWwudF310KWmPQ9ZDQ51LiT0Xoh1S_o,52705
+sutro-0.1.19.dist-info/METADATA,sha256=6Irt5RX_DaIyRC6ig4XvHr0p0CKnFK0nmWA23oi1mCA,669
+sutro-0.1.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sutro-0.1.19.dist-info/entry_points.txt,sha256=eXvr4dvMV4UmZgR0zmrY8KOmNpo64cJkhNDywiadRFM,40
+sutro-0.1.19.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sutro-0.1.19.dist-info/RECORD,,

sutro-0.1.17.dist-info/RECORD

@@ -1,8 +0,0 @@
-sutro/__init__.py,sha256=yUiVwcZ8QamSqDdRHgzoANyTZ-x3cPzlt2Fs5OllR_w,402
-sutro/cli.py,sha256=6Qy9Vwaaho92HeO8YA_z1De4zp1dEFkSX3bEnLvdbkE,13203
-sutro/sdk.py,sha256=z3cvU9zyAi7EJF_rFIb9mQMjE9zuh576-p8KuP_F1PM,46500
-sutro-0.1.17.dist-info/METADATA,sha256=UDHp7-xr8tdKmDl_GbxxRbVVC6nx62yr6-nka4vBwOY,669
-sutro-0.1.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sutro-0.1.17.dist-info/entry_points.txt,sha256=eXvr4dvMV4UmZgR0zmrY8KOmNpo64cJkhNDywiadRFM,40
-sutro-0.1.17.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sutro-0.1.17.dist-info/RECORD,,