# flake8: noqa
# isort: skip_file

# __pytorch_save_start__
import os
import tempfile

import numpy as np
import torch
import torch.nn as nn
from torch.optim import Adam

import ray.train.torch
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer


def train_func(config):
    n = 100
    # create a toy dataset
    # data   : X - dim = (n, 4)
    # target : Y - dim = (n, 1)
    X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
    Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
    # toy neural network : 1-layer
    # Wrap the model in DDP
    model = ray.train.torch.prepare_model(nn.Linear(4, 1))
    criterion = nn.MSELoss()

    optimizer = Adam(model.parameters(), lr=3e-4)
    for epoch in range(config["num_epochs"]):
        y = model.forward(X)
        loss = criterion(y, Y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        metrics = {"loss": loss.item()}

        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            checkpoint = None

            should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
            # In standard DDP training, where the model is the same across all ranks,
            # only the global rank 0 worker needs to save and report the checkpoint
            if train.get_context().get_world_rank() == 0 and should_checkpoint:
                torch.save(
                    model.module.state_dict(),  # NOTE: Unwrap the model.
                    os.path.join(temp_checkpoint_dir, "model.pt"),
                )
                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)

            train.report(metrics, checkpoint=checkpoint)


trainer = TorchTrainer(
    train_func,
    train_loop_config={"num_epochs": 5},
    scaling_config=ScalingConfig(num_workers=2),
)
result = trainer.fit()
# __pytorch_save_end__

# __pytorch_restore_start__
import os
import tempfile

import numpy as np
import torch
import torch.nn as nn
from torch.optim import Adam

import ray.train.torch
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer


def train_func(config):
    n = 100
    # create a toy dataset
    # data   : X - dim = (n, 4)
    # target : Y - dim = (n, 1)
    X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
    Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
    # toy neural network : 1-layer
    model = nn.Linear(4, 1)
    optimizer = Adam(model.parameters(), lr=3e-4)
    criterion = nn.MSELoss()

    # Wrap the model in DDP and move it to GPU.
    model = ray.train.torch.prepare_model(model)

    # ====== Resume training state from the checkpoint. ======
    start_epoch = 0
    checkpoint = train.get_checkpoint()
    if checkpoint:
        with checkpoint.as_directory() as checkpoint_dir:
            model_state_dict = torch.load(
                os.path.join(checkpoint_dir, "model.pt"),
                # map_location=...,  # Load onto a different device if needed.
            )
            model.module.load_state_dict(model_state_dict)
            optimizer.load_state_dict(
                torch.load(os.path.join(checkpoint_dir, "optimizer.pt"))
            )
            start_epoch = (
                torch.load(os.path.join(checkpoint_dir, "extra_state.pt"))["epoch"] + 1
            )
    # ========================================================

    for epoch in range(start_epoch, config["num_epochs"]):
        y = model.forward(X)
        loss = criterion(y, Y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        metrics = {"loss": loss.item()}

        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            checkpoint = None

            should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
            # In standard DDP training, where the model is the same across all ranks,
            # only the global rank 0 worker needs to save and report the checkpoint
            if train.get_context().get_world_rank() == 0 and should_checkpoint:
                # === Make sure to save all state needed for resuming training ===
                torch.save(
                    model.module.state_dict(),  # NOTE: Unwrap the model.
                    os.path.join(temp_checkpoint_dir, "model.pt"),
                )
                torch.save(
                    optimizer.state_dict(),
                    os.path.join(temp_checkpoint_dir, "optimizer.pt"),
                )
                torch.save(
                    {"epoch": epoch},
                    os.path.join(temp_checkpoint_dir, "extra_state.pt"),
                )
                # ================================================================
                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)

            train.report(metrics, checkpoint=checkpoint)

        if epoch == 1:
            raise RuntimeError("Intentional error to showcase restoration!")


trainer = TorchTrainer(
    train_func,
    train_loop_config={"num_epochs": 5},
    scaling_config=ScalingConfig(num_workers=2),
    run_config=train.RunConfig(failure_config=train.FailureConfig(max_failures=1)),
)
result = trainer.fit()
# __pytorch_restore_end__

