github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/ml/inference/pytorch_inference_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# pytype: skip-file

import os
import shutil
import tempfile
import unittest
from collections import OrderedDict

import numpy as np
import pytest

import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to

# Protect against environments where pytorch library is not available.
# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
try:
  import torch
  from apache_beam.ml.inference.base import PredictionResult
  from apache_beam.ml.inference.base import RunInference
  from apache_beam.ml.inference import pytorch_inference
  from apache_beam.ml.inference.pytorch_inference import default_keyed_tensor_inference_fn
  from apache_beam.ml.inference.pytorch_inference import default_tensor_inference_fn
  from apache_beam.ml.inference.pytorch_inference import make_keyed_tensor_model_fn
  from apache_beam.ml.inference.pytorch_inference import make_tensor_model_fn
  from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor
  from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor
except ImportError:
  raise unittest.SkipTest('PyTorch dependencies are not installed')

try:
  from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
except ImportError:
  GCSFileSystem = None  # type: ignore

TWO_FEATURES_EXAMPLES = [
    torch.from_numpy(np.array([1, 5], dtype="float32")),
    torch.from_numpy(np.array([3, 10], dtype="float32")),
    torch.from_numpy(np.array([-14, 0], dtype="float32")),
    torch.from_numpy(np.array([0.5, 0.5], dtype="float32")),
]

TWO_FEATURES_PREDICTIONS = [
    PredictionResult(ex, pred) for ex,
    pred in zip(
        TWO_FEATURES_EXAMPLES,
        torch.Tensor(
            [f1 * 2.0 + f2 * 3 + 0.5
             for f1, f2 in TWO_FEATURES_EXAMPLES]).reshape(-1, 1))
]

TWO_FEATURES_DICT_OUT_PREDICTIONS = [
    PredictionResult(
        p.example, {
            "output1": p.inference, "output2": p.inference
        }) for p in TWO_FEATURES_PREDICTIONS
]

KEYED_TORCH_EXAMPLES = [
    {
        'k1': torch.from_numpy(np.array([1], dtype="float32")),
        'k2': torch.from_numpy(np.array([1.5], dtype="float32"))
    },
    {
        'k1': torch.from_numpy(np.array([5], dtype="float32")),
        'k2': torch.from_numpy(np.array([5.5], dtype="float32"))
    },
    {
        'k1': torch.from_numpy(np.array([-3], dtype="float32")),
        'k2': torch.from_numpy(np.array([-3.5], dtype="float32"))
    },
    {
        'k1': torch.from_numpy(np.array([10.0], dtype="float32")),
        'k2': torch.from_numpy(np.array([10.5], dtype="float32"))
    },
]

KEYED_TORCH_PREDICTIONS = [
    PredictionResult(ex, pred) for ex,
    pred in zip(
        KEYED_TORCH_EXAMPLES,
        torch.Tensor([(example['k1'] * 2.0 + 0.5) + (example['k2'] * 2.0 + 0.5)
                      for example in KEYED_TORCH_EXAMPLES]).reshape(-1, 1))
]

KEYED_TORCH_HELPER_PREDICTIONS = [
    PredictionResult(ex, pred) for ex,
    pred in zip(
        KEYED_TORCH_EXAMPLES,
        torch.Tensor([(example['k1'] * 2.0 + 0.5) +
                      (example['k2'] * 2.0 + 0.5) + 0.5
                      for example in KEYED_TORCH_EXAMPLES]).reshape(-1, 1))
]

KEYED_TORCH_DICT_OUT_PREDICTIONS = [
    PredictionResult(
        p.example, {
            "output1": p.inference, "output2": p.inference
        }) for p in KEYED_TORCH_PREDICTIONS
]


class TestPytorchModelHandlerForInferenceOnly(PytorchModelHandlerTensor):
  def __init__(self, device, *, inference_fn=default_tensor_inference_fn):
    self._device = device
    self._inference_fn = inference_fn
    self._state_dict_path = None
    self._torch_script_model_path = None


class TestPytorchModelHandlerKeyedTensorForInferenceOnly(
    PytorchModelHandlerKeyedTensor):
  def __init__(self, device, *, inference_fn=default_keyed_tensor_inference_fn):
    self._device = device
    self._inference_fn = inference_fn
    self._state_dict_path = None
    self._torch_script_model_path = None


