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

chalk/client/client_grpc.py

@@ -74,6 +74,10 @@ from chalk._gen.chalk.server.v1.model_registry_pb2 import (
 )
 from chalk._gen.chalk.server.v1.model_registry_pb2_grpc import ModelRegistryServiceStub
 from chalk._gen.chalk.server.v1.offline_queries_pb2_grpc import OfflineQueryMetadataServiceStub
+from chalk._gen.chalk.server.v1.scheduled_query_pb2_grpc import ScheduledQueryServiceStub
+from chalk._gen.chalk.server.v1.scheduled_query_run_pb2 import GetScheduledQueryRunsRequest
+from chalk._gen.chalk.server.v1.scheduler_pb2 import ManualTriggerScheduledQueryRequest
+from chalk._gen.chalk.server.v1.scheduler_pb2_grpc import SchedulerServiceStub
 from chalk._gen.chalk.server.v1.script_tasks_pb2 import CreateScriptTaskRequest, CreateScriptTaskResponse
 from chalk._gen.chalk.server.v1.script_tasks_pb2_grpc import ScriptTaskServiceStub
 from chalk._gen.chalk.server.v1.team_pb2 import (
@@ -83,6 +87,7 @@ from chalk._gen.chalk.server.v1.team_pb2 import (
     ListServiceTokensResponse,
 )
 from chalk._gen.chalk.server.v1.team_pb2_grpc import TeamServiceStub
+from chalk._gen.chalk.streaming.v1.simple_streaming_service_pb2_grpc import SimpleStreamingServiceStub
 from chalk.client import ChalkAuthException, FeatureReference
 from chalk.client.client_impl import _validate_context_dict  # pyright: ignore[reportPrivateUsage]
 from chalk.client.models import (
@@ -92,6 +97,9 @@ from chalk.client.models import (
     CreateBranchResponse,
     GetRegisteredModelResponse,
     GetRegisteredModelVersionResponse,
+)
+from chalk.client.models import ManualTriggerScheduledQueryResponse as ManualTriggerScheduledQueryResponseDataclass
+from chalk.client.models import (
     ModelUploadUrlResponse,
     OnlineQuery,
     OnlineQueryResponse,
@@ -99,6 +107,9 @@ from chalk.client.models import (
     RegisterModelResponse,
     RegisterModelVersionResponse,
     ResourceRequests,
+    ScheduledQueryRun,
+    StreamResolverTestResponse,
+    StreamResolverTestStatus,
     UploadFeaturesResponse,
 )
 from chalk.client.serialization.model_serialization import ModelSerializer
@@ -109,6 +120,7 @@ from chalk.features._encoding.inputs import GRPC_ENCODE_OPTIONS, InputEncodeOpti
 from chalk.features._encoding.json import FeatureEncodingOptions
 from chalk.features._encoding.outputs import encode_outputs
 from chalk.features.feature_set import is_feature_set_class
+from chalk.features.resolver import Resolver
 from chalk.features.tag import DeploymentId
 from chalk.importer import CHALK_IMPORT_FLAG
 from chalk.ml import LocalSourceConfig, ModelEncoding, ModelRunCriterion, ModelType, SourceConfig
@@ -118,13 +130,15 @@ from chalk.parsed._proto.utils import datetime_to_proto_timestamp, value_to_prot
 from chalk.utils import df_utils
 from chalk.utils.df_utils import record_batch_to_arrow_ipc
 from chalk.utils.grpc import AuthenticatedChalkClientInterceptor, TokenRefresher, UnauthenticatedChalkClientInterceptor
-from chalk.utils.tracing import add_trace_headers
+from chalk.utils.tracing import add_trace_headers, safe_trace
 
 if TYPE_CHECKING:
     from pyarrow import RecordBatch, Table
+    from pydantic import BaseModel
 
     from chalk._gen.chalk.server.v1.builder_pb2 import StartBranchResponse
     from chalk._gen.chalk.server.v1.builder_pb2_grpc import BuilderServiceStub
+    from chalk.client import ChalkError
 
 CHALK_GRPC_TRACE_ID_HEADER: str = "x-chalk-trace-id"
 
@@ -212,29 +226,26 @@ def _parse_uri_for_engine(query_server_uri: str) -> ParsedUri:
     return ParsedUri(uri_without_scheme=uri_without_scheme, use_tls=use_tls)
 
 
-channel_options: List[tuple[str, str | int]] = [
-    ("grpc.max_send_message_length", 1024 * 1024 * 100),  # 100MB
-    ("grpc.max_receive_message_length", 1024 * 1024 * 100),  # 100MB
+default_channel_options: Dict[str, str | int] = {
+    "grpc.max_send_message_length": 1024 * 1024 * 100,  # 100MB
+    "grpc.max_receive_message_length": 1024 * 1024 * 100,  # 100MB
     # https://grpc.io/docs/guides/performance/#python
-    (grpc.experimental.ChannelOptions.SingleThreadedUnaryStream, 1),
-    (
-        "grpc.service_config",
-        json.dumps(
-            {
-                "methodConfig": [
-                    {
-                        "name": [{}],
-                        "maxAttempts": 5,
-                        "initialBackoff": "0.1s",
-                        "maxBackoff": "1s",
-                        "backoffMultiplier": 2,
-                        "retryableStatusCodes": ["UNAVAILABLE"],
-                    }
-                ]
-            }
-        ),
+    grpc.experimental.ChannelOptions.SingleThreadedUnaryStream: 1,
+    "grpc.service_config": json.dumps(
+        {
+            "methodConfig": [
+                {
+                    "name": [{}],
+                    "maxAttempts": 5,
+                    "initialBackoff": "0.1s",
+                    "maxBackoff": "1s",
+                    "backoffMultiplier": 2,
+                    "retryableStatusCodes": ["UNAVAILABLE"],
+                }
+            ]
+        }
     ),
-]
+}
 
 
 T = TypeVar("T")
@@ -282,6 +293,22 @@ class StubProvider:
             )
         return OfflineQueryMetadataServiceStub(self._server_channel)
 
+    @cached_property
+    def scheduled_query_stub(self) -> SchedulerServiceStub:
+        if self._server_channel is None:
+            raise ValueError(
+                "The GRPC engine service is not available. If you would like to set up a GRPC service, please contact Chalk."
+            )
+        return SchedulerServiceStub(self._server_channel)
+
+    @cached_property
+    def scheduled_query_run_stub(self) -> ScheduledQueryServiceStub:
+        if self._server_channel is None:
+            raise ValueError(
+                "The GRPC engine service is not available. If you would like to set up a GRPC service, please contact Chalk."
+            )
+        return ScheduledQueryServiceStub(self._server_channel)
+
     @cached_property
     def sql_stub(self) -> SqlServiceStub:
         if self._engine_channel is None:
@@ -298,6 +325,14 @@ class StubProvider:
             )
         return DataFrameServiceStub(self._engine_channel)
 
+    @cached_property
+    def streaming_stub(self) -> SimpleStreamingServiceStub:
+        if self._engine_channel is None:
+            raise ValueError(
+                "The GRPC engine service is not available. If you would like to set up a GRPC service, please contact Chalk."
+            )
+        return SimpleStreamingServiceStub(self._engine_channel)
+
     @cached_property
     def model_stub(self) -> ModelRegistryServiceStub:
         if self._server_channel is None:
@@ -337,17 +372,25 @@ class StubProvider:
         deployment_tag: str | None = None,
         skip_api_server: bool = False,
         additional_headers: List[tuple[str, str]] | None = None,
+        channel_options: List[tuple[str, str | int]] | None = None,
     ):
         super().__init__()
         additional_headers_nonempty: List[tuple[str, str]] = [] if additional_headers is None else additional_headers
         token_refresher: TokenRefresher | None = None
+        channel_options_merged: Dict[str, str | int] = default_channel_options.copy()
+        if channel_options:
+            channel_options_merged.update(dict(channel_options))
         if skip_api_server:
             # Omits the auth handshake with the API server. Primarily for internal use/testing -- if used in production,
             # this client will simply fail to connect. If True then query_server must be provided & point to
             # `localhost/127.0.0.1`.
             if query_server is None:
                 raise ValueError("If skipping API server auth, query_server URI must be provided.")
-            elif not (query_server.startswith("localhost") or query_server.startswith("127.0.0.1")):
+            parsed_uri = _parse_uri_for_engine(query_server)
+            if not (
+                parsed_uri.uri_without_scheme.startswith("localhost")
+                or parsed_uri.uri_without_scheme.startswith("127.0.0.1")
+            ):
                 warnings.warn(
                     "Skipping API server auth should only be enabled if query_server URI is localhost. It will fail to authenticate against a production engine."
                 )
@@ -363,13 +406,13 @@ class StubProvider:
             _unauthenticated_server_channel: grpc.Channel = (
                 grpc.insecure_channel(
                     target=server_host,
-                    options=channel_options,
+                    options=list(channel_options_merged.items()),
                 )
                 if server_host.startswith("localhost") or server_host.startswith("127.0.0.1")
                 else grpc.secure_channel(
                     target=server_host,
                     credentials=grpc.ssl_channel_credentials(),
-                    options=channel_options,
+                    options=list(channel_options_merged.items()),
                 )
             )
 
@@ -443,12 +486,12 @@ class StubProvider:
                     grpc.secure_channel(
                         target=parsed_uri.uri_without_scheme,
                         credentials=grpc.ssl_channel_credentials(),
-                        options=channel_options,
+                        options=list(channel_options_merged.items()),
                     )
                     if parsed_uri.use_tls
                     else grpc.insecure_channel(
                         target=parsed_uri.uri_without_scheme,
-                        options=channel_options,
+                        options=list(channel_options_merged.items()),
                     )
                 ),
                 *interceptors,
