fix

Author: chiajunglien

dags/tpu_observability/utils/jobset_util.py

@@ -916,7 +916,7 @@ def wait_for_jobset_metrics(
   types returned by the Prometheus exporter.
 
   Args:
-    metric_name (ReplicatedJobStatus): The status of the replicated job to
+    metric_name (JobSetHealthiness): The status of the replicated job to
       monitor.
     expected_value (any): The target value to wait for. If set to
       "USE_CONFIG_REPLICAS", the sensor will dynamically retrieve the