def _compare_prediction_result(x, y):
  if isinstance(x.example, dict):
    example_equals = all(
        torch.equal(x, y) for x,
        y in zip(x.example.values(), y.example.values()))
  else:
    example_equals = torch.equal(x.example, y.example)
  if not example_equals:
    return False

  if isinstance(x.inference, dict):
    return all(
        torch.equal(x, y) for x,
        y in zip(x.inference.values(), y.inference.values()))

  return torch.equal(x.inference, y.inference)


def custom_tensor_inference_fn(
    batch, model, device, inference_args, model_id=None):
  predictions = [
      PredictionResult(ex, pred) for ex,
      pred in zip(
          batch,
          torch.Tensor([item * 2.0 + 1.5 for item in batch]).reshape(-1, 1))
  ]
  return predictions


class PytorchLinearRegression(torch.nn.Module):
  def __init__(self, input_dim, output_dim):
    super().__init__()
    self.linear = torch.nn.Linear(input_dim, output_dim)

  def forward(self, x):
    out = self.linear(x)
    return out

  def generate(self, x):
    out = self.linear(x) + 0.5
    return out


class PytorchLinearRegressionDict(torch.nn.Module):
  def __init__(self, input_dim, output_dim):
    super().__init__()
    self.linear = torch.nn.Linear(input_dim, output_dim)

  def forward(self, x):
    out = self.linear(x)
    return {'output1': out, 'output2': out}


class PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs(torch.nn.Module):
  """
  A linear model with batched keyed inputs and non-batchable extra args.

  Note: k1 and k2 are batchable examples passed in as a dict from str to tensor.
  prediction_param_array, prediction_param_bool are non-batchable extra args
  (typically model-related info) used to configure the model before its predict
  call is invoked
  """
  def __init__(self, input_dim, output_dim):
    super().__init__()
    self.linear = torch.nn.Linear(input_dim, output_dim)

  def forward(self, k1, k2, prediction_param_array, prediction_param_bool):
    if not prediction_param_bool:
      raise ValueError("Expected prediction_param_bool to be True")
    if not torch.all(prediction_param_array):
      raise ValueError("Expected prediction_param_array to be all True")
    out = self.linear(k1) + self.linear(k2)
    return out


