.. _faq_training:

Training Models and Generating Predictions
===========================================

Here, we discuss the final step of the Kumo workflow: training a model, and
generating predictions. The primary interface used here is
:class:`~kumoai.trainer.Trainer`, which you may already be familiar with if you
have worked with other machine learning libraries (*e.g.*
`scikit-learn <https://scikit-learn.org/stable/index.html>`__) before. A
:class:`~kumoai.trainer.Trainer` has two important methods:

* :meth:`~kumoai.trainer.Trainer.fit`, which takes a
  :class:`~kumoai.graph.Graph` and :class:`~kumoai.pquery.TrainingTable`
  (or a :class:`~kumoai.pquery.TrainingTableJob`, if the training table
  was generated in a non-blocking manner), and trains a model on this graph
  and training table.
* :meth:`~kumoai.trainer.Trainer.predict`, which takes a
  :class:`~kumoai.graph.Graph` and :class:`~kumoai.pquery.PredictionTable`
  (or a :class:`~kumoai.pquery.PredictionTableJob`, if the prediction table
  was generated in a non-blocking manner), a job ID corresponding to a trained
  model, and other parameters detailing where to output the predictions. It
  generates predictions for each entity in the prediction table, and writes
  the outputs to the specified output connector.

.. note::

    Training a model is fully customizable, with a detailed suite of
    `model plan <https://docs.kumo.ai/docs/advanced-operations>`__ options. For a guide on tuning your model for optimal
    performance, see `here <https://docs.kumo.ai/docs/debugging-poor-model-performance>`__.

.. note::

    You can view all your launched jobs in the Kumo UI, at the URL
    ``https://<customer_id>.kumoai.cloud/jobs``. Jobs are keyed by their
    unique job ID, and contain all specified job tags as well.

.. contents:: FAQ
    :local:

How do I create a Trainer? What's a model plan?
-----------------------------------------------

Creating a :class:`~kumoai.trainer.Trainer` object requires a model plan, which
defines the search space to be used when exploring model configurations for
model training.

You can suggest a model plan for your predictive query with
:meth:`~kumoai.pquery.PredictiveQuery.suggest_model_plan`, which will produce
an object of type :class:`~kumoai.trainer.ModelPlan`:

.. code-block:: python

    pquery = kumoai.PredictiveQuery(graph=..., query="...")
    model_plan = pquery.suggest_model_plan()

    print(model_plan)

The model plan can be edited with full granularity; see
`here <https://docs.kumo.ai/docs/advanced-operations>`__ for documentation,
and the :class:`~kumoai.trainer.ModelPlan` object for the exposed customizable
attributes.

Once you have customized your model plan to your liking, you can create a
:class:`~kumoai.trainer.Trainer` by simply passing the model plan in:

.. code-block:: python

    trainer = kumoai.Trainer(mdoel_plan)

That's all!

How do I train a model?
-----------------------

Training a model amounts to calling :meth:`~kumoai.trainer.Trainer.fit`, which
accepts the following arguments:

* A :class:`~kumoai.graph.Graph`, which defines the data that the model will
  be trained on. Note if you have already called
  :meth:`~kumoai.graph.Graph.snapshot`, this snapshot of the data will be
  used when training your model.
* A :class:`~kumoai.pquery.TrainingTable` or
  :class:`~kumoai.pquery.TrainingTableJob`, generated by
  :meth:`~kumoai.pquery.PredictiveQuery.generate_training_table`. This
  defines the training examples that will be used by the model; if a
  :class:`~kumoai.pquery.TrainingTableJob` is passed, its execution
  will be sequenced before training by the Kumo platform.
* :obj:`non_blocking`, which can be set to ``True`` if you would like to
  schedule training and return immediately, or ``False`` if you would like
  to wait for training to complete.
* :obj:`custom_tags`, which define a custom mapping of key/value tags that
  you can use to label your training job.

Training will raise return a :class:`~kumoai.trainer.TrainingJobResult` if
``non_blocking=False`` and training completes successfully, or a
:class:`~kumoai.trainer.TrainingJob` if ``non_blocking=True``. Each
training job is associated with a unique Job ID, starting with
``trainingjob-``.

An example invocation of :meth:`~kumoai.trainer.Trainer.fit` is as follows:

.. code-block:: python

    graph = kumoai.Graph(...)
    pquery = kumoai.PredictiveQuery(graph=graph, query="...")

    # Generate the training table, but do not wait for its completion; just
    # schedule it using `non_blocking=True`:
    training_table_plan = pquery.suggest_training_table_plan()
    training_table = pquery.generate_training_table(
        training_table_plan, non_blocking=True)

    # Create a trainer with a suggested model plan:
    model_plan = pquery.suggest_model_plan()
    trainer = kumoai.Trainer(model_plan)

    # Schedule a training job (`non_blocking=True`) given on the defined graph
    # and training table future:
    training_job_future = trainer.fit(
        graph=graph,
        train_table=training_table,
        non_blocking=True,
        custom_tags={'author': 'trial'},  # any custom key/value pairs
    )

    # Print the training job ID:
    print(f"Training job ID: {training_job_future.id}")

    # Attach to the training job to watch its status and see logs (you can
    # detach anytime without canceling the job):
    training_job_future.attach()


