kumoai 2.13.0.dev202512091732__cp311-cp311-macosx_11_0_arm64.whl → 2.14.0.dev202601051732__cp311-cp311-macosx_11_0_arm64.whl

kumoai/experimental/rfm/sagemaker.py

@@ -1,6 +1,6 @@
 import base64
 import json
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 import requests
 
@@ -48,8 +48,8 @@ class KumoClient_SageMakerAdapter(KumoClient):
 
         # Recording buffers.
         self._recording_active = False
-        self._recorded_reqs: List[Dict[str, Any]] = []
-        self._recorded_resps: List[Dict[str, Any]] = []
+        self._recorded_reqs: list[dict[str, Any]] = []
+        self._recorded_resps: list[dict[str, Any]] = []
 
     def authenticate(self) -> None:
         # TODO(siyang): call /ping to verify?
@@ -92,7 +92,7 @@ class KumoClient_SageMakerAdapter(KumoClient):
         self._recorded_reqs.clear()
         self._recorded_resps.clear()
 
-    def end_recording(self) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]:
+    def end_recording(self) -> list[tuple[dict[str, Any], dict[str, Any]]]:
         """Stop recording and return recorded requests/responses."""
         assert self._recording_active
         self._recording_active = False

kumoai/experimental/rfm/task_table.py