@pytest.mark.uses_pytorch
class PytorchRunInferenceTest(unittest.TestCase):
  def test_run_inference_single_tensor_feature(self):
    examples = [
        torch.from_numpy(np.array([1], dtype="float32")),
        torch.from_numpy(np.array([5], dtype="float32")),
        torch.from_numpy(np.array([-3], dtype="float32")),
        torch.from_numpy(np.array([10.0], dtype="float32")),
    ]
    expected_predictions = [
        PredictionResult(ex, pred) for ex,
        pred in zip(
            examples,
            torch.Tensor([example * 2.0 + 0.5
                          for example in examples]).reshape(-1, 1))
    ]

    model = PytorchLinearRegression(input_dim=1, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerForInferenceOnly(
        torch.device('cpu'))
    predictions = inference_runner.run_inference(examples, model)
    for actual, expected in zip(predictions, expected_predictions):
      self.assertEqual(actual, expected)

  def test_run_inference_multiple_tensor_features(self):
    model = PytorchLinearRegression(input_dim=2, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerForInferenceOnly(
        torch.device('cpu'))
    predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, model)
    for actual, expected in zip(predictions, TWO_FEATURES_PREDICTIONS):
      self.assertEqual(actual, expected)

  def test_run_inference_multiple_tensor_features_dict_output(self):
    model = PytorchLinearRegressionDict(input_dim=2, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerForInferenceOnly(
        torch.device('cpu'))
    predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, model)
    for actual, expected in zip(predictions, TWO_FEATURES_DICT_OUT_PREDICTIONS):
      self.assertEqual(actual, expected)

  def test_run_inference_custom(self):
    examples = [
        torch.from_numpy(np.array([1], dtype="float32")),
        torch.from_numpy(np.array([5], dtype="float32")),
        torch.from_numpy(np.array([-3], dtype="float32")),
        torch.from_numpy(np.array([10.0], dtype="float32")),
    ]
    expected_predictions = [
        PredictionResult(ex, pred) for ex,
        pred in zip(
            examples,
            torch.Tensor([example * 2.0 + 1.5
                          for example in examples]).reshape(-1, 1))
    ]

    model = PytorchLinearRegression(input_dim=1, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerForInferenceOnly(
        torch.device('cpu'), inference_fn=custom_tensor_inference_fn)
    predictions = inference_runner.run_inference(examples, model)
    for actual, expected in zip(predictions, expected_predictions):
      self.assertEqual(actual, expected)

  def test_run_inference_keyed(self):
    """
    This tests for inputs that are passed as a dictionary from key to tensor
    instead of a standard non-keyed tensor example.

    Example:
    Typical input format is
    input = torch.tensor([1, 2, 3])

    But Pytorch syntax allows inputs to have the form
    input = {
      'k1' : torch.tensor([1, 2, 3]),
      'k2' : torch.tensor([4, 5, 6])
    }
    """
    class PytorchLinearRegressionMultipleArgs(torch.nn.Module):
      def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = torch.nn.Linear(input_dim, output_dim)

      def forward(self, k1, k2):
        out = self.linear(k1) + self.linear(k2)
        return out

    model = PytorchLinearRegressionMultipleArgs(input_dim=1, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly(
        torch.device('cpu'))
    predictions = inference_runner.run_inference(KEYED_TORCH_EXAMPLES, model)
    for actual, expected in zip(predictions, KEYED_TORCH_PREDICTIONS):
      self.assertTrue(_compare_prediction_result(actual, expected))

  def test_run_inference_keyed_dict_output(self):
    class PytorchLinearRegressionMultipleArgsDict(torch.nn.Module):
      def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = torch.nn.Linear(input_dim, output_dim)

      def forward(self, k1, k2):
        out = self.linear(k1) + self.linear(k2)
        return {'output1': out, 'output2': out}

    model = PytorchLinearRegressionMultipleArgsDict(input_dim=1, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly(
        torch.device('cpu'))
    predictions = inference_runner.run_inference(KEYED_TORCH_EXAMPLES, model)
    for actual, expected in zip(predictions, KEYED_TORCH_DICT_OUT_PREDICTIONS):
      self.assertTrue(_compare_prediction_result(actual, expected))

  def test_inference_runner_inference_args(self):
    """
    This tests for non-batchable input arguments. Since we do the batching
    for the user, we have to distinguish between the inputs that should be
    batched and the ones that should not be batched.
    """
    inference_args = {
        'prediction_param_array': torch.from_numpy(
            np.array([1, 2], dtype="float32")),
        'prediction_param_bool': True
    }

    model = PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs(
        input_dim=1, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly(
        torch.device('cpu'))
    predictions = inference_runner.run_inference(
        batch=KEYED_TORCH_EXAMPLES, model=model, inference_args=inference_args)
    for actual, expected in zip(predictions, KEYED_TORCH_PREDICTIONS):
      self.assertEqual(actual, expected)

  def test_run_inference_helper(self):
    examples = [
        torch.from_numpy(np.array([1], dtype="float32")),
        torch.from_numpy(np.array([5], dtype="float32")),
        torch.from_numpy(np.array([-3], dtype="float32")),
        torch.from_numpy(np.array([10.0], dtype="float32")),
    ]
    expected_predictions = [
        PredictionResult(ex, pred) for ex,
        pred in zip(
            examples,
            torch.Tensor([example * 2.0 + 1.0
                          for example in examples]).reshape(-1, 1))
    ]

    gen_fn = make_tensor_model_fn('generate')

    model = PytorchLinearRegression(input_dim=1, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    inference_runner = TestPytorchModelHandlerForInferenceOnly(
        torch.device('cpu'), inference_fn=gen_fn)
    predictions = inference_runner.run_inference(examples, model)
    for actual, expected in zip(predictions, expected_predictions):
      self.assertEqual(actual, expected)

  def test_run_inference_keyed_helper(self):
    """
    This tests for inputs that are passed as a dictionary from key to tensor
    instead of a standard non-keyed tensor example.

    Example:
    Typical input format is
    input = torch.tensor([1, 2, 3])

    But Pytorch syntax allows inputs to have the form
    input = {
      'k1' : torch.tensor([1, 2, 3]),
      'k2' : torch.tensor([4, 5, 6])
    }
    """
    class PytorchLinearRegressionMultipleArgs(torch.nn.Module):
      def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = torch.nn.Linear(input_dim, output_dim)

      def forward(self, k1, k2):
        out = self.linear(k1) + self.linear(k2)
        return out

      def generate(self, k1, k2):
        out = self.linear(k1) + self.linear(k2) + 0.5
        return out

    model = PytorchLinearRegressionMultipleArgs(input_dim=1, output_dim=1)
    model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                     ('linear.bias', torch.Tensor([0.5]))]))
    model.eval()

    gen_fn = make_keyed_tensor_model_fn('generate')

    inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly(
        torch.device('cpu'), inference_fn=gen_fn)
    predictions = inference_runner.run_inference(KEYED_TORCH_EXAMPLES, model)
    for actual, expected in zip(predictions, KEYED_TORCH_HELPER_PREDICTIONS):
      self.assertTrue(_compare_prediction_result(actual, expected))

  def test_num_bytes(self):
    inference_runner = TestPytorchModelHandlerForInferenceOnly(
        torch.device('cpu'))
    examples = torch.from_numpy(
        np.array([1, 5, 3, 10, -14, 0, 0.5, 0.5],
                 dtype="float32")).reshape(-1, 2)
    self.assertEqual((examples[0].element_size()) * 8,
                     inference_runner.get_num_bytes(examples))

  def test_namespace(self):
    inference_runner = TestPytorchModelHandlerForInferenceOnly(
        torch.device('cpu'))
    self.assertEqual('BeamML_PyTorch', inference_runner.get_metrics_namespace())


@pytest.mark.uses_pytorch
class PytorchRunInferencePipelineTest(unittest.TestCase):
  def setUp(self):
    self.tmpdir = tempfile.mkdtemp()

  def tearDown(self):
    shutil.rmtree(self.tmpdir)

  def test_pipeline_local_model_simple(self):
    with TestPipeline() as pipeline:
      state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
                                ('linear.bias', torch.Tensor([0.5]))])
      path = os.path.join(self.tmpdir, 'my_state_dict_path')
      torch.save(state_dict, path)

      model_handler = PytorchModelHandlerTensor(
          state_dict_path=path,
          model_class=PytorchLinearRegression,
          model_params={
              'input_dim': 2, 'output_dim': 1
          })

      pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES)
      predictions = pcoll | RunInference(model_handler)
      assert_that(
          predictions,
          equal_to(
              TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result))

  def test_pipeline_local_model_extra_inference_args(self):
    with TestPipeline() as pipeline:
      inference_args = {
          'prediction_param_array': torch.from_numpy(
              np.array([1, 2], dtype="float32")),
          'prediction_param_bool': True
      }

      state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                                ('linear.bias', torch.Tensor([0.5]))])
      path = os.path.join(self.tmpdir, 'my_state_dict_path')
      torch.save(state_dict, path)

      model_handler = PytorchModelHandlerKeyedTensor(
          state_dict_path=path,
          model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs,
          model_params={
              'input_dim': 1, 'output_dim': 1
          })

      pcoll = pipeline | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES)
      inference_args_side_input = (
          pipeline | 'create side' >> beam.Create(inference_args))
      predictions = pcoll | RunInference(
          model_handler=model_handler,
          inference_args=beam.pvalue.AsDict(inference_args_side_input))
      assert_that(
          predictions,
          equal_to(
              KEYED_TORCH_PREDICTIONS, equals_fn=_compare_prediction_result))

  def test_pipeline_local_model_extra_inference_args_batching_args(self):
    with TestPipeline() as pipeline:
      inference_args = {
          'prediction_param_array': torch.from_numpy(
              np.array([1, 2], dtype="float32")),
          'prediction_param_bool': True
      }

      state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                                ('linear.bias', torch.Tensor([0.5]))])
      path = os.path.join(self.tmpdir, 'my_state_dict_path')
      torch.save(state_dict, path)

      def batch_validator_keyed_tensor_inference_fn(
          batch,
          model,
          device,
          inference_args,
          model_id,
      ):
        if len(batch) != 2:
          raise Exception(
              f'Expected batch of size 2, received batch of size {len(batch)}')
        return default_keyed_tensor_inference_fn(
            batch, model, device, inference_args, model_id)

      model_handler = PytorchModelHandlerKeyedTensor(
          state_dict_path=path,
          model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs,
          model_params={
              'input_dim': 1, 'output_dim': 1
          },
          inference_fn=batch_validator_keyed_tensor_inference_fn,
          min_batch_size=2,
          max_batch_size=2)

      pcoll = pipeline | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES)
      inference_args_side_input = (
          pipeline | 'create side' >> beam.Create(inference_args))
      predictions = pcoll | RunInference(
          model_handler=model_handler,
          inference_args=beam.pvalue.AsDict(inference_args_side_input))
      assert_that(
          predictions,
          equal_to(
              KEYED_TORCH_PREDICTIONS, equals_fn=_compare_prediction_result))

  @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed')
  def test_pipeline_gcs_model(self):
    with TestPipeline() as pipeline:
      examples = torch.from_numpy(
          np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1))
      expected_predictions = [
          PredictionResult(ex, pred) for ex,
          pred in zip(
              examples,
              torch.Tensor([example * 2.0 + 0.5
                            for example in examples]).reshape(-1, 1))
      ]

      gs_pth = 'gs://apache-beam-ml/models/' \
          'pytorch_lin_reg_model_2x+0.5_state_dict.pth'
      model_handler = PytorchModelHandlerTensor(
          state_dict_path=gs_pth,
          model_class=PytorchLinearRegression,
          model_params={
              'input_dim': 1, 'output_dim': 1
          })

      pcoll = pipeline | 'start' >> beam.Create(examples)
      predictions = pcoll | RunInference(model_handler)
      assert_that(
          predictions,
          equal_to(expected_predictions, equals_fn=_compare_prediction_result))

  @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed')
  def test_pipeline_gcs_model_control_batching(self):
    with TestPipeline() as pipeline:
      examples = torch.from_numpy(
          np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1))
      expected_predictions = [
          PredictionResult(ex, pred) for ex,
          pred in zip(
              examples,
              torch.Tensor([example * 2.0 + 0.5
                            for example in examples]).reshape(-1, 1))
      ]

      def batch_validator_tensor_inference_fn(
          batch,
          model,
          device,
          inference_args,
          model_id,
      ):
        if len(batch) != 2:
          raise Exception(
              f'Expected batch of size 2, received batch of size {len(batch)}')
        return default_tensor_inference_fn(
            batch, model, device, inference_args, model_id)


      gs_pth = 'gs://apache-beam-ml/models/' \
          'pytorch_lin_reg_model_2x+0.5_state_dict.pth'
      model_handler = PytorchModelHandlerTensor(
          state_dict_path=gs_pth,
          model_class=PytorchLinearRegression,
          model_params={
              'input_dim': 1, 'output_dim': 1
          },
          inference_fn=batch_validator_tensor_inference_fn,
          min_batch_size=2,
          max_batch_size=2)

      pcoll = pipeline | 'start' >> beam.Create(examples)
      predictions = pcoll | RunInference(model_handler)
      assert_that(
          predictions,
          equal_to(expected_predictions, equals_fn=_compare_prediction_result))

  def test_invalid_input_type(self):
    with self.assertRaisesRegex(TypeError, "expected Tensor as element"):
      with TestPipeline() as pipeline:
        examples = np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1)

        state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                                  ('linear.bias', torch.Tensor([0.5]))])
        path = os.path.join(self.tmpdir, 'my_state_dict_path')
        torch.save(state_dict, path)

        model_handler = PytorchModelHandlerTensor(
            state_dict_path=path,
            model_class=PytorchLinearRegression,
            model_params={
                'input_dim': 1, 'output_dim': 1
            })

        pcoll = pipeline | 'start' >> beam.Create(examples)
        # pylint: disable=expression-not-assigned
        pcoll | RunInference(model_handler)

  def test_gpu_auto_convert_to_cpu(self):
    """
    This tests the scenario in which the user defines `device='GPU'` for the
    PytorchModelHandlerX, but runs the pipeline on a machine without GPU, we
    automatically detect this discrepancy and do automatic conversion to CPU.
    A warning is also logged to inform the user.
    """
    with self.assertLogs() as log:
      with TestPipeline() as pipeline:
        examples = torch.from_numpy(
            np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1))

        state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                                  ('linear.bias', torch.Tensor([0.5]))])
        path = os.path.join(self.tmpdir, 'my_state_dict_path')
        torch.save(state_dict, path)

        model_handler = PytorchModelHandlerTensor(
            state_dict_path=path,
            model_class=PytorchLinearRegression,
            model_params={
                'input_dim': 1, 'output_dim': 1
            },
            device='GPU')
        # Upon initialization, device is cuda
        self.assertEqual(model_handler._device, torch.device('cuda'))

        pcoll = pipeline | 'start' >> beam.Create(examples)
        # pylint: disable=expression-not-assigned
        pcoll | RunInference(model_handler)

        # During model loading, device converted to cuda
        self.assertEqual(model_handler._device, torch.device('cuda'))

      self.assertIn("INFO:root:Device is set to CUDA", log.output)
      self.assertIn(
          "WARNING:root:Model handler specified a 'GPU' device, but GPUs " \
          "are not available. Switching to CPU.",
          log.output)

  def test_load_torch_script_model(self):
    torch_model = PytorchLinearRegression(2, 1)
    torch_script_model = torch.jit.script(torch_model)

    torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt')

    torch.jit.save(torch_script_model, torch_script_path)

    model_handler = PytorchModelHandlerTensor(
        torch_script_model_path=torch_script_path)

    torch_script_model = model_handler.load_model()

    self.assertTrue(isinstance(torch_script_model, torch.jit.ScriptModule))

  def test_inference_torch_script_model(self):
    torch_model = PytorchLinearRegression(2, 1)
    torch_model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
                     ('linear.bias', torch.Tensor([0.5]))]))

    torch_script_model = torch.jit.script(torch_model)

    torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt')

    torch.jit.save(torch_script_model, torch_script_path)

    model_handler = PytorchModelHandlerTensor(
        torch_script_model_path=torch_script_path)

    with TestPipeline() as pipeline:
      pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES)
      predictions = pcoll | RunInference(model_handler)
      assert_that(
          predictions,
          equal_to(
              TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result))

  def test_torch_model_class_none(self):
    torch_model = PytorchLinearRegression(2, 1)
    torch_path = os.path.join(self.tmpdir, 'torch_model.pt')

    torch.save(torch_model, torch_path)

    with self.assertRaisesRegex(
        RuntimeError,
        "A state_dict_path has been supplied to the model "
        "handler, but the required model_class is missing. "
        "Please provide the model_class in order to"):
      _ = PytorchModelHandlerTensor(state_dict_path=torch_path)

    with self.assertRaisesRegex(
        RuntimeError,
        "A state_dict_path has been supplied to the model "
        "handler, but the required model_class is missing. "
        "Please provide the model_class in order to"):
      _ = (PytorchModelHandlerKeyedTensor(state_dict_path=torch_path))

  def test_torch_model_state_dict_none(self):
    with self.assertRaisesRegex(
        RuntimeError,
        "A model_class has been supplied to the model "
        "handler, but the required state_dict_path is missing. "
        "Please provide the state_dict_path in order to"):
      _ = PytorchModelHandlerTensor(model_class=PytorchLinearRegression)

    with self.assertRaisesRegex(
        RuntimeError,
        "A model_class has been supplied to the model "
        "handler, but the required state_dict_path is missing. "
        "Please provide the state_dict_path in order to"):
      _ = PytorchModelHandlerKeyedTensor(model_class=PytorchLinearRegression)

  def test_specify_torch_script_path_and_state_dict_path(self):
    torch_model = PytorchLinearRegression(2, 1)
    torch_path = os.path.join(self.tmpdir, 'torch_model.pt')

    torch.save(torch_model, torch_path)
    torch_script_model = torch.jit.script(torch_model)

    torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt')

    torch.jit.save(torch_script_model, torch_script_path)
    with self.assertRaisesRegex(
        RuntimeError, "Please specify either torch_script_model_path or "):
      _ = PytorchModelHandlerTensor(
          state_dict_path=torch_path,
          model_class=PytorchLinearRegression,
          torch_script_model_path=torch_script_path)

  def test_prediction_result_model_id_with_torch_script_model(self):
    torch_model = PytorchLinearRegression(2, 1)
    torch_script_model = torch.jit.script(torch_model)
    torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt')
    torch.jit.save(torch_script_model, torch_script_path)

    model_handler = PytorchModelHandlerTensor(
        torch_script_model_path=torch_script_path)

    def check_torch_script_model_id(element):
      assert ('torch_script_model.pt' in element.model_id) is True

    with TestPipeline() as pipeline:
      pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES)
      predictions = pcoll | RunInference(model_handler)
      _ = predictions | beam.Map(check_torch_script_model_id)

  def test_prediction_result_model_id_with_torch_model(self):
    # weights associated with PytorchLinearRegression class
    state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
                              ('linear.bias', torch.Tensor([0.5]))])
    torch_path = os.path.join(self.tmpdir, 'torch_model.pt')
    torch.save(state_dict, torch_path)

    model_handler = PytorchModelHandlerTensor(
        state_dict_path=torch_path,
        model_class=PytorchLinearRegression,
        model_params={
            'input_dim': 2, 'output_dim': 1
        })

    def check_torch_script_model_id(element):
      assert ('torch_model.pt' in element.model_id) is True

    with TestPipeline() as pipeline:
      pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES)
      predictions = pcoll | RunInference(model_handler)
      _ = predictions | beam.Map(check_torch_script_model_id)

  def test_env_vars_set_correctly_tensor_handler(self):
    torch_model = PytorchLinearRegression(2, 1)
    torch_model.load_state_dict(
        OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
                     ('linear.bias', torch.Tensor([0.5]))]))

    torch_script_model = torch.jit.script(torch_model)

    torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt')

    torch.jit.save(torch_script_model, torch_script_path)

    handler_with_vars = PytorchModelHandlerTensor(
        torch_script_model_path=torch_script_path, env_vars={'FOO': 'bar'})
    os.environ.pop('FOO', None)
    self.assertFalse('FOO' in os.environ)
    with TestPipeline() as pipeline:
      _ = (
          pipeline
          | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES)
          | RunInference(handler_with_vars))
      pipeline.run()
      self.assertTrue('FOO' in os.environ)
      self.assertTrue((os.environ['FOO']) == 'bar')

  def test_env_vars_set_correctly_keyed_tensor_handler(self):
    os.environ.pop('FOO', None)
    self.assertFalse('FOO' in os.environ)
    with TestPipeline() as pipeline:
      inference_args = {
          'prediction_param_array': torch.from_numpy(
              np.array([1, 2], dtype="float32")),
          'prediction_param_bool': True
      }

      state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
                                ('linear.bias', torch.Tensor([0.5]))])
      path = os.path.join(self.tmpdir, 'my_state_dict_path')
      torch.save(state_dict, path)

      handler_with_vars = PytorchModelHandlerKeyedTensor(
          env_vars={'FOO': 'bar'},
          state_dict_path=path,
          model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs,
          model_params={
              'input_dim': 1, 'output_dim': 1
          })
      inference_args_side_input = (
          pipeline | 'create side' >> beam.Create(inference_args))

      _ = (
          pipeline
          | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES)
          | RunInference(
              model_handler=handler_with_vars,
              inference_args=beam.pvalue.AsDict(inference_args_side_input)))
      pipeline.run()
      self.assertTrue('FOO' in os.environ)
      self.assertTrue((os.environ['FOO']) == 'bar')


