chalkpy 2.96.7__tar.gz → 2.96.8__tar.gz

{chalkpy-2.96.7 → chalkpy-2.96.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkpy
-Version: 2.96.7
+Version: 2.96.8
 Summary: Python SDK for Chalk
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

chalkpy-2.96.8/chalk/_version.py

@@ -0,0 +1 @@
+__version__ = "2.96.8"

{chalkpy-2.96.7 → chalkpy-2.96.8}/chalk/client/client.py

@@ -33,6 +33,7 @@ from chalk.client.models import (
     GetRegisteredModelResponse,
     GetRegisteredModelVersionResponse,
     ManualTriggerScheduledQueryResponse,
+    OfflineQueryDeadlineOptions,
     OfflineQueryInputUri,
     OnlineQuery,
     OnlineQueryContext,
@@ -927,7 +928,7 @@ class ChalkClient:
         store_offline: bool = False,
         num_shards: int | None = None,
         num_workers: int | None = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
@@ -1043,8 +1044,8 @@ class ChalkClient:
             If specified, the query will be run asynchronously across a maximum `num_workers` pod workers at any time.
             This parameter is useful if you have a large number of shards and would like to limit the number of pods running at once.
         completion_deadline
-            If specified, shards must complete within 'completion_deadline' duration, or they will be terminated.
-            Terminated shards can be tried.
+            If specified as a timedelta, applies a completion deadline to each shard; each shard's query will fail (allowing retries) if it does not complete within the duration.
+            If specified as an OfflineQueryDeadlineOptions, allows more fine-grained control of shard- or query-level deadlines, with options to retry on failure or not.
         max_retries
             If specified, failed offline query shards will be retried. The retry budget is shared across all shards.
             By default, max_retries=num_shards/

{chalkpy-2.96.7 → chalkpy-2.96.8}/chalk/client/client_async.py

@@ -12,6 +12,7 @@ from chalk.client.models import (
     BulkOnlineQueryResult,
     FeatureReference,
     FeatureStatisticsResponse,
+    OfflineQueryDeadlineOptions,
     OfflineQueryInputUri,
     OnlineQuery,
     OnlineQueryContext,
@@ -647,7 +648,7 @@ class AsyncChalkClient:
         store_offline: bool = False,
         num_shards: int | None = None,
         num_workers: int | None = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
@@ -750,6 +751,9 @@ class AsyncChalkClient:
         num_workers
             If specified, the query will be run asynchronously across a maximum `num_workers` pod workers at any time.
             This parameter is useful if you have a large number of shards and would like to limit the number of pods running at once.
+        completion_deadline
+            If specified as a timedelta, applies a completion deadline to each shard; each shard's query will fail (allowing retries) if it does not complete within the duration.
+            If specified as an OfflineQueryDeadlineOptions, allows more fine-grained control of shard- or query-level deadlines, with options to retry on failure or not.
         query_name
             The name of the query to execute. If provided, will create a new named query or fill in missing parameters from a preexisting execution.
         query_name_version

{chalkpy-2.96.7 → chalkpy-2.96.8}/chalk/client/client_impl.py

@@ -103,6 +103,7 @@ from chalk.client.models import (
     MultiUploadFeaturesRequest,
     MultiUploadFeaturesResponse,
     OfflineQueryContext,
+    OfflineQueryDeadlineOptions,
     OfflineQueryInput,
     OfflineQueryInputSql,
     OfflineQueryInputUri,
@@ -2229,7 +2230,7 @@ https://docs.chalk.ai/cli/apply
         store_offline: bool = False,
         num_shards: int | None = None,
         num_workers: int | None = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
@@ -3607,7 +3608,7 @@ https://docs.chalk.ai/cli/apply
         num_shards: int | None = None,
         num_workers: int | None = None,
         feature_for_lower_upper_bound: Optional[str] = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         optional_output_expressions: Optional[List[str]] = None,
         required_output_expressions: Optional[List[str]] = None,
@@ -3651,6 +3652,13 @@ https://docs.chalk.ai/cli/apply
         upper_bound_str = process_bound(upper_bound)
         if branch is ...:
             branch = self._branch
+
+        retyped_completion_deadline: Union[None, str, OfflineQueryDeadlineOptions] = None
+        if isinstance(completion_deadline, OfflineQueryDeadlineOptions):
+            retyped_completion_deadline = completion_deadline.with_chalk_durations()
+        elif isinstance(completion_deadline, timedelta):
+            retyped_completion_deadline = timedelta_to_duration(completion_deadline)
+
         req = CreateOfflineQueryJobRequest(
             output=optional_output,
             output_expressions=optional_output_expressions or [],
@@ -3683,7 +3691,7 @@ https://docs.chalk.ai/cli/apply
             num_shards=num_shards,
             num_workers=num_workers,
             feature_for_lower_upper_bound=feature_for_lower_upper_bound,
-            completion_deadline=timedelta_to_duration(completion_deadline) if completion_deadline is not None else None,
+            completion_deadline=retyped_completion_deadline,
             max_retries=max_retries,
             use_job_queue=use_job_queue,
             overlay_graph=_get_overlay_graph_b64(),

{chalkpy-2.96.7 → chalkpy-2.96.8}/chalk/client/models.py

@@ -20,6 +20,7 @@ from chalk.features.tag import EnvironmentId
 from chalk.prompts import Prompt
 from chalk.queries.query_context import ContextJsonDict
 from chalk.utils.df_utils import read_parquet
+from chalk.utils.duration import timedelta_to_duration
 from chalk.utils.missing_dependency import missing_dependency_exception
 
 if TYPE_CHECKING:
@@ -820,6 +821,48 @@ class ResourceRequests(BaseModel):
     """Resource group to use for this job. If not specified, the default resource group will be used."""
 
 
+class OfflineQueryDeadlineOptions(BaseModel):
+    """
+    Specification for setting deadlines for shards of the query or the entire query itself.
+    """
+
+    shard_deadline: Union[timedelta, str, None] = None
+    """
+    Maximum amount of time a query shard can work before being failed.
+    """
+
+    retry_on_shard_deadline: Optional[bool] = None
+    """
+    Whether to retry when the per-shard deadline is triggered. Will default to true.
+    """
+
+    query_deadline: Union[timedelta, str, None] = None
+    """
+    Maximum amount of time that the entire query can work before being failed.
+    """
+
+    retry_on_query_deadline: Optional[bool] = None
+    """
+    Whether to retry when the entire query's deadline is triggered. Will default to false.
+    """
+
+    def with_chalk_durations(self) -> OfflineQueryDeadlineOptions:
+        return OfflineQueryDeadlineOptions(
+            shard_deadline=(
+                timedelta_to_duration(self.shard_deadline)
+                if isinstance(self.shard_deadline, timedelta)
+                else self.shard_deadline
+            ),
+            retry_on_shard_deadline=self.retry_on_shard_deadline,
+            query_deadline=(
+                timedelta_to_duration(self.query_deadline)
+                if isinstance(self.query_deadline, timedelta)
+                else self.query_deadline
+            ),
+            retry_on_query_deadline=self.retry_on_query_deadline,
+        )
+
+
 class CreateOfflineQueryJobRequest(BaseModel):
     output: List[str]
     """A list of output feature root fqns to query"""
@@ -902,7 +945,7 @@ class CreateOfflineQueryJobRequest(BaseModel):
     num_workers: Optional[int] = None
     feature_for_lower_upper_bound: Optional[str] = None
 
-    completion_deadline: Optional[str] = None
+    completion_deadline: Union[None, str, OfflineQueryDeadlineOptions] = None
     max_retries: Optional[int] = None
 
     use_job_queue: bool = False

chalkpy-2.96.7/chalk/_version.py

@@ -1 +0,0 @@
-__version__ = "2.96.7"