Repository: mlflow/mlflow
Pull request: Replace `dev/check-notebooks.sh` with `clint` rule for trace UI detection
This PR replaces the standalone `dev/check-notebooks.sh` script with a proper `clint` rule that integrates seamlessly with the existing linting infrastructure.

## Changes Made

1. **Added new rule**: `ForbiddenTraceUIInNotebook` (MLF0022) in `dev/clint/src/clint/rules.py` that detects the presence of MLflow Trace UI iframe content in notebooks
2. **Updated linter**: Modified `lint_file()` function in `dev/clint/src/clint/linter.py` to check for the forbidden trace UI pattern in `.ipynb` files
3. **Removed shell script**: Deleted `dev/check-notebooks.sh` 
4. **Updated pre-commit config**: Removed the `check-notebooks` hook from `.pre-commit-config.yaml` since `clint` already processes `.ipynb` files

## Functionality

The new rule detects notebooks containing `static-files/lib/notebook-trace-renderer/index.html` and provides the same helpful error message as the original script:

```
Found the MLflow Trace UI iframe in the notebook. The trace UI in cell outputs will not render correctly in previews or the website. Please run `mlflow.tracing.disable_notebook_display()` and rerun the cell to remove the iframe.
```

## Testing

The implementation correctly handles:
- Notebooks with trace UI content (triggers MLF0022 violation)
- Clean notebooks (no violations)
- Edge cases like notebooks without cells
- Integration with existing clint rules and pre-commit hooks

## Benefits

- **Consistency**: Uses the same linting infrastructure as other rules
- **Maintainability**: Single codebase instead of separate shell script
- **Integration**: Works with existing `clint` ignore mechanisms and tooling
- **Performance**: No additional pre-commit hook needed

