code: def test_trainer_with_ray_data(
      ray_start_6_cpus_2_gpus, strategy_name, accelerator
  ):
      """Test Data integration with ddp and fsdp."""
  
      if accelerator == "cpu" and strategy_name == "fsdp":
          return
  
      num_epochs = 4
      batch_size = 8
      num_workers = 2
      dataset_size = 256
  
      strategy_map = {
          "ddp": RayDDPStrategy(),
          "fsdp": RayFSDPStrategy(),
      }
  
      dataset = np.random.rand(dataset_size, 32).astype(np.float32)
      train_dataset = ray.data.from_numpy(dataset)
      val_dataset = ray.data.from_numpy(dataset)
  
      def train_loop():
          model = LinearModule(
              input_dim=32, output_dim=4, strategy=strategy_name
          )
  
          strategy = strategy_map[strategy_name]
  
          trainer = pl.Trainer(
              max_epochs=num_epochs,
              devices="auto",
              accelerator=accelerator,
              strategy=strategy,
              plugins=[RayLightningEnvironment()],
              callbacks=[RayTrainReportCallback()],
          )
  
          train_data_iterable = ray.train.get_dataset_shard(
              "train"
          ).iter_torch_batches(batch_size=batch_size)
          val_data_iterable = ray.train.get_dataset_shard(
              "val"
          ).iter_torch_batches(batch_size=batch_size)
  
          trainer.fit(
              model,
              train_dataloaders=train_data_iterable,
              val_dataloaders=val_data_iterable,
          )
  
      trainer = TorchTrainer(
          train_loop_per_worker=train_loop,
          scaling_config=ScalingConfig(
              num_workers=2, use_gpu=(accelerator == "gpu")
          ),
          datasets={"train": train_dataset, "val": val_dataset},
      )
  
      results = trainer.fit()
      assert results.metrics["epoch"] == num_epochs - 1
      assert (
          results.metrics["step"]
          == num_epochs * dataset_size / num_workers / batch_size
      )
      assert "loss" in results.metrics
      assert "val_loss" in results.metrics
  
file_path: ray/python/ray/train/tests/test_torch_lightning_train.py
imports: import os
  import numpy
  import pytest
  import ray
  from ray.train import ScalingConfig
  from ray.train.lightning import RayDDPStrategy, RayDeepSpeedStrategy, RayFSDPStrategy, RayLightningEnvironment, RayTrainReportCallback
  from ray.train.lightning._lightning_utils import import_lightning
  from ray.train.tests.lightning_test_utils import DummyDataModule, LinearModule
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 4 for calling a global function with example usage: test_trainer_with_ray_data(ray_start_6_cpus_2_gpus, strategy_name, accelerator) without return types