chalkpy 2.96.7__py3-none-any.whl → 2.96.8__py3-none-any.whl

chalk/_version.py

@@ -1 +1 @@
-__version__ = "2.96.7"
+__version__ = "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/

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

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(),

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.dist-info → chalkpy-2.96.8.dist-info}/METADATA

@@ -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.7.dist-info → chalkpy-2.96.8.dist-info}/RECORD

@@ -1,5 +1,5 @@
 chalk/__init__.py,sha256=vKsx9-cl5kImlVWGHVRYO6bweBm79NAzGs3l36u71wM,2657
-chalk/_version.py,sha256=3_AK7WV-Cb97v3VLeqbWqWc5ckXI1N8Wm6F3FsWkhQI,23
+chalk/_version.py,sha256=XQBn2m_TF6iSSjpgoPtUzJ93WpjTyMNEkfPlAWxDliQ,23
 chalk/cli.py,sha256=ckqqfOI-A2mT23-rnZzDMmblYj-2x1VBX8ebHlIEn9A,5873
 chalk/importer.py,sha256=m4lMn1lSYj_euDq8CS7LYTBnek9JOcjGJf9-82dJHbA,64441
 chalk/prompts.py,sha256=2H9UomLAamdfRTNUdKs9i3VTpiossuyRhntqsAXUhhg,16117
@@ -612,14 +612,14 @@ chalk/_validation/validation.py,sha256=9cCMfZa9-1wxkXLme_ylmD5vIA1qExJD6aqbYvbmK
 chalk/byte_transmit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/byte_transmit/model.py,sha256=LFX8pj9X_CWXeap7fDnMl9YmXsYTgq7jBAbEWkxoYoE,13048
 chalk/client/__init__.py,sha256=wu3WQVzE5gRj6noQQDOdYJUgDaz_9QtbjXH4KuoIlXQ,1782
-chalk/client/client.py,sha256=uu0S3bG0_6HG9S2seIH1oaww6Jz684fG9fIm3x2MwvI,103648
-chalk/client/client_async.py,sha256=I-1igsjYgK6qFh-KN524QBSvtxc8Yntltgt5fRHpQvA,51114
+chalk/client/client.py,sha256=59UYbIq7KHwNm1ZFw5TmCV2JvbJxn45Ar_AjfSkuzQ8,103907
+chalk/client/client_async.py,sha256=wD38RIkwVLaKYKj7K1bWDXnVZkLFYHsbY_5VhtflMvo,51559
 chalk/client/client_async_impl.py,sha256=ZphhgTB49JBWHCGXe-dI0wWWKc9zPcOczy02q_gFy50,6925
 chalk/client/client_grpc.py,sha256=skpSHYCyE8hg0T3Bcl5R8MMi92QjrFtxQRu_NxCSVlw,106952
-chalk/client/client_impl.py,sha256=rHkVYn0jawWwlFYPLseGuSm8sHwuLlMGY9uf1KSa8wY,211960
+chalk/client/client_impl.py,sha256=3NDl1Vo7YJ0grraxJeN6vjFj1hBpU7YT_ikolK31D6k,212395
 chalk/client/dataset.py,sha256=LneWwaAOHCjtj7gaJjsSeVNruj-QJ51hjRi62zrFNVE,77561
 chalk/client/exc.py,sha256=kZJ80YbSeSRDmTLTh240j_eRdJFZBa7IaDsNSRoDroU,4145
-chalk/client/models.py,sha256=X3sq-kv4F601ABUcwE3Y7X-yF4sTHwbCk1HfX7xnhO0,64019
+chalk/client/models.py,sha256=YwXywN6OSPOda1DKJuPVe6hgsHzf8ewvX2dYOBOZk7I,65534
 chalk/client/response.py,sha256=m8sQCOj7YVv3mZSZMIC1rIMzFMQ9rfMdBRLg5NRmOOE,53257
 chalk/client/_internal_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/client/_internal_models/check.py,sha256=3Xfo4Ws4rvwjeVg0-5-kejfRfRBJeqHmnRhW-WEz784,917
@@ -827,8 +827,8 @@ chalk/utils/tracing.py,sha256=NiiM-9dbuJhSCv6R1npR1uYNKWlkqTR6Ygm0Voi2NrY,13078
 chalk/utils/weak_set_by_identity.py,sha256=VmikA_laYwFeOphCwXJIuyOIkrdlQe0bSzaXq7onoQw,953
 chalk/utils/pydanticutil/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/utils/pydanticutil/pydantic_compat.py,sha256=O575lLYJ5GvZC4HMzR9yATxf9XwjC6NrDUXbNwZidlE,3031
-chalkpy-2.96.7.dist-info/METADATA,sha256=QWqtSzsETPKu6gG6oHi0s_nhv7SfLSuca2-SjNFKAKU,27754
-chalkpy-2.96.7.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-chalkpy-2.96.7.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
-chalkpy-2.96.7.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
-chalkpy-2.96.7.dist-info/RECORD,,
+chalkpy-2.96.8.dist-info/METADATA,sha256=wLMi2G7i8tczLoHoAAYiMWG9g2oujf4DoUmq3hqkP98,27754
+chalkpy-2.96.8.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+chalkpy-2.96.8.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
+chalkpy-2.96.8.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
+chalkpy-2.96.8.dist-info/RECORD,,