@@ -463,6 +506,7 @@ class StubRefresher:
         deployment_tag: str | None = None,
         skip_api_server: bool = False,
         additional_headers: List[tuple[str, str]] | None = None,
+        channel_options: List[tuple[str, str | int]] | None = None,
     ):
         super().__init__()
         self._token_config = token_config
@@ -470,6 +514,7 @@ class StubRefresher:
         self._deployment_tag = deployment_tag
         self._skip_api_server = skip_api_server
         self._additional_headers = additional_headers
+        self._channel_options = channel_options
         self._stub = self._refresh_stub()
 
     def _refresh_stub(self) -> StubProvider:
@@ -479,6 +524,7 @@ class StubRefresher:
             deployment_tag=self._deployment_tag,
             skip_api_server=self._skip_api_server,
             additional_headers=self._additional_headers,
+            channel_options=self._channel_options,
         )
         return self._stub
 
@@ -515,6 +561,12 @@ class StubRefresher:
     def call_offline_query_stub(self, fn: Callable[[OfflineQueryMetadataServiceStub], T]) -> T:
         return self._retry_callable(fn, lambda: self._stub.offline_query_stub)
 
+    def call_scheduled_query_stub(self, fn: Callable[[SchedulerServiceStub], T]) -> T:
+        return self._retry_callable(fn, lambda: self._stub.scheduled_query_stub)
+
+    def call_scheduled_query_run_stub(self, fn: Callable[[ScheduledQueryServiceStub], T]) -> T:
+        return self._retry_callable(fn, lambda: self._stub.scheduled_query_run_stub)
+
     def call_sql_stub(self, fn: Callable[[SqlServiceStub], T]) -> T:
         return self._retry_callable(fn, lambda: self._stub.sql_stub)
 