# __checkpoint_from_single_worker_start__
import tempfile

from ray import train


def train_fn(config):
    ...

    metrics = {...}
    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
        checkpoint = None

        # Only the global rank 0 worker saves and reports the checkpoint
        if train.get_context().get_world_rank() == 0:
            ...  # Save checkpoint to temp_checkpoint_dir

            checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)

        train.report(metrics, checkpoint=checkpoint)
# __checkpoint_from_single_worker_end__


# __lightning_save_example_start__
import lightning.pytorch as pl

from ray import train
from ray.train.lightning import RayTrainReportCallback
from ray.train.torch import TorchTrainer


class MyLightningModule(pl.LightningModule):
    # ...

    def on_validation_epoch_end(self):
        ...
        mean_acc = calculate_accuracy()
        self.log("mean_accuracy", mean_acc, sync_dist=True)


def train_func():
    ...
    model = MyLightningModule(...)
    datamodule = MyLightningDataModule(...)

    trainer = pl.Trainer(
        # ...
        callbacks=[RayTrainReportCallback()]
    )
    trainer.fit(model, datamodule=datamodule)


ray_trainer = TorchTrainer(
    train_func,
    scaling_config=train.ScalingConfig(num_workers=2),
    run_config=train.RunConfig(
        checkpoint_config=train.CheckpointConfig(
            num_to_keep=2,
            checkpoint_score_attribute="mean_accuracy",
            checkpoint_score_order="max",
        ),
    ),
)
# __lightning_save_example_end__


# __lightning_custom_save_example_start__
import os
from tempfile import TemporaryDirectory

from lightning.pytorch.callbacks import Callback

import ray
import ray.train
from ray.train import Checkpoint


class CustomRayTrainReportCallback(Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        should_checkpoint = trainer.current_epoch % 3 == 0

        with TemporaryDirectory() as tmpdir:
            # Fetch metrics from `self.log(..)` in the LightningModule
            metrics = trainer.callback_metrics
            metrics = {k: v.item() for k, v in metrics.items()}

            # Add customized metrics
            metrics["epoch"] = trainer.current_epoch
            metrics["custom_metric"] = 123

            checkpoint = None
            global_rank = ray.train.get_context().get_world_rank() == 0
            if global_rank == 0 and should_checkpoint:
                # Save model checkpoint file to tmpdir
                ckpt_path = os.path.join(tmpdir, "ckpt.pt")
                trainer.save_checkpoint(ckpt_path, weights_only=False)

                checkpoint = Checkpoint.from_directory(tmpdir)

            # Report to train session
            ray.train.report(metrics=metrics, checkpoint=checkpoint)
# __lightning_custom_save_example_end__

# __lightning_restore_example_start__
import os

from ray import train
from ray.train import Checkpoint
from ray.train.torch import TorchTrainer
from ray.train.lightning import RayTrainReportCallback


def train_func():
    model = MyLightningModule(...)
    datamodule = MyLightningDataModule(...)
    trainer = pl.Trainer(..., callbacks=[RayTrainReportCallback()])

    checkpoint = train.get_checkpoint()
    if checkpoint:
        with checkpoint.as_directory() as ckpt_dir:
            ckpt_path = os.path.join(ckpt_dir, RayTrainReportCallback.CHECKPOINT_NAME)
            trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path)
    else:
        trainer.fit(model, datamodule=datamodule)


ray_trainer = TorchTrainer(
    train_func,
    scaling_config=train.ScalingConfig(num_workers=2),
    run_config=train.RunConfig(
        checkpoint_config=train.CheckpointConfig(num_to_keep=2),
    ),
)
# __lightning_restore_example_end__


# __transformers_save_example_start__
from transformers import TrainingArguments

from ray import train
from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
from ray.train.torch import TorchTrainer


