import pytest
import torch

import ray
from ray.ml.predictors.integrations.torch import TorchPredictor
from ray.ml.train.integrations.torch import TorchTrainer
from ray import train
from ray.ml.examples.pytorch.torch_linear_example import train_func as linear_train_func


@pytest.fixture
def ray_start_4_cpus():
    address_info = ray.init(num_cpus=4)
    yield address_info
    # The code after the yield will run as teardown code.
    ray.shutdown()


@pytest.mark.parametrize("num_workers", [1, 2])
def test_torch_linear(ray_start_4_cpus, num_workers):
    def train_func(config):
        result = linear_train_func(config)
        assert len(result) == epochs
        assert result[-1]["loss"] < result[0]["loss"]

    num_workers = num_workers
    epochs = 3
    scaling_config = {"num_workers": num_workers}
    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        train_loop_config=config,
        scaling_config=scaling_config,
    )
    trainer.fit()


def test_torch_e2e(ray_start_4_cpus):
    def train_func():
        model = torch.nn.Linear(1, 1)
        train.save_checkpoint(model=model)

    scaling_config = {"num_workers": 2}
    trainer = TorchTrainer(
        train_loop_per_worker=train_func, scaling_config=scaling_config
    )
    result = trainer.fit()

    predict_dataset = ray.data.range(3)

    class TorchScorer:
        def __init__(self):
            self.pred = TorchPredictor.from_checkpoint(result.checkpoint)

        def __call__(self, x):
            return self.pred.predict(x, dtype=torch.float)

    predictions = predict_dataset.map_batches(
        TorchScorer, batch_format="pandas", compute="actors"
    )
    assert predictions.count() == 3


def test_torch_e2e_state_dict(ray_start_4_cpus):
    def train_func():
        model = torch.nn.Linear(1, 1).state_dict()
        train.save_checkpoint(model=model)

    scaling_config = {"num_workers": 2}
    trainer = TorchTrainer(
        train_loop_per_worker=train_func, scaling_config=scaling_config
    )
    result = trainer.fit()

    # If loading from a state dict, a model definition must be passed in.
    with pytest.raises(ValueError):
        TorchPredictor.from_checkpoint(result.checkpoint)

    class TorchScorer:
        def __init__(self):
            self.pred = TorchPredictor.from_checkpoint(
                result.checkpoint, model=torch.nn.Linear(1, 1)
            )

        def __call__(self, x):
            return self.pred.predict(x, dtype=torch.float)

    predict_dataset = ray.data.range(3)
    predictions = predict_dataset.map_batches(
        TorchScorer, batch_format="pandas", compute="actors"
    )
    assert predictions.count() == 3


if __name__ == "__main__":
    import pytest
    import sys

    sys.exit(pytest.main(["-v", "-x", __file__]))

[(19, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""num_workers"""', '[1, 2]'], {}), False, 'import pytest\n'), (13, 'ray.init', 'ray.init', ([], {'num_cpus': '(4)'}), False, 'import ray\n'), (16, 'ray.shutdown', 'ray.shutdown', ([], {}), False, 'import ray\n'), (30, 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', ([], {'train_loop_per_worker': 'train_func', 'train_loop_config': 'config', 'scaling_config': 'scaling_config'}), False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), (44, 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', ([], {'train_loop_per_worker': 'train_func', 'scaling_config': 'scaling_config'}), False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), (49, 'ray.data.range', 'ray.data.range', (['(3)'], {}), False, 'import ray\n'), (70, 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', ([], {'train_loop_per_worker': 'train_func', 'scaling_config': 'scaling_config'}), False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), (88, 'ray.data.range', 'ray.data.range', (['(3)'], {}), False, 'import ray\n'), (22, 'ray.ml.examples.pytorch.torch_linear_example.train_func', 'linear_train_func', (['config'], {}), True, 'from ray.ml.examples.pytorch.torch_linear_example import train_func as linear_train_func\n'), (40, 'torch.nn.Linear', 'torch.nn.Linear', (['(1)', '(1)'], {}), False, 'import torch\n'), (41, 'ray.train.save_checkpoint', 'train.save_checkpoint', ([], {'model': 'model'}), False, 'from ray import train\n'), (67, 'ray.train.save_checkpoint', 'train.save_checkpoint', ([], {'model': 'model'}), False, 'from ray import train\n'), (76, 'pytest.raises', 'pytest.raises', (['ValueError'], {}), False, 'import pytest\n'), (77, 'ray.ml.predictors.integrations.torch.TorchPredictor.from_checkpoint', 'TorchPredictor.from_checkpoint', (['result.checkpoint'], {}), False, 'from ray.ml.predictors.integrations.torch import TorchPredictor\n'), (99, 'pytest.main', 'pytest.main', (["['-v', '-x', __file__]"], {}), False, 'import pytest\n'), (53, 'ray.ml.predictors.integrations.torch.TorchPredictor.from_checkpoint', 'TorchPredictor.from_checkpoint', (['result.checkpoint'], {}), False, 'from ray.ml.predictors.integrations.torch import TorchPredictor\n'), (66, 'torch.nn.Linear', 'torch.nn.Linear', (['(1)', '(1)'], {}), False, 'import torch\n'), (82, 'torch.nn.Linear', 'torch.nn.Linear', (['(1)', '(1)'], {}), False, 'import torch\n')]