refactor test_utils filenames (#3445)

Summary:
Pull Request resolved: #3445

# context
* configurable components are essential for test and benchmark
* rename those files as *_config such as pipeline_config, table_config

Reviewed By: spmex

Differential Revision: D84227634

fbshipit-source-id: 8e7025af797d45ed8abad34fa6e4923dfc5343e6

Author: TroyGarden

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -44,15 +44,15 @@
 from torchrec.distributed.comm import get_local_size
 from torchrec.distributed.embedding_types import EmbeddingComputeKernel
 from torchrec.distributed.planner import Topology
+from torchrec.distributed.test_utils.model_input import ModelInput
 
 from torchrec.distributed.test_utils.multi_process import (
     MultiProcessContext,
     run_multi_process_func,
 )
+from torchrec.distributed.test_utils.pipeline_config import PipelineConfig
 from torchrec.distributed.test_utils.table_config import EmbeddingTablesConfig
-from torchrec.distributed.test_utils.test_input import ModelInput
 from torchrec.distributed.test_utils.test_model import TestOverArchLarge
-from torchrec.distributed.test_utils.train_pipeline import PipelineConfig
 from torchrec.distributed.train_pipeline import TrainPipeline
 from torchrec.distributed.types import ShardingType
 from torchrec.modules.embedding_configs import EmbeddingBagConfig

torchrec/distributed/benchmark/benchmark_utils.py

@@ -32,7 +32,7 @@
 from torchrec.distributed.planner.constants import NUM_POOLINGS, POOLING_FACTOR
 from torchrec.distributed.planner.planners import HeteroEmbeddingShardingPlanner
 from torchrec.distributed.planner.types import ParameterConstraints
-from torchrec.distributed.test_utils.test_input import ModelInput
+from torchrec.distributed.test_utils.model_input import ModelInput
 from torchrec.distributed.test_utils.test_model import (
     TestEBCSharder,
     TestSparseNN,

torchrec/distributed/test_utils/test_input.py renamed to torchrec/distributed/test_utils/model_input.py

@@ -40,12 +40,12 @@
     get_module_to_default_sharders,
     table_wise,
 )
+from torchrec.distributed.test_utils.model_input import ModelInput
 
 from torchrec.distributed.test_utils.multi_process import (
     MultiProcessContext,
     MultiProcessTestBase,
 )
-from torchrec.distributed.test_utils.test_input import ModelInput
 from torchrec.distributed.test_utils.test_model_parallel import ModelParallelTestShared
 from torchrec.distributed.test_utils.test_sharding import (
     copy_state_dict,

torchrec/distributed/test_utils/train_pipeline.py renamed to torchrec/distributed/test_utils/pipeline_config.py

@@ -11,7 +11,7 @@
 
 import torch
 from torch import nn
-from torchrec.distributed.test_utils.test_input import ModelInput
+from torchrec.distributed.test_utils.model_input import ModelInput
 from torchrec.modules.deepfm import DeepFM, FactorizationMachine
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor

torchrec/distributed/tests/test_dynamic_sharding.py

@@ -12,7 +12,7 @@
 import torch
 from torch import nn
 from torchrec.datasets.utils import Batch
-from torchrec.distributed.test_utils.test_input import ModelInput
+from torchrec.distributed.test_utils.model_input import ModelInput
 from torchrec.modules.crossnet import LowRankCrossNet
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.mlp import MLP

torchrec/models/deepfm.py

@@ -14,7 +14,7 @@
 import torch
 from parameterized import parameterized
 from torch.testing import FileCheck  # @manual
-from torchrec.distributed.test_utils.test_input import ModelInput
+from torchrec.distributed.test_utils.model_input import ModelInput
 from torchrec.fx import symbolic_trace, Tracer
 from torchrec.models.deepfm import (
     DenseArch,

torchrec/models/dlrm.py

@@ -16,7 +16,7 @@
 from torch import nn
 from torch.testing import FileCheck  # @manual
 from torchrec.datasets.utils import Batch
-from torchrec.distributed.test_utils.test_input import ModelInput
+from torchrec.distributed.test_utils.model_input import ModelInput
 from torchrec.fx import symbolic_trace
 from torchrec.ir.serializer import JsonSerializer
 from torchrec.ir.utils import decapsulate_ir_modules, encapsulate_ir_modules