def train_func(config):
    ...

    # Configure logging, saving, evaluation strategies as usual.
    args = TrainingArguments(
        ...,
        eval_strategy="epoch",
        save_strategy="epoch",
        logging_strategy="step",
    )

    trainer = transformers.Trainer(args, ...)

    # Add a report callback to transformers Trainer
    # =============================================
    trainer.add_callback(RayTrainReportCallback())
    trainer = prepare_trainer(trainer)

    trainer.train()


ray_trainer = TorchTrainer(
    train_func,
    run_config=train.RunConfig(
        checkpoint_config=train.CheckpointConfig(
            num_to_keep=3,
            checkpoint_score_attribute="eval_loss",  # The monitoring metric
            checkpoint_score_order="min",
        )
    ),
)
# __transformers_save_example_end__


# __transformers_custom_save_example_start__
from ray import train

from transformers.trainer_callback import TrainerCallback


class MyTrainReportCallback(TrainerCallback):
    def __init__(self):
        super().__init__()
        self.metrics = {}

    def on_log(self, args, state, control, model=None, logs=None, **kwargs):
        """Log is called on evaluation step and logging step."""
        self.metrics.update(logs)

    def on_save(self, args, state, control, **kwargs):
        """Event called after a checkpoint save."""

        checkpoint = None
        if train.get_context().get_world_rank() == 0:
            # Build a Ray Train Checkpoint from the latest checkpoint
            checkpoint_path = transformers.trainer.get_last_checkpoint(args.output_dir)
            checkpoint = Checkpoint.from_directory(checkpoint_path)

        # Report to Ray Train with up-to-date metrics
        ray.train.report(metrics=self.metrics, checkpoint=checkpoint)

        # Clear the metrics buffer
        self.metrics = {}
# __transformers_custom_save_example_end__


# __distributed_checkpointing_start__
from ray import train
from ray.train import Checkpoint
from ray.train.torch import TorchTrainer


def train_func(config):
    ...

    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
        rank = train.get_context().get_world_rank()
        torch.save(
            ...,
            os.path.join(temp_checkpoint_dir, f"model-rank={rank}.pt"),
        )
        checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)

        train.report(metrics, checkpoint=checkpoint)


trainer = TorchTrainer(
    train_func,
    scaling_config=train.ScalingConfig(num_workers=2),
    run_config=train.RunConfig(storage_path="s3://bucket/"),
)
# The checkpoint in cloud storage will contain: model-rank=0.pt, model-rank=1.pt
# __distributed_checkpointing_end__

# __inspect_checkpoint_example_start__
from pathlib import Path

from ray.train import Checkpoint

# For demonstration, create a locally available directory with a `model.pt` file.
example_checkpoint_dir = Path("/tmp/test-checkpoint")
example_checkpoint_dir.mkdir()
example_checkpoint_dir.joinpath("model.pt").touch()

# Create the checkpoint, which is a reference to the directory.
checkpoint = Checkpoint.from_directory(example_checkpoint_dir)

# Inspect the checkpoint's contents with either `as_directory` or `to_directory`:
with checkpoint.as_directory() as checkpoint_dir:
    assert Path(checkpoint_dir).joinpath("model.pt").exists()

checkpoint_dir = checkpoint.to_directory()
assert Path(checkpoint_dir).joinpath("model.pt").exists()
# __inspect_checkpoint_example_end__

# __inspect_transformers_checkpoint_example_start__
# After training finished
checkpoint = result.checkpoint
with checkpoint.as_directory() as checkpoint_dir:
    hf_checkpoint_path = f"{checkpoint_dir}/checkpoint/"
# __inspect_transformers_checkpoint_example_end__

# __inspect_lightning_checkpoint_example_start__
# After training finished
checkpoint = result.checkpoint
with checkpoint.as_directory() as checkpoint_dir:
    lightning_checkpoint_path = f"{checkpoint_dir}/checkpoint.ckpt"
# __inspect_lightning_checkpoint_example_end__

# __checkpoint_upload_mode_sync_start__
def train_fn(config):
    ...
    metrics = {...}
    with tempfile.TemporaryDirectory() as tmpdir:
        ...  # Save checkpoint to tmpdir
        checkpoint = Checkpoint.from_directory(tmpdir)
        train.report(
            metrics,
            checkpoint=checkpoint,
            checkpoint_upload_mode=train.CheckpointUploadMode.SYNC,
        )