@@ -536,6 +588,9 @@ class StubRefresher:
     def call_job_queue_stub(self, fn: Callable[[DataPlaneJobQueueServiceStub], T]) -> T:
         return self._retry_callable(fn, lambda: self._stub.job_queue_stub)
 
+    def call_streaming_stub(self, fn: Callable[[SimpleStreamingServiceStub], T]) -> T:
+        return self._retry_callable(fn, lambda: self._stub.streaming_stub)
+
     @property
     def log_stub(self) -> LogSearchServiceStub:
         return self._stub.log_stub
@@ -561,6 +616,7 @@ class ChalkGRPCClient:
         additional_headers: List[tuple[str, str]] | None = None,
         query_server: str | None = None,
         input_compression: typing.Literal["lz4", "zstd", "uncompressed"] = "lz4",
+        channel_options: List[Tuple[str, str | int]] | None = None,
         **kwargs: Any,
     ):
         """Create a `ChalkGRPCClient` with the given credentials.
@@ -616,6 +672,7 @@ class ChalkGRPCClient:
             deployment_tag=deployment_tag,
             additional_headers=additional_headers,
             skip_api_server=kwargs.get("_skip_api_server", False),
+            channel_options=channel_options,
         )
 
     _INPUT_ENCODE_OPTIONS = GRPC_ENCODE_OPTIONS
@@ -796,6 +853,7 @@ class ChalkGRPCClient:
             request_timeout=request_timeout,
             headers=headers,
             query_context=_validate_context_dict(query_context),
+            trace=trace,
         )
         return OnlineQueryConverter.online_query_bulk_response_decode_to_single(bulk_response)
 
@@ -823,41 +881,42 @@ class ChalkGRPCClient:
         query_context: Mapping[str, Union[str, int, float, bool, None]] | None = None,
         trace: bool = False,
     ) -> online_query_pb2.OnlineQueryBulkResponse:
-        request = self._make_query_bulk_request(
-            input={k: [v] for k, v in input.items()},
-            output=output,
-            now=[now] if now is not None else [],
-            staleness=staleness or {},
-            tags=tags or (),
-            correlation_id=correlation_id,
-            query_name=query_name,
-            query_name_version=query_name_version,
-            include_meta=include_meta,
-            meta=meta or {},
-            explain=explain,
-            store_plan_stages=store_plan_stages,
-            value_metrics_tag_by_features=value_metrics_tag_by_features,
-            encoding_options=encoding_options,
-            required_resolver_tags=required_resolver_tags or (),
-            planner_options=planner_options or {},
-            query_context=query_context,
-        )
-        if trace:
-            extra_headers: dict[str, str] = {}
-            extra_headers = add_trace_headers(extra_headers)
-            headers = _merge_headers(extra_headers, headers)
-        metadata = _canonicalize_headers(headers)
-        return self._stub_refresher.call_query_stub(
-            lambda x: x.OnlineQueryBulk(
-                request,
-                timeout=request_timeout,
-                metadata=metadata,
+        with safe_trace("_online_query_grpc_request"):
+            request = self._make_query_bulk_request(
+                input={k: [v] for k, v in input.items()},
+                output=output,
+                now=[now] if now is not None else [],
+                staleness=staleness or {},
+                tags=tags or (),
+                correlation_id=correlation_id,
+                query_name=query_name,
+                query_name_version=query_name_version,
+                include_meta=include_meta,
+                meta=meta or {},
+                explain=explain,
+                store_plan_stages=store_plan_stages,
+                value_metrics_tag_by_features=value_metrics_tag_by_features,
+                encoding_options=encoding_options,
+                required_resolver_tags=required_resolver_tags or (),
+                planner_options=planner_options or {},
+                query_context=query_context,
+            )
+            if trace:
+                extra_headers: dict[str, str] = {}
+                extra_headers = add_trace_headers(extra_headers)
+                headers = _merge_headers(extra_headers, headers)
+            metadata = _canonicalize_headers(headers)
+            return self._stub_refresher.call_query_stub(
+                lambda x: x.OnlineQueryBulk(
+                    request,
+                    timeout=request_timeout,
+                    metadata=metadata,
+                )
             )
-        )
 
     def online_query_bulk(
         self,
-        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame],
+        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame, None] = None,
         output: Sequence[FeatureReference] = (),
         now: Optional[Sequence[dt.datetime]] = None,
         staleness: Optional[Mapping[FeatureReference, str]] = None,
@@ -876,9 +935,17 @@ class ChalkGRPCClient:
         request_timeout: Optional[float] = None,
         headers: Mapping[str, str | bytes] | Sequence[tuple[str, str | bytes]] | None = None,
         query_context: Mapping[str, Union[str, int, float, bool, None]] | str | None = None,
+        *,
+        input_sql: str | None = None,
     ) -> BulkOnlineQueryResult:
+        if input is None and input_sql is None:
+            raise TypeError("One of `input` or `input_sql` is required")
+        if input is not None and input_sql is not None:
+            raise TypeError("`input` and `input_sql` are mutually exclusive")
+
         response, call = self._online_query_bulk_grpc_request(
             input=input,
+            input_sql=input_sql,
             output=output,
             now=now,
             staleness=staleness,
@@ -905,7 +972,8 @@ class ChalkGRPCClient:
     def _online_query_bulk_grpc_request(
         self,
         *,
-        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame],
+        input: Union[Mapping[FeatureReference, Sequence[Any]], DataFrame, None] = None,
+        input_sql: str | None = None,
         output: Sequence[FeatureReference] = (),
         now: Optional[Sequence[dt.datetime]] = None,
         staleness: Optional[Mapping[FeatureReference, str]] = None,
@@ -926,8 +994,10 @@ class ChalkGRPCClient:
         query_context: Mapping[str, Union[str, int, float, bool, None]] | None = None,
     ) -> Tuple[online_query_pb2.OnlineQueryBulkResponse, grpc.Call]:
         """Returns the raw GRPC response and metadata"""
+
         request = self._make_query_bulk_request(
             input=input,
+            input_sql=input_sql,
             output=output,
             now=now or (),
             staleness=staleness or {},
@@ -1116,7 +1186,9 @@ class ChalkGRPCClient:
 
     def _make_query_bulk_request(
         self,
-        input: Mapping[FeatureReference, Sequence[Any]] | DataFrame,
+        *,
+        input: Mapping[FeatureReference, Sequence[Any]] | DataFrame | None = None,
+        input_sql: str | None = None,
         output: Sequence[FeatureReference],
         now: Sequence[dt.datetime],
         staleness: Mapping[FeatureReference, str],
@@ -1134,9 +1206,19 @@ class ChalkGRPCClient:
         planner_options: Mapping[str, str | int | bool],
         query_context: Mapping[str, Union[str, int, float, bool, None]] | str | None,
     ) -> online_query_pb2.OnlineQueryBulkRequest:
-        inputs_bytes = get_features_feather_bytes(
-            input, self._INPUT_ENCODE_OPTIONS, compression=self._input_compression
-        )
+        if input is None and input_sql is None:
+            raise TypeError("One of `input` or `input_sql` is required")
+        if input is not None and input_sql is not None:
+            raise TypeError("`input` and `input_sql` are mutually exclusive")
+
+        inputs_feather: bytes | None
+        if input is None:
+            inputs_feather = None
+        else:
+            inputs_feather = get_features_feather_bytes(
+                input, self._INPUT_ENCODE_OPTIONS, compression=self._input_compression
+            )
+
         encoded_outputs = encode_outputs(output)
         outputs = encoded_outputs.string_outputs
         # Currently assume every feature tag is just a fqn instead of a more complex expr.
@@ -1165,7 +1247,8 @@ class ChalkGRPCClient:
         query_context = _validate_context_dict(query_context)
         query_context_proto = {k: value_to_proto(v) for k, v in query_context.items()} if query_context else None
         return online_query_pb2.OnlineQueryBulkRequest(
-            inputs_feather=inputs_bytes,
+            inputs_feather=inputs_feather,
+            inputs_sql=input_sql,
             outputs=[online_query_pb2.OutputExpr(feature_fqn=o) for o in outputs]
             + [online_query_pb2.OutputExpr(feature_expression=o) for o in encoded_outputs.feature_expressions_proto],
             now=now_proto,
@@ -1193,6 +1276,101 @@ class ChalkGRPCClient:
             body_type=online_query_pb2.FEATHER_BODY_TYPE_RECORD_BATCHES,
         )
 
+    def run_scheduled_query(
+        self,
+        name: str,
+        planner_options: Optional[Mapping[str, Any]],
+        incremental_resolvers: Optional[Sequence[str]],
+        max_samples: Optional[int],
+        env_overrides: Optional[Mapping[str, str]],
+    ) -> ManualTriggerScheduledQueryResponseDataclass:
+        """
+        Manually trigger a scheduled query request.
+
+        Parameters
+        ----------
+        name
+            The name of the scheduled query to be triggered.
+        incremental_resolvers
+            If set to None, Chalk will incrementalize resolvers in the query's root namespaces.
+            If set to a list of resolvers, this set will be used for incrementalization.
+            Incremental resolvers must return a feature time in its output, and must return a `DataFrame`.
+            Most commonly, this will be the name of a SQL file resolver. Chalk will ingest all new data
+            from these resolvers and propagate changes to values in the root namespace.
+        max_samples
+            The maximum number of samples to compute.
+        env_overrides:
+            A dictionary of environment values to override during this specific triggered query.
+
+        Other Parameters
+        ----------------
+        planner_options
+            A dictionary of options to pass to the planner.
+            These are typically provided by Chalk Support for specific use cases.
+
+        Returns
+        -------
+        ManualTriggerScheduledQueryResponse
+            A response message containing metadata around the triggered run.
+
+        Examples
+        --------
+        >>> from chalk.client.client_grpc import ChalkGRPCClient
+        >>> ChalkGRPCClient().run_scheduled_query(
+        ...     name="my_scheduled_query",
+        ... )
+        """
+        proto_resp = self._stub_refresher.call_scheduled_query_stub(
+            lambda x: x.ManualTriggerScheduledQuery(
+                request=ManualTriggerScheduledQueryRequest(
+                    cron_query_name=name,
+                    planner_options=planner_options or {},
+                    incremental_resolvers=incremental_resolvers or (),
+                    max_samples=max_samples,
+                    env_overrides=env_overrides or {},
+                ),
+            )
+        )
+        return ManualTriggerScheduledQueryResponseDataclass.from_proto(proto_resp)
+
+    def get_scheduled_query_run_history(
+        self,
+        name: str,
+        limit: int = 10,
+    ) -> List[ScheduledQueryRun]:
+        """
+        Get the run history for a scheduled query.
+
+        Parameters
+        ----------
+        name
+            The name of the scheduled query.
+        limit
+            The maximum number of runs to return. Defaults to 10.
+
+        Returns
+        -------
+        list[ScheduledQueryRun]
+            A response message containing the list of scheduled query runs.
+
+        Examples
+        --------
+        >>> from chalk.client.client_grpc import ChalkGRPCClient
+        >>> ChalkGRPCClient().get_scheduled_query_run_history(
+        ...     name="my_scheduled_query",
+        ...     limit=20,
+        ... )
+        """
+        proto_resp = self._stub_refresher.call_scheduled_query_run_stub(
+            lambda x: x.GetScheduledQueryRuns(
+                GetScheduledQueryRunsRequest(
+                    cron_name=name,
+                    limit=limit,
+                )
+            )
+        )
+        return [ScheduledQueryRun.from_proto(run) for run in proto_resp.runs]
+
     def get_graph(self, deployment: DeploymentId | None = None) -> Graph:
         """Get the graph for a given deployment.
 
