Back out "Convert directory fbcode/torchrec to use the Ruff Formatter"

Summary:
Original commit changeset: ee300de21222

Original Phabricator Diff: D66013071

bypass-github-export-checks

Reviewed By: aporialiao

Differential Revision: D66198773

fbshipit-source-id: 4a8e5a124937a8329d7ed39444d3fdc5e4f2d10c

Author: Thomas Polasek

benchmarks/ebc_benchmarks.py

@@ -163,6 +163,7 @@ def get_fused_ebc_uvm_time(
     location: EmbeddingLocation,
     epochs: int = 100,
 ) -> Tuple[float, float]:
+
     fused_ebc = FusedEmbeddingBagCollection(
         tables=embedding_bag_configs,
         optimizer_type=torch.optim.SGD,
@@ -194,6 +195,7 @@ def get_ebc_comparison(
     device: torch.device,
     epochs: int = 100,
 ) -> Tuple[float, float, float, float, float]:
+
     # Simple EBC module wrapping a list of nn.EmbeddingBag
     ebc = EmbeddingBagCollection(
         tables=embedding_bag_configs,

benchmarks/ebc_benchmarks_utils.py

@@ -26,6 +26,7 @@ def get_random_dataset(
     embedding_bag_configs: List[EmbeddingBagConfig],
     pooling_factors: Optional[Dict[str, int]] = None,
 ) -> IterableDataset[Batch]:
+
     if pooling_factors is None:
         pooling_factors = {}
 
@@ -56,6 +57,7 @@ def train_one_epoch(
     dataset: IterableDataset[Batch],
     device: torch.device,
 ) -> float:
+
     start_time = time.perf_counter()
 
     for data in dataset:
@@ -80,6 +82,7 @@ def train_one_epoch_fused_optimizer(
     dataset: IterableDataset[Batch],
     device: torch.device,
 ) -> float:
+
     start_time = time.perf_counter()
 
     for data in dataset:
@@ -103,6 +106,7 @@ def train(
     device: torch.device,
     epochs: int = 100,
 ) -> Tuple[float, float]:
+
     training_time = []
     for _ in range(epochs):
         if optimizer:

examples/bert4rec/bert4rec_main.py

@@ -35,6 +35,7 @@
 
 # OSS import
 try:
+
     # pyre-ignore[21]
     # @manual=//torchrec/github/examples/bert4rec:bert4rec_metrics
     from bert4rec_metrics import recalls_and_ndcgs_for_ks

examples/golden_training/train_dlrm_data_parallel.py

@@ -160,7 +160,7 @@ def train(
     )
 
     def dense_filter(
-        named_parameters: Iterator[Tuple[str, nn.Parameter]],
+        named_parameters: Iterator[Tuple[str, nn.Parameter]]
     ) -> Iterator[Tuple[str, nn.Parameter]]:
         for fqn, param in named_parameters:
             if "sparse" not in fqn:

examples/retrieval/two_tower_retrieval.py

@@ -27,6 +27,7 @@
 
 # OSS import
 try:
+
     # pyre-ignore[21]
     # @manual=//torchrec/github/examples/retrieval:knn_index
     from knn_index import get_index

tools/lint/black_linter.py

@@ -179,9 +179,7 @@ def main() -> None:
         level=(
             logging.NOTSET
             if args.verbose
-            else logging.DEBUG
-            if len(args.filenames) < 1000
-            else logging.INFO
+            else logging.DEBUG if len(args.filenames) < 1000 else logging.INFO
         ),
         stream=sys.stderr,
     )

torchrec/datasets/criteo.py

@@ -351,9 +351,8 @@ def get_file_row_ranges_and_remainder(
 
             # If the ranges overlap.
             if rank_left_g <= file_right_g and rank_right_g >= file_left_g:
-                overlap_left_g, overlap_right_g = (
-                    max(rank_left_g, file_left_g),
-                    min(rank_right_g, file_right_g),
+                overlap_left_g, overlap_right_g = max(rank_left_g, file_left_g), min(
+                    rank_right_g, file_right_g
                 )
 
                 # Convert overlap in global numbers to (local) numbers specific to the

torchrec/datasets/random.py

@@ -33,6 +33,7 @@ def __init__(
         *,
         min_ids_per_features: Optional[List[int]] = None,
     ) -> None:
+
         self.keys = keys
         self.keys_length: int = len(keys)
         self.batch_size = batch_size
@@ -75,6 +76,7 @@ def __next__(self) -> Batch:
         return batch
 
     def _generate_batch(self) -> Batch:
+
         values = []
         lengths = []
         for key_idx, _ in enumerate(self.keys):

torchrec/datasets/test_utils/criteo_test_utils.py

@@ -103,6 +103,7 @@ def _create_dataset_npys(
         labels: Optional[np.ndarray] = None,
     ) -> Generator[Tuple[str, ...], None, None]:
         with tempfile.TemporaryDirectory() as tmpdir:
+
             if filenames is None:
                 filenames = [filename]
 

torchrec/distributed/batched_embedding_kernel.py

@@ -785,7 +785,9 @@ def purge(self) -> None:
     def named_split_embedding_weights(
         self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
-        assert remove_duplicate, "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
         for config, param in zip(
             self._config.embedding_tables,
             self.emb_module.split_embedding_weights(),
@@ -897,7 +899,9 @@ def named_parameters(
     def named_split_embedding_weights(
         self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
-        assert remove_duplicate, "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
         for config, tensor in zip(
             self._config.embedding_tables,
             self.split_embedding_weights(),
@@ -1078,9 +1082,8 @@ def named_parameters(
         combined_key = "/".join(
             [config.name for config in self._config.embedding_tables]
         )
-        yield (
-            append_prefix(prefix, f"{combined_key}.weight"),
-            cast(nn.Parameter, self._emb_module.weights),
+        yield append_prefix(prefix, f"{combined_key}.weight"), cast(
+            nn.Parameter, self._emb_module.weights
         )
 
 
@@ -1098,8 +1101,7 @@ def __init__(
         self._pg = pg
 
         self._pooling: PoolingMode = pooling_type_to_pooling_mode(
-            config.pooling,
-            sharding_type,  # pyre-ignore[6]
+            config.pooling, sharding_type  # pyre-ignore[6]
         )
 
         self._local_rows: List[int] = []
@@ -1218,7 +1220,9 @@ def purge(self) -> None:
     def named_split_embedding_weights(
         self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, torch.Tensor]]:
-        assert remove_duplicate, "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
         for config, tensor in zip(
             self._config.embedding_tables,
             self.emb_module.split_embedding_weights(),
@@ -1358,7 +1362,9 @@ def named_parameters(
     def named_split_embedding_weights(
         self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
     ) -> Iterator[Tuple[str, PartiallyMaterializedTensor]]:
-        assert remove_duplicate, "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in BaseBatchedEmbedding.named_split_embedding_weights"
         for config, tensor in zip(
             self._config.embedding_tables,
             self.split_embedding_weights(),
@@ -1561,7 +1567,6 @@ def named_parameters(
         combined_key = "/".join(
             [config.name for config in self._config.embedding_tables]
         )
-        yield (
-            append_prefix(prefix, f"{combined_key}.weight"),
-            cast(nn.Parameter, self._emb_module.weights),
+        yield append_prefix(prefix, f"{combined_key}.weight"), cast(
+            nn.Parameter, self._emb_module.weights
         )

torchrec/distributed/benchmark/benchmark_inference.py

@@ -250,7 +250,9 @@ def main() -> None:
             mb = int(float(num * dim) / 1024 / 1024)
             tables_info += f"\nTABLE[{i}][{num:9}, {dim:4}] u8: {mb:6}Mb"
 
-        report: str = f"REPORT BENCHMARK {datetime_sfx} world_size:{args.world_size} batch_size:{args.batch_size}\n"
+        report: str = (
+            f"REPORT BENCHMARK {datetime_sfx} world_size:{args.world_size} batch_size:{args.batch_size}\n"
+        )
         report += f"Module: {module_name}\n"
         report += tables_info
         report += "\n"

torchrec/distributed/benchmark/benchmark_train.py

@@ -157,7 +157,9 @@ def main() -> None:
             tables_info += f"\nTABLE[{i}][{num:9}, {dim:4}] {mb:6}Mb"
 
         ### Benchmark no VBE
-        report: str = f"REPORT BENCHMARK {datetime_sfx} world_size:{args.world_size} batch_size:{args.batch_size}\n"
+        report: str = (
+            f"REPORT BENCHMARK {datetime_sfx} world_size:{args.world_size} batch_size:{args.batch_size}\n"
+        )
         report += f"Module: {module_name}\n"
         report += tables_info
         report += "\n"
@@ -179,7 +181,9 @@ def main() -> None:
         )
 
         ### Benchmark with VBE
-        report: str = f"REPORT BENCHMARK (VBE) {datetime_sfx} world_size:{args.world_size} batch_size:{args.batch_size}\n"
+        report: str = (
+            f"REPORT BENCHMARK (VBE) {datetime_sfx} world_size:{args.world_size} batch_size:{args.batch_size}\n"
+        )
         report += f"Module: {module_name} (VBE)\n"
         report += tables_info
         report += "\n"

torchrec/distributed/benchmark/benchmark_utils.py

@@ -128,7 +128,6 @@ def __str__(self) -> str:
 @dataclass
 class BenchmarkResult:
     "Class for holding results of benchmark runs"
-
     short_name: str
     elapsed_time: torch.Tensor  # milliseconds
     mem_stats: List[MemoryStats]  # memory stats per rank
@@ -555,7 +554,9 @@ def fx_script_module(eager_module: torch.nn.Module) -> torch.nn.Module:
         return fx_script_module(
             # pyre-fixme[6]: For 1st argument expected `Module` but got
             #  `Optional[Module]`.
-            sharded_module if not benchmark_unsharded_module else module
+            sharded_module
+            if not benchmark_unsharded_module
+            else module
         )
     else:
         # pyre-fixme[7]: Expected `Module` but got `Optional[Module]`.
@@ -966,6 +967,7 @@ def multi_process_benchmark(
     # pyre-ignore
     **kwargs,
 ) -> BenchmarkResult:
+
     def setUp() -> None:
         if "MASTER_ADDR" not in os.environ:
             os.environ["MASTER_ADDR"] = str("localhost")

torchrec/distributed/comm_ops.py

@@ -477,6 +477,7 @@ def variable_batch_alltoall_pooled(
     group: Optional[dist.ProcessGroup] = None,
     codecs: Optional[QuantizedCommCodecs] = None,
 ) -> Awaitable[Tensor]:
+
     if group is None:
         group = dist.distributed_c10d._get_default_group()
 

torchrec/distributed/composable/tests/test_embedding.py

@@ -210,6 +210,7 @@ def test_sharding_ebc(
         use_apply_optimizer_in_backward: bool,
         use_index_dedup: bool,
     ) -> None:
+
         WORLD_SIZE = 2
 
         embedding_config = [

torchrec/distributed/composable/tests/test_embeddingbag.py

@@ -292,6 +292,7 @@ def test_sharding_ebc(
         sharding_type: str,
         use_apply_optimizer_in_backward: bool,
     ) -> None:
+
         # TODO DistributedDataParallel needs full support of registering fused optims before we can enable this.
         assume(
             not (

torchrec/distributed/embedding.py

@@ -167,6 +167,7 @@ def create_sharding_infos_by_sharding(
     table_name_to_parameter_sharding: Dict[str, ParameterSharding],
     fused_params: Optional[Dict[str, Any]],
 ) -> Dict[str, List[EmbeddingShardingInfo]]:
+
     if fused_params is None:
         fused_params = {}
 
@@ -248,6 +249,7 @@ def create_sharding_infos_by_sharding_device_group(
     table_name_to_parameter_sharding: Dict[str, ParameterSharding],
     fused_params: Optional[Dict[str, Any]],
 ) -> Dict[Tuple[str, str], List[EmbeddingShardingInfo]]:
+
     if fused_params is None:
         fused_params = {}
 

torchrec/distributed/embedding_kernel.py

@@ -105,15 +105,7 @@ def get_key_from_embedding_table(embedding_table: ShardedEmbeddingTable) -> str:
 
         assert embedding_table.local_rows == param.size(  # pyre-ignore[16]
             0
-        ), (
-            # pyre-fixme[16]: Item `Tuple` of `PartiallyMaterializedTensor | Tensor
-            #  | Module | Tuple[Tensor, Optional[Tensor], Optional[Tensor]]` has no
-            #  attribute `size`.
-            # pyre-fixme[16]: Item `Tuple` of `PartiallyMaterializedTensor | Tensor
-            #  | Module | Tuple[Tensor, Optional[Tensor], Optional[Tensor]]` has no
-            #  attribute `shape`.
-            f"{embedding_table.local_rows=}, {param.size(0)=}, {param.shape=}"
-        )
+        ), f"{embedding_table.local_rows=}, {param.size(0)=}, {param.shape=}"  # pyre-ignore[16]
 
         if qscale is not None:
             assert embedding_table.local_cols == param.size(1)  # pyre-ignore[16]

torchrec/distributed/embedding_tower_sharding.py

@@ -237,6 +237,7 @@ def input_dist(
         features: KeyedJaggedTensor,
         optional_features: Optional[KeyedJaggedTensor] = None,
     ) -> Awaitable[Awaitable[KJTList]]:
+
         # optional_features are populated only if both kjt and weighted kjt present in tower
         if self._wkjt_feature_names and self._kjt_feature_names:
             kjt_features = features
@@ -505,7 +506,9 @@ def __init__(
                 if lt_tables.issubset(pt_tables):
                     found_physical_tower = True
                     break
-            assert found_physical_tower, f"tables in a logical tower must be in the same physical tower, logical tower tables: {lt_tables}, tables_per_pt: {tables_per_pt}"
+            assert (
+                found_physical_tower
+            ), f"tables in a logical tower must be in the same physical tower, logical tower tables: {lt_tables}, tables_per_pt: {tables_per_pt}"
 
         logical_to_physical_order: List[List[int]] = [
             [] for _ in range(self._cross_pg_world_size)
@@ -604,6 +607,7 @@ def _create_input_dist(
         kjt_feature_names: List[str],
         wkjt_feature_names: List[str],
     ) -> None:
+
         if self._kjt_feature_names != kjt_feature_names:
             self._has_kjt_features_permute = True
             for f in self._kjt_feature_names:
@@ -940,6 +944,7 @@ def __init__(
         fused_params: Optional[Dict[str, Any]] = None,
         qcomm_codecs_registry: Optional[Dict[str, QuantizedCommCodecs]] = None,
     ) -> None:
+
         super().__init__(
             fused_params=fused_params, qcomm_codecs_registry=qcomm_codecs_registry
         )
@@ -955,6 +960,7 @@ def shard(
         device: Optional[torch.device] = None,
         module_fqn: Optional[str] = None,
     ) -> ShardedEmbeddingTowerCollection:
+
         return ShardedEmbeddingTowerCollection(
             module=module,
             table_name_to_parameter_sharding=params,

torchrec/distributed/embedding_types.py

@@ -514,6 +514,7 @@ def sharding_types(self, compute_device_type: str) -> List[str]:
         return types
 
     def shardable_parameters(self, module: M) -> Dict[str, nn.Parameter]:
+
         shardable_params: Dict[str, nn.Parameter] = {}
         for name, param in module.state_dict().items():
             if name.endswith(".weight"):

torchrec/distributed/embeddingbag.py

@@ -208,6 +208,7 @@ def create_sharding_infos_by_sharding(
     fused_params: Optional[Dict[str, Any]],
     suffix: Optional[str] = "weight",
 ) -> Dict[str, List[EmbeddingShardingInfo]]:
+
     if fused_params is None:
         fused_params = {}
 
@@ -312,6 +313,7 @@ def create_sharding_infos_by_sharding_device_group(
     fused_params: Optional[Dict[str, Any]],
     suffix: Optional[str] = "weight",
 ) -> Dict[Tuple[str, str], List[EmbeddingShardingInfo]]:
+
     if fused_params is None:
         fused_params = {}
 

torchrec/distributed/fp_embeddingbag.py

@@ -115,6 +115,7 @@ def compute(
         ctx: EmbeddingBagCollectionContext,
         dist_input: KJTList,
     ) -> List[torch.Tensor]:
+
         fp_features = self.apply_feature_processors_to_kjt_list(dist_input)
         return self._embedding_bag_collection.compute(ctx, fp_features)
 
@@ -185,6 +186,7 @@ def shard(
         device: Optional[torch.device] = None,
         module_fqn: Optional[str] = None,
     ) -> ShardedFeatureProcessedEmbeddingBagCollection:
+
         if device is None:
             device = torch.device("cuda")