# __checkpoint_upload_mode_sync_end__

# __checkpoint_upload_mode_async_start__
def train_fn(config):
    ...
    metrics = {...}
    tmpdir = tempfile.mkdtemp()
    ...  # Save checkpoint to tmpdir
    checkpoint = Checkpoint.from_directory(tmpdir)
    train.report(
        metrics,
        checkpoint=checkpoint,
        checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
    )
# __checkpoint_upload_mode_async_end__

# __checkpoint_upload_mode_no_upload_start__
from s3torchconnector.dcp import S3StorageWriter
from torch.distributed.checkpoint.state_dict_saver import save
from torch.distributed.checkpoint.state_dict import get_state_dict


def train_fn(config):
    ...
    for epoch in range(config["num_epochs"]):
        # Directly upload checkpoint to s3 with Torch
        model, optimizer = ...
        storage_context = ray.train.get_context().get_storage()
        checkpoint_path = (
            f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
        )
        storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
        model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
        save(
            {"model": model_dict, "opt": opt_dict},
            storage_writer=storage_writer,
        )

        # Report that checkpoint to Ray Train
        metrics = {...}
        checkpoint = Checkpoint(checkpoint_path)
        train.report(
            metrics,
            checkpoint=checkpoint,
            checkpoint_upload_mode=train.CheckpointUploadMode.NO_UPLOAD,
        )
# __checkpoint_upload_mode_no_upload_end__


# __checkpoint_upload_fn_start__
from torch.distributed.checkpoint.state_dict_saver import async_save
from s3torchconnector.dcp import S3StorageWriter
from torch.distributed.checkpoint.state_dict import get_state_dict

from ray import train
from ray.train import Checkpoint


def train_fn(config):
    ...
    for epoch in config["num_epochs"]:
        # Start async checkpoint upload to s3 with Torch
        model, optimizer = ...
        storage_context = train.get_context().get_storage()
        checkpoint_path = (
            f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
        )
        storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
        model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
        ckpt_ref = async_save(
            {"model": model_dict, "opt": opt_dict},
            storage_writer=storage_writer,
        )

        def wait_async_save(checkpoint, checkpoint_dir_name):
            # This function waits for checkpoint to be finalized before returning it as is
            ckpt_ref.result()
            return checkpoint

        # Ray Train kicks off a thread that waits for the async checkpoint upload to complete
        # before reporting the checkpoint
        metrics = {...}
        checkpoint = Checkpoint(checkpoint_path)
        train.report(
            metrics=metrics,
            checkpoint=checkpoint,
            checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
            checkpoint_upload_fn=wait_async_save,
            # As uploading into the experiment directory then don't delete the checkpoint after upload is complete
            delete_local_checkpoint_after_upload=False,
        )

trainer = TorchTrainer(
   train_fn,
   train_loop_config={"num_epochs": 3},
   scaling_config=train.ScalingConfig(num_workers=2, use_gpu=True),
   # we need a cpu backend for async_save and a gpu backend for training
   torch_config=train.torch.TorchConfig(backend="cpu:gloo,cuda:nccl"),
   run_config=train.RunConfig(storage_path="s3://bucket/")
)
# __checkpoint_upload_fn_end__

# __get_all_reported_checkpoints_example_start__
import ray.train
from ray.train import CheckpointConsistencyMode

def train_fn():
    for epoch in range(2):
        metrics = {"train/loss": 0.1}
        checkpoint = ...
        ray.train.report(
            metrics,
            checkpoint=checkpoint,
            validation=...,
        )

    # Get committed checkpoints which may still have ongoing validations.
    committed_checkpoints = ray.train.get_all_reported_checkpoints(
        consistency_mode=CheckpointConsistencyMode.COMMITTED)

    # Wait for all pending validations to finish to access reported checkpoints
    # with validation metrics attached.
    validated_checkpoints = ray.train.get_all_reported_checkpoints()
    ...
# __get_all_reported_checkpoints_example_end__