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, 1, 19, 47), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(19, 25, 19, 38): '"""num_workers"""', (19, 40, 19, 46): '[1, 2]'}, {}), "('num_workers', [1, 2])", False, 'import pytest\n'), ((13, 19, 13, 39), 'ray.init', 'ray.init', (), '', False, 'import ray\n'), ((16, 4, 16, 18), 'ray.shutdown', 'ray.shutdown', ({}, {}), '()', False, 'import ray\n'), ((30, 14, 34, 5), 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', (), '', False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), ((44, 14, 46, 5), 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', (), '', False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), ((49, 22, 49, 39), 'ray.data.range', 'ray.data.range', ({(49, 37, 49, 38): '3'}, {}), '(3)', False, 'import ray\n'), ((70, 14, 72, 5), 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', (), '', False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), ((88, 22, 88, 39), 'ray.data.range', 'ray.data.range', ({(88, 37, 88, 38): '3'}, {}), '(3)', False, 'import ray\n'), ((22, 17, 22, 42), 'ray.ml.examples.pytorch.torch_linear_example.train_func', 'linear_train_func', ({(22, 35, 22, 41): 'config'}, {}), '(config)', True, 'from ray.ml.examples.pytorch.torch_linear_example import train_func as linear_train_func\n'), ((40, 16, 40, 37), 'torch.nn.Linear', 'torch.nn.Linear', ({(40, 32, 40, 33): '1', (40, 35, 40, 36): '1'}, {}), '(1, 1)', False, 'import torch\n'), ((41, 8, 41, 42), 'ray.train.save_checkpoint', 'train.save_checkpoint', (), '', False, 'from ray import train\n'), ((67, 8, 67, 42), 'ray.train.save_checkpoint', 'train.save_checkpoint', (), '', False, 'from ray import train\n'), ((76, 9, 76, 34), 'pytest.raises', 'pytest.raises', ({(76, 23, 76, 33): 'ValueError'}, {}), '(ValueError)', False, 'import pytest\n'), ((77, 8, 77, 57), 'ray.ml.predictors.integrations.torch.TorchPredictor.from_checkpoint', 'TorchPredictor.from_checkpoint', ({(77, 39, 77, 56): 'result.checkpoint'}, {}), '(result.checkpoint)', False, 'from ray.ml.predictors.integrations.torch import TorchPredictor\n'), ((99, 13, 99, 48), 'pytest.main', 'pytest.main', ({(99, 25, 99, 47): "['-v', '-x', __file__]"}, {}), "(['-v', '-x', __file__])", False, 'import pytest\n'), ((53, 24, 53, 73), 'ray.ml.predictors.integrations.torch.TorchPredictor.from_checkpoint', 'TorchPredictor.from_checkpoint', ({(53, 55, 53, 72): 'result.checkpoint'}, {}), '(result.checkpoint)', False, 'from ray.ml.predictors.integrations.torch import TorchPredictor\n'), ((66, 16, 66, 37), 'torch.nn.Linear', 'torch.nn.Linear', ({(66, 32, 66, 33): '1', (66, 35, 66, 36): '1'}, {}), '(1, 1)', False, 'import torch\n'), ((82, 41, 82, 62), 'torch.nn.Linear', 'torch.nn.Linear', ({(82, 57, 82, 58): '1', (82, 60, 82, 61): '1'}, {}), '(1, 1)', False, 'import torch\n')]