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__]))

[((456, 502), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""num_workers"""', '[1, 2]'], {}), "('num_workers', [1, 2])\n", (479, 502), False, 'import pytest\n'), ((332, 352), 'ray.init', 'ray.init', ([], {'num_cpus': '(4)'}), '(num_cpus=4)\n', (340, 352), False, 'import ray\n'), ((438, 452), 'ray.shutdown', 'ray.shutdown', ([], {}), '()\n', (450, 452), False, 'import ray\n'), ((908, 1015), 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', ([], {'train_loop_per_worker': 'train_func', 'train_loop_config': 'config', 'scaling_config': 'scaling_config'}), '(train_loop_per_worker=train_func, train_loop_config=config,\n    scaling_config=scaling_config)\n', (920, 1015), False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), ((1259, 1336), 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', ([], {'train_loop_per_worker': 'train_func', 'scaling_config': 'scaling_config'}), '(train_loop_per_worker=train_func, scaling_config=scaling_config)\n', (1271, 1336), False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), ((1401, 1418), 'ray.data.range', 'ray.data.range', (['(3)'], {}), '(3)\n', (1415, 1418), False, 'import ray\n'), ((2009, 2086), 'ray.ml.train.integrations.torch.TorchTrainer', 'TorchTrainer', ([], {'train_loop_per_worker': 'train_func', 'scaling_config': 'scaling_config'}), '(train_loop_per_worker=train_func, scaling_config=scaling_config)\n', (2021, 2086), False, 'from ray.ml.train.integrations.torch import TorchTrainer\n'), ((2596, 2613), 'ray.data.range', 'ray.data.range', (['(3)'], {}), '(3)\n', (2610, 2613), False, 'import ray\n'), ((602, 627), 'ray.ml.examples.pytorch.torch_linear_example.train_func', 'linear_train_func', (['config'], {}), '(config)\n', (619, 627), True, 'from ray.ml.examples.pytorch.torch_linear_example import train_func as linear_train_func\n'), ((1139, 1160), 'torch.nn.Linear', 'torch.nn.Linear', (['(1)', '(1)'], {}), '(1, 1)\n', (1154, 1160), False, 'import torch\n'), ((1169, 1203), 'ray.train.save_checkpoint', 'train.save_checkpoint', ([], {'model': 'model'}), '(model=model)\n', (1190, 1203), False, 'from ray import train\n'), ((1919, 1953), 'ray.train.save_checkpoint', 'train.save_checkpoint', ([], {'model': 'model'}), '(model=model)\n', (1940, 1953), False, 'from ray import train\n'), ((2212, 2237), 'pytest.raises', 'pytest.raises', (['ValueError'], {}), '(ValueError)\n', (2225, 2237), False, 'import pytest\n'), ((2247, 2296), 'ray.ml.predictors.integrations.torch.TorchPredictor.from_checkpoint', 'TorchPredictor.from_checkpoint', (['result.checkpoint'], {}), '(result.checkpoint)\n', (2277, 2296), False, 'from ray.ml.predictors.integrations.torch import TorchPredictor\n'), ((2840, 2875), 'pytest.main', 'pytest.main', (["['-v', '-x', __file__]"], {}), "(['-v', '-x', __file__])\n", (2851, 2875), False, 'import pytest\n'), ((1495, 1544), 'ray.ml.predictors.integrations.torch.TorchPredictor.from_checkpoint', 'TorchPredictor.from_checkpoint', (['result.checkpoint'], {}), '(result.checkpoint)\n', (1525, 1544), False, 'from ray.ml.predictors.integrations.torch import TorchPredictor\n'), ((1876, 1897), 'torch.nn.Linear', 'torch.nn.Linear', (['(1)', '(1)'], {}), '(1, 1)\n', (1891, 1897), False, 'import torch\n'), ((2446, 2467), 'torch.nn.Linear', 'torch.nn.Linear', (['(1)', '(1)'], {}), '(1, 1)\n', (2461, 2467), False, 'import torch\n')]