latest composer-workload-simulator, and ruff formatting (#1480)

Co-authored-by: Andrew Gold <[email protected]>

Author: CYarros10

Diff for: tools/cloud-composer-stress-testing/cloud-composer-dag-generator/dag_generator/__main__.py

@@ -21,7 +21,7 @@
  b) run $python main.py
  c) move dags folder generated to the dag buckets in composer like:
         gsutil cp -r out  gs://BUCKET_NAME/dags
-NOTE: the "number_of_operators_defined" variable in the configuration file 
+NOTE: the "number_of_operators_defined" variable in the configuration file
         (config.json) allows to create up to 5 differents kind of task,
         none has complex functionallity:
         a) bash_operator_echo
@@ -32,22 +32,23 @@
 """
 
 import json
+import math
 import random
+
 import modules.initDag
 import modules.operators
-import math
 
 
 def get_config():
     """module to read configs"""
-    f = open('config.json', "r")
+    f = open("config.json", "r")
     data = json.loads(f.read())
     f.close()
     return data
 
 
 def get_init_content(i):
-    """Initialise test DAG with headers """
+    """Initialise test DAG with headers"""
     modules.initDag.get_init_dag(i)
 
 
@@ -60,70 +61,74 @@ def get_task_dag(min_number_of_task_in_dag):
     file.close()
     return data
 
-# build the dags 
+
+# build the dags
 def main():
     """main function to create test DAGs"""
     # read config file
     data = get_config()
 
-    number_of_dags_to_generate = data['number_of_dags_to_generate']
-    min_number_of_task_in_dag = data['min_number_of_task_in_dag']
-    max_number_of_task_in_dag = data['max_number_of_task_in_dag']
-    task_min_time_in_sec = data['task_min_time_in_sec']
-    task_max_time_in_sec = data['task_max_time_in_sec']
-    percentage_of_job_in_parallel = data['percentage_of_job_in_parallel']
-    number_of_operators_defined = data['number_of_operators_defined']
-    file_index = data['file_start_index']
-    schedules = data['schedules']
+    number_of_dags_to_generate = data["number_of_dags_to_generate"]
+    min_number_of_task_in_dag = data["min_number_of_task_in_dag"]
+    max_number_of_task_in_dag = data["max_number_of_task_in_dag"]
+    task_min_time_in_sec = data["task_min_time_in_sec"]
+    task_max_time_in_sec = data["task_max_time_in_sec"]
+    percentage_of_job_in_parallel = data["percentage_of_job_in_parallel"]
+    number_of_operators_defined = data["number_of_operators_defined"]
+    file_index = data["file_start_index"]
+    schedules = data["schedules"]
 
     # creatting DAG's files
     for i in range(number_of_dags_to_generate):
         task_list = []
         dagf = open(f"out/dagFile_{file_index+i}.py", "w+")
         dagf.write(
             modules.initDag.get_init_dag(
-                file_index + i,
-                schedules[random.randrange(0,
-                                           len(schedules) - 1)]))
+                file_index + i, schedules[random.randrange(0, len(schedules) - 1)]
+            )
+        )
         dagf.write(modules.operators.start_task())
         dagf.write(modules.operators.stop_task())
         for task_index in range(
-                random.randrange(min_number_of_task_in_dag,
-                                 max_number_of_task_in_dag)):
+            random.randrange(min_number_of_task_in_dag, max_number_of_task_in_dag)
+        ):
             task_list.append("task_{index}".format(index=task_index))
-            if (task_index % number_of_operators_defined == 0):
+            if task_index % number_of_operators_defined == 0:
                 dagf.write(modules.operators.bash_operator_echo(task_index))
-            elif (task_index % number_of_operators_defined == 1):
+            elif task_index % number_of_operators_defined == 1:
                 dagf.write(
                     modules.operators.bash_operator_sleep(
                         task_index,
-                        random.randrange(task_min_time_in_sec,
-                                         task_max_time_in_sec)))
-            elif (task_index % number_of_operators_defined == 2):
+                        random.randrange(task_min_time_in_sec, task_max_time_in_sec),
+                    )
+                )
+            elif task_index % number_of_operators_defined == 2:
                 dagf.write(
                     modules.operators.python_operator_task_sleep(
                         task_index,
-                        random.randrange(task_min_time_in_sec,
-                                         task_max_time_in_sec)))
-            elif (task_index % number_of_operators_defined == 3):
-                dagf.write(
-                    modules.operators.bash_operator_task_ping(task_index))
+                        random.randrange(task_min_time_in_sec, task_max_time_in_sec),
+                    )
+                )
+            elif task_index % number_of_operators_defined == 3:
+                dagf.write(modules.operators.bash_operator_task_ping(task_index))
             else:
-                dagf.write(
-                    modules.operators.python_operator_task_print(task_index))
-        no_tasks_in_parallel = math.ceil(percentage_of_job_in_parallel / 100 *
-                                         len(task_list))
+                dagf.write(modules.operators.python_operator_task_print(task_index))
+        no_tasks_in_parallel = math.ceil(
+            percentage_of_job_in_parallel / 100 * len(task_list)
+        )
         parallel_tasks = []
-        if (no_tasks_in_parallel > 1):
+        if no_tasks_in_parallel > 1:
             for parallel_task_index in range(no_tasks_in_parallel):
                 parallel_tasks.append(task_list.pop())
-            task_list.insert(random.randrange(1,
-                                              len(task_list) - 2),
-                             "[{task}]".format(task=",".join(parallel_tasks)))
-        dagf.write("\n\tchain(start_task,{tasks},stop_task)".format(
-            tasks=",".join(task_list)))
+            task_list.insert(
+                random.randrange(1, len(task_list) - 2),
+                "[{task}]".format(task=",".join(parallel_tasks)),
+            )
+        dagf.write(
+            "\n\tchain(start_task,{tasks},stop_task)".format(tasks=",".join(task_list))
+        )
         dagf.close()
 
 
 if __name__ == "__main__":
-    main()
+    main()

Diff for: tools/cloud-composer-stress-testing/cloud-composer-dag-generator/dag_generator/modules/initDag.py

@@ -18,20 +18,22 @@
 
 def get_init_dag(dag_number, schedule):
     today = datetime.datetime.now()
-    if (schedule == "min"):
+    if schedule == "min":
         dag_schedule = "{min} * * * *".format(min=random.randrange(1, 10))
-    elif (schedule == "hour"):
-        dag_schedule = "{min} {hour} * * *".format(min=random.randrange(0, 59),
-                                                   hour=random.randrange(0, 23))
-    elif (schedule == "everyhalfhour"):
+    elif schedule == "hour":
+        dag_schedule = "{min} {hour} * * *".format(
+            min=random.randrange(0, 59), hour=random.randrange(0, 23)
+        )
+    elif schedule == "everyhalfhour":
         dag_schedule = "*/30 * * * *"
-    elif (schedule == "everyhour"):
+    elif schedule == "everyhour":
         dag_schedule = "0 * * * *"
     else:
         dag_schedule = "{min} {hour} {day} * *".format(
             min=random.randrange(0, 59),
             hour=random.randrange(0, 23),
-            day=random.randrange(1, 28))
+            day=random.randrange(1, 28),
+        )
 
     lines = """import time
 from datetime import datetime