@@ -1486,53 +1664,53 @@ class ChalkGRPCClient:
         source_config: Optional[SourceConfig] = None,
         dependencies: Optional[List[str]] = None,
     ) -> RegisterModelVersionResponse:
-        """
-        Register a model in the Chalk model registry.
+        """Register a model in the Chalk model registry.
 
         Parameters
         ----------
-        name : str
-           Unique name for the model
-        aliases : list of str, optional
-           List of version aliases (e.g., ["v1.0", "latest"])
-        model : object, optional
-           Python model object (for object-based registration)
-        model_paths : list of str, optional
-           Paths to model files (for file-based registration)
-        additional_files : List[str], optional
-           Additional files needed for inference (tokenizers, configs, etc.)
-        model_type : ModelType, optional
-           Type of model framework
-        model_encoding : ModelEncoding, optional
-           Serialization format
-        input_schema : dict, list, or Any
-           Definition of the input schema. Can be:
-           - dict: Dictionary mapping column names to dtypes for tabular data
-           - list: List of (shape, dtype) tuples for tensor data
-        output_schema : dict, list, or Any
-           Definition of the output schema. Can be:
-           - dict: Dictionary mapping column names to dtypes for tabular data
-           - list: List of (shape, dtype) tuples for tensor data
-        metadata : dict, optional
-           Additional metadata dictionary containing framework info,
-           training details, performance metrics, etc.
-        input_features : FeatureReference, str, optional
+        name
+            Unique name for the model.
+        aliases
+            List of version aliases (e.g., `["v1.0", "latest"]`).
+        model
+            Python model object (for object-based registration).
+        model_paths
+            Paths to model files (for file-based registration).
+        additional_files
+            Additional files needed for inference (tokenizers, configs, etc.)
+        model_type
+            Type of model framework.
+        model_encoding
+            Serialization format.
+        input_schema
+            Definition of the input schema. Can be:
+            - `dict`: Dictionary mapping column names to dtypes for tabular data
+            - `list`: List of `(shape, dtype)` tuples for tensor data
+        output_schema
+            Definition of the output schema. Can be:
+            - `dict`: Dictionary mapping column names to dtypes for tabular data
+            - `list`: List of `(shape, dtype)` tuples for tensor data
+        metadata
+            Additional metadata dictionary containing framework info,
+            training details, performance metrics, etc.
+        input_features
             The features to be used as inputs to the model.
             For example, `[User.message]`. Features can also be expressed as snakecased strings,
