Fix: Add validation of success status

Author: chiajunglien

dags/tpu_observability/jobset_healthiness_validation.py

@@ -38,6 +38,7 @@
 SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
 
 FAIL_WORKLOAD = "python3 -c 'import logging; import sys; logging.error(\"Simulating Failure\"); sys.exit(1)'"
+SUCCESS_WORKLOAD = "python3 -c 'import logging; import sys; logging.info(\"Simulating Success\"); sys.exit(0)'"
 
 # Keyword arguments are generated dynamically at runtime (pylint does not
 # know this signature).
@@ -164,6 +165,33 @@
           jobset_config=jobset_config,
       )
 
+      with TaskGroup(group_id="inject_and_validate_success") as success_test:
+        cleanup_for_success = jobset.end_workload.override(
+            task_id="cleanup_before_success_injection"
+        )(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+        )
+
+        start_success_job = jobset.run_workload.override(
+            task_id="start_success_job"
+        )(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+            workload_type=SUCCESS_WORKLOAD,
+        )
+
+        validate_succeeded_metric = jobset.wait_for_jobset_metrics.override(
+            task_id="wait_for_succeeded_count"
+        )(
+            metric_name=ReplicatedJobStatus.SUCCEEDED,
+            expected_value="USE_CONFIG_REPLICAS",
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+        )
+
+        chain(cleanup_for_success, start_success_job, validate_succeeded_metric)
+
       with TaskGroup(group_id="inject_and_validate_failure") as failure_test:
         cleanup_for_failure = jobset.end_workload.override(
             task_id="cleanup_before_failure_injection"
@@ -214,6 +242,7 @@
           suspend_action,
           validate_suspended,
           resume_action,
+          success_test,
           failure_test,
           cleanup_workload,
           cleanup_node_pool,