chore: import upstream snapshot with attribution

release/train_tests/benchmark/train_benchmark.py

@@ -0,0 +1,123 @@
+import json
+import logging
+import os
+import pprint
+import time
+
+import ray.train
+from ray._private.test_utils import safe_write_to_results_json
+from ray.train.torch import TorchTrainer
+from ray.train.v2._internal.util import date_str
+
+from config import BenchmarkConfig, cli_to_config
+from benchmark_factory import BenchmarkFactory
+from ray_dataloader_factory import RayDataLoaderFactory
+
+
+logger = logging.getLogger(__name__)
+
+
+METRICS_OUTPUT_PATH = "/mnt/cluster_storage/train_benchmark_metrics.json"
+
+
+def train_fn_per_worker(config):
+    factory: BenchmarkFactory = config["factory"]
+
+    if factory.benchmark_config.task == "recsys":
+        from recsys.torchrec_runner import TorchRecRunner
+
+        runner = TorchRecRunner(factory)
+    else:
+        from runner import VanillaTorchRunner
+
+        runner = VanillaTorchRunner(factory)
+
+    runner.run()
+
+    metrics = runner.get_metrics(
+        dataset_creation_time=config.get("dataset_creation_time", 0)
+    )
+    if ray.train.get_context().get_world_rank() == 0:
+        with open(METRICS_OUTPUT_PATH, "w") as f:
+            json.dump(metrics, f)
+
+
+def get_datasets_and_data_config(factory: BenchmarkFactory):
+    dataloader_factory = factory.get_dataloader_factory()
+    if isinstance(dataloader_factory, RayDataLoaderFactory):
+        datasets = dataloader_factory.get_ray_datasets()
+        data_config = dataloader_factory.get_ray_data_config()
+    else:
+        datasets = {}
+        data_config = None
+
+    return datasets, data_config
+
+
+def main():
+    start_time = time.perf_counter()
+    logging.basicConfig(level=logging.INFO)
+
+    benchmark_config: BenchmarkConfig = cli_to_config()
+    logger.info(
+        "\nBenchmark config:\n" + pprint.pformat(benchmark_config.__dict__, indent=2)
+    )
+
+    if benchmark_config.task == "image_classification":
+        from image_classification.factory import ImageClassificationFactory
+
+        factory = ImageClassificationFactory(benchmark_config)
+    elif benchmark_config.task == "recsys":
+        from recsys.recsys_factory import RecsysFactory
+
+        factory = RecsysFactory(benchmark_config)
+    else:
+        raise ValueError(f"Unknown task: {benchmark_config.task}")
+
+    datasets, data_config = get_datasets_and_data_config(factory)
+
+    dataset_creation_time = time.perf_counter() - start_time
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn_per_worker,
+        train_loop_config={
+            "factory": factory,
+            "dataset_creation_time": dataset_creation_time,
+        },
+        scaling_config=ray.train.ScalingConfig(
+            num_workers=benchmark_config.num_workers,
+            use_gpu=not benchmark_config.mock_gpu,
+            resources_per_worker={"MOCK_GPU": 1} if benchmark_config.mock_gpu else None,
+        ),
+        run_config=ray.train.RunConfig(
+            storage_path=f"{os.environ['ANYSCALE_ARTIFACT_STORAGE']}/train_benchmark/",
+            name=f"{benchmark_config.task}-{date_str(include_ms=True)}",
+            failure_config=ray.train.FailureConfig(
+                max_failures=benchmark_config.max_failures
+            ),
+        ),
+        datasets=datasets,
+        dataset_config=data_config,
+    )
+
+    trainer.fit()
+
+    end_time = time.perf_counter()
+
+    with open(METRICS_OUTPUT_PATH, "r") as f:
+        metrics = json.load(f)
+
+    final_metrics_str = (
+        f"\nTotal training time: {end_time - start_time} seconds\n"
+        "Final metrics:\n" + "-" * 80 + "\n" + pprint.pformat(metrics) + "\n" + "-" * 80
+    )
+    logger.info(final_metrics_str)
+
+    # Write metrics as a release test result.
+    safe_write_to_results_json(metrics)
+
+
+if __name__ == "__main__":
+    # Workers need to access the working directory module.
+    ray.init(runtime_env={"working_dir": os.path.dirname(__file__)})
+    main()