-            e.g. `["user.message"]`
-        output_features : FeatureReference, str, optional
+            e.g. `["user.message"]`.
+        output_features
             The features to be used as outputs to the model.
             For example, `[User.is_spam]`. Features can also be expressed as snakecased strings,
-            e.g. `["user.is_spam"]`
-        source_config : LocalSourceConfig, S3SourceConfig, HFSourceConfig, optional
+            e.g. `["user.is_spam"]`.
+        source_config
             Config to pass credentials to access files from a remote source.
-        dependencies : List[str], optional
+        dependencies
             List of package dependencies needed to run this model.
-            e.g. ["torch==2.7.1", "numpy==1.26.4"]
+            e.g. `["torch==2.7.1", "numpy==1.26.4"]`.
+
         Returns
         -------
         ModelVersion
-           The registered model version object
+            The registered model version object.
 
         Examples
         --------
@@ -1613,9 +1791,29 @@ class ChalkGRPCClient:
                             "Failed to register model. Please specify a model encoding if using model_paths."
                         )
 
+                # Auto-convert ONNX list schemas to dict format if needed
+                if model_type == ModelType.ONNX:
+                    input_schema = model_serializer.convert_onnx_list_schema_to_dict(input_schema, model, is_input=True)
+                    output_schema = model_serializer.convert_onnx_list_schema_to_dict(
+                        output_schema, model, is_input=False
+                    )
+
                 input_model_schema = model_serializer.convert_schema(input_schema)
                 output_model_schema = model_serializer.convert_schema(output_schema)
 