@@ -0,0 +1,290 @@
+import copy
+from collections.abc import Sequence
+
+import pandas as pd
+from kumoapi.task import TaskType
+from kumoapi.typing import Dtype, Stype
+from typing_extensions import Self
+
+from kumoai.experimental.rfm.base import Column
+from kumoai.experimental.rfm.infer import contains_timestamp, infer_dtype
+
+
+class TaskTable:
+    r"""A :class:`TaskTable` fully specifies the task, *i.e.* its context and
+    prediction examples with entity IDs, targets and timestamps.
+
+    Args:
+        task_type: The task type.
+        context_df: The data frame holding context examples.
+        pred_df: The data frame holding prediction examples.
+        entity_table_name: The entity table to predict for. For link prediction
+            tasks, needs to hold both entity and target table names.
+        entity_column: The name of the entity column.
+        target_column: The name of the target column.
+        time_column: The name of the time column to use as anchor time. If
+            ``TaskTable.ENTITY_TIME``, use the timestamp of the entity table
+            as anchor time.
+    """
+    ENTITY_TIME = '__entity_time__'
+
+    def __init__(
+        self,
+        task_type: TaskType,
+        context_df: pd.DataFrame,
+        pred_df: pd.DataFrame,
+        entity_table_name: str | Sequence[str],
+        entity_column: str,
+        target_column: str,
+        time_column: str | None = None,
+    ) -> None:
+
+        task_type = TaskType(task_type)
+        if task_type not in {  # Currently supported task types:
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+                TaskType.REGRESSION,
+                TaskType.TEMPORAL_LINK_PREDICTION,
+        }:
+            raise ValueError  # TODO
+        self._task_type = task_type
+
+        # TODO Binary classification and regression checks
+
+        # TODO Check dfs (unify from local table)
+        if context_df.empty:
+            raise ValueError("No context examples given")
+        self._context_df = context_df.copy(deep=False)
+
+        if pred_df.empty:
+            raise ValueError("Provide at least one entity to predict for")
+        self._pred_df = pred_df.copy(deep=False)
+
+        self._dtype_dict: dict[str, Dtype] = {
+            column_name: infer_dtype(context_df[column_name])
+            for column_name in context_df.columns
+        }
+
+        self._entity_table_names: tuple[str] | tuple[str, str]
+        if isinstance(entity_table_name, str):
+            self._entity_table_names = (entity_table_name, )
+        elif len(entity_table_name) == 1:
+            self._entity_table_names = (entity_table_name[0], )
+        elif len(entity_table_name) == 2:
+            self._entity_table_names = (
+                entity_table_name[0],
+                entity_table_name[1],
+            )
+        else:
+            raise ValueError  # TODO
+
+        self._entity_column: str = ''
+        self._target_column: str = ''
+        self._time_column: str | None = None
+
+        self.entity_column = entity_column
+        self.target_column = target_column
+        if time_column is not None:
+            self.time_column = time_column
+
+    @property
+    def num_context_examples(self) -> int:
+        return len(self._context_df)
+
+    @property
+    def num_prediction_examples(self) -> int:
+        return len(self._pred_df)
+
+    @property
+    def task_type(self) -> TaskType:
+        r"""The task type."""
+        return self._task_type
+
+    def narrow_context(self, start: int, length: int) -> Self:
+        r"""Returns a new :class:`TaskTable` that holds a narrowed version of
+        context examples.
+
+        Args:
+            start: Index of the prediction examples to start narrowing.
+            length: Length of the prediction examples.
+        """
+        out = copy.copy(self)
+        df = out._context_df.iloc[start:start + length].reset_index(drop=True)
+        out._context_df = df
+        return out
+
+    def narrow_prediction(self, start: int, length: int) -> Self:
+        r"""Returns a new :class:`TaskTable` that holds a narrowed version of
+        prediction examples.
+
+        Args:
+            start: Index of the prediction examples to start narrowing.
+            length: Length of the prediction examples.
+        """
+        out = copy.copy(self)
+        df = out._pred_df.iloc[start:start + length].reset_index(drop=True)
+        out._pred_df = df
+        return out
+
+    # Entity column ###########################################################
+
+    @property
+    def entity_table_name(self) -> str:
+        return self._entity_table_names[0]
+
+    @property
+    def entity_table_names(self) -> tuple[str] | tuple[str, str]:
+        return self._entity_table_names
+
+    @property
+    def entity_column(self) -> Column:
+        return Column(
+            name=self._entity_column,
+            expr=None,
+            dtype=self._dtype_dict[self._entity_column],
+            stype=Stype.ID,
+        )
+
+    @entity_column.setter
+    def entity_column(self, name: str) -> None:
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not Stype.ID.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._entity_column = name
+
+    # Target column ###########################################################
+
+    @property
+    def evaluate(self) -> bool:
+        r"""Returns ``True`` if this task can be used for model evaluation."""
+        return self._target_column in self._pred_df
+
+    @property
+    def _target_stype(self) -> Stype:
+        if self.task_type in {
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+        }:
+            return Stype.categorical
+        if self.task_type in {TaskType.REGRESSION}:
+            return Stype.numerical
+        if self.task_type.is_link_pred:
+            return Stype.multicategorical
+        raise ValueError
+
+    @property
+    def target_column(self) -> Column:
+        return Column(
+            name=self._target_column,
+            expr=None,
+            dtype=self._dtype_dict[self._target_column],
+            stype=self._target_stype,
+        )
+
+    @target_column.setter
+    def target_column(self, name: str) -> None:
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if not self._target_stype.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._target_column = name
+
+    # Time column #############################################################
+
+    def has_time_column(self) -> bool:
+        r"""Returns ``True`` if this task has a time column; ``False``
+        otherwise.
+        """
+        return self._time_column not in {None, self.ENTITY_TIME}
+
+    @property
+    def use_entity_time(self) -> bool:
+        r"""Whether to use the timestamp of the entity table as anchor time."""
+        return self._time_column == self.ENTITY_TIME
+
+    @property
+    def time_column(self) -> Column | None:
+        r"""The time column of this task.
+
+        The getter returns the time column of this task, or ``None`` if no
+        such time column is present.
+
+        The setter sets a column as a time column for this task, and raises a
+        :class:`ValueError` if the time column has a non-timestamp compatible
+        data type or if the column name does not match a column in the data
+        frame.
+        """
+        if not self.has_time_column():
+            return None
+        assert self._time_column is not None
+        return Column(
+            name=self._time_column,
+            expr=None,
+            dtype=self._dtype_dict[self._time_column],
+            stype=Stype.timestamp,
+        )
+
+    @time_column.setter
+    def time_column(self, name: str | None) -> None:
+        if name is None or name == self.ENTITY_TIME:
+            self._time_column = name
+            return
+
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not contains_timestamp(
+                ser=self._context_df[name],
+                column_name=name,
+                dtype=self._dtype_dict[name],
+        ):
+            raise ValueError  # TODO
+
+        self._time_column = name
+
+    # Metadata ################################################################
+
+    @property
+    def metadata(self) -> pd.DataFrame:
+        raise NotImplementedError
+
+    def print_metadata(self) -> None:
+        raise NotImplementedError
+
+    # Python builtins #########################################################
+
+    def __hash__(self) -> int:
+        return hash((
+            self.task_type,
+            self.entity_table_names,
+            self._entity_column,
+            self._target_column,
+            self._time_column,
+        ))
+
+    def __repr__(self) -> str:
+        if self.task_type.is_link_pred:
+            entity_table_repr = f'entity_table_names={self.entity_table_names}'
+        else:
+            entity_table_repr = f'entity_table_name={self.entity_table_name}'
+
+        if self.use_entity_time:
+            time_repr = 'use_entity_time=True'
+        else:
+            time_repr = f'time_column={self._time_column}'
+
+        return (f'{self.__class__.__name__}(\n'
+                f'  task_type={self.task_type},\n'
+                f'  num_context_examples={self.num_context_examples},\n'
+                f'  num_prediction_examples={self.num_prediction_examples},\n'
+                f'  {entity_table_repr},\n'
+                f'  entity_column={self._entity_column},\n'
+                f'  target_column={self._target_column},\n'
+                f'  {time_repr},\n'
+                f')')