How do I view the metrics and artifacts of a trained model?
-----------------------------------------------------------

Recall that a trained model is represented by a
:class:`~kumoai.trainer.TrainingJobResult` object; if you have a
:class:`~kumoai.trainer.TrainingJob`, you need to await its completion
by calling :meth:`~kumoai.trainer.TrainingJob.result` before proceeding.

A :class:`~kumoai.trainer.TrainingJobResult` exposes numerous methods to help
analyze the performance of a trained model, including
:meth:`~kumoai.trainer.TrainingJobResult.metrics` and
:meth:`~kumoai.trainer.TrainingJobResult.holdout_df`. A full set of
visualizations, performance graphs, and explainability can all be accessed
at the URL specified by :py:attr:`~kumoai.trainer.TrainingJobResult.tracking_url`.

How do I generate predictions?
------------------------------

Predicting on a trained model amounts to calling
:meth:`~kumoai.trainer.Trainer.predict`, which accepts the following arguments:

* A :class:`~kumoai.graph.Graph`, which defines the data that the model will
  use to make predictions on. Note if you have already called
  :meth:`~kumoai.graph.Graph.snapshot`, this snapshot of the data will be
  used when generating predictions.
* A :class:`~kumoai.pquery.PredictionTable` or
  :class:`~kumoai.pquery.PredictionTableJob`, generated by
  :meth:`~kumoai.pquery.PredictiveQuery.generate_prediction_table` or supplied
  via a custom path. This defines the prediction examples that will be used by
  the model; if a :class:`~kumoai.pquery.PredictionTableJob` is passed, its
  execution will be sequenced before prediction by the Kumo platform.
* :obj:`training_job_id`, which defines the job ID of the training job whose
  model will be used for making predictions.
* :obj:`non_blocking`, which can be set to ``True`` if you would like to
  schedule prediction and return immediately, or ``False`` if you would like
  to wait for prediction to complete.
* :obj:`custom_tags`, which define a custom mapping of key/value tags that
  you can use to label your training job.
* *additional arguments documented in* :meth:`~kumoai.trainer.Trainer.predict`
  that can be used to specify where predictions should be output to.

Prediction will raise return a :class:`~kumoai.trainer.TrainingJobResult` if
``non_blocking=False`` and training completes successfully, or a
:class:`~kumoai.trainer.TrainingJob` if ``non_blocking=True``. Each
batch prediction job is associated with a unique Job ID, starting with
``bp-job-``.


An example invocation of :meth:`~kumoai.trainer.Trainer.predict` is as follows:

.. code-block:: python

    # Assume we have a completed training job id:
    completed_job_id = "<completed_training_job_id>"

    # Output connector:
    output_connector = ...  # any Kumo Connector

    # Load the trainer and predictive query from a completed training job:
    trainer = kumoai.Trainer.load(completed_job_id)
    pquery = kumoai.PredictiveQuery.load_from_training_job(completed_job_id)

    # Generate the prediction table, but do not wait for its completion; just
    # schedule it using `non_blocking=True`:
    prediction_table_plan = pquery.suggest_prediction_table_plan()
    prediction_table = pquery.generate_prediction_table(
        prediction_table_plan, non_blocking=True)

    # Schedule a prediction job (`non_blocking=True`) given on the defined
    # graph and prediction table future:

    # For v1.4 and above:
    from kumoai.artifact_export.config import OutputConfig
    # For v1.3 and below (backward compatible):
    # from kumoai.trainer.config import OutputConfig

    prediction_job_future = trainer.predict(
        graph=graph,
        prediction_table=prediction_table,
        training_job_id=completed_job_id,
        non_blocking=True,
        custom_tags={'author': 'trial'},  # any custom key/value pairs
        output_config=OutputConfig(
            output_types={'predictions', 'embeddings'},
            output_connector=output_connector,
            output_table_name='kumo_predictions',
        ),
    )

    # Print the prediction job ID:
    print(f"Prediction job ID: {prediction_job_future.id}")

    # Attach to the prediction job to watch its status and see logs (you can
    # detach anytime without canceling the job):
    prediction_job_future.attach()


How do I poll a training or prediction job's status?
----------------------------------------------------------

Any job scheduled with ``non_blocking=True`` will be represented as a
Future object, that has various methods to poll the scheduled job for its
status or completion. Common patterns include:

* Querying ``future.status()`` for the status of the scheduled job in a loop
* Calling ``future.attach()`` to attach to the future and print logs
  periodically; when the future is complete, this method will return the
  resolved output (*e.g.* :class:`~kumoai.trainer.TrainingJob` becomes
  :class:`~kumoai.trainer.TrainingJobResult`)
* Calling ``future.result()`` will block until the future is complete,
  and return the resolved output.