+                # Final validation: ONNX models must use tabular schemas
+                if model_type == ModelType.ONNX:
+                    if input_model_schema is not None and not input_model_schema.HasField("tabular"):
+                        raise ValueError(
+                            "ONNX models must be registered with tabular input schema (dict format). "
+                            + "Use dict format like {'input': Tensor[...]} instead of list format."
+                        )
+                    if output_model_schema is not None and not output_model_schema.HasField("tabular"):
+                        raise ValueError(
+                            "ONNX models must be registered with tabular output schema (dict format). "
+                            + "Use dict format like {'output': Vector[...]} instead of list format."
+                        )
+
                 all_files_to_process, model_file_names = model_file_uploader.prepare_file_mapping(
                     model_paths, additional_files
                 )
@@ -1758,6 +1956,7 @@ class ChalkGRPCClient:
         name: str,
         model_artifact_id: Optional[str] = None,
         run_id: Optional[str] = None,
+        run_name: Optional[str] = None,
         criterion: Optional[ModelRunCriterion] = None,
         aliases: Optional[List[str]] = None,
     ) -> RegisterModelVersionResponse:
@@ -1772,6 +1971,8 @@ class ChalkGRPCClient:
             Artifact UUID to promote to a model version.
         run_id: str, optional
             run id that produce the artifact to promote.
+        run_name: str, optional
+            run name used in the checkpointer for artifact to promote.
         criterion: ModelRunCriterion, optional
             criterion on which to select the artifact from the training run.
             If none provided, the latest artifact in the run will be selected.