@@ -52,9 +54,11 @@ def get_init_dag(dag_number, schedule):
     default_args=default_args,
     catchup=False
 ) as dag:
-    """.format(dag_id=dag_number,
-               start_year=today.year,
-               schedule=dag_schedule,
-               start_month=today.month,
-               start_day=today.day)
+    """.format(
+        dag_id=dag_number,
+        start_year=today.year,
+        schedule=dag_schedule,
+        start_month=today.month,
+        start_day=today.day,
+    )
     return lines

Diff for: tools/cloud-composer-stress-testing/cloud-composer-workload-simulator/.gitignore

@@ -1,3 +1,3 @@
 **/.DS_Store
 taskflow_collections/__pycache__
-utils/*
+.ruff_cache

Diff for: tools/cloud-composer-stress-testing/cloud-composer-workload-simulator/README.md

@@ -61,6 +61,7 @@ Sizing a Composer Environment is an estimated process. There are many different
 experiment_id: experiment_1
 number_of_dags: 10
 tasks_per_dag: 3
+paused: 0.5 # weighted chance for dag to be paused
 
 # Schedules and weights
 schedules:
@@ -97,7 +98,7 @@ default_settings:
   deferrable: true
   retries: 1
   catchup: false
-  is_paused_upon_creation: false
+  is_paused_upon_creation: false # true will pause all dags
   execution_timeout: 30
   sla: 25
   project_id: your-project

Diff for: tools/cloud-composer-stress-testing/cloud-composer-workload-simulator/configs/custom.yaml

@@ -15,6 +15,7 @@
 experiment_id: custom_experiment_1
 number_of_dags: 1
 min_tasks_per_dag: 1
+paused: 0.5
 
 # Schedules and weights
 schedules:
@@ -35,13 +36,13 @@ default_settings:
   project_id: your-project
   region: your-region
   composer_environment: your-environment
-  deferrable: true
+  deferrable: false
   retries: 1
-  retry_delay: 2 # minutes
+  retry_delay: 2  # minutes
   catchup: false
-  is_paused_upon_creation: false
-  dagrun_timeout: 60 # minutes
-  execution_timeout: 30 # minutes
-  sla: 25 # minutes
-  mode: poke
-  poke_interval: 120 # seconds
+  is_paused_upon_creation: false # True will override and set all dags to paused
+  dagrun_timeout: 60  # minutes 
+  execution_timeout: 30  # minutes
+  sla: 25  # minutes
+  mode: reschedule
+  poke_interval: 60  # seconds

Diff for: tools/cloud-composer-stress-testing/cloud-composer-workload-simulator/configs/sample.yaml

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-experiment_id: experiment_4
+experiment_id: sample
 number_of_dags: 1
 min_tasks_per_dag: 10
+paused: 0.5
 
 # Schedules and weights
 schedules:
@@ -52,13 +53,13 @@ default_settings:
   project_id: your-project
   region: your-region
   composer_environment: your-environment
-  deferrable: true
+  deferrable: false
   retries: 1
-  retry_delay: 2 # minutes
+  retry_delay: 2  # minutes
   catchup: false
-  is_paused_upon_creation: false
-  dagrun_timeout: 60 # minutes
-  execution_timeout: 30 # minutes
-  sla: 25 # minutes
-  mode: poke
-  poke_interval: 120 # seconds
+  is_paused_upon_creation: false # True will override and set all dags to paused
+  dagrun_timeout: 60  # minutes 
+  execution_timeout: 30  # minutes
+  sla: 25  # minutes
+  mode: reschedule
+  poke_interval: 60  # seconds