diff --git a/scripts/ci/prek/check_metrics_synced_with_the_registry.py b/scripts/ci/prek/check_metrics_synced_with_the_registry.py index 140e6ab9e91ed..c680b56661de3 100644 --- a/scripts/ci/prek/check_metrics_synced_with_the_registry.py +++ b/scripts/ci/prek/check_metrics_synced_with_the_registry.py @@ -97,16 +97,25 @@ def normalize_metric_name(registry_metric_name: str) -> str: return re.sub(r"\{[^}]+\}", "*", registry_metric_name) -# Sentinel returned when a dynamic metric name is partially matched based on a common prefix. +# Sentinel returned when a dynamic metric name is partially matched based on its static part. # For dynamic metric names that include variables, the check can't find an exact match with a registry -# entry or its type. So, a partially matched prefix is good enough and type checking is skipped. +# entry or its type. So, a partial match is good enough and type checking is skipped. _PREFIX_MATCHED = "__prefix_matched__" def find_prefix_matched_registry_entries(metric_name: str, metrics_registry: dict[str, dict]) -> list[str]: - """Return the registry entry names whose name matches the static prefix of a dynamic metric name.""" + """Return the registry entry names whose name matches the static part of a dynamic metric name. + + ``{stats_prefix}.cache_hit`` has no static prefix, so it matches on its static suffix, which + must start with a dot to only match whole segments. + """ base = metric_name.split("{")[0].rstrip(".") - return [name for name in metrics_registry if name == base or name.startswith(base + ".")] + if base: + return [name for name in metrics_registry if name == base or name.startswith(base + ".")] + suffix = metric_name.rsplit("}", 1)[-1] + if not suffix.startswith("."): + return [] + return [name for name in metrics_registry if name.endswith(suffix)] def find_registry_match(metric_name: str, metrics_registry: dict[str, dict]) -> str | None: @@ -127,11 +136,8 @@ def find_registry_match(metric_name: str, metrics_registry: dict[str, dict]) -> # Dynamic metric name. if "{" in metric_name and find_prefix_matched_registry_entries(metric_name, metrics_registry): - # Metric prefix matches the prefix of a dynamic registry entry. - # If the static part before the first variable, matches an exact registry entry name, - # or a dotted-prefix of one, then the name is considered covered and - # _PREFIX_MATCHED is returned. The type check must be skipped because - # the resulting metric name with all variables expanded, cannot be determined. + # The type check must be skipped because the name with all variables expanded + # cannot be determined. return _PREFIX_MATCHED # All checks for matching failed. diff --git a/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py b/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py index e9b5e2ebe2979..c6b895987db59 100644 --- a/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py +++ b/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py @@ -119,6 +119,7 @@ def test_normalize_metric_name(metric_name, expected_result): pytest.param("dag.{x}.{y}.duration", "task.duration", id="legacy_name_match_different_structure"), pytest.param("ti.{state}", _PREFIX_MATCHED, id="prefix_match_returns_sentinel"), pytest.param("dagrun.duration.{state}", _PREFIX_MATCHED, id="prefix_match_dotted_base"), + pytest.param("{stats_prefix}.open_slots", _PREFIX_MATCHED, id="suffix_match_returns_sentinel"), pytest.param("non.existent.{var}", None, id="dynamic_metric_no_prefix_match_returns_none"), pytest.param("non.existent", None, id="static_metric_not_in_registry_returns_none"), ], @@ -226,6 +227,12 @@ def test_extract_metric_names_from_ast_node(code: str, expected_result): ), pytest.param("dagrun.duration.{state}", ["dagrun.duration.success"], id="dotted_base_prefix"), pytest.param("non.existent.{var}", [], id="no_prefix_match_returns_empty_list"), + pytest.param( + "{stats_prefix}.open_slots", + ["pool.open_slots", "executor.open_slots"], + id="leading_variable_matches_on_static_suffix", + ), + pytest.param("{stats_prefix}pen_slots", [], id="suffix_not_starting_at_a_dot_returns_empty_list"), ], ) def test_find_prefix_matched_registry_entries(metric_name, expected_result): @@ -277,6 +284,18 @@ def test_find_prefix_matched_registry_entries(metric_name, expected_result): ], id="legacy_name_structure_match_marks_entry_used", ), + pytest.param( + {"{stats_prefix}.open_slots"}, + [ + "dagrun.duration.success", + "scheduler.heartbeat", + "task.duration", + "ti.queued", + "ti.scheduled", + "ti.start.{dag_id}.{task_id}", + ], + id="suffix_match_marks_all_matching_entries_used", + ), ], ) def test_compute_unused_registry_entries(code_metric_names, expected_unused):