code: def test_train_with_custom_data_module(self):
      if skip_lightning_tests:
          self.skipTest(
              "Spark PyTorch Lightning tests conflict with Tensorflow 2.5.x: "
              "https://github.com/horovod/horovod/pull/3263"
          )
  
      from horovod.spark.data_loaders.pytorch_data_loaders import (
          PytorchAsyncDataLoader,
      )
  
      class CustomDataModule(pl.LightningDataModule):
          """Custom DataModule for Lightning Estimator, using PytorchAsyncDataLoader"""
  
          def __init__(
              self,
              train_dir: str,
              val_dir: str,
              has_val: bool = True,
              train_batch_size: int = 32,
              val_batch_size: int = 32,
              shuffle: bool = True,
              num_reader_epochs=None,
              cur_shard: int = 0,
              shard_count: int = 1,
              schema_fields=None,
              storage_options=None,
              steps_per_epoch_train: int = 1,
              steps_per_epoch_val: int = 1,
              verbose=True,
              **kwargs,
          ):
              super().__init__()
              self.train_dir = train_dir
              self.val_dir = val_dir
              self.has_val = has_val
              self.train_batch_size = train_batch_size
              self.val_batch_size = val_batch_size
              self.shuffle = shuffle
              self.num_reader_epochs = num_reader_epochs
              self.cur_shard = cur_shard
              self.shard_count = shard_count
              self.schema_fields = schema_fields
              self.storage_options = storage_options
              self.steps_per_epoch_train = steps_per_epoch_train
              self.steps_per_epoch_val = steps_per_epoch_val
              self.verbose = verbose
  
          def setup(self, stage=None):
              # Assign train/val datasets for use in dataloaders
              from petastorm import make_batch_reader
  
              if stage == "fit" or stage is None:
                  self.train_reader = make_batch_reader(
                      self.train_dir,
                      num_epochs=self.num_reader_epochs,
                      cur_shard=self.cur_shard,
                      shard_count=self.shard_count,
                      shuffle_rows=self.shuffle,
                      shuffle_row_groups=self.shuffle,
                      hdfs_driver="libhdfs",
                      schema_fields=self.schema_fields,
                      storage_options=self.storage_options,
                  )
                  if self.has_val:
                      self.val_reader = make_batch_reader(
                          self.val_dir,
                          num_epochs=self.num_reader_epochs,
                          cur_shard=self.cur_shard,
                          shard_count=self.shard_count,
                          hdfs_driver="libhdfs",
                          schema_fields=self.schema_fields,
                          storage_options=self.storage_options,
                      )
  
          def teardown(self, stage=None):
              if stage == "fit" or stage is None:
                  if self.verbose:
                      print("Tear down petastorm readers")
                  self.train_reader.stop()
                  self.train_reader.join()
                  if self.has_val:
                      self.val_reader.stop()
                      self.val_reader.join()
  
          def train_dataloader(self):
              if self.verbose:
                  print("Setup train dataloader")
              kwargs = dict(
                  reader=self.train_reader,
                  batch_size=self.train_batch_size,
                  name="train dataloader",
                  shuffling_queue_capacity=0,
                  limit_step_per_epoch=self.steps_per_epoch_train,
                  verbose=self.verbose,
              )
              return PytorchAsyncDataLoader(**kwargs)
  
          def val_dataloader(self):
              if not self.has_val:
                  return None
              if self.verbose:
                  print("setup val dataloader")
              kwargs = dict(
                  reader=self.val_reader,
                  batch_size=self.val_batch_size,
                  name="val dataloader",
                  shuffling_queue_capacity=0,
                  limit_step_per_epoch=self.steps_per_epoch_val,
                  verbose=self.verbose,
              )
              return PytorchAsyncDataLoader(**kwargs)
  
      with spark_session("test_fit_model") as spark:
          df = create_noisy_xor_data(spark)
          model = create_xor_model()
  
          with local_store() as store:
              torch_estimator = hvd_spark.TorchEstimator(
                  num_proc=2,
                  store=store,
                  model=model,
                  input_shapes=[[-1, 2]],
                  feature_cols=["features"],
                  label_cols=["y"],
                  validation=0.2,
                  batch_size=4,
                  epochs=2,
                  data_module=CustomDataModule,
                  verbose=2,
                  trainer_args={"num_sanity_val_steps": 0},
              )
  
              torch_model = torch_estimator.fit(df)
  
              # TODO: Find a way to pass log metrics from remote, and assert base on the logger.
              trained_model = torch_model.getModel()
              pred = trained_model(
                  torch.ones([1, 2], dtype=torch.int32)
              )
              assert len(pred) == 1
              assert pred.dtype == torch.float32
  
imports: import logging
  import os
  import sys
  import unittest
  import warnings
  import mock
  from unittest.mock import Mock
  import pytest
  import numpy
  from packaging import version
  import torch
  import torch.nn
  from torch.nn import functional
  from pyspark.ml.linalg import VectorUDT
  from pyspark.sql.types import FloatType, IntegerType
  import pytorch_lightning
  import horovod
  import horovod.torch
  from horovod.torch.elastic import run
  import horovod.spark.lightning
  from horovod.spark.lightning import remote
  from horovod.spark.lightning.estimator import EstimatorParams, _torch_param_serialize, MIN_PL_VERSION
  from horovod.spark.lightning.legacy import to_lightning_module
  from horovod.common.util import gloo_built
  from horovod.spark.common import constants, util
  from common import tempdir, spawn, is_built
  from spark_common import CallbackBackend, create_noisy_xor_data, create_noisy_xor_data_with_val, create_xor_data, local_store, spark_session
label_desc: Use image node_id 22 for calling the SparkLightningTests obj's underlying member method code with example usage: obj.test_train_with_custom_data_module() and returns: PytorchAsyncDataLoader, PytorchAsyncDataLoader, None