fix

Author: chiajunglien

dags/tpu_observability/utils/jobset_util.py

@@ -1025,13 +1025,10 @@ def wait_for_jobset_metrics(
   Args:
     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
-      expected count from the jobset_config.
+    expected_value (int): The target value to wait for.
     node_pool (Info): An instance of the Info class containing GKE cluster
       metadata.
-    jobset_config (JobSet): An instance of the JobSet class representing
-      the target jobset configuration.
+    jobset_name: The name of the JobSet.
     start_time (TimeUtil, optional): The UTC timestamp to start polling
       from. If not provided, defaults to 60 minutes before the current time.