# __validation_fn_simple_start__
import os
import torch

import ray.train
import ray.data

# Define Ray Data validation dataset outside validation function because it is not json serializable
validation_dataset = ...


def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
    # Load the checkpoint
    model = ...
    with checkpoint.as_directory() as checkpoint_dir:
        model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
        model.load_state_dict(model_state_dict)
    model.eval()

    # Perform validation on the data
    total_accuracy = 0
    with torch.no_grad():
        for batch in validation_dataset.iter_torch_batches(batch_size=128):
            images, labels = batch["image"], batch["label"]
            outputs = model(images)
            total_accuracy += (outputs.argmax(1) == labels).sum().item()
    return {"score": total_accuracy / len(validation_dataset)}
# __validation_fn_simple_end__

# __validation_fn_torch_trainer_start__
import torchmetrics
from torch.nn import CrossEntropyLoss

import ray.train.torch
from ray.data import ExecutionOptions


def eval_only_train_fn(config_dict: dict) -> dict:
    # Load the checkpoint
    model = ...
    with config_dict["checkpoint"].as_directory() as checkpoint_dir:
        model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
        model.load_state_dict(model_state_dict)
    model.cuda().eval()

    # Set up metrics and data loaders
    criterion = CrossEntropyLoss()
    mean_valid_loss = torchmetrics.MeanMetric().cuda()
    test_data_shard = ray.train.get_dataset_shard("validation")
    test_dataloader = test_data_shard.iter_torch_batches(batch_size=128)

    # Compute metric and return it directly from the train function
    with torch.no_grad():
        for batch in test_dataloader:
            images, labels = batch["image"], batch["label"]
            outputs = model(images)
            loss = criterion(outputs, labels)
            mean_valid_loss(loss)
    return {"score": mean_valid_loss.compute().item()}


def validation_fn(checkpoint: ray.train.Checkpoint, train_run_name: str, epoch: int) -> dict:
    trainer = ray.train.torch.TorchTrainer(
        eval_only_train_fn,
        train_loop_config={"checkpoint": checkpoint},
        scaling_config=ray.train.ScalingConfig(
            num_workers=2, use_gpu=True, accelerator_type="A10G"
        ),
        # Give unique name to validation run so it does not attempt to load placeholder checkpoint.
        # Also allows you to better associate training runs with validation runs.
        run_config=ray.train.RunConfig(
            name=f"{train_run_name}_validation_epoch_{epoch}"
        ),
        # Use weaker GPUs for validation
        datasets={"validation": validation_dataset},
        # Pin to the "validation" subcluster so it doesn't compete with
        # training. See https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
        dataset_config=ray.train.DataConfig(
            execution_options={
                "validation": ExecutionOptions(
                    label_selector={"ray-subcluster": "validation"}
                ),
            },
        ),
    )
    result = trainer.fit()
    # return_value holds the value returned by train function of worker 0
    return result.return_value
# __validation_fn_torch_trainer_end__

# __validation_fn_map_batches_start__
import ray.data


class Predictor:
    def __init__(self, checkpoint: ray.train.Checkpoint):
        self.model = ...
        with checkpoint.as_directory() as checkpoint_dir:
            model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
            self.model.load_state_dict(model_state_dict)
        self.model.cuda().eval()

    def __call__(self, batch: dict) -> dict:
        image = torch.as_tensor(batch["image"], dtype=torch.float32, device="cuda")
        label = torch.as_tensor(batch["label"], dtype=torch.float32, device="cuda")
        pred = self.model(image)
        return {"res": (pred.argmax(1) == label).cpu().numpy()}


# Construct ``validation_dataset`` under a DataContext copy pinned to the
# "validation" subcluster. ``Dataset.context`` is a deep copy of the
# current context taken at construction, so the selector is baked in and
# every downstream operator (including the ``map_batches`` below) inherits
# it — no in-function mutation needed. See
# https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "validation"}
with ray.data.DataContext.current(ctx):
    validation_dataset = ray.data.read_parquet(...)


def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
    # Set name to avoid confusion; default name is "Dataset"
    validation_dataset.set_name("validation")
    eval_res = validation_dataset.map_batches(
        Predictor,
        batch_size=128,
        num_gpus=1,
        fn_constructor_kwargs={"checkpoint": checkpoint},
        concurrency=2,
    )
    mean = eval_res.mean(["res"])
    return {
        "score": mean,
    }
# __validation_fn_map_batches_end__

# __validation_fn_report_start__
import tempfile

from ray.data import ExecutionOptions
from ray.train import ValidationConfig, ValidationTaskConfig


def train_func(config: dict) -> None:
    ...
    epochs = ...
    model = ...
    rank = ray.train.get_context().get_world_rank()
    for epoch in epochs:
        ...  # training step
        if rank == 0:
            training_metrics = {"loss": ..., "epoch": epoch}
            local_checkpoint_dir = tempfile.mkdtemp()
            torch.save(
                model.module.state_dict(),
                os.path.join(local_checkpoint_dir, "model.pt"),
            )
            ray.train.report(
                training_metrics,
                checkpoint=ray.train.Checkpoint.from_directory(local_checkpoint_dir),
                checkpoint_upload_mode=ray.train.CheckpointUploadMode.ASYNC,
                validation=ValidationTaskConfig(fn_kwargs={
                    "train_run_name": ray.train.get_context().get_experiment_name(),
                    "epoch": epoch,
                }),
            )
        else:
            ray.train.report({}, None)


