!pip install -qU "ray[tune]" sklearn xgboost_ray wandb

"""
Explanation: Logging results and uploading models to Weights & Biases
In this example, we train a simple XGBoost model and log the training
results to Weights & Biases. We also save the resulting model checkpoints
as artifacts.
Let's start with installing our dependencies:
End of explanation
"""


import ray

from ray.air import RunConfig
from ray.air.result import Result
from ray.train.xgboost import XGBoostTrainer
from ray.tune.integration.wandb import WandbLoggerCallback
from sklearn.datasets import load_breast_cancer

"""
Explanation: Then we need some imports:
End of explanation
"""


def get_train_dataset() -> ray.data.Dataset:
    """Return the "Breast cancer" dataset as a Ray dataset."""
    data_raw = load_breast_cancer(as_frame=True)
    df = data_raw["data"]
    df["target"] = data_raw["target"]
    return ray.data.from_pandas(df)

"""
Explanation: We define a simple function that returns our training dataset as a Ray Dataset:
End of explanation
"""


def train_model(train_dataset: ray.data.Dataset, wandb_project: str) -> Result:
    """Train a simple XGBoost model and return the result."""
    trainer = XGBoostTrainer(
        scaling_config={"num_workers": 2},
        params={"tree_method": "auto"},
        label_column="target",
        datasets={"train": train_dataset},
        num_boost_round=10,
        run_config=RunConfig(
            callbacks=[
                # This is the part needed to enable logging to Weights & Biases.
                # It assumes you've logged in before, e.g. with `wandb login`.
                WandbLoggerCallback(
                    project=wandb_project,
                    save_checkpoints=True,
                )
            ]
        ),
    )
    result = trainer.fit()
    return result

"""
Explanation: Now we define a simple training function. All the magic happens within the WandbLoggerCallback:
python
WandbLoggerCallback(
    project=wandb_project,
    save_checkpoints=True,
)
It will automatically log all results to Weights & Biases and upload the checkpoints as artifacts. It assumes you're logged in into Wandb via an API key or wandb login.
End of explanation
"""


wandb_project = "ray_air_example"

train_dataset = get_train_dataset()
result = train_model(train_dataset=train_dataset, wandb_project=wandb_project)

"""
Explanation: Let's kick off a run:
End of explanation
"""


