diff --git a/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py b/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py index 9c01c72c7..aa12bfe73 100644 --- a/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py +++ b/envs/monkey_zoo/blackbox/utils/gcp_machine_handlers.py @@ -45,24 +45,16 @@ def start_machines(machine_list): """ LOGGER.info("Setting up all GCP machines...") try: - arglist = [] - for zone in machine_list: - arglist.append((MACHINE_STARTING_COMMAND, machine_list, zone)) - with Pool(2) as pool: - pool.map(run_gcp_command, arglist) - LOGGER.info("GCP machines successfully started.") + run_gcp_pool(MACHINE_STARTING_COMMAND, machine_list) + LOGGER.info("GCP machines successfully started.") except Exception as e: LOGGER.error("GCP Handler failed to start GCP machines: %s" % e) def stop_machines(machine_list): try: - arglist = [] - for zone in machine_list: - arglist.append((MACHINE_STOPPING_COMMAND, machine_list, zone)) - with Pool(2) as pool: - pool.map(run_gcp_command, arglist) - LOGGER.info("GCP machines stopped successfully.") + run_gcp_pool(MACHINE_STOPPING_COMMAND, machine_list) + LOGGER.info("GCP machines stopped successfully.") except Exception as e: LOGGER.error("GCP Handler failed to stop network machines: %s" % e) @@ -78,6 +70,13 @@ def get_set_project_command(project): def run_gcp_command(arglist): gcp_cmd, machine_list, zone = arglist subprocess.call( # noqa DUO116 - (gcp_cmd % (" ".join(machine_list[zone]), zone)), + (gcp_cmd % (" ".join(machine_list), zone)), shell=True, ) + + +def run_gcp_pool(gcp_command, machine_list): + arglist = [(gcp_command, machine_list[zone], zone) for zone in machine_list] + with Pool(2) as pool: + pool.map(run_gcp_command, arglist) +