kumoai 2.14.0.dev202512211732__cp313-cp313-win_amd64.whl → 2.15.0.dev202601151732__cp313-cp313-win_amd64.whl

kumoai/experimental/rfm/task_table.py

@@ -0,0 +1,292 @@
+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
+
+        self._query: str = ''  # A description of the task, e.g., for XAI.
+
+    @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/training_table.py

@@ -199,6 +199,7 @@ class TrainingTable:
         self,
         source_table_type: SourceTableType,
         train_table_mod: TrainingTableSpec,
+        extensive_validation: bool = False,
     ) -> None:
         r"""Validates the modified training table.
 
@@ -206,6 +207,8 @@ class TrainingTable:
             source_table_type: The source table to be used as the modified
                 training table.
             train_table_mod: The modification specification.
+            extensive_validation: Enable extensive validation for custom
+                table.
 
         Raises:
             ValueError: If the modified training table is invalid.
@@ -215,7 +218,8 @@ class TrainingTable:
             global_state.client.generate_train_table_job_api)
         response = api.validate_custom_train_table(self.job_id,
                                                    source_table_type,
-                                                   train_table_mod)
+                                                   train_table_mod,
+                                                   extensive_validation)
         if not response.ok:
             raise ValueError("Invalid weighted train table",
                              response.error_message)
@@ -225,6 +229,7 @@ class TrainingTable:
         source_table: SourceTable,
         train_table_mod: TrainingTableSpec,
         validate: bool = True,
+        extensive_validation: bool = False,
     ) -> Self:
         r"""Sets the `source_table` as the modified training table.
 
@@ -243,6 +248,9 @@ class TrainingTable:
             train_table_mod: The modification specification.
             validate: Whether to validate the modified training table. This can
                 be slow for large tables.
+            extensive_validation: Whether to validate number of rows in
+                existing and modified training table.
+                It can be slow for large tables.
         """
         if isinstance(source_table.connector, S3Connector):
             # Special handling for s3 as `source_table._to_api_source_table`
@@ -252,7 +260,13 @@ class TrainingTable:
         else:
             source_table_type = source_table._to_api_source_table()
         if validate:
-            self.validate_custom_table(source_table_type, train_table_mod)
+            if extensive_validation:
+                logger.warning(
+                    "You have opted in to perform extensive validation on"
+                    " your custom training table."
+                    " This operation can be slow for large tables.")
+            self.validate_custom_table(source_table_type, train_table_mod,
+                                       extensive_validation)
         self._custom_train_table = CustomTrainingTable(
             source_table=source_table_type, table_mod_spec=train_table_mod,
             validated=validate)

kumoai/testing/snow.py

@@ -10,7 +10,7 @@ def connect(
     id: str,
     account: str,
     user: str,
-    warehouse: str,
+    warehouse: str | None = None,
     database: str | None = None,
     schema: str | None = None,
 ) -> Connection:
@@ -42,8 +42,8 @@ def connect(
     return _connect(
         account=account,
         user=user,
-        warehouse='WH_XS',
-        database='KUMO',
+        warehouse=warehouse or 'WH_XS',
+        database=database or '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/display.py

@@ -0,0 +1,87 @@
+from collections.abc import Sequence
+
+import pandas as pd
+from rich import box
+from rich.console import Console
+from rich.table import Table
+from rich.text import Text
+
+from kumoai import in_notebook, in_snowflake_notebook
+
+
+def message(msg: str) -> None:
+    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.replace("`", "'"))
+
+
+def title(msg: str) -> None:
+    if in_notebook():
+        message(f"### {msg}")
+    else:
+        msg = msg.replace("`", "'")
+        Console().print(f"[bold]{msg}[/bold]", highlight=False)
+
+
+def italic(msg: str) -> None:
+    if in_notebook():
+        message(f"*{msg}*")
+    else:
+        msg = msg.replace("`", "'")
+        Console().print(
+            f"[italic]{msg}[/italic]",
+            highlight=False,
+            style='dim',
+        )
+
+
+def unordered_list(items: Sequence[str]) -> None:
+    if in_notebook():
+        msg = '\n'.join([f"- {item}" for item in items])
+        message(msg)
+    else:
+        text = Text('\n').join(
+            Text.assemble(
+                Text(' • ', style='yellow'),
+                Text(item.replace('`', '')),
+            ) for item in items)
+        Console().print(text, highlight=False)
+
+
+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:
+        Console().print(to_rich_table(df))
+
+
+def to_rich_table(df: pd.DataFrame) -> Table:
+    table = Table(box=box.ROUNDED)
+    for column in df.columns:
+        table.add_column(str(column))
+    for _, row in df.iterrows():
+        values: list[str | Text] = []
+        for value in row:
+            if str(value) == 'True':
+                values.append('✅')
+            elif str(value) in {'False', '-'}:
+                values.append(Text('-', style='dim'))
+            else:
+                values.append(str(value))
+        table.add_row(*values)
+    return table

kumoai/utils/progress_logger.py

@@ -24,6 +24,7 @@ class ProgressLogger:
     def __init__(self, msg: str, verbose: bool = True) -> None:
         self.msg = msg
         self.verbose = verbose
+        self.depth = 0
 
         self.logs: list[str] = []
 
@@ -55,10 +56,13 @@ class ProgressLogger:
         pass
 
     def __enter__(self) -> Self:
-        self.start_time = time.perf_counter()
+        self.depth += 1
+        if self.depth == 1:
+            self.start_time = time.perf_counter()
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.depth -= 1
         self.end_time = time.perf_counter()
 
     def __repr__(self) -> str:
@@ -123,6 +127,9 @@ class RichProgressLogger(ProgressLogger):
 
         super().__enter__()
 
+        if self.depth > 1:
+            return self
+
         if not in_notebook():  # Render progress bar in TUI.
             sys.stdout.write("\x1b]9;4;3\x07")
             sys.stdout.flush()
@@ -142,6 +149,9 @@ class RichProgressLogger(ProgressLogger):
 
         super().__exit__(exc_type, exc_val, exc_tb)
 
+        if self.depth > 1:
+            return
+
         if exc_type is not None:
             self._exception = True
 
@@ -213,6 +223,9 @@ class StreamlitProgressLogger(ProgressLogger):
 
         import streamlit as st
 
+        if self.depth > 1:
+            return self
+
         # Adjust layout for prettier output:
         st.markdown(STREAMLIT_CSS, unsafe_allow_html=True)
 
@@ -253,7 +266,7 @@ class StreamlitProgressLogger(ProgressLogger):
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         super().__exit__(exc_type, exc_val, exc_tb)
 
-        if not self.verbose or self._status is None:
+        if not self.verbose or self._status is None or self.depth > 1:
             return
 
         label = f'{self._sanitize_text(self.msg)} ({self.duration:.2f}s)'