@@ -1789,11 +1990,15 @@ class ChalkGRPCClient:
         ... )
         """
         if model_artifact_id is not None:
-            if run_id is not None or criterion is not None:
-                raise ValueError("Please specify only one of 'model_artifact_id' or (run_name, run criterion)")
+            if run_id is not None or criterion is not None or run_name is not None:
+                raise ValueError(
+                    "Please specify only one of 'model_artifact_id', (run_id, run criterion), (run_name, run criterion)"
+                )
         else:
-            if run_id is None and criterion is None:
-                raise ValueError("Please specify at least one of 'model_artifact_id' or (run_name, run criterion)")
+            if run_name is None and run_id is None:
+                raise ValueError(
+                    "Please specify only one of 'model_artifact_id', (run_id, run criterion), (run_name, run criterion)"
+                )
 
         try:
             resp: CreateModelVersionFromArtifactResponse = self._stub_refresher.call_model_stub(
@@ -1802,7 +2007,8 @@ class ChalkGRPCClient:
                         model_name=name,
                         model_artifact_id=model_artifact_id,
                         training_run=ModelSerializer.convert_run_criterion_to_proto(
-                            run_name=run_id,
+                            run_id=run_id,
+                            run_name=run_name,
                             criterion=criterion,
                         ),
                         aliases=aliases,
@@ -2011,3 +2217,232 @@ class ChalkGRPCClient:
             poll_interval=poll_interval,
             output_callback=output_callback,
         )
+
+    def test_streaming_resolver(
+        self,
+        resolver: str | Resolver,
+        message_bodies: "list[str | bytes | BaseModel] | None" = None,
+        message_keys: list[str | None] | None = None,
+        message_timestamps: list[str | dt.datetime] | None = None,
+        message_filepath: str | None = None,
+        request_timeout: Optional[float] = None,
+    ) -> "StreamResolverTestResponse":
+        """Test a streaming resolver with supplied messages.
+
+        This method tests streaming resolvers using the gRPC TestStreamingResolver endpoint.
+        It supports both deployed resolvers (by FQN) and static/undeployed resolvers
+        (automatically serialized from Resolver objects).
+
+        Parameters
+        ----------
+        resolver : str | Resolver
+            The streaming resolver or its string name. If a StreamResolver object with
+            feature_expressions is provided, it will be automatically serialized for testing.
+        message_bodies : list[str | bytes | BaseModel], optional
+            The message bodies to process. Can be JSON strings, raw bytes,
+            or Pydantic models (will be serialized to JSON).
+            Either message_bodies or message_filepath must be provided.
+        message_keys : list[str | None], optional
+            Optional keys for each message. If not provided, all keys will be None.
+            Must match length of message_bodies if provided.
+        message_timestamps : list[str | datetime], optional
+            Optional timestamps for each message. If not provided, current time
+            will be used. Must match length of message_bodies if provided.
+        message_filepath : str, optional
+            A filepath from which test messages will be ingested.
+            This file should be newline delimited JSON with format:
+            {"message_key": "my-key", "message_body": {"field1": "value1"}}
+            Each line may optionally contain a "message_timestamp" field.
+            Either message_bodies or message_filepath must be provided.
+        request_timeout : float, optional
+            Request timeout in seconds.
+
+        Returns
+        -------
+        StreamResolverTestResponse
+            Response containing:
+            - status: SUCCESS or FAILURE
+            - data_uri: Optional signed URL to parquet file with results
+            - errors: List of ChalkError objects
+            - message: Human-readable message
+
+        Examples
+        --------
+        >>> from chalk.client.client_grpc import ChalkGRPCClient
+        >>> client = ChalkGRPCClient()
+        >>> response = client.test_streaming_resolver(
+        ...     resolver="my_module.my_stream_resolver",
+        ...     message_bodies=[
+        ...         '{"user_id": 1, "event": "login"}',
+        ...         '{"user_id": 2, "event": "logout"}',
+        ...     ],
+        ...     message_keys=["user_1", "user_2"],
+        ... )
+        >>> print(f"Status: {response.status}")
+        >>> if response.data_uri:
+        ...     print(f"Results at: {response.data_uri}")
+        """
+        import base64
+        import json
+        from uuid import uuid4
+
+        import pyarrow as pa
+
+        from chalk._gen.chalk.streaming.v1.simple_streaming_service_pb2 import TestStreamingResolverRequest
+        from chalk.utils.pydanticutil.pydantic_compat import get_pydantic_model_json, is_pydantic_basemodel_instance
+
+        # Determine if resolver is static and needs serialization
+        resolver_fqn: str | None = None
+        static_stream_resolver_b64: str | None = None
+
+        if isinstance(resolver, str):
+            resolver_fqn = resolver
+        else:
+            from chalk.features.resolver import StreamResolver
+
+            resolver_fqn = resolver.fqn
+
+            if isinstance(resolver, StreamResolver) and resolver.feature_expressions:
+                from chalk.parsed.to_proto import ToProtoConverter
+
+                proto_resolver = ToProtoConverter.convert_stream_resolver(resolver)
+                static_stream_resolver_b64 = base64.b64encode(
+                    proto_resolver.SerializeToString(deterministic=True)
+                ).decode("utf-8")
+
+        # Load from file if provided
+        if message_filepath is not None:
+            if message_bodies is not None:
+                raise ValueError("Cannot provide both message_filepath and message_bodies")
+
+            loaded_bodies: list[Any] = []
+            loaded_keys: list[str | None] = []
+            loaded_timestamps: list[str | None] = []
+
+            with open(message_filepath, "r") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    msg = json.loads(line)
+                    loaded_bodies.append(msg.get("message_body", msg))
+                    loaded_keys.append(msg.get("message_key"))
+                    loaded_timestamps.append(msg.get("message_timestamp"))
+
+            message_bodies = loaded_bodies
+            if message_keys is None and any(k is not None for k in loaded_keys):
+                message_keys = loaded_keys
+            if message_timestamps is None and any(t is not None for t in loaded_timestamps):
+                # Cast needed: loaded_timestamps is list[str | None] from JSON,
+                # but message_timestamps is list[str | datetime] - strings will be parsed later
+                message_timestamps = typing.cast(list[str | dt.datetime], loaded_timestamps)
+
+        # Validate inputs
+        if message_bodies is None:
+            raise ValueError("Either message_bodies or message_filepath must be provided")
+
+        num_messages = len(message_bodies)
+        if num_messages == 0:
+            raise ValueError("message_bodies cannot be empty")
+
+        if message_keys is not None and len(message_keys) != num_messages:
+            raise ValueError(
+                f"message_keys length ({len(message_keys)}) must match message_bodies length ({num_messages})"
+            )
+
+        if message_timestamps is not None and len(message_timestamps) != num_messages:
+            raise ValueError(
+                f"message_timestamps length ({len(message_timestamps)}) must match message_bodies length ({num_messages})"
+            )
+
+        # Generate defaults
+        message_ids = [str(uuid4()) for _ in range(num_messages)]
+
+        if message_keys is None:
+            message_keys = typing.cast(list[str | None], [None] * num_messages)
+
+        if message_timestamps is None:
+            message_timestamps = typing.cast(list[str | dt.datetime], [dt.datetime.now()] * num_messages)
+
+        # Convert message bodies to bytes
+        processed_bodies: list[bytes] = []
+        for body in message_bodies:
+            if isinstance(body, bytes):
+                processed_bodies.append(body)
+            elif isinstance(body, str):
+                processed_bodies.append(body.encode("utf-8"))
+            elif is_pydantic_basemodel_instance(body):
+                # Use utility function that handles both Pydantic v1 and v2
+                processed_bodies.append(get_pydantic_model_json(body).encode("utf-8"))
+            else:
+                # Try JSON serialization for dict-like objects
+                processed_bodies.append(json.dumps(body).encode("utf-8"))
+
+        # Convert timestamps to unix timestamps in milliseconds (int64)
+        # At this point message_timestamps is guaranteed to be non-None due to the default assignment above
+        assert message_timestamps is not None
+        processed_timestamps: list[int] = []
+        for ts in message_timestamps:
+            if isinstance(ts, str):
+                # Parse ISO format string
+                parsed = dt.datetime.fromisoformat(ts.replace("Z", "+00:00"))
+                processed_timestamps.append(int(parsed.timestamp() * 1000))  # milliseconds
+            else:
+                # Type narrowing: ts must be dt.datetime here
+                processed_timestamps.append(int(ts.timestamp() * 1000))  # milliseconds
+
+        # Create Arrow table
+        table = pa.table(
+            {
+                "message_id": message_ids,
+                "message_key": message_keys,
+                "message_data": processed_bodies,
+                "publish_timestamp": processed_timestamps,
+            }
+        )
+
+        # Serialize to Arrow IPC format
+        sink = pa.BufferOutputStream()
+        with pa.ipc.new_stream(sink, table.schema) as writer:
+            writer.write_table(table)
+        input_data = sink.getvalue().to_pybytes()
+
+        # Create gRPC request
+        request = TestStreamingResolverRequest(
+            resolver_fqn=resolver_fqn or "",
+            input_data=input_data,
+            operation_id=None,
+            debug=True,
+        )
+
+        if static_stream_resolver_b64:
+            request.static_stream_resolver_b64 = static_stream_resolver_b64
+
+        # Call new TestStreamingResolver endpoint
+        proto_response = self._stub_refresher.call_streaming_stub(
+            lambda x: x.TestStreamingResolver(
+                request,
+                timeout=request_timeout,
+            )
+        )
+
+        # Convert proto response to StreamResolverTestResponse
+        from chalk._gen.chalk.streaming.v1.simple_streaming_service_pb2 import TEST_STREAM_RESOLVER_STATUS_SUCCESS
+
+        status = (
+            StreamResolverTestStatus.SUCCESS
+            if proto_response.status == TEST_STREAM_RESOLVER_STATUS_SUCCESS
+            else StreamResolverTestStatus.FAILURE
+        )
+
+        # Convert proto errors to ChalkError objects
+        errors_list: list[ChalkError] = []
+        if proto_response.errors:
+            errors_list = [ChalkErrorConverter.chalk_error_decode(err) for err in proto_response.errors]
+
+        return StreamResolverTestResponse(
+            status=status,
+            data_uri=proto_response.data_uri if proto_response.HasField("data_uri") else None,
+            errors=errors_list if errors_list else None,
+            message=proto_response.message if proto_response.message else None,
+        )