@pytest.mark.uses_pytorch
class PytorchInferenceTestWithMocks(unittest.TestCase):
  def setUp(self):
    self._load_model = pytorch_inference._load_model
    pytorch_inference._load_model = unittest.mock.MagicMock(
        return_value=("model", "device"))
    self.tmpdir = tempfile.mkdtemp()
    self.state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
                                   ('linear.bias', torch.Tensor([0.5]))])
    self.torch_path = os.path.join(self.tmpdir, 'torch_model.pt')
    torch.save(self.state_dict, self.torch_path)
    self.model_params = {'input_dim': 2, 'output_dim': 1}

  def tearDown(self):
    pytorch_inference._load_model = self._load_model
    shutil.rmtree(self.tmpdir)

  def test_load_model_args_tensor(self):
    load_model_args = {'weights_only': True}
    model_handler = PytorchModelHandlerTensor(
        state_dict_path=self.torch_path,
        model_class=PytorchLinearRegression,
        model_params=self.model_params,
        load_model_args=load_model_args)
    model_handler.load_model()
    pytorch_inference._load_model.assert_called_with(
        model_class=PytorchLinearRegression,
        state_dict_path=self.torch_path,
        device=torch.device('cpu'),
        model_params=self.model_params,
        torch_script_model_path=None,
        load_model_args=load_model_args)

  def test_load_model_args_keyed_tensor(self):
    load_model_args = {'weights_only': True}
    model_handler = PytorchModelHandlerKeyedTensor(
        state_dict_path=self.torch_path,
        model_class=PytorchLinearRegression,
        model_params=self.model_params,
        load_model_args=load_model_args)
    model_handler.load_model()
    pytorch_inference._load_model.assert_called_with(
        model_class=PytorchLinearRegression,
        state_dict_path=self.torch_path,
        device=torch.device('cpu'),
        model_params=self.model_params,
        torch_script_model_path=None,
        load_model_args=load_model_args)


if __name__ == '__main__':
  unittest.main()