kumoai/pquery/predictive_query.py

@@ -370,9 +370,11 @@ class PredictiveQuery:
         train_table_job_api = global_state.client.generate_train_table_job_api
         job_id: GenerateTrainTableJobID = train_table_job_api.create(
             GenerateTrainTableRequest(
-                dict(custom_tags), pq_id, plan,
-                graph_snapshot_id=self.graph.snapshot(
-                    non_blocking=non_blocking)))
+                dict(custom_tags),
+                pq_id,
+                plan,
+                None,
+            ))
 
         self._train_table = TrainingTableJob(job_id=job_id)
         if non_blocking:
@@ -451,9 +453,11 @@ class PredictiveQuery:
         bp_table_api = global_state.client.generate_prediction_table_job_api
         job_id: GeneratePredictionTableJobID = bp_table_api.create(
             GeneratePredictionTableRequest(
-                dict(custom_tags), pq_id, plan,
-                graph_snapshot_id=self.graph.snapshot(
-                    non_blocking=non_blocking)))
+                dict(custom_tags),
+                pq_id,
+                plan,
+                None,
+            ))
 
         self._prediction_table = PredictionTableJob(job_id=job_id)
         if non_blocking:

kumoai/testing/snow.py

@@ -0,0 +1,50 @@
+import json
+import os
+
+from kumoai.experimental.rfm.backend.snow import Connection
+from kumoai.experimental.rfm.backend.snow import connect as _connect
+
+
+def connect(
+    region: str,
+    id: str,
+    account: str,
+    user: str,
+    warehouse: str,
+    database: str | None = None,
+    schema: str | None = None,
+) -> Connection:
+
+    kwargs = dict(password=os.getenv('SNOWFLAKE_PASSWORD'))
+    if kwargs['password'] is None:
+        import boto3
+        from cryptography.hazmat.primitives import serialization
+
+        client = boto3.client(
+            service_name='secretsmanager',
+            region_name=region,
+        )
+        secret_id = (f'arn:aws:secretsmanager:{region}:{id}:secret:'
+                     f'{account}.snowflakecomputing.com')
+        response = client.get_secret_value(SecretId=secret_id)['SecretString']
+        secret = json.loads(response)
+
+        private_key = serialization.load_pem_private_key(
+            secret['kumo_user_secretkey'].encode(),
+            password=None,
+        )
+        kwargs['private_key'] = private_key.private_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
+    return _connect(
+        account=account,
+        user=user,
+        warehouse='WH_XS',
+        database='KUMO',
+        schema=schema,
+        session_parameters=dict(CLIENT_TELEMETRY_ENABLED=False),
+        **kwargs,
+    )

