From 4c9056e4b0772bdc09dc382abad4a53fcbed9a21 Mon Sep 17 00:00:00 2001 From: Andreas Motl Date: Tue, 5 Mar 2024 19:27:34 +0100 Subject: [PATCH] Autoscaling: Refactor code to top-level code environment For making `autoscale.py` a real program, and good Python citizen, it is important to consider your file might be considered as a module, and _imported_. In this case, you don't want to have your main code executed. https://docs.python.org/3/library/__main__.html In this case, the program would have caused any sort of test runner, which scans directory trees for functions to executed, by, well, importing all modules, to eventually run into an eternal loop. When putting code into the top-level code environment, as exercised on behalf of the `__name__ == '__main__'` code guard, it will safely execute as a Python program, but not when used as a Python module. In this spirit, a Python file can be both a program and a module at the same time. --- topic/autoscaling/autoscale.py | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/topic/autoscaling/autoscale.py b/topic/autoscaling/autoscale.py index 4ba14613..a0c061b2 100644 --- a/topic/autoscaling/autoscale.py +++ b/topic/autoscaling/autoscale.py @@ -131,16 +131,22 @@ def scale_cluster(num, cluster_status, max_num_shard): DELAY_SECONDS = 5 MAX_NUM_SHARDS = 30 -# Main loop to monitor and adjust cluster size based on shard count -while True: - try: - status = get_cluster_status() # Fetch current cluster status - number_shards = num_shards() # Calculate average shard count - if number_shards is not None: - logging.info(f"Current avg number of shards: {number_shards}") - scale_cluster(number_shards, status, MAX_NUM_SHARDS) - else: - logging.error("Failed to retrieve shard metrics.") - except Exception as e: - logging.error(f"An error occurred: {e}") - time.sleep(DELAY_SECONDS) + +def main(): + # Main loop to monitor and adjust cluster size based on shard count + while True: + try: + status = get_cluster_status() # Fetch current cluster status + number_shards = num_shards() # Calculate average shard count + if number_shards is not None: + logging.info(f"Current avg number of shards: {number_shards}") + scale_cluster(number_shards, status, MAX_NUM_SHARDS) + else: + logging.error("Failed to retrieve shard metrics.") + except Exception as e: + logging.error(f"An error occurred: {e}") + time.sleep(DELAY_SECONDS) + + +if __name__ == "__main__": + main()