Merge branch 'main' into tchow/emr-submission

Author: tchow-zlai

cloud_gcp/iceberg-bigquery-1.11.0-SNAPSHOT.jar

@@ -233,15 +233,7 @@ def Aggregation(
     if isinstance(operation, tuple):
         operation, arg_map = operation[0], operation[1]
 
-    def normalize(w: Union[common.Window, str]) -> common.Window:
-        if isinstance(w, str):
-            return window_utils._from_str(w)
-        elif isinstance(w, common.Window):
-            return w
-        else:
-            raise Exception("window should be either a string like '7d', '24h', or a Window type")
-
-    norm_windows = [normalize(w) for w in windows] if windows else None
+    norm_windows = [window_utils.normalize_window(w) for w in windows] if windows else None
 
     agg = ttypes.Aggregation(input_column, operation, arg_map, norm_windows, buckets)
 

python/src/ai/chronon/group_by.py

@@ -1,10 +1,12 @@
 import inspect
 from dataclasses import dataclass
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import gen_thrift.api.ttypes as ttypes
+import gen_thrift.common.ttypes as common
 
 from ai.chronon import utils
+from ai.chronon import windows as window_utils
 from ai.chronon.data_types import DataType, FieldsType
 from ai.chronon.utils import ANY_SOURCE_TYPE, normalize_source
 
@@ -14,6 +16,17 @@ class ModelBackend:
     SAGEMAKER = ttypes.ModelBackend.SageMaker
 
 
+class DeploymentStrategyType:
+    # deploys the model in a blue-green fashion (~2x capacity) to another endpoint and gradually ramps traffic
+    BLUE_GREEN = ttypes.DeploymentStrategyType.BLUE_GREEN
+
+    # deploys the model in a rolling manner by gradually scaling down existing instances and scaling up new instances
+    ROLLING = ttypes.DeploymentStrategyType.ROLLING
+
+    # deploys the model immediately to the endpoint without any traffic ramping
+    IMMEDIATE = ttypes.DeploymentStrategyType.IMMEDIATE
+
+
 @dataclass
 class ResourceConfig:
     min_replica_count: Optional[int] = None
@@ -38,20 +51,150 @@ def to_thrift(self):
         resource_config_thrift = None
         if self.resource_config:
             resource_config_thrift = self.resource_config.to_thrift()
-        
+
         return ttypes.InferenceSpec(
             modelBackend=self.model_backend,
             modelBackendParams=self.model_backend_params,
             resourceConfig=resource_config_thrift,
         )
 
 
+@dataclass
+class TrainingSpec:
+    training_data_source: Optional[ANY_SOURCE_TYPE] = None
+    training_data_window: Optional[Union[common.Window, str]] = None
+    schedule: Optional[str] = None
+    image: Optional[str] = None
+    python_module: Optional[str] = None
+    resource_config: Optional[ResourceConfig] = None
+    job_configs: Optional[Dict[str, str]] = None
+
+    def to_thrift(self):
+        resource_config_thrift = None
+        if self.resource_config:
+            resource_config_thrift = self.resource_config.to_thrift()
+
+        training_data_source_thrift = None
+        if self.training_data_source:
+            training_data_source_thrift = normalize_source(self.training_data_source)
+
+        # Normalize window - convert string like "30d" or "24h" to common.Window
+        training_data_window_thrift = None
+        if self.training_data_window:
+            training_data_window_thrift = window_utils.normalize_window(self.training_data_window)
+
+        return ttypes.TrainingSpec(
+            trainingDataSource=training_data_source_thrift,
+            trainingDataWindow=training_data_window_thrift,
+            schedule=self.schedule,
+            image=self.image,
+            pythonModule=self.python_module,
+            resourceConfig=resource_config_thrift,
+            jobConfigs=self.job_configs,
+        )
+
+
+@dataclass
+class ServingContainerConfig:
+    image: Optional[str] = None
+    serving_health_route: Optional[str] = None
+    serving_predict_route: Optional[str] = None
+    serving_container_env_vars: Optional[Dict[str, str]] = None
+
+    def to_thrift(self):
+        return ttypes.ServingContainerConfig(
+            image=self.image,
+            servingHealthRoute=self.serving_health_route,
+            servingPredictRoute=self.serving_predict_route,
+            servingContainerEnvVars=self.serving_container_env_vars,
+        )
+
+
+@dataclass
+class EndpointConfig:
+    endpoint_name: Optional[str] = None
+    additional_configs: Optional[Dict[str, str]] = None
+
+    def to_thrift(self):
+        return ttypes.EndpointConfig(
+            endpointName=self.endpoint_name,
+            additionalConfigs=self.additional_configs,
+        )
+
+
+@dataclass
+class Metric:
+    name: Optional[str] = None
+    threshold: Optional[float] = None
+
+    def to_thrift(self):
+        return ttypes.Metric(
+            name=self.name,
+            threshold=self.threshold,
+        )
+
+
+@dataclass
+class RolloutStrategy:
+    rollout_type: Optional[DeploymentStrategyType] = None
+    validation_traffic_percent_ramps: Optional[List[int]] = None
+    validation_traffic_duration_mins: Optional[List[int]] = None
+    rollout_metric_thresholds: Optional[List[Metric]] = None
+
+    def to_thrift(self):
+        rollout_metric_thresholds_thrift = None
+        if self.rollout_metric_thresholds:
+            rollout_metric_thresholds_thrift = [metric.to_thrift() for metric in self.rollout_metric_thresholds]
+
+        return ttypes.RolloutStrategy(
+            rolloutType=self.rollout_type,
+            validationTrafficPercentRamps=self.validation_traffic_percent_ramps,
+            validationTrafficDurationMins=self.validation_traffic_duration_mins,
+            rolloutMetricThresholds=rollout_metric_thresholds_thrift,
+        )
+
+
+@dataclass
+class DeploymentSpec:
+    container_config: Optional[ServingContainerConfig] = None
+    endpoint_config: Optional[EndpointConfig] = None
+    resource_config: Optional[ResourceConfig] = None
+    rollout_strategy: Optional[RolloutStrategy] = None
+
+    def to_thrift(self):
+        container_config_thrift = None
+        if self.container_config:
+            container_config_thrift = self.container_config.to_thrift()
+
+        endpoint_config_thrift = None
+        if self.endpoint_config:
+            endpoint_config_thrift = self.endpoint_config.to_thrift()
+
+        resource_config_thrift = None
+        if self.resource_config:
+            resource_config_thrift = self.resource_config.to_thrift()
+
+        rollout_strategy_thrift = None
+        if self.rollout_strategy:
+            rollout_strategy_thrift = self.rollout_strategy.to_thrift()
+
+        return ttypes.DeploymentSpec(
+            containerConfig=container_config_thrift,
+            endpointConfig=endpoint_config_thrift,
+            resourceConfig=resource_config_thrift,
+            rolloutStrategy=rollout_strategy_thrift,
+        )
+
+
 def Model(
     version: str,
     inference_spec: Optional[InferenceSpec] = None,
     input_mapping: Optional[Dict[str, str]] = None,
     output_mapping: Optional[Dict[str, str]] = None,
     value_fields: Optional[FieldsType] = None,
+    model_artifact_base_uri: Optional[str] = None,
+    training_conf: Optional[TrainingSpec] = None,
+    deployment_conf: Optional[DeploymentSpec] = None,
     output_namespace: Optional[str] = None,
     table_properties: Optional[Dict[str, str]] = None,
     tags: Optional[Dict[str, str]] = None,
@@ -76,6 +219,15 @@ def Model(
         If provided, creates a STRUCT schema that will be set as the model's valueSchema.
         Example: [('score', DataType.DOUBLE), ('category', DataType.STRING)]
     :type value_fields: FieldsType
+    :param model_artifact_base_uri:
+        Base URI where trained model artifacts are stored
+    :type model_artifact_base_uri: str
+    :param training_conf:
+        Configs related to orchestrating model training jobs
+    :type training_conf: TrainingSpec
+    :param deployment_conf:
+        Configs related to orchestrating model deployment
+    :type deployment_conf: DeploymentSpec
     :param output_namespace:
         Namespace for the model output
     :type output_namespace: str
@@ -114,14 +266,27 @@ def Model(
     if value_fields:
         schema_name = "model_value_schema"
         value_schema = DataType.STRUCT(schema_name, *value_fields)
-    
+
+    # Convert training_conf to thrift if provided
+    training_conf_thrift = None
+    if training_conf:
+        training_conf_thrift = training_conf.to_thrift()
+
+    # Convert deployment_conf to thrift if provided
+    deployment_conf_thrift = None
+    if deployment_conf:
+        deployment_conf_thrift = deployment_conf.to_thrift()
+
     # Create and return the Model object
     model = ttypes.Model(
         metaData=meta_data,
         inferenceSpec=inference_spec_thrift,
         inputMapping=input_mapping,
         outputMapping=output_mapping,
         valueSchema=value_schema,
+        modelArtifactBaseUri=model_artifact_base_uri,
+        trainingConf=training_conf_thrift,
+        deploymentConf=deployment_conf_thrift,
     )
 
     return model

python/src/ai/chronon/model.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 import gen_thrift.common.ttypes as common
 
 
@@ -46,3 +48,30 @@ def _from_str(s: str) -> common.Window:
         if "invalid literal for int()" in str(e):
             raise ValueError(f"Invalid numeric value in duration: {value}") from e
         raise e from None
+
+
+def normalize_window(w: Union[common.Window, str]) -> common.Window:
+    """
+    Normalizes a window specification to a common.Window object.
+
+    Accepts either a Window object directly or a string like "30d" or "24h".
+    This is used across the codebase (e.g., in GroupBy aggregations and TrainingSpec).
+
+    Args:
+        w: Either a common.Window object or a string like "7d", "24h"
+
+    Returns:
+        common.Window: The normalized window object
+
+    Raises:
+        TypeError: If the input is neither a string nor a Window object
+    """
+    if isinstance(w, str):
+        return _from_str(w)
+    elif isinstance(w, common.Window):
+        return w
+    else:
+        raise TypeError(
+            f"Window should be either a string like '7d', '24h', or a Window type, "
+            f"got {type(w).__name__}"
+        )