From b85be760c15c84a20e66ea59315316e594f7c522 Mon Sep 17 00:00:00 2001 From: QuanyiLi Date: Mon, 8 May 2023 10:02:13 +0100 Subject: [PATCH] auto reduce worker num --- scenarionet/tests/test_combine_dataset.py | 3 +-- scenarionet/verifier/utils.py | 17 ++++++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/scenarionet/tests/test_combine_dataset.py b/scenarionet/tests/test_combine_dataset.py index b14bbb8..07228c7 100644 --- a/scenarionet/tests/test_combine_dataset.py +++ b/scenarionet/tests/test_combine_dataset.py @@ -18,9 +18,8 @@ def test_combine_multiple_dataset(): summary, sorted_scenarios, mapping = read_dataset_summary(dataset_path) for scenario_file in sorted_scenarios: read_scenario(os.path.join(dataset_path, mapping[scenario_file], scenario_file)) - num_worker = 4 if len(summary) > 4 else 1 success, result = verify_loading_into_metadrive( - dataset_path, result_save_dir="test_dataset", steps_to_run=1000, num_workers=num_worker + dataset_path, result_save_dir="test_dataset", steps_to_run=1000, num_workers=4 ) assert success diff --git a/scenarionet/verifier/utils.py b/scenarionet/verifier/utils.py index 54fb23c..3d34e3f 100644 --- a/scenarionet/verifier/utils.py +++ b/scenarionet/verifier/utils.py @@ -16,8 +16,12 @@ def verify_loading_into_metadrive(dataset_path, result_save_dir, steps_to_run=10 assert os.path.exists(result_save_dir ) and os.path.isdir(result_save_dir), "Argument result_save_dir must be an existing dir" num_scenario = get_number_of_scenarios(dataset_path) - argument_list = [] + if num_scenario < num_workers: + # single process + logger.info("Use one worker, as num_scenario < num_workers:") + num_workers = 1 + argument_list = [] func = partial(loading_wrapper, dataset_path=dataset_path, steps_to_run=steps_to_run) num_scenario_each_worker = int(num_scenario // num_workers) @@ -41,18 +45,17 @@ def verify_loading_into_metadrive(dataset_path, result_save_dir, steps_to_run=10 json.dump(logs, f, indent=4) if result: - print("All scenarios can be loaded successfully!") + logger.info("All scenarios can be loaded successfully!") else: - print("Fail to load all scenarios, see log for more details! Number of failed scenarios: {}".format(len(logs))) + logger.info( + "Fail to load all scenarios, see log for more details! Number of failed scenarios: {}".format(len(logs))) return result, logs def loading_into_metadrive(start_scenario_index, num_scenario, dataset_path, steps_to_run): - print( + logger.info( "================ Begin Scenario Loading Verification for scenario {}-{} ================ \n".format( - start_scenario_index, num_scenario + start_scenario_index - ) - ) + start_scenario_index, num_scenario + start_scenario_index)) success = True env = ScenarioEnv( {