chalkpy 2.90.1__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/ml/model_reference.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import inspect
 import os
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 from chalk.ml.model_version import ModelVersion
 from chalk.ml.utils import (
@@ -15,6 +16,9 @@ from chalk.ml.utils import (
 from chalk.utils.object_inspect import get_source_object_starting
 from chalk.utils.source_parsing import should_skip_source_code_parsing
 
+if TYPE_CHECKING:
+    from chalk.features.resolver import ResourceHint
+
 
 class ModelReference:
     def __init__(
@@ -24,6 +28,8 @@ class ModelReference:
         version: int | None = None,
         alias: str | None = None,
         as_of_date: datetime | None = None,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
     ):
         """Specifies the model version that should be loaded into the deployment.
 
@@ -74,6 +80,8 @@ class ModelReference:
         self.as_of_date = as_of_date
         self.alias = alias
         self.identifier = identifier
+        self.resource_hint = resource_hint
+        self.resource_group = resource_group
 
         self.filename = filename
         self.source_line_start = source_line_start
@@ -110,6 +118,8 @@ class ModelReference:
                 model_class=ModelClass(model_artifact_metadata.spec.model_class)
                 if model_artifact_metadata.spec.model_class
                 else None,
+                resource_hint=resource_hint,
+                resource_group=resource_group,
             )
 
             from chalk.features.hooks import before_all
@@ -117,14 +127,22 @@ class ModelReference:
             def hook():
                 mv.load_model()
 
-            before_all(hook)
+            before_all(hook, resource_hint=resource_hint, resource_group=resource_group)
 
             self.model_version = mv
         else:
-            self.model_version = ModelVersion(name=name, identifier=identifier)
+            self.model_version = ModelVersion(
+                name=name, identifier=identifier, resource_hint=resource_hint, resource_group=resource_group
+            )
 
     @classmethod
-    def as_of(cls, name: str, when: datetime) -> ModelVersion:
+    def as_of(
+        cls,
+        name: str,
+        when: datetime,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
+    ) -> ModelVersion:
         """Creates a ModelReference for a specific point in time.
 
         Parameters
@@ -133,6 +151,11 @@ class ModelReference:
             The name of the model.
         when
             The datetime to use for creating the model version identifier.
+        resource_hint
+            Whether this model loading is bound by CPU, I/O, or GPU.
+        resource_group
+            The resource group for the model: this is used to isolate execution
+            onto a separate pod (or set of nodes), such as on a GPU-enabled node.
 
         Returns
         -------
@@ -144,13 +167,20 @@ class ModelReference:
         >>> import datetime
         >>> timestamp = datetime.datetime(2023, 10, 15, 14, 30, 0)
         >>> model = ModelReference.as_of("fraud_model", timestamp)
+        >>> model = ModelReference.as_of("fraud_model", timestamp, resource_hint="gpu", resource_group="gpu-group")
         """
 
-        mr = ModelReference(name=name, as_of_date=when)
+        mr = ModelReference(name=name, as_of_date=when, resource_hint=resource_hint, resource_group=resource_group)
         return mr.model_version
 
     @classmethod
-    def from_version(cls, name: str, version: int) -> ModelVersion:
+    def from_version(
+        cls,
+        name: str,
+        version: int,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
+    ) -> ModelVersion:
         """Creates a ModelReference using a numeric version identifier.
 
         Parameters
@@ -159,6 +189,11 @@ class ModelReference:
             The name of the model.
         version
             The version number. Must be a non-negative integer.
+        resource_hint
+            Whether this model loading is bound by CPU, I/O, or GPU.
+        resource_group
+            The resource group for the model: this is used to isolate execution
+            onto a separate pod (or set of nodes), such as on a GPU-enabled node.
 
         Returns
         -------
@@ -173,15 +208,22 @@ class ModelReference:
         Examples
         --------
         >>> model = ModelReference.from_version("fraud_model", 1)
+        >>> model = ModelReference.from_version("fraud_model", 1, resource_hint="gpu", resource_group="gpu-group")
         """
         if version < 0:
             raise ValueError("Version number must be a non-negative integer.")
 
-        mr = ModelReference(name=name, version=version)
+        mr = ModelReference(name=name, version=version, resource_hint=resource_hint, resource_group=resource_group)
         return mr.model_version
 
     @classmethod
-    def from_alias(cls, name: str, alias: str) -> ModelVersion:
+    def from_alias(
+        cls,
+        name: str,
+        alias: str,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
+    ) -> ModelVersion:
         """Creates a ModelReference using an alias identifier.
 
         Parameters
@@ -190,6 +232,11 @@ class ModelReference:
             The name of the model.
         alias
             The alias string. Must be non-empty.
+        resource_hint
+            Whether this model loading is bound by CPU, I/O, or GPU.
+        resource_group
+            The resource group for the model: this is used to isolate execution
+            onto a separate pod (or set of nodes), such as on a GPU-enabled node.
 
         Returns
         -------
@@ -204,11 +251,12 @@ class ModelReference:
         Examples
         --------
         >>> model = ModelReference.from_alias("fraud_model", "latest")
+        >>> model = ModelReference.from_alias("fraud_model", "latest", resource_hint="gpu", resource_group="gpu-group")
         """
         if not alias:
             raise ValueError("Alias must be a non-empty string.")
 
-        mr = ModelReference(name=name, alias=alias)
+        mr = ModelReference(name=name, alias=alias, resource_hint=resource_hint, resource_group=resource_group)
         return mr.model_version
 
 

chalk/ml/model_version.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from chalk.ml.model_hooks import MODEL_REGISTRY
 from chalk.ml.utils import ModelClass, ModelEncoding, ModelType
 
+if TYPE_CHECKING:
+    from chalk.features.resolver import ResourceHint
+
 
 class ModelVersion:
     def __init__(
@@ -20,6 +23,8 @@ class ModelVersion:
         model_encoding: ModelEncoding | None = None,
         model_class: ModelClass | None = None,
         filename: str | None = None,
+        resource_hint: "ResourceHint | None" = None,
+        resource_group: str | None = None,
     ):
         """Specifies the model version that should be loaded into the deployment.
 
@@ -41,6 +46,8 @@ class ModelVersion:
         self.model_encoding = model_encoding
         self.model_class = model_class
         self.filename = filename
+        self.resource_hint: "ResourceHint | None" = resource_hint
+        self.resource_group = resource_group
 
         self._model = None
         self._predictor = None
@@ -58,26 +65,15 @@ class ModelVersion:
                 model_type=self.model_type, encoding=self.model_encoding, model_class=self.model_class
             )
             if model is not None and self.filename is not None:
-                self._model = model.load_model(self.filename)
+                self._model = model.load_model(self.filename, resource_hint=self.resource_hint)
             else:
                 raise ValueError(
                     f"No load function defined for type {self.model_type}, encoding {self.model_encoding}, and class {self.model_class}"
                 )
 
     def predict(self, X: Any):
-        """Loads the model from the specified filename using the appropriate hook."""
-
-        if self._predictor is None:
-            if self.model_type is None or self.model_encoding is None:
-                raise ValueError("Model type and encoding must be specified to use predict.")
-            self._predictor = MODEL_REGISTRY.get(
-                model_type=self.model_type, encoding=self.model_encoding, model_class=self.model_class
-            )
-            if self._predictor is None:
-                raise ValueError(
-                    f"No predict function defined for type {self.model_type} and extension {self.model_encoding}"
-                )
-        return self._predictor.predict(self.model, X)
+        """Runs prediction using the loaded model."""
+        return self.predictor.predict(self.model, X)
 
     @property
     def model(self) -> Any:
@@ -86,3 +82,16 @@ class ModelVersion:
             self.load_model()
 
         return self._model
+
+    @property
+    def predictor(self) -> Any:
+        """Returns the predictor instance, initializing it if needed."""
+        if self._predictor is None:
+            if self.model_type is None or self.model_encoding is None:
+                raise ValueError("Model type and encoding must be specified to use predictor.")
+            self._predictor = MODEL_REGISTRY.get(
+                model_type=self.model_type, encoding=self.model_encoding, model_class=self.model_class
+            )
+            if self._predictor is None:
+                raise ValueError(f"No predictor defined for type {self.model_type} and encoding {self.model_encoding}")
+        return self._predictor

chalk/ml/utils.py

@@ -400,6 +400,26 @@ class ModelAttributeExtractor:
 
     @staticmethod
     def infer_model_type(model: Any) -> Tuple[Optional[ModelType], Optional[ModelClass]]:
+        # ONNX - check early since ONNX models are commonly wrapped
+        try:
+            import onnx  # pyright: ignore[reportMissingImports]
+
+            if isinstance(model, onnx.ModelProto):
+                return ModelType.ONNX, None
+            # Check if model has a wrapped ONNX ModelProto (e.g., model._model)
+            if hasattr(model, "_model") and isinstance(model._model, onnx.ModelProto):
+                return ModelType.ONNX, None
+        except ImportError:
+            pass
+
+        try:
+            import onnxruntime  # pyright: ignore[reportMissingImports]
+
+            if isinstance(model, onnxruntime.InferenceSession):
+                return ModelType.ONNX, None
+        except ImportError:
+            pass
+
         # PYTORCH
         try:
             import torch.nn as nn  # pyright: ignore[reportMissingImports]
@@ -458,23 +478,6 @@ class ModelAttributeExtractor:
         except ImportError:
             pass
 
-        # ONNX - do we need this one?
-        try:
-            import onnx  # pyright: ignore[reportMissingImports]
-
-            if isinstance(model, onnx.ModelProto):
-                return ModelType.ONNX, None
-        except ImportError:
-            pass
-
-        try:
-            import onnxruntime  # pyright: ignore[reportMissingImports]
-
-            if isinstance(model, onnxruntime.InferenceSession):
-                return ModelType.ONNX, None
-        except ImportError:
-            pass
-
         # SKLEARN
         try:
             import sklearn.base  # pyright: ignore[reportMissingImports]

chalk/operators/_utils.py

@@ -9,6 +9,7 @@ import pyarrow
 from chalk import DataFrame, Features, StaticOperator
 from chalk._gen.chalk.expression.v1 import expression_pb2 as expr_pb
 from chalk.client import ChalkError, ChalkException, ErrorCode, ErrorCodeCategory
+from chalk.df.LazyFramePlaceholder import LazyFramePlaceholder
 from chalk.features.feature_field import Feature
 
 
@@ -79,7 +80,7 @@ def static_resolver_to_operator(
     fn: Callable,
     inputs: Sequence[Union[Feature, type[DataFrame]]],
     output: Optional[type[Features]],
-) -> StaticOperator | DfPlaceholder | ChalkDataFrame:
+) -> StaticOperator | DfPlaceholder | ChalkDataFrame | LazyFramePlaceholder:
     if output is None:
         raise _GetStaticOperatorError(
             resolver_fqn=fqn,
@@ -96,8 +97,14 @@ def static_resolver_to_operator(
             message="Static resolver must take no arguments and have exactly one DataFrame output",
             underlying_exception=None,
         )
+
     try:
-        placeholder_inputs = [DfPlaceholder(schema_dict=schema_for_input(input_type)) for input_type in inputs]
+        placeholder_inputs = [
+            LazyFramePlaceholder.named_table(
+                name=f"resolver_df_input_{input_index}", schema=pyarrow.schema(schema_for_input(input_type))
+            )
+            for input_index, input_type in enumerate(inputs)
+        ]
         static_operator = fn(*placeholder_inputs)
     except Exception as e:
         # Weird hacky way to return a placeholder even if the resolver fails.
@@ -108,7 +115,7 @@ def static_resolver_to_operator(
         )
     else:
         if (
-            not isinstance(static_operator, (StaticOperator, DfPlaceholder))
+            not isinstance(static_operator, (StaticOperator, DfPlaceholder, LazyFramePlaceholder))
             and not static_operator.__class__.__name__ == "ChalkDataFrame"
             and not static_operator.__class__.__name__ == "LazyFrame"
             and not (

chalk/parsed/_proto/export.py

@@ -27,6 +27,7 @@ from chalk.parsed._proto.utils import (
     convert_failed_import_to_gql,
     convert_failed_import_to_proto,
     datetime_to_proto_timestamp,
+    timedelta_to_proto_duration,
 )
 from chalk.parsed._proto.validation import validate_artifacts
 from chalk.parsed.to_proto import ToProtoConverter
@@ -145,6 +146,24 @@ def export_from_registry() -> export_pb.Export:
     """
     failed_protos: List[export_pb.FailedImport] = []
 
+    # Validate registries BEFORE conversion to catch errors early
+    # This ensures parity with GQL validation path
+    from chalk.parsed.validation_from_registries import validate_all_from_registries
+
+    try:
+        validate_all_from_registries(
+            features_registry=FeatureSetBase.registry,
+            resolver_registry=RESOLVER_REGISTRY,
+        )
+    except Exception as e:
+        # If validation fails, add to failed_protos but continue
+        # to allow other validation to complete
+        from chalk._lsp.error_builder import LSPErrorBuilder
+
+        if not LSPErrorBuilder.promote_exception(e):
+            # Not an LSP error, so log it as a failed import
+            failed_protos.append(build_failed_import(e, "validation"))
+
     graph_res = ToProtoConverter.convert_graph(
         features_registry=FeatureSetBase.registry,
         resolver_registry=RESOLVER_REGISTRY.get_all_resolvers(),
@@ -193,6 +212,9 @@ def export_from_registry() -> export_pb.Export:
                 file_name=cron.filename,
                 resource_group=cron.resource_group,
                 planner_options=cron.planner_options,
+                completion_deadline=timedelta_to_proto_duration(cron.completion_deadline)
+                if cron.completion_deadline is not None
+                else cron.completion_deadline,
             )
         )
 

chalk/parsed/duplicate_input_gql.py

@@ -279,6 +279,9 @@ class UpsertCronQueryGQL:
     upperBoundStr: Optional[str] = None
     resourceGroup: Optional[str] = None
     plannerOptions: Optional[Dict[str, str]] = None
+    completionDeadline: Optional[str] = None
+    numShards: Optional[int] = None
+    numWorkers: Optional[int] = None
 
 
 @dataclasses_json.dataclass_json

chalk/parsed/json_conversions.py

@@ -428,6 +428,9 @@ def convert_type_to_gql(
             incrementalSources=None if t.incremental_resolvers is None else list(t.incremental_resolvers),
             resourceGroup=t.resource_group,
             plannerOptions=t.planner_options,
+            completionDeadline=None if t.completion_deadline is None else timedelta_to_duration(t.completion_deadline),
+            numShards=t.num_shards,
+            numWorkers=t.num_workers,
         )
 
     if isinstance(t, NamedQuery):
@@ -539,20 +542,23 @@ def convert_type_to_gql(
             )
 
         elif t.join is not None:
-            # If a has_one/has_many has an incorrect type annotation
-            builder = t.features_cls.__chalk_error_builder__
-            builder.add_diagnostic(
-                range=builder.annotation_range(t.attribute_name),
-                message=(
-                    f"The attribute '{t.features_cls.__name__}.{t.attribute_name}' "
-                    f"has a join filter ({t.join}) but its type annotation is not a feature class or "
-                    f"DataFrame ({t.typ})."
-                ),
-                label="Incorrect join type annotation",
-                raise_error=TypeError,
-                code="34",
-                code_href="https://docs.chalk.ai/docs/has-many",
-            )
+            # Check if user tried to use DataFrame (even if validation failed)
+            # Use is_dataframe_annotation() to detect DataFrame types without triggering validation errors
+            if not t.typ.is_dataframe_annotation():
+                # If a has_one/has_many has an incorrect type annotation
+                builder = t.features_cls.__chalk_error_builder__
+                builder.add_diagnostic(
+                    range=builder.annotation_range(t.attribute_name),
+                    message=(
+                        f"The attribute '{t.features_cls.__name__}.{t.attribute_name}' "
+                        f"has a join filter ({t.join}) but its type annotation is not a feature class or "
+                        f"DataFrame ({t.typ})."
+                    ),
+                    label="Incorrect join type annotation",
+                    raise_error=TypeError,
+                    code="34",
+                    code_href="https://docs.chalk.ai/docs/has-many",
+                )
 
         elif t.is_feature_time:
             feature_time_kind_gql = UpsertFeatureTimeKindGQL()

chalk/parsed/to_proto.py

@@ -21,6 +21,7 @@ from chalk._gen.chalk.graph.v1 import graph_pb2 as pb
 from chalk._gen.chalk.graph.v2 import sources_pb2 as sources_pb
 from chalk._gen.chalk.lsp.v1.lsp_pb2 import Location, Position, Range
 from chalk._validation.feature_validation import FeatureValidation
+from chalk.df.LazyFramePlaceholder import LazyFramePlaceholder
 from chalk.features import (
     CacheStrategy,
     Feature,
@@ -899,7 +900,7 @@ class ToProtoConverter:
                     else None,
                     backfill_schedule=mat.backfill_schedule,
                     approx_top_k_arg_k=aggregation_kwargs.get("k")
-                    if mat.aggregation in ("approx_top_k", "min_by_n", "max_by_n")
+                    if mat.aggregation in ("approx_top_k", "approx_percentile", "min_by_n", "max_by_n")
                     else None,
                 ),
                 tags=f.tags,
@@ -995,7 +996,7 @@ class ToProtoConverter:
                             else None,
                             continuous_resolver=wmp.continuous_resolver,
                             approx_top_k_arg_k=aggregation_kwargs.get("k")
-                            if wmp.aggregation in ("approx_top_k", "min_by_n", "max_by_n")
+                            if wmp.aggregation in ("approx_top_k", "approx_percentile", "min_by_n", "max_by_n")
                             else None,
                         )
                         if wmp is not None
@@ -1025,6 +1026,9 @@ class ToProtoConverter:
                 expression=ToProtoConverter.convert_underscore(f.underscore_expression)
                 if f.underscore_expression is not None
                 else None,
+                offline_expression=ToProtoConverter.convert_underscore(f.offline_underscore_expression)
+                if f.offline_underscore_expression is not None
+                else None,
                 expression_definition_location=ToProtoConverter.convert_expression_definition_location(
                     f.underscore_expression
                 )
@@ -1147,9 +1151,13 @@ class ToProtoConverter:
             raise ValueError(f"Unsupported resource hint: {r.resource_hint}")
 
         static_operation = None
+        static_operation_dataframe = None
         if r.static:
             static_operator = static_resolver_to_operator(fqn=r.fqn, fn=r.fn, inputs=r.inputs, output=r.output)
-            static_operation = static_operator._to_proto()  # pyright: ignore[reportPrivateUsage]
+            if isinstance(static_operator, LazyFramePlaceholder):
+                static_operation_dataframe = static_operator._to_proto()  # pyright: ignore[reportPrivateUsage]
+            else:
+                static_operation = static_operator._to_proto()  # pyright: ignore[reportPrivateUsage]
 
         function_reference_proto = ToProtoConverter.create_function_reference(
             r.fn,
@@ -1158,7 +1166,9 @@ class ToProtoConverter:
             filename=r.filename,
             source_line=r.source_line,
         )
-
+        postprocessing_underscore_expr: expr_pb.LogicalExprNode | None = None
+        if isinstance(r.postprocessing, Underscore):
+            postprocessing_underscore_expr = r.postprocessing._to_proto()  # pyright: ignore[reportPrivateUsage]
         return pb.Resolver(
             fqn=r.fqn,
             kind=(
@@ -1186,9 +1196,11 @@ class ToProtoConverter:
             unique_on=tuple(x.root_fqn for x in r.unique_on) if r.unique_on is not None else (),
             partitioned_by=(x.root_fqn for x in r.partitioned_by) if r.partitioned_by is not None else (),
             static_operation=static_operation,
+            static_operation_dataframe=static_operation_dataframe,
             sql_settings=ToProtoConverter.convert_sql_settings(r.sql_settings) if r.sql_settings else None,
             output_row_order=r.output_row_order,
             venv=r.venv,
+            underscore_expr=postprocessing_underscore_expr,
         )
 
     @staticmethod

chalk/parsed/user_types_to_json.py

@@ -110,6 +110,22 @@ def project_settings_to_gql(config: ProjectSettings) -> ProjectSettingsGQL:
 
 
 def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGraphGQL:
+    # Validate registries BEFORE conversion to catch errors early
+    # This ensures parity with Proto validation path
+    from chalk.parsed.validation_from_registries import validate_all_from_registries
+
+    try:
+        validate_all_from_registries(
+            features_registry=FeatureSetBase.registry,
+            resolver_registry=RESOLVER_REGISTRY,
+        )
+    except Exception as e:
+        # If validation fails, add to failed but continue
+        # to allow other validation to complete
+        if not LSPErrorBuilder.promote_exception(e):
+            # Not an LSP error, so log it as a failed import
+            failed.append(build_failed_import(e, "validation"))
+
     features = []
     feature_classes: list[FeatureClassGQL] = []
     for x in FeatureSetBase.registry.values():
@@ -173,6 +189,14 @@ def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGr
     named_queries: list[UpsertNamedQueryGQL] = []
     for named_query in NAMED_QUERY_REGISTRY.values():
         if named_query.filename is None or _is_relative_to(Path(named_query.filename), scope_to):
+            # only try to convert if there are no errors
+            if not named_query.errors:
+                try:
+                    named_queries.append(convert_type_to_gql(named_query, path_prefix=path_prefix_to_remove))
+                except Exception as e:
+                    failed.append(build_failed_import(e, f"named query '{named_query.name}'"))
+
+            # named_query.errors can go from empty to non-empty after conversion
             if named_query.errors:
                 code_object_diagnostics[named_query.filename or ""].extend(
                     [
@@ -185,15 +209,17 @@ def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGr
                         for error in named_query.errors
                     ]
                 )
-                continue
-            try:
-                named_queries.append(convert_type_to_gql(named_query, path_prefix=path_prefix_to_remove))
-            except Exception as e:
-                failed.append(build_failed_import(e, f"named query '{named_query.name}'"))
 
     model_references: list[UpsertModelReferenceGQL] = []
     for mr in MODEL_REFERENCE_REGISTRY.values():
         if mr.filename is None or _is_relative_to(Path(mr.filename), scope_to):
+            # similar logic to named_queries above
+            if not mr.errors:
+                try:
+                    model_references.append(convert_type_to_gql(mr, path_prefix=path_prefix_to_remove))
+                except Exception as e:
+                    failed.append(build_failed_import(e, f"model reference '{mr.name}'"))
+
             if mr.errors:
                 code_object_diagnostics[mr.filename or ""].extend(
                     [
@@ -206,11 +232,6 @@ def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGr
                         for error in mr.errors
                     ]
                 )
-                continue
-            try:
-                model_references.append(convert_type_to_gql(mr, path_prefix=path_prefix_to_remove))
-            except Exception as e:
-                failed.append(build_failed_import(e, f"model reference '{mr.name}'"))
 
     # online store configs
     for osc in ONLINE_STORE_CONFIG_REGISTRY.values():