Fixes #16008.

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more [Copilot coding agent tips](https://gh.io/copilot-coding-agent-tips) in the docs.

Problem statement:
Remove `dev/check-notebooks.sh` and use `clint` instead
`clint` is a custom linter used in this repository to check python scripts/notebooks. We should use it to check notebooks instead of the `dev/check-notebooks.sh` script.

Instructions:

1. Add a new rule in `dev/clint/src/clint/rules.py`.
2. Update `_lint_cell` in `dev/clint/src/clint/linter.py` to use the new rule.
3. Remove the `check-notebooks` hook in `.pre-commit-config.yaml`.
4. Remove `dev/check-notebooks.sh`.

Testing:

Create a dummy notebook that violates the new rule, and ensure `clint` catches it.

Hints from issue discussion:

Commit message: Refactor _has_trace_ui_content to reduce nesting

Co-authored-by: harupy <17039389+harupy@users.noreply.github.com>

Patch to review:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 9442f05520350..45aa0f372631e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -49,14 +49,6 @@ repos:
         stages: [pre-commit]
         require_serial: true
 
-      - id: check-notebooks
-        name: check-notebooks
-        entry: dev/check-notebooks.sh
-        language: system
-        files: '\.ipynb$'
-        stages: [pre-commit]
-        require_serial: true
-
       - id: clint
         name: clint
         entry: clint
diff --git a/dev/check-notebooks.sh b/dev/check-notebooks.sh
deleted file mode 100755
index 89fc4c959cdc3..0000000000000
--- a/dev/check-notebooks.sh
+++ /dev/null
@@ -1,13 +0,0 @@
-#!/usr/bin/env bash
-
-TRACE_UI_SRC="static-files/lib/notebook-trace-renderer/index.html"
-
-# If offending files are found, list them and throw an exception
-if grep -In $TRACE_UI_SRC --include="*.ipynb" "$@"; then
-    echo -e "\nError: Found the MLflow Trace UI iframe in the above notebooks."
-    echo "The trace UI in cell outputs will not render correctly in previews or the website"
-    echo "Please run \`mlflow.tracing.disable_notebook_display()\` and rerun the cell to remove the iframe."
-    exit 1
-else
-  exit 0
-fi
diff --git a/dev/clint/src/clint/linter.py b/dev/clint/src/clint/linter.py
index 65c5cd6cb6d2e..8ff002f5a42f0 100644
--- a/dev/clint/src/clint/linter.py
+++ b/dev/clint/src/clint/linter.py
@@ -674,21 +674,55 @@ def post_visit(self) -> None:
                     self._check(loc, rules.LazyModule())
 
 
+def _has_trace_ui_content(output: dict[str, Any]) -> bool:
+    """Check if an output contains MLflow trace UI content."""
+    data = output.get("data")
+    if not data:
+        return False
+
+    for content_type, content in data.items():
+        if not isinstance(content, (str, list)):
+            continue
+
+        content_str = "".join(content) if isinstance(content, list) else content
+        if "static-files/lib/notebook-trace-renderer/index.html" in content_str:
+            return True
+
+    return False
+
+
 def _lint_cell(path: Path, config: Config, cell: dict[str, Any], index: int) -> list[Violation]:
+    violations = []
     type_ = cell.get("cell_type")
+
+    # Check for forbidden trace UI iframe in cell outputs
+    outputs = cell.get("outputs")
+    if outputs:
+        for output in outputs:
+            if _has_trace_ui_content(output):
+                violations.append(
+                    Violation(
+                        rules.ForbiddenTraceUIInNotebook(),
+                        path,
+                        Location(0, 0),
+                        cell=index,
+                    )
+                )
+                break
+
     if type_ != "code":
-        return []
+        return violations
 
     src = "\n".join(cell.get("source", []))
     try:
         tree = ast.parse(src)
     except SyntaxError:
         # Ignore non-python cells such as `!pip install ...`
-        return []
+        return violations
 
     linter = Linter(path=path, config=config, ignore=ignore_map(src), cell=index)
     linter.visit(tree)
-    violations = linter.violations
+    violations.extend(linter.violations)
 
     if not src.strip():
         violations.append(
@@ -705,11 +739,11 @@ def _lint_cell(path: Path, config: Config, cell: dict[str, Any], index: int) ->
 def lint_file(path: Path, config: Config) -> list[Violation]:
     code = path.read_text()
     if path.suffix == ".ipynb":
+        violations = []
         if cells := json.loads(code).get("cells"):
-            violations = []
             for idx, cell in enumerate(cells, start=1):
                 violations.extend(_lint_cell(path, config, cell, idx))
-            return violations
+        return violations
     elif path.suffix in {".rst", ".md", ".mdx"}:
         violations = []
         code_blocks = (
diff --git a/dev/clint/src/clint/rules.py b/dev/clint/src/clint/rules.py
index b4fa3248e70a5..fefb22d1d4e20 100644
--- a/dev/clint/src/clint/rules.py
+++ b/dev/clint/src/clint/rules.py
@@ -351,3 +351,16 @@ def _message(self) -> str:
         return (
             "Usage of `set_active_model` is not allowed in mlflow, use `_set_active_model` instead."
         )
+
+
+class ForbiddenTraceUIInNotebook(Rule):
+    def _id(self) -> str:
+        return "MLF0022"
+
+    def _message(self) -> str:
+        return (
+            "Found the MLflow Trace UI iframe in the notebook. "
+            "The trace UI in cell outputs will not render correctly in previews or the website. "
+            "Please run `mlflow.tracing.disable_notebook_display()` and rerun the cell "
+            "to remove the iframe."
+        )


Code review comments:
diff_hunk: @@ -674,21 +674,55 @@ def post_visit(self) -> None:
                     self._check(loc, rules.LazyModule())
 
 
+def _has_trace_ui_content(output: dict[str, Any]) -> bool:
+    """Check if an output contains MLflow trace UI content."""
+    data = output.get("data")
+    if not data:
+        return False
+
+    for content_type, content in data.items():
+        if not isinstance(content, (str, list)):
+            continue
+
+        content_str = "".join(content) if isinstance(content, list) else content
+        if "static-files/lib/notebook-trace-renderer/index.html" in content_str:
+            return True
+
+    return False
+
+
 def _lint_cell(path: Path, config: Config, cell: dict[str, Any], index: int) -> list[Violation]:
+    violations = []
     type_ = cell.get("cell_type")
+
+    # Check for forbidden trace UI iframe in cell outputs
+    outputs = cell.get("outputs")
+    if outputs:
path: dev/clint/src/clint/linter.py
text: @user1:
```suggestion
    # Check for forbidden trace UI iframe in cell outputs
    if outputs := cell.get("outputs"):
```

@author:
Applied the walrus operator for cleaner, more Pythonic code. Commit: dc3fba7