kumoai/trainer/distilled_trainer.py

@@ -0,0 +1,175 @@
+import logging
+from typing import Literal, Mapping, Optional, Union, overload
+
+from kumoapi.distilled_model_plan import DistilledModelPlan
+from kumoapi.jobs import DistillationJobRequest, DistillationJobResource
+
+from kumoai import global_state
+from kumoai.client.jobs import TrainingJobID
+from kumoai.graph import Graph
+from kumoai.pquery.training_table import TrainingTable, TrainingTableJob
+from kumoai.trainer.job import TrainingJob, TrainingJobResult
+
+logger = logging.getLogger(__name__)
+
+
+class DistillationTrainer:
+    r"""A trainer supports creating a Kumo machine learning model
+    for use in an online serving endpoint. The distllation process involes
+    training a shallow model on a :class:`~kumoai.pquery.PredictiveQuery` using
+    the embeddings generated by a base model :args:`base_training_job_id`.
+
+    Args:
+        model_plan: The distilled model plan to use for the distillation process.
+        base_training_job_id: The ID of the base training job to use for the distillation process.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        model_plan: DistilledModelPlan,
+        base_training_job_id: TrainingJobID,
+    ) -> None:
+        self.model_plan: DistilledModelPlan = model_plan
+        self.base_training_job_id: TrainingJobID = base_training_job_id
+
+        # Cached from backend:
+        self._training_job_id: Optional[TrainingJobID] = None
+
+    # Metadata ################################################################
+
+    @property
+    def is_trained(self) -> bool:
+        r"""Returns ``True`` if this trainer instance has successfully been
+        trained (and is therefore ready for prediction); ``False`` otherwise.
+        """
+        raise NotImplementedError(
+            "Checking if a distilled trainer is trained is not "
+            "implemented yet.")
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+    ) -> TrainingJobResult:
+        pass
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: Literal[False],
+    ) -> TrainingJobResult:
+        pass
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: Literal[True],
+    ) -> TrainingJob:
+        pass
+
+    @overload
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: bool,
+    ) -> Union[TrainingJob, TrainingJobResult]:
+        pass
+
+    def fit(
+        self,
+        graph: Graph,
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: bool = False,
+        custom_tags: Mapping[str, str] = {},
+    ) -> Union[TrainingJob, TrainingJobResult]:
+        r"""Fits a model to the specified graph and training table, with the
+        strategy defined by :class:`DistilledTrainer`'s :obj:`model_plan`.
+
+        Args:
+            graph: The :class:`~kumoai.graph.Graph` object that represents the
+                tables and relationships that Kumo will learn from.
+            train_table: The :class:`~kumoai.pquery.TrainingTable`, or
+                in-progress :class:`~kumoai.pquery.TrainingTableJob`, that
+                represents the training data produced by a
+                :class:`~kumoai.pquery.PredictiveQuery` on :obj:`graph`.
+            non_blocking: Whether this operation should return immediately
+                after launching the training job, or await completion of the
+                training job.
+            custom_tags: Additional, customer defined k-v tags to be associated
+                with the job to be launched. Job tags are useful for grouping
+                and searching jobs.
+
+        Returns:
+            Union[TrainingJobResult, TrainingJob]:
+                If ``non_blocking=False``, returns a training job object. If
+                ``non_blocking=True``, returns a training job future object.
+        """
+        # TODO(manan, siyang): remove soon:
+        job_id = train_table.job_id
+        assert job_id is not None
+
+        train_table_job_api = global_state.client.generate_train_table_job_api
+        pq_id = train_table_job_api.get(job_id).config.pquery_id
+        assert pq_id is not None
+
+        custom_table = None
+        if isinstance(train_table, TrainingTable):
+            custom_table = train_table._custom_train_table
+
+        # NOTE the backend implementation currently handles sequentialization
+        # between a training table future and a training job; that is, if the
+        # training table future is still executing, the backend will wait on
+        # the job ID completion before executing a training job. This preserves
+        # semantics for both futures, ensures that Kumo works as expected if
+        # used only via REST API, and allows us to avoid chaining calllbacks
+        # in an ugly way here:
+        api = global_state.client.distillation_job_api
+        self._training_job_id = api.create(
+            DistillationJobRequest(
+                dict(custom_tags),
+                pquery_id=pq_id,
+                base_training_job_id=self.base_training_job_id,
+                distilled_model_plan=self.model_plan,
+                graph_snapshot_id=graph.snapshot(non_blocking=non_blocking),
+                train_table_job_id=job_id,
+                custom_train_table=custom_table,
+            ))
+
+        out = TrainingJob(job_id=self._training_job_id)
+        if non_blocking:
+            return out
+        return out.attach()
+
+    @classmethod
+    def _load_from_job(
+        cls,
+        job: DistillationJobResource,
+    ) -> 'DistillationTrainer':
+        trainer = cls(job.config.distilled_model_plan,
+                      job.config.base_training_job_id)
+        trainer._training_job_id = job.job_id
+        return trainer
+
+    @classmethod
+    def load(cls, job_id: TrainingJobID) -> 'DistillationTrainer':
+        r"""Creates a :class:`~kumoai.trainer.Trainer` instance from a training
+        job ID.
+        """
+        raise NotImplementedError(
+            "Loading a distilled trainer from a job ID is not implemented yet."
+        )
+
+    @classmethod
+    def load_from_tags(cls, tags: Mapping[str, str]) -> 'DistillationTrainer':
+        raise NotImplementedError(
+            "Loading a distilled trainer from tags is not implemented yet.")

kumoai/utils/__init__.py

@@ -1,10 +1,11 @@
-from .progress_logger import ProgressLogger, InteractiveProgressLogger
+from .sql import quote_ident
+from .progress_logger import ProgressLogger
 from .forecasting import ForecastVisualizer
 from .datasets import from_relbench
 
 __all__ = [
+    'quote_ident',
     'ProgressLogger',
-    'InteractiveProgressLogger',
     'ForecastVisualizer',
     'from_relbench',
 ]

kumoai/utils/display.py

@@ -0,0 +1,51 @@
+from collections.abc import Sequence
+
+import pandas as pd
+
+from kumoai import in_notebook, in_snowflake_notebook
+
+
+def message(msg: str) -> None:
+    msg = msg.replace("`", "'") if not in_notebook() else msg
+
+    if in_snowflake_notebook():
+        import streamlit as st
+        st.markdown(msg)
+    elif in_notebook():
+        from IPython.display import Markdown, display
+        display(Markdown(msg))
+    else:
+        print(msg)
+
+
+def title(msg: str) -> None:
+    message(f"### {msg}" if in_notebook() else f"{msg}:")
+
+
+def italic(msg: str) -> None:
+    message(f"*{msg}*" if in_notebook() else msg)
+
+
+def unordered_list(items: Sequence[str]) -> None:
+    if in_notebook():
+        msg = '\n'.join([f"- {item}" for item in items])
+    else:
+        msg = '\n'.join([f"• {item.replace('`', '')}" for item in items])
+    message(msg)
+
+
+def dataframe(df: pd.DataFrame) -> None:
+    if in_snowflake_notebook():
+        import streamlit as st
+        st.dataframe(df, hide_index=True)
+    elif in_notebook():
+        from IPython.display import display
+        try:
+            if hasattr(df.style, 'hide'):
+                display(df.style.hide(axis='index'))  # pandas=2
+            else:
+                display(df.style.hide_index())  # pandas<1.3
+        except ImportError:
+            print(df.to_string(index=False))  # missing jinja2
+    else:
+        print(df.to_string(index=False))