def run_trainer() -> ray.train.Result:
    # 1) Construction-time tasks (parquet schema inference, file listing)
    # read the current DataContext. Pin them to "training" with a copy of
    # the DataContext applied via the DataContext.current() context
    # manager — scoped to the `with` block so it doesn't leak. See
    # https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
    ctx = ray.data.DataContext.get_current().copy()
    ctx.execution_options.label_selector = {"ray-subcluster": "training"}
    with ray.data.DataContext.current(ctx):
        train_dataset = ray.data.read_parquet(...)

    trainer = ray.train.torch.TorchTrainer(
        train_func,
        validation_config=ValidationConfig(fn=validation_fn),
        # Pass training dataset in datasets arg to split it across training workers
        datasets={"train": train_dataset},
        # 2) DataConfig.execution_options REPLACES ds.context.execution_options
        # wholesale at training start, dropping anything not re-specified
        # (including label_selector). Restate the selector here so per-worker
        # ingest stays pinned to "training".
        dataset_config=ray.train.DataConfig(
            datasets_to_split=["train"],
            execution_options={
                "train": ExecutionOptions(
                    label_selector={"ray-subcluster": "training"}
                ),
            },
        ),
        scaling_config=ray.train.ScalingConfig(
            num_workers=2,
            use_gpu=True,
            # Use powerful GPUs for training
            accelerator_type="A100",
        ),
    )
    return trainer.fit()
# __validation_fn_report_end__

# __exp_tracking_same_run_wandb_start__
import wandb
import ray.train
from ray.train import ValidationConfig, ValidationTaskConfig


entity = "my_entity"
project = "my_project"
num_epochs = ...


def validation_fn(checkpoint: ray.train.Checkpoint, wandb_run_id: str, val_step: int) -> dict:
    wandb.init(
        entity=entity,
        project=project,
        settings=wandb.Settings(mode="shared", x_primary=False),
        id=wandb_run_id,
    )
    score = ...
    wandb.log({"validation/loss": score, "val_step": val_step})
    wandb.finish()  # flush the metrics
    return {"validation/loss": score}


def train_func():
    if ray.train.get_context().get_world_rank() == 0:
        run = wandb.init(
            entity=entity,
            project=project,
            settings=wandb.Settings(mode="shared", x_primary=True,)
        )
        wandb.define_metric("val_step", hidden=True)
        wandb.define_metric("train_step", hidden=True)
        wandb.define_metric("validation/loss", step_metric="val_step")
        wandb.define_metric("train/loss", step_metric="train_step")

    for epoch in range(num_epochs):
        loss = ...
        if ray.train.get_context().get_world_rank() == 0:
            wandb.log({"train/loss": loss, "train_step": epoch})
            checkpoint = ...
            ray.train.report(
                {"train/loss": loss},
                checkpoint=checkpoint,
                validation=ValidationTaskConfig(
                    fn_kwargs={"wandb_run_id": run.id, "val_step": epoch}
                ),
            )
        else:
            ray.train.report({}, None)

    if ray.train.get_context().get_world_rank() == 0:
        wandb.finish()


# __exp_tracking_same_run_wandb_end__

# __exp_tracking_same_run_mlflow_start__
import mlflow
from mlflow.tracking import MlflowClient
import ray.train
from ray.train import ValidationConfig, ValidationTaskConfig


tracking_uri = "my_uri"
experiment_name = "my_experiment"
num_epochs = ...

def validation_fn(
    checkpoint: ray.train.Checkpoint, mlflow_run_id: str, val_step: int
) -> dict:
    client = MlflowClient(tracking_uri=tracking_uri)
    score = ...
    client.log_metric(mlflow_run_id, "val_score", score, step=val_step)
    return {"val_score": score}


def train_func():
    if ray.train.get_context().get_world_rank() == 0:
        client = MlflowClient(tracking_uri=tracking_uri)
        experiment = client.get_experiment_by_name(experiment_name)
        run = client.create_run(experiment_id=experiment.experiment_id)

    for epoch in range(num_epochs):
        loss = ...
        if ray.train.get_context().get_world_rank() == 0:
            client.log_metric(run.info.run_id, "train_loss", loss, step=epoch)
            checkpoint = ...
            ray.train.report(
                {"train_loss": loss},
                checkpoint=checkpoint,
                validation=ValidationTaskConfig(
                    fn_kwargs={"mlflow_run_id": run.info.run_id, "val_step": epoch}
                ),
            )
        else:
            ray.train.report({}, None)

    if ray.train.get_context().get_world_rank() == 0:
        client.set_terminated(run.info.run_id)

# __exp_tracking_same_run_mlflow_end__