kumoai 2.13.0.dev202511131731__cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/pquery/predictive_query.py

@@ -0,0 +1,637 @@
+import logging
+from typing import List, Literal, Mapping, Optional, Tuple, Union, overload
+
+from kumoapi.jobs import (
+    GeneratePredictionTableRequest,
+    GenerateTrainTableRequest,
+)
+from kumoapi.model_plan import (
+    InferredType,
+    PredictionTableGenerationPlan,
+    RunMode,
+    SuggestModelPlanRequest,
+    TrainingTableGenerationPlan,
+)
+from kumoapi.pquery import PQueryResource
+from kumoapi.task import TaskType
+from kumoapi.train import TrainingTableSpec
+from typing_extensions import Self
+
+from kumoai import global_state
+from kumoai.client.jobs import (
+    GeneratePredictionTableJobID,
+    GenerateTrainTableJobID,
+    TrainingJobAPI,
+)
+from kumoai.graph import Graph
+from kumoai.pquery.prediction_table import PredictionTable, PredictionTableJob
+from kumoai.pquery.training_table import TrainingTable, TrainingTableJob
+from kumoai.trainer import (
+    BaselineTrainer,
+    ModelPlan,
+    Trainer,
+    TrainingJob,
+    TrainingJobResult,
+)
+from kumoai.trainer.job import BaselineJob, BaselineJobResult
+
+logger = logging.getLogger(__name__)
+
+PredictiveQueryID = str
+
+
+class PredictiveQuery:
+    r"""The Kumo predictive query is a declarative syntax for describing a
+    machine learning task. Predictive queries are written using the predictive
+    query language (PQL), a concise SQL-like syntax that allows you to define a
+    model for a new business problem.
+
+    A predictive query object can be created from a
+    :class:`~kumoai.graph.Graph` and a query string. For information on the
+    construction of a query string, please visit the Kumo
+    `documentation <https://docs.kumo.ai/docs/pquery-structure/>`__.
+
+    .. code-block:: python
+
+        import kumoai
+
+        # See `Graph` documentation for more information:
+        graph = kumoai.Graph(...)
+
+        # Create a predictive query representing a machine learning problem
+        # over this Graph:
+        pquery = kumoai.PredictiveQuery(
+            graph=graph,
+            query=(
+                "PREDICT MAX(transaction.Quantity, 0, 30) "
+                "FOR EACH customer.CustomerID"
+            ),
+        )
+
+        # Validate the predictive query configuration, for syntax and
+        # correctness:
+        pquery.validate(verbose=True)
+
+        # Get the machine learning task type corresponding to this predictive
+        # query (e.g. binary classification, regression, link prediction, etc.)
+        print(pquery.get_task_type())
+
+        # Suggest a training table generation plan and use it to generate a
+        # training table from this query, to be used in `Trainer.fit`:
+        training_table_plan = pquery.suggest_training_table_plan()
+        training_table = pquery.generate_training_table(training_table_plan)
+
+        # Suggest a prediction table generation plan and use it to generate a
+        # prediction table from this query, to be used in `Trainer.predict`:
+        pred_table_plan = pquery.suggest_prediction_table_plan()
+        pred_table = pquery.generate_prediction_table(pred_table_plan)
+
+    Args:
+        graph: The :class:`~kumoai.graph.Graph` object which the predictive
+            query is defined over.
+        query: A string representation of the predictive query.
+    """
+    def __init__(
+        self,
+        graph: Graph,
+        query: str,
+    ) -> None:
+        self.graph = graph
+        self.query = query
+
+        # A predictive query owns a trainer object, which is used internally
+        # to support `fit` and `predict` directly on this object. A user can
+        # also inspect the training table, prediction table, and trainer
+        # objects, but cannot set them; any advanced configuration must be
+        # done directly via `Trainer`:
+        self._train_table: Optional[Union[TrainingTable,
+                                          TrainingTableJob]] = None
+        self._prediction_table: Optional[Union[PredictionTable,
+                                               PredictionTableJob]] = None
+
+    # Metadata ################################################################
+
+    @property
+    def id(self) -> str:
+        r"""Returns the unique ID for this predictive query, determined from
+        its schema and the schema of its associated graph. Two queries that
+        differ either in their syntax or in their graph will have different
+        ids.
+        """
+        return self.save()
+
+    @property
+    def train_table(self) -> Union[TrainingTable, TrainingTableJob]:
+        r"""Returns the training table that was last generated by this
+        predictive query. If the predictive query has not yet generated a
+        training table, raises a :class:`ValueError`.
+
+        Note that the training table may be of type
+        :class:`~kumoai.pquery.TrainingTable` or
+        :class:`~kumoai.pquery.TrainingTableJob`, depending on whether the
+        training table was generated with or without waiting for its
+        completion, respectively.
+        """
+        if not self._train_table:
+            raise ValueError(
+                "This predictive query has not yet generated a training "
+                "table. Please call `generate_training_table` to generate "
+                "a training table before proceeding.")
+
+        return self._train_table
+
+    @property
+    def prediction_table(self) -> Union[PredictionTable, PredictionTableJob]:
+        r"""Returns the prediction table that was last generated by this
+        predictive query. If the predictive query has not yet generated a
+        prediction table, raises a :class:`ValueError`.
+
+        Note that the prediction table may be of type
+        :class:`~kumoai.pquery.PredictionTable` or
+        :class:`~kumoai.pquery.PredictionTableJob`, depending on whether the
+        prediction table was generated with or without waiting for its
+        completion, respectively.
+        """
+        if not self._prediction_table:
+            raise ValueError(
+                "This predictive query has not yet generated a prediction "
+                "table. Please call `generate_prediction_table` to generate a "
+                "prediction table before proceeding.")
+
+        return self._prediction_table
+
+    def get_task_type(self) -> TaskType:
+        r"""Returns the task type of this predictive query. The task type of
+        the query corresponds to the machine learning problem that this query
+        translates to in the Kumo platform; for more information about possible
+        task types, please visit the Kumo `documentation
+        <https://docs.kumo.ai/docs/task-types/>`__.
+        """
+        try:
+            self.validate(verbose=False)
+        except ValueError as e:
+            raise ValueError(
+                f"Predictive query {self.query} is improperly configured, so "
+                f"a task type cannot be obtained. Please ensure your query "
+                f"has a valid configuration before proceeding. You can use "
+                f"the `validate` method to verify the validity of your query."
+            ) from e
+        task_type, _ = global_state.client.pquery_api.infer_task_type(
+            pquery_string=self.query, graph_id=self.graph.id)
+        return task_type
+
+    def validate(self, verbose: bool = True) -> Self:
+        r"""Validates the syntax of this predictive query, ensuring that
+        the query is formulated correctly in Kumo's Predictive Query Language
+        and that the query makes semantic sense (defines a suitable predictive
+        problem) on this :class:`~kumoai.graph.Graph`.
+
+        Args:
+            verbose: Whether to log non-error output of this validation.
+
+        Raises:
+            ValueError:
+                if validation fails.
+
+        Example:
+            >>> import kumoai
+            >>> query = kumoai.PredictiveQuery(...)  # doctest: +SKIP
+            >>> query.validate()  # doctest: +SKIP
+            ValidationResponse(warnings=[], errors=[])
+        """
+        self.graph.save()  # Need a valid graph ID; also validates graph.
+
+        resp = global_state.client.pquery_api.validate(
+            self._to_api_pquery_resource())
+
+        if not resp.ok:
+            raise ValueError(resp.error_message())
+        if verbose:
+            if resp.empty():
+                logger.info("Query %s is configured correctly.", self.query)
+            else:
+                logger.warning(resp.message())
+
+        return self
+
+    # Persistence #############################################################
+
+    def _to_api_pquery_resource(
+        self,
+        name: Optional[str] = None,
+    ) -> PQueryResource:
+        return PQueryResource(
+            name=name,
+            query_string=self.query,
+            graph=self.graph._to_api_graph_definition(),
+            desc="",
+        )
+
+    def save(self, name: Optional[str] = None) -> PredictiveQueryID:
+        r"""Saves a predictive query to Kumo, returning a unique ID for this
+        query. If a name is provided, saves it as a named, re-usable template.
+
+        Args:
+            name: Optional name for the template. If provided, saves the
+                query as a named template. If the name already exists,
+                that template will be overwritten.
+
+        Example:
+            >>> import kumoai
+            >>> query = kumoai.PredictiveQuery(...)  # doctest: +SKIP
+            >>> query.save()  # doctest: +SKIP
+            pquery-xxx
+            >>> query.save("my_template")  # doctest: +SKIP
+            my_template
+        """
+        try:
+            self.validate(verbose=False)
+        except ValueError as e:
+            raise ValueError(
+                f"Predictive query {self.query} is improperly configured, so "
+                f"it cannot be saved. Please ensure your query "
+                f"has a valid configuration before proceeding. You can use "
+                f"the `validate` method to verify the validity of your query."
+            ) from e
+
+        if name is not None:
+            template_resource = global_state.client.pquery_api.get_if_exists(
+                name)
+            if template_resource is not None:
+                template_string = template_resource.query_string
+                logger.warning(
+                    ("Predictive query template %s already exists, with "
+                     "query string %s. This template will be overridden with "
+                     "configuration %s."), name, template_string, self.query)
+
+        self.graph.save()
+        return global_state.client.pquery_api.create(
+            pquery=self._to_api_pquery_resource(name))
+
+    @classmethod
+    def load(cls, pq_id_or_template: str) -> 'PredictiveQuery':
+        r"""Loads a predictive query from either a predictive query ID or a
+        named template. Returns a :class:`~kumoai.pquery.PredictiveQuery`
+        object that contains the loaded query along with its associated graph,
+        tables, etc.
+        """
+        api = global_state.client.pquery_api
+        res = api.get_if_exists(pq_id_or_template)
+        if not res:
+            raise ValueError(
+                f"Predictive query {pq_id_or_template} was not found.")
+        return cls(
+            graph=Graph._from_api_graph_definition(res.graph),
+            query=res.query_string,
+        )
+
+    @classmethod
+    def load_from_training_job(cls, training_job_id: str) -> 'PredictiveQuery':
+        r"""Loads a predictive query from a training job, regardless of the
+        training job's status. Returns a
+        :class:`~kumoai.pquery.PredictiveQuery` object that contains the loaded
+        query along with its associated graph, tables, etc.
+        """
+        train_api: TrainingJobAPI = global_state.client.training_job_api
+        job = train_api.get(training_job_id)
+        id_or_name = job.config.pquery_id
+        return PredictiveQuery.load(pq_id_or_template=id_or_name)
+
+    # Training & Prediction Table Generation ##################################
+
+    @overload
+    def generate_training_table(
+        self,
+        plan: Optional[TrainingTableGenerationPlan] = None,
+    ) -> TrainingTable:
+        pass
+
+    @overload
+    def generate_training_table(
+        self,
+        plan: Optional[TrainingTableGenerationPlan] = None,
+        *,
+        non_blocking: Literal[False],
+    ) -> TrainingTable:
+        pass
+
+    @overload
+    def generate_training_table(
+        self,
+        plan: Optional[TrainingTableGenerationPlan] = None,
+        *,
+        non_blocking: Literal[True],
+    ) -> TrainingTableJob:
+        pass
+
+    @overload
+    def generate_training_table(
+        self,
+        plan: Optional[TrainingTableGenerationPlan] = None,
+        *,
+        non_blocking: bool,
+    ) -> Union[TrainingTable, TrainingTableJob]:
+        pass
+
+    def generate_training_table(
+        self,
+        plan: Optional[TrainingTableGenerationPlan] = None,
+        *,
+        non_blocking: bool = False,
+        custom_tags: Mapping[str, str] = {},
+    ) -> Union[TrainingTable, TrainingTableJob]:
+        r"""Generates a training table from the specified :attr:`query`
+        string.
+
+        Args:
+            plan: A specification of the parameters for training table
+                generation. If not provided, will use an intelligently
+                generated default plan based on the query and graph. This plan
+                is equivalent to the plan inferred with
+                ``suggest_training_table_plan(run_mode=RunMode.NORMAL)``.
+            non_blocking: Whether this operation should return immediately
+                after launching the training table generation job, or await
+                completion of the generated training table.
+            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[TrainingTable, TrainingTableJob]:
+                If ``non_blocking=False``, returns a training table object. If
+                ``non_blocking=True``, returns a training table future object.
+        """
+        pq_id = self.save()
+
+        # TODO(manan): improve this...
+        if not plan:
+            plan = self.suggest_training_table_plan()
+
+        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)))
+
+        self._train_table = TrainingTableJob(job_id=job_id)
+        if non_blocking:
+            return self._train_table
+        self._train_table = self._train_table.attach()
+        return self._train_table
+
+    @overload
+    def generate_prediction_table(
+        self,
+        plan: Optional[PredictionTableGenerationPlan] = None,
+    ) -> PredictionTable:
+        pass
+
+    @overload
+    def generate_prediction_table(
+        self,
+        plan: Optional[PredictionTableGenerationPlan] = None,
+        *,
+        non_blocking: Literal[False],
+    ) -> PredictionTable:
+        pass
+
+    @overload
+    def generate_prediction_table(
+        self,
+        plan: Optional[PredictionTableGenerationPlan] = None,
+        *,
+        non_blocking: Literal[True],
+    ) -> PredictionTableJob:
+        pass
+
+    @overload
+    def generate_prediction_table(
+        self,
+        plan: Optional[PredictionTableGenerationPlan] = None,
+        *,
+        non_blocking: bool,
+    ) -> Union[PredictionTable, PredictionTableJob]:
+        pass
+
+    def generate_prediction_table(
+        self,
+        plan: Optional[PredictionTableGenerationPlan] = None,
+        *,
+        non_blocking: bool = False,
+        custom_tags: Mapping[str, str] = {},
+    ) -> Union[PredictionTable, PredictionTableJob]:
+        r"""Generates a prediction table from the predictive query
+        :attr:`query` string.
+
+        Args:
+            plan: A specification of the parameters for prediction table
+                generation. If not provided, will use an intelligently
+                generated default plan based on the query and graph. This plan
+                is equivalent to the plan inferred with
+                ``suggest_prediction_table_plan(run_mode=RunMode.NORMAL)``.
+            non_blocking: Whether this operation should return immediately
+                after launching the prediction table generation job, or await
+                completion of the generated prediction table.
+            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[PredictionTable, PredictionTableJob]:
+                If ``non_blocking=False``, returns a prediction table object.
+                If ``non_blocking=True``, returns a prediction table future
+                object.
+        """
+        pq_id = self.save()
+
+        if not plan:
+            plan = self.suggest_prediction_table_plan()
+
+        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)))
+
+        self._prediction_table = PredictionTableJob(job_id=job_id)
+        if non_blocking:
+            return self._prediction_table
+        self._prediction_table = self._prediction_table.result()
+        return self._prediction_table
+
+    # Training & Prediction ###################################################
+
+    def suggest_training_table_plan(
+        self,
+        run_mode: RunMode = RunMode.FAST,
+    ) -> TrainingTableGenerationPlan:
+        r"""Suggests a training table generation plan given the predictive
+        query and graph. This training table generation plan can be used to
+        alter the approach Kumo uses to generate the training table for your
+        predictive query.
+
+        Args:
+            run_mode: A representation of how quickly you would like your
+                predictive query to complete. Faster run modes correspond to
+                lower training times, at the cost of potentially lower
+                performance.
+        """
+        self.graph.save()
+        req = SuggestModelPlanRequest(
+            query_string=self.query,
+            graph_id=self.graph.id,
+            run_mode=run_mode,
+        )
+        return global_state.client.pquery_api.suggest_training_table_plan(req)
+
+    def suggest_prediction_table_plan(self, ) -> PredictionTableGenerationPlan:
+        r"""Suggests a prediction table generation plan given the predictive
+        query and graph. This prediction table generation plan can be used to
+        alter the approach Kumo uses to generate the prediction table for your
+        predictive query.
+        """
+        return PredictionTableGenerationPlan(anchor_time=InferredType.VALUE)
+
+    def suggest_model_plan(
+        self,
+        run_mode: RunMode = RunMode.FAST,
+        train_table_spec: Optional[TrainingTableSpec] = None,
+    ) -> ModelPlan:
+        r"""Suggests a modeling plan given the predictive query and graph. This
+        model plan can be used to alter the approach Kumo uses to train your
+        machine learning model.
+
+        Args:
+            run_mode: A representation of how quickly you would like your
+                predictive query to complete. Faster run modes correspond to
+                lower training times, at the cost of potentially lower
+                performance.
+            train_table_spec: Needed if the original train table has been
+                modified by adding a weight column.
+        """
+        self.graph.save()
+        req = SuggestModelPlanRequest(
+            query_string=self.query,
+            graph_id=self.graph.id,
+            run_mode=run_mode,
+            train_table_spec=train_table_spec,
+        )
+        return global_state.client.pquery_api.suggest_model_plan(req)
+
+    @overload
+    def fit(
+        self,
+        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
+        model_plan: Optional[ModelPlan] = None,
+    ) -> Tuple[Trainer, TrainingJobResult]:
+        pass
+
+    @overload
+    def fit(
+        self,
+        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
+        model_plan: Optional[ModelPlan] = None,
+        *,
+        non_blocking: Literal[False],
+    ) -> Tuple[Trainer, TrainingJobResult]:
+        pass
+
+    @overload
+    def fit(
+        self,
+        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
+        model_plan: Optional[ModelPlan] = None,
+        *,
+        non_blocking: Literal[True],
+    ) -> Tuple[Trainer, TrainingJob]:
+        pass
+
+    @overload
+    def fit(
+        self,
+        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
+        model_plan: Optional[ModelPlan] = None,
+        *,
+        non_blocking: bool,
+    ) -> Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]:
+        pass
+
+    def fit(
+        self,
+        training_table_plan: Optional[TrainingTableGenerationPlan] = None,
+        model_plan: Optional[ModelPlan] = None,
+        *,
+        non_blocking: bool = False,
+    ) -> Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]:
+        r"""Trains a Kumo model on this predictive query, given optional
+        additional specifications of the training table generation plan and
+        the model plan.
+
+        Args:
+            training_table_plan: A specification of the parameters for training
+                table generation. If not provided, will use an intelligently
+                generated default plan based on the query and graph. This plan
+                is equivalent to the plan inferred with
+                ``suggest_training_table_plan(run_mode=RunMode.NORMAL)``.
+            model_plan: A specification of the parameters for model training.
+                If not provided, will use an intelligently generated default
+                plan based on the query and graph. This plan
+                is equivalent to the plan inferred with
+                ``suggest_model_plan(run_mode=RunMode.NORMAL)``.
+            non_blocking: Whether this operation should return immediately
+                after launching the training job, or await completion of the
+                training job.
+
+        Returns:
+            Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]:
+                A tuple with two elements. The first element is the trainer
+                object used to launch the training job. The second element
+                is either a training job object (if ``non_blocking=True``)
+                or a training job future object (if ``non_blocking=False``).
+        """
+        # If we have already generated the training table, use it with Trainer:
+        if self._train_table is None:
+            # Nonblocking generate:
+            self._train_table = self.generate_training_table(
+                training_table_plan, non_blocking=True)
+
+        # TODO(manan): what if `self._train_table` represents a failed job?
+        model_plan = model_plan or self.suggest_model_plan()
+        trainer = Trainer(model_plan)
+        return (trainer,
+                trainer.fit(self.graph, self.train_table,
+                            non_blocking=non_blocking))
+
+    def generate_baseline(
+        self,
+        metrics: List[str],
+        train_table: Union[TrainingTable, TrainingTableJob],
+        *,
+        non_blocking: bool = False,
+    ) -> Union[BaselineJob, BaselineJobResult]:
+        r"""Runs a baseline model on this predictive query, given metrics and
+        optional additional specifications of the training table generation
+        plan.
+
+        Args:
+            metrics (List[str]): A list to metrics that baseline model will be
+                evaluated on.
+            train_table (Union[TrainingTable, TrainingTableJob]): 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 (bool): Whether this operation should
+                return immediately after launching the baseline job, or await
+                completion of the baseline job. Defaults to False.
+
+        Returns:
+            Union[BaselineJob, BaselineJobResult]:  either a baseline job
+                object (if ``non_blocking=True``) or a baseline job future
+                object (if ``non_blocking=False``).
+        """ # noqa
+        baseline_trainer = BaselineTrainer(metrics)
+        return baseline_trainer.run(self.graph, train_table,
+                                    non_blocking=non_blocking)