chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,10 @@
|
||||
FROM python:3.10-slim
|
||||
|
||||
WORKDIR /mlflow
|
||||
|
||||
RUN pip install protobuf packaging
|
||||
|
||||
# Copy only necessary files to compile protos
|
||||
COPY mlflow/protos ./mlflow/protos/
|
||||
COPY dev ./dev/
|
||||
COPY tests/protos ./tests/protos/
|
||||
@@ -0,0 +1,8 @@
|
||||
**/*_pb2.py
|
||||
**/*.pyc
|
||||
**/__pycache__/
|
||||
*.class
|
||||
*.jar
|
||||
**/build
|
||||
**/dist
|
||||
**/*.egg-info
|
||||
@@ -0,0 +1,17 @@
|
||||
## MLflow Dev Scripts
|
||||
|
||||
This directory contains automation scripts for MLflow developers and the build infrastructure.
|
||||
|
||||
## Job Statuses
|
||||
|
||||
[](https://github.com/mlflow/dev/actions/workflows/examples.yml?query=workflow%3AExamples+event%3Aschedule)
|
||||
[](https://github.com/mlflow/dev/actions/workflows/cross-version-tests.yml?query=workflow%3A%22Cross+version+tests%22+event%3Aschedule)
|
||||
[](https://github.com/mlflow/dev/actions/workflows/xtest-viz.yml)
|
||||
[](https://github.com/mlflow/dev/actions/workflows/r.yml?query=workflow%3AR+event%3Aschedule)
|
||||
[](https://github.com/mlflow/dev/actions/workflows/test-requirements.yml?query=workflow%3A%22Test+requirements%22+event%3Aschedule)
|
||||
[](https://github.com/mlflow/mlflow/actions/workflows/push-images.yml?query=event%3Arelease)
|
||||
[](https://github.com/mlflow/dev/actions/workflows/slow-tests.yml?query=event%3Aschedule)
|
||||
[](https://github.com/mlflow/mlflow-website/actions/workflows/e2e.yml?query=event%3Aschedule)
|
||||
[](https://github.com/mlflow/mlflow/actions/workflows/update-model-catalog.yml?query=event%3Aschedule)
|
||||
[](https://github.com/mlflow/mlflow/actions/workflows/gateway-benchmark.yml?query=event%3Aschedule)
|
||||
[](https://github.com/mlflow/mlflow/actions/workflows/tracing-benchmark.yml)
|
||||
@@ -0,0 +1,145 @@
|
||||
# MLflow AI Gateway Benchmark
|
||||
|
||||
Measures the **proxy overhead** of the MLflow tracking-server-backed AI Gateway under
|
||||
concurrent load. A fake OpenAI server simulates the upstream provider at a fixed latency,
|
||||
so results reflect pure MLflow processing time rather than provider variance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Python 3.10+ with [`uv`](https://docs.astral.sh/uv/) — all scripts must be run via `uv run`, which handles dependency installation automatically via inline script metadata
|
||||
- Docker (required for `--database postgres` and `multi` mode)
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
cd dev/benchmarks/gateway
|
||||
|
||||
# 4 instances behind nginx (default, requires Docker)
|
||||
uv run run.py
|
||||
|
||||
# Single instance, SQLite (no Docker needed)
|
||||
uv run run.py --instances 1
|
||||
|
||||
# Single instance, PostgreSQL
|
||||
uv run run.py --instances 1 --database postgres
|
||||
|
||||
# Scale up
|
||||
uv run run.py --instances 8 --workers 8
|
||||
|
||||
# Benchmark an existing endpoint directly (skips all setup)
|
||||
uv run run.py --url http://your-server/gateway/my-endpoint/mlflow/invocations
|
||||
|
||||
# Basic-auth enabled (starts MLflow with --app-name=basic-auth,
|
||||
# sends Authorization: Basic on every request)
|
||||
uv run run.py --instances 1 --auth
|
||||
|
||||
```
|
||||
|
||||
## What is measured
|
||||
|
||||
Latency is measured **client-side** using `time.perf_counter()` around each `aiohttp` request.
|
||||
Each sample covers the full round-trip: client serialization → loopback → full server processing → response deserialization. Only HTTP 200 responses count toward latency stats; errors are tracked separately.
|
||||
|
||||
Connection pooling and HTTP keep-alive are enabled, so TCP handshake cost is amortized after the warmup phase.
|
||||
|
||||
### What is NOT measured
|
||||
|
||||
| Factor | In this benchmark | In production |
|
||||
| ------------------ | ---------------------------------------------- | --------------------------- |
|
||||
| Network latency | ~0 ms (loopback) | 1–100 ms per hop |
|
||||
| TLS/SSL | None (plain HTTP) | ~5–20 ms per new connection |
|
||||
| Provider inference | Fixed fake delay (`--fake-delay-ms`) | Variable (50 ms – 60 s+) |
|
||||
| Authentication | Off by default; basic-auth opt-in via `--auth` | Token validation, RBAC |
|
||||
|
||||
## What MLflow does per request
|
||||
|
||||
Each invocation through the tracking-server gateway runs these steps:
|
||||
|
||||
```
|
||||
1. Config resolution (DB-backed, cached after first hit)
|
||||
2. Secret decryption (cached, 60 s TTL)
|
||||
3. Provider instantiation
|
||||
4. Tracing (if usage_tracking=True)
|
||||
5. HTTP call to LLM API
|
||||
```
|
||||
|
||||
Steps 1 (config resolution) and 4 (tracing) have historically been the dominant bottlenecks.
|
||||
Config caching (enabled by default) eliminates most of step 1's cost. Tracing overhead
|
||||
depends on the span processor in use.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Single instance (`--instances 1`)
|
||||
|
||||
```
|
||||
benchmark.py ──aiohttp──▶ MLflow server (:5731) ──▶ fake_server.py (:9137)
|
||||
│
|
||||
SQLite or PostgreSQL
|
||||
```
|
||||
|
||||
### Multi-instance (`--instances N`, default)
|
||||
|
||||
```
|
||||
benchmark.py ──aiohttp──▶ nginx LB (:5731) ──round-robin──▶ MLflow :5800
|
||||
MLflow :5801
|
||||
MLflow :580N
|
||||
│
|
||||
fake_server.py (:9137)
|
||||
PostgreSQL (Docker)
|
||||
```
|
||||
|
||||
MLflow instances are started **sequentially** (instance 0 first) to let it initialize the
|
||||
DB schema before the others join. All instances share one PostgreSQL database.
|
||||
|
||||
## Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
| ----------------------------- | -------------- | ----------------------------------------------------------------------- |
|
||||
| `--url URL` | — | Benchmark this URL directly, skip all setup |
|
||||
| `--instances N` | 4 | MLflow instances. Use 1 for single-instance (no nginx, optional SQLite) |
|
||||
| `--workers N` | 4 | MLflow worker processes per instance |
|
||||
| `--database sqlite\|postgres` | `sqlite` | Database to use — only applies when `--instances 1` |
|
||||
| `--no-usage-tracking` | — | Disable usage tracking (tracing) on the endpoint |
|
||||
| `--port N` | 5731 | Port to benchmark (MLflow port for single, nginx LB port for multi) |
|
||||
| `--base-port N` | 5800 | First MLflow instance port in multi mode (rest are +1, +2, …) |
|
||||
| `--fake-server-port N` | 9137 | Fake OpenAI server port |
|
||||
| `--requests N` | 2000 | Requests per run |
|
||||
| `--max-concurrent N` | 50 | Max concurrent requests |
|
||||
| `--runs N` | 3 | Number of benchmark runs |
|
||||
| `--fake-delay-ms N` | 50 | Simulated provider latency in ms |
|
||||
| `--min-rps N` | — | Fail (exit 1) if average throughput falls below N req/s |
|
||||
| `--max-p50-ms N` | — | Fail (exit 1) if average P50 latency exceeds N ms (CI threshold) |
|
||||
| `--max-p99-ms N` | — | Fail (exit 1) if average P99 latency exceeds N ms (CI threshold) |
|
||||
| `--auth` | off | Start MLflow with `--app-name=basic-auth`; send Basic auth on requests |
|
||||
| `--auth-username USER` | `admin` | Basic-auth username (matches `mlflow/server/auth/basic_auth.ini`) |
|
||||
| `--auth-password PASS` | `password1234` | Basic-auth password (matches `mlflow/server/auth/basic_auth.ini`) |
|
||||
|
||||
All flags can also be set via environment variables (same name, uppercased):
|
||||
`INSTANCES`, `WORKERS_PER_INSTANCE`, `REQUESTS`, `MAX_CONCURRENT`, `RUNS`,
|
||||
`FAKE_RESPONSE_DELAY_MS`, `MLFLOW_PORT`, `BASE_PORT`, `FAKE_SERVER_PORT`,
|
||||
`AUTH`, `AUTH_USERNAME`, `AUTH_PASSWORD`.
|
||||
|
||||
To avoid conflicts with a local PostgreSQL instance, override the port via `GATEWAY_BENCH_POSTGRES_PORT` (default: 5432).
|
||||
|
||||
## Known limitations
|
||||
|
||||
- **Loopback only** — all processes run on the same machine. Results don't include real
|
||||
network latency between client, gateway, and provider.
|
||||
- **No TLS** — MLflow is started with `--disable-security-middleware`. Production deployments
|
||||
add TLS termination overhead.
|
||||
- **Fixed provider latency** — `fake_server.py` always responds in exactly `--fake-delay-ms`.
|
||||
Real providers have high variance (P99 often 5–10× P50).
|
||||
- **Basic-auth is opt-in, no RBAC** — `--auth` enables `basic-auth` with the default
|
||||
admin user (full permissions), which measures the cost of HTTP Basic authentication
|
||||
and user lookup but not fine-grained RBAC checks against non-admin users.
|
||||
- **Single machine resource contention** — with multiple instances, all MLflow instances, nginx,
|
||||
PostgreSQL, and the benchmark client share CPU/memory. On a server with dedicated resources
|
||||
per instance, throughput will be higher.
|
||||
|
||||
## Files
|
||||
|
||||
| File | Purpose |
|
||||
| ---------------- | ------------------------------------------------------------------------------ |
|
||||
| `run.py` | Main entry point — orchestrates servers, Docker, endpoint setup, and benchmark |
|
||||
| `benchmark.py` | Async HTTP benchmark client (standalone or imported by `run.py`) |
|
||||
| `fake_server.py` | Fake OpenAI-compatible server for controlled latency simulation |
|
||||
@@ -0,0 +1,362 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["aiohttp>=3.13.3,<4", "rich>=14.3.3,<15"]
|
||||
# ///
|
||||
"""Async HTTP benchmark client for the MLflow AI Gateway.
|
||||
|
||||
Can be imported by run.py or used standalone:
|
||||
uv run benchmark.py --url http://127.0.0.1:5731/gateway/benchmark-chat/mlflow/invocations
|
||||
uv run benchmark.py --url http://... --requests 5000 --max-concurrent 100 --runs 3
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import math
|
||||
import statistics
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
from rich.console import Console # type: ignore[import-not-found]
|
||||
from rich.progress import ( # type: ignore[import-not-found]
|
||||
BarColumn,
|
||||
MofNCompleteColumn,
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TaskID,
|
||||
TextColumn,
|
||||
TimeElapsedColumn,
|
||||
)
|
||||
from rich.table import Table # type: ignore[import-not-found]
|
||||
|
||||
console = Console()
|
||||
|
||||
_BODY = {
|
||||
"messages": [{"role": "user", "content": "benchmark request"}],
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 50,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunResult:
|
||||
latencies_ms: list[float] = field(default_factory=list)
|
||||
failures: dict[str, int] = field(default_factory=dict)
|
||||
wall_time: float = 0.0
|
||||
|
||||
@property
|
||||
def n_success(self) -> int:
|
||||
return len(self.latencies_ms)
|
||||
|
||||
@property
|
||||
def n_failures(self) -> int:
|
||||
return sum(self.failures.values())
|
||||
|
||||
@property
|
||||
def throughput(self) -> float:
|
||||
return self.n_success / self.wall_time if self.wall_time > 0 else 0.0
|
||||
|
||||
def percentile(self, p: float) -> float:
|
||||
if not self.latencies_ms:
|
||||
return 0.0
|
||||
s = sorted(self.latencies_ms)
|
||||
idx = max(0, math.ceil(p / 100 * len(s)) - 1)
|
||||
return s[idx]
|
||||
|
||||
|
||||
async def _send(
|
||||
session: aiohttp.ClientSession,
|
||||
url: str,
|
||||
sem: asyncio.Semaphore,
|
||||
auth: aiohttp.BasicAuth | None = None,
|
||||
) -> tuple[float, str | None]:
|
||||
async with sem:
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
async with session.post(url, json=_BODY, auth=auth) as resp:
|
||||
await resp.read()
|
||||
ms = (time.perf_counter() - t0) * 1000
|
||||
if resp.status == 200:
|
||||
return ms, None
|
||||
return ms, f"HTTP {resp.status}"
|
||||
except Exception as e:
|
||||
return (time.perf_counter() - t0) * 1000, type(e).__name__
|
||||
|
||||
|
||||
async def _run_once(
|
||||
url: str,
|
||||
n: int,
|
||||
max_concurrent: int,
|
||||
progress: Progress,
|
||||
task_id: TaskID,
|
||||
auth: aiohttp.BasicAuth | None = None,
|
||||
) -> RunResult:
|
||||
sem = asyncio.Semaphore(max_concurrent)
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=max(max_concurrent * 2, 200),
|
||||
limit_per_host=max(max_concurrent, 200),
|
||||
force_close=False,
|
||||
enable_cleanup_closed=True,
|
||||
)
|
||||
result = RunResult()
|
||||
total_time = 0.0
|
||||
max_time = 0.0
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
t0 = time.perf_counter()
|
||||
for coro in asyncio.as_completed([_send(session, url, sem, auth) for _ in range(n)]):
|
||||
ms, error = await coro
|
||||
if error:
|
||||
result.failures[error] = result.failures.get(error, 0) + 1
|
||||
else:
|
||||
result.latencies_ms.append(ms)
|
||||
total_time += ms
|
||||
if ms > max_time:
|
||||
max_time = ms
|
||||
|
||||
n_ok = result.n_success
|
||||
n_fail = result.n_failures
|
||||
mean = total_time / n_ok if n_ok else 0.0
|
||||
fail_part = f"[red]✗{n_fail}[/red] " if n_fail else ""
|
||||
live = f"{fail_part}✓{n_ok} mean={mean:.0f}ms max={max_time:.0f}ms"
|
||||
progress.update(task_id, advance=1, live=live)
|
||||
|
||||
result.wall_time = time.perf_counter() - t0
|
||||
return result
|
||||
|
||||
|
||||
async def _warmup(
|
||||
url: str, n: int, max_concurrent: int, auth: aiohttp.BasicAuth | None = None
|
||||
) -> None:
|
||||
sem = asyncio.Semaphore(max_concurrent)
|
||||
connector = aiohttp.TCPConnector(limit=max(max_concurrent * 2, 200))
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
await asyncio.gather(*[_send(session, url, sem, auth) for _ in range(n)])
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
url: str,
|
||||
n_requests: int = 2000,
|
||||
max_concurrent: int = 50,
|
||||
runs: int = 3,
|
||||
auth: aiohttp.BasicAuth | None = None,
|
||||
) -> list[RunResult]:
|
||||
warmup_n = min(max(50, max_concurrent), n_requests)
|
||||
console.print(f" [dim]Warming up ({warmup_n} requests)...[/dim]")
|
||||
asyncio.run(_warmup(url, warmup_n, max_concurrent, auth))
|
||||
|
||||
results = []
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
MofNCompleteColumn(),
|
||||
TimeElapsedColumn(),
|
||||
TextColumn(" {task.fields[live]}"),
|
||||
console=console,
|
||||
) as progress:
|
||||
for i in range(runs):
|
||||
task_id = progress.add_task(f" Run {i + 1}/{runs}", total=n_requests, live="")
|
||||
results.append(
|
||||
asyncio.run(_run_once(url, n_requests, max_concurrent, progress, task_id, auth))
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def results_to_dict(results: list[RunResult]) -> dict[str, Any]:
|
||||
runs = [
|
||||
{
|
||||
"n_success": r.n_success,
|
||||
"n_failures": r.n_failures,
|
||||
"failures": r.failures,
|
||||
"wall_time_s": r.wall_time,
|
||||
"mean_ms": statistics.mean(r.latencies_ms) if r.latencies_ms else 0.0,
|
||||
"p50_ms": r.percentile(50),
|
||||
"p95_ms": r.percentile(95),
|
||||
"p99_ms": r.percentile(99),
|
||||
"max_ms": max(r.latencies_ms) if r.latencies_ms else 0.0,
|
||||
"rps": r.throughput,
|
||||
}
|
||||
for r in results
|
||||
]
|
||||
summary: dict[str, Any] = (
|
||||
{
|
||||
"avg_mean_ms": statistics.mean(
|
||||
statistics.mean(r.latencies_ms) if r.latencies_ms else 0.0 for r in results
|
||||
),
|
||||
"avg_p50_ms": statistics.mean(r.percentile(50) for r in results),
|
||||
"avg_p99_ms": statistics.mean(r.percentile(99) for r in results),
|
||||
"avg_rps": statistics.mean(r.throughput for r in results),
|
||||
}
|
||||
if results
|
||||
else {}
|
||||
)
|
||||
return {"runs": runs, "summary": summary}
|
||||
|
||||
|
||||
def print_results(results: list[RunResult]) -> None:
|
||||
table = Table(show_header=True, header_style="bold cyan", box=None, padding=(0, 2))
|
||||
table.add_column("Run", style="dim", width=5)
|
||||
table.add_column("Mean ms", justify="right")
|
||||
table.add_column("P50 ms", justify="right")
|
||||
table.add_column("P95 ms", justify="right")
|
||||
table.add_column("P99 ms", justify="right")
|
||||
table.add_column("Max ms", justify="right")
|
||||
table.add_column("Req/s", justify="right")
|
||||
table.add_column("Failures", justify="right")
|
||||
|
||||
means = []
|
||||
p50s = []
|
||||
p95s = []
|
||||
p99s = []
|
||||
maxes = []
|
||||
throughputs = []
|
||||
for i, r in enumerate(results):
|
||||
mean = statistics.mean(r.latencies_ms) if r.latencies_ms else 0.0
|
||||
p50 = r.percentile(50)
|
||||
p95 = r.percentile(95)
|
||||
p99 = r.percentile(99)
|
||||
mx = max(r.latencies_ms) if r.latencies_ms else 0.0
|
||||
means.append(mean)
|
||||
p50s.append(p50)
|
||||
p95s.append(p95)
|
||||
p99s.append(p99)
|
||||
maxes.append(mx)
|
||||
throughputs.append(r.throughput)
|
||||
fail_str = f"[red]{r.n_failures}[/red]" if r.n_failures else "0"
|
||||
table.add_row(
|
||||
str(i + 1),
|
||||
f"{mean:.1f}",
|
||||
f"{p50:.1f}",
|
||||
f"{p95:.1f}",
|
||||
f"{p99:.1f}",
|
||||
f"{mx:.1f}",
|
||||
f"{r.throughput:.0f}",
|
||||
fail_str,
|
||||
)
|
||||
|
||||
if len(results) > 1:
|
||||
table.add_section()
|
||||
table.add_row(
|
||||
"[bold]avg[/bold]",
|
||||
f"[bold]{statistics.mean(means):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(p50s):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(p95s):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(p99s):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(maxes):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(throughputs):.0f}[/bold]",
|
||||
"",
|
||||
)
|
||||
|
||||
console.print()
|
||||
console.print(table)
|
||||
|
||||
combined: dict[str, int] = {}
|
||||
for r in results:
|
||||
for k, v in r.failures.items():
|
||||
combined[k] = combined.get(k, 0) + v
|
||||
if combined:
|
||||
console.print()
|
||||
console.print("[red]Failure breakdown:[/red]")
|
||||
for reason, count in sorted(combined.items(), key=lambda x: -x[1]):
|
||||
console.print(f" {reason}: {count}")
|
||||
|
||||
|
||||
def check_thresholds(
|
||||
results: list[RunResult],
|
||||
min_rps: float | None = None,
|
||||
max_p50_ms: float | None = None,
|
||||
max_p99_ms: float | None = None,
|
||||
) -> bool:
|
||||
"""Check results against performance thresholds. Returns True if all pass."""
|
||||
avg_rps = statistics.mean(r.throughput for r in results)
|
||||
avg_p50 = statistics.mean(r.percentile(50) for r in results)
|
||||
avg_p99 = statistics.mean(r.percentile(99) for r in results)
|
||||
passed = True
|
||||
|
||||
if min_rps is not None and avg_rps < min_rps:
|
||||
console.print(
|
||||
f"\n[red]THRESHOLD FAILED:[/red] avg throughput {avg_rps:.0f} req/s"
|
||||
f" < minimum {min_rps:.0f} req/s"
|
||||
)
|
||||
passed = False
|
||||
|
||||
if max_p50_ms is not None and avg_p50 > max_p50_ms:
|
||||
console.print(
|
||||
f"\n[red]THRESHOLD FAILED:[/red] avg P50 {avg_p50:.1f} ms > maximum {max_p50_ms:.1f} ms"
|
||||
)
|
||||
passed = False
|
||||
|
||||
if max_p99_ms is not None and avg_p99 > max_p99_ms:
|
||||
console.print(
|
||||
f"\n[red]THRESHOLD FAILED:[/red] avg P99 {avg_p99:.1f} ms > maximum {max_p99_ms:.1f} ms"
|
||||
)
|
||||
passed = False
|
||||
|
||||
if passed and (min_rps is not None or max_p50_ms is not None or max_p99_ms is not None):
|
||||
console.print("\n[green]All thresholds passed.[/green]")
|
||||
|
||||
return passed
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Async HTTP benchmark client for MLflow Gateway")
|
||||
parser.add_argument("--url", required=True, help="Gateway invocation URL")
|
||||
parser.add_argument("--requests", type=int, default=2000)
|
||||
parser.add_argument("--max-concurrent", type=int, default=50)
|
||||
parser.add_argument("--runs", type=int, default=3)
|
||||
parser.add_argument(
|
||||
"--min-rps",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Fail (exit 1) if average throughput falls below N req/s",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p50-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Fail (exit 1) if average P50 latency exceeds N ms",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p99-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Fail (exit 1) if average P99 latency exceeds N ms",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-username",
|
||||
default=None,
|
||||
help="Basic auth username. If set together with --auth-password, sent on every request.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-password",
|
||||
default=None,
|
||||
help="Basic auth password. If set together with --auth-username, sent on every request.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
auth = (
|
||||
aiohttp.BasicAuth(args.auth_username, args.auth_password)
|
||||
if args.auth_username and args.auth_password
|
||||
else None
|
||||
)
|
||||
|
||||
console.print(f"\n[bold]Benchmarking[/bold] {args.url}")
|
||||
console.print(
|
||||
f" {args.requests} requests · {args.max_concurrent} concurrent · {args.runs} runs\n"
|
||||
)
|
||||
results = run_benchmark(args.url, args.requests, args.max_concurrent, args.runs, auth)
|
||||
print_results(results)
|
||||
|
||||
if not check_thresholds(
|
||||
results, min_rps=args.min_rps, max_p50_ms=args.max_p50_ms, max_p99_ms=args.max_p99_ms
|
||||
):
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,68 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["fastapi>=0.115.0,<1", "uvicorn[standard]>=0.30.0,<1"]
|
||||
# ///
|
||||
"""Fake OpenAI-compatible server for benchmarking.
|
||||
|
||||
Returns synthetic responses after a configurable delay so benchmarks measure
|
||||
MLflow overhead rather than provider latency.
|
||||
|
||||
Run standalone:
|
||||
uv run fake_server.py
|
||||
PORT=9200 uv run fake_server.py
|
||||
|
||||
Or with multiple workers (as launched by run.py):
|
||||
uvicorn fake_server:app --workers 8 --port 9137
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
DELAY_MS = int(os.environ.get("FAKE_RESPONSE_DELAY_MS", "50"))
|
||||
|
||||
|
||||
class ChatRequest(BaseModel):
|
||||
model: str = "gpt-4o-mini"
|
||||
messages: list[dict[str, str]] = Field(min_length=1)
|
||||
stream: bool = False
|
||||
temperature: float = 1.0
|
||||
max_tokens: int = 50
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat_completions(req: ChatRequest) -> dict[str, Any]:
|
||||
await asyncio.sleep(DELAY_MS / 1000)
|
||||
return {
|
||||
"id": "chatcmpl-fake",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": req.model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Hello!"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict[str, str]:
|
||||
# Polled by run.py's _wait_for_port to detect when the server is ready.
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
port = int(os.environ.get("PORT", "9137"))
|
||||
host = os.environ.get("FAKE_SERVER_HOST", "127.0.0.1")
|
||||
uvicorn.run("fake_server:app", host=host, port=port, log_level="warning")
|
||||
@@ -0,0 +1,793 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["aiohttp>=3.13.3,<4", "psycopg2-binary>=2.9,<3", "rich>=14.3.3,<15"]
|
||||
# ///
|
||||
"""MLflow AI Gateway benchmark runner.
|
||||
|
||||
Orchestrates fake OpenAI server, MLflow server(s), optional PostgreSQL and
|
||||
nginx (via Docker), then runs the async benchmark client.
|
||||
|
||||
Usage:
|
||||
uv run run.py # 4 instances, PostgreSQL, nginx (Docker)
|
||||
uv run run.py --instances 1 # single instance, SQLite, no Docker
|
||||
uv run run.py --instances 1 --database postgres
|
||||
uv run run.py --instances 8 --workers 8
|
||||
uv run run.py --url http://... # benchmark an existing endpoint directly
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
import aiohttp # type: ignore[import-not-found]
|
||||
import benchmark as bm # local module; path inserted above
|
||||
from rich.console import Console # type: ignore[import-not-found]
|
||||
from rich.panel import Panel # type: ignore[import-not-found]
|
||||
from rich.progress import ( # type: ignore[import-not-found]
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TextColumn,
|
||||
TimeElapsedColumn,
|
||||
)
|
||||
|
||||
SCRIPT_DIR = Path(__file__).parent
|
||||
|
||||
FAKE_SERVER_PORT = 9137
|
||||
FAKE_SERVER_WORKERS = 8
|
||||
MLFLOW_PORT = 5731
|
||||
INSTANCE_BASE_PORT = 5800
|
||||
POSTGRES_PORT = int(os.environ.get("GATEWAY_BENCH_POSTGRES_PORT", "5432"))
|
||||
POSTGRES_PASSWORD = "benchmarkpass"
|
||||
ENDPOINT_NAME = "benchmark-chat"
|
||||
|
||||
_API_SECRET_CREATE = "gateway/secrets/create"
|
||||
_API_MODEL_DEF_CREATE = "gateway/model-definitions/create"
|
||||
_API_ENDPOINT_CREATE = "gateway/endpoints/create"
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def _uv_prefix() -> list[str]:
|
||||
"""Return uv run prefix when inside the mlflow repo, else empty list."""
|
||||
in_repo = (
|
||||
shutil.which("uv")
|
||||
and subprocess.run(
|
||||
["git", "rev-parse", "HEAD"], cwd=SCRIPT_DIR, capture_output=True
|
||||
).returncode
|
||||
== 0
|
||||
)
|
||||
return ["uv", "run", "--no-build-isolation", "--extra", "gateway"] if in_repo else []
|
||||
|
||||
|
||||
def _subprocess_env() -> dict[str, str]:
|
||||
return os.environ | {"OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES"}
|
||||
|
||||
|
||||
def _wait_for_port(port: int, label: str, log_file: Path | None = None, timeout: int = 30) -> None:
|
||||
url = f"http://127.0.0.1:{port}/health"
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
) as progress:
|
||||
progress.add_task(f" Waiting for {label}...", total=None)
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=1):
|
||||
break
|
||||
except Exception:
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
console.print(f" [red]✗ {label} failed to start within {timeout}s[/red]")
|
||||
if log_file and log_file.exists():
|
||||
console.print(" [yellow]Last 20 lines of log:[/yellow]")
|
||||
for line in log_file.read_text().splitlines()[-20:]:
|
||||
console.print(f" [dim]{line}[/dim]")
|
||||
sys.exit(1)
|
||||
console.print(f" [green]✓[/green] {label} ready")
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_fake_server(
|
||||
work_dir: str, port: int = FAKE_SERVER_PORT, workers: int = FAKE_SERVER_WORKERS
|
||||
) -> Generator[None, None, None]:
|
||||
prefix = _uv_prefix()
|
||||
log_file = Path(work_dir) / "fake_server.log"
|
||||
with (
|
||||
log_file.open("w") as f,
|
||||
subprocess.Popen(
|
||||
[
|
||||
*prefix,
|
||||
"uvicorn",
|
||||
"fake_server:app",
|
||||
"--workers",
|
||||
str(workers),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(port),
|
||||
"--log-level",
|
||||
"warning",
|
||||
],
|
||||
cwd=SCRIPT_DIR,
|
||||
stdout=f,
|
||||
stderr=f,
|
||||
env=_subprocess_env(),
|
||||
) as proc,
|
||||
):
|
||||
_wait_for_port(port, "fake OpenAI server", log_file)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
proc.terminate()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_mlflow(
|
||||
work_dir: str,
|
||||
port: int,
|
||||
workers: int,
|
||||
backend_uri: str,
|
||||
label: str = "MLflow server",
|
||||
host: str = "127.0.0.1",
|
||||
auth: bool = False,
|
||||
) -> Generator[None, None, None]:
|
||||
prefix = _uv_prefix()
|
||||
# basic-auth requires the `auth` extra (Flask-WTF) at runtime.
|
||||
if auth and prefix:
|
||||
prefix = [*prefix, "--extra", "auth"]
|
||||
# psycopg2-binary lives in the `db` extra.
|
||||
if backend_uri.startswith("postgresql") and prefix:
|
||||
prefix = [*prefix, "--extra", "db"]
|
||||
log_file = Path(work_dir) / f"mlflow-{port}.log"
|
||||
cmd = [
|
||||
*prefix,
|
||||
"mlflow",
|
||||
"server",
|
||||
"--backend-store-uri",
|
||||
backend_uri,
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
"--workers",
|
||||
str(workers),
|
||||
"--disable-security-middleware",
|
||||
]
|
||||
if auth:
|
||||
cmd += ["--app-name", "basic-auth"]
|
||||
with (
|
||||
log_file.open("w") as f,
|
||||
subprocess.Popen(cmd, cwd=SCRIPT_DIR, stdout=f, stderr=f, env=_subprocess_env()) as proc,
|
||||
):
|
||||
_wait_for_port(port, label, log_file)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
proc.terminate()
|
||||
|
||||
|
||||
def _check_docker() -> None:
|
||||
try:
|
||||
result = subprocess.run(["docker", "info"], capture_output=True)
|
||||
except FileNotFoundError:
|
||||
console.print(
|
||||
"[red]Docker is not installed. Install it at https://docs.docker.com/get-docker/[/red]"
|
||||
)
|
||||
sys.exit(1)
|
||||
if result.returncode != 0:
|
||||
console.print("[red]Docker daemon is not running. Please start Docker and try again.[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_postgres(container_name: str = "benchmark-postgres") -> Generator[str, None, None]:
|
||||
"""Start a PostgreSQL Docker container. Yields the connection URI."""
|
||||
subprocess.run(["docker", "rm", "-f", container_name], capture_output=True)
|
||||
|
||||
with subprocess.Popen(
|
||||
[
|
||||
"docker",
|
||||
"run",
|
||||
"--rm",
|
||||
"--name",
|
||||
container_name,
|
||||
"-e",
|
||||
f"POSTGRES_PASSWORD={POSTGRES_PASSWORD}",
|
||||
"-e",
|
||||
"POSTGRES_DB=mlflow",
|
||||
"-p",
|
||||
f"127.0.0.1:{POSTGRES_PORT}:5432",
|
||||
"postgres:16-alpine",
|
||||
"-c",
|
||||
"max_connections=500",
|
||||
],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
):
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
) as progress:
|
||||
progress.add_task(" Starting PostgreSQL...", total=None)
|
||||
deadline = time.monotonic() + 30
|
||||
while time.monotonic() < deadline:
|
||||
if (
|
||||
subprocess.run(
|
||||
["docker", "exec", container_name, "pg_isready", "-U", "postgres"],
|
||||
capture_output=True,
|
||||
).returncode
|
||||
== 0
|
||||
):
|
||||
break
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
console.print(" [red]✗ PostgreSQL failed to start within 30s[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
console.print(" [green]✓[/green] PostgreSQL ready")
|
||||
try:
|
||||
yield f"postgresql://postgres:{POSTGRES_PASSWORD}@127.0.0.1:{POSTGRES_PORT}/mlflow"
|
||||
finally:
|
||||
subprocess.run(["docker", "kill", container_name], capture_output=True)
|
||||
|
||||
|
||||
def _basic_auth_header(creds: tuple[str, str] | None) -> dict[str, str]:
|
||||
if creds is None:
|
||||
return {}
|
||||
token = base64.b64encode(f"{creds[0]}:{creds[1]}".encode()).decode()
|
||||
return {"Authorization": f"Basic {token}"}
|
||||
|
||||
|
||||
def _api_post(
|
||||
tracking_uri: str,
|
||||
path: str,
|
||||
body: dict[str, Any],
|
||||
creds: tuple[str, str] | None = None,
|
||||
) -> Any:
|
||||
url = f"{tracking_uri.rstrip('/')}/api/3.0/mlflow/{path}"
|
||||
headers = {"Content-Type": "application/json", **_basic_auth_header(creds)}
|
||||
req = urllib.request.Request(url, data=json.dumps(body).encode(), headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
return json.loads(resp.read())
|
||||
except urllib.error.HTTPError as e:
|
||||
console.print(f" [red]API error {e.code} at {url}: {e.read().decode()}[/red]")
|
||||
sys.exit(1)
|
||||
except urllib.error.URLError as e:
|
||||
console.print(f" [red]API error at {url}: {e.reason}[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def _setup_endpoint(
|
||||
tracking_uri: str,
|
||||
fake_server_url: str,
|
||||
endpoint_name: str,
|
||||
usage_tracking: bool,
|
||||
creds: tuple[str, str] | None = None,
|
||||
) -> str:
|
||||
"""Create secret → model definition → endpoint. Returns the invocation URL."""
|
||||
console.print(" Creating secret...")
|
||||
secret_id = _api_post(
|
||||
tracking_uri,
|
||||
_API_SECRET_CREATE,
|
||||
{
|
||||
"secret_name": "benchmark-secret",
|
||||
"secret_value": {"api_key": "fake-benchmark-key"},
|
||||
"provider": "openai",
|
||||
"auth_config": {"api_base": fake_server_url},
|
||||
},
|
||||
creds,
|
||||
)["secret"]["secret_id"]
|
||||
|
||||
console.print(" Creating model definition...")
|
||||
model_def_id = _api_post(
|
||||
tracking_uri,
|
||||
_API_MODEL_DEF_CREATE,
|
||||
{
|
||||
"name": "benchmark-model",
|
||||
"secret_id": secret_id,
|
||||
"provider": "openai",
|
||||
"model_name": "gpt-4o-mini",
|
||||
},
|
||||
creds,
|
||||
)["model_definition"]["model_definition_id"]
|
||||
|
||||
console.print(f" Creating endpoint '{endpoint_name}' (usage_tracking={usage_tracking})...")
|
||||
_api_post(
|
||||
tracking_uri,
|
||||
_API_ENDPOINT_CREATE,
|
||||
{
|
||||
"name": endpoint_name,
|
||||
"model_configs": [
|
||||
{"model_definition_id": model_def_id, "linkage_type": "PRIMARY", "weight": 1.0}
|
||||
],
|
||||
"usage_tracking": usage_tracking,
|
||||
},
|
||||
creds,
|
||||
)
|
||||
|
||||
invoke_url = f"{tracking_uri.rstrip('/')}/gateway/{endpoint_name}/mlflow/invocations"
|
||||
console.print(f" [green]✓[/green] Endpoint ready: [cyan]{invoke_url}[/cyan]")
|
||||
return invoke_url
|
||||
|
||||
|
||||
def _sanity_check(url: str, creds: tuple[str, str] | None = None) -> None:
|
||||
console.print(" Sending sanity-check request...")
|
||||
body = json.dumps({"messages": [{"role": "user", "content": "test"}]}).encode()
|
||||
headers = {"Content-Type": "application/json", **_basic_auth_header(creds)}
|
||||
req = urllib.request.Request(url, data=body, headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
if resp.status != 200:
|
||||
console.print(f" [red]✗ Sanity check failed: HTTP {resp.status}[/red]")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
console.print(f" [red]✗ Sanity check failed: {e}[/red]")
|
||||
sys.exit(1)
|
||||
console.print(" [green]✓[/green] Sanity check passed")
|
||||
|
||||
|
||||
def _run_benchmark(
|
||||
url: str,
|
||||
n_requests: int,
|
||||
max_concurrent: int,
|
||||
runs: int,
|
||||
min_rps: float | None = None,
|
||||
max_p50_ms: float | None = None,
|
||||
max_p99_ms: float | None = None,
|
||||
output: Path | None = None,
|
||||
creds: tuple[str, str] | None = None,
|
||||
) -> None:
|
||||
auth = aiohttp.BasicAuth(*creds) if creds else None
|
||||
results = bm.run_benchmark(url, n_requests, max_concurrent, runs, auth)
|
||||
bm.print_results(results)
|
||||
if output is not None:
|
||||
output.write_text(json.dumps(bm.results_to_dict(results), indent=2))
|
||||
console.print(f" Results saved to [cyan]{output}[/cyan]")
|
||||
if not bm.check_thresholds(
|
||||
results, min_rps=min_rps, max_p50_ms=max_p50_ms, max_p99_ms=max_p99_ms
|
||||
):
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_nginx(
|
||||
work_dir: str, instance_ports: list[int], port: int, container_name: str = "benchmark-nginx"
|
||||
) -> Generator[None, None, None]:
|
||||
nginx_dir = Path(work_dir) / "nginx"
|
||||
conf_d = nginx_dir / "conf.d"
|
||||
conf_d.mkdir(parents=True)
|
||||
|
||||
upstream_lines = "\n".join(f" server host.docker.internal:{p};" for p in instance_ports)
|
||||
(conf_d / "mlflow.conf").write_text(
|
||||
f"upstream mlflow_backends {{\n"
|
||||
f"{upstream_lines}\n"
|
||||
f" keepalive 512;\n"
|
||||
f" keepalive_requests 100000;\n"
|
||||
f" keepalive_timeout 60s;\n"
|
||||
f"}}\n"
|
||||
f"server {{\n"
|
||||
f" listen {port} reuseport backlog=65535;\n"
|
||||
f" location / {{\n"
|
||||
f" proxy_pass http://mlflow_backends;\n"
|
||||
f" proxy_http_version 1.1;\n"
|
||||
f' proxy_set_header Connection "";\n'
|
||||
f" proxy_set_header Host $host;\n"
|
||||
f" proxy_set_header X-Real-IP $remote_addr;\n"
|
||||
f" proxy_connect_timeout 5s;\n"
|
||||
f" proxy_send_timeout 60s;\n"
|
||||
f" proxy_read_timeout 60s;\n"
|
||||
f" }}\n"
|
||||
f"}}\n"
|
||||
)
|
||||
(nginx_dir / "nginx.conf").write_text(
|
||||
"worker_processes auto;\n"
|
||||
"worker_rlimit_nofile 65535;\n"
|
||||
"events {\n"
|
||||
" worker_connections 16384;\n"
|
||||
" use epoll;\n"
|
||||
" multi_accept on;\n"
|
||||
"}\n"
|
||||
"http {\n"
|
||||
" access_log off;\n"
|
||||
" tcp_nodelay on;\n"
|
||||
" keepalive_timeout 65;\n"
|
||||
" keepalive_requests 100000;\n"
|
||||
" reset_timedout_connection on;\n"
|
||||
" include /etc/nginx/conf.d/*.conf;\n"
|
||||
"}\n"
|
||||
)
|
||||
|
||||
subprocess.run(["docker", "rm", "-f", container_name], capture_output=True)
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
) as progress:
|
||||
progress.add_task(" Starting nginx...", total=None)
|
||||
subprocess.run(
|
||||
[
|
||||
"docker",
|
||||
"run",
|
||||
"--rm",
|
||||
"-d",
|
||||
"--name",
|
||||
container_name,
|
||||
"--add-host=host.docker.internal:host-gateway",
|
||||
"--ulimit",
|
||||
"nofile=65535:65535",
|
||||
"-v",
|
||||
f"{nginx_dir / 'nginx.conf'}:/etc/nginx/nginx.conf:ro",
|
||||
"-v",
|
||||
f"{conf_d}:/etc/nginx/conf.d:ro",
|
||||
"-p",
|
||||
f"127.0.0.1:{port}:{port}",
|
||||
"nginx:alpine",
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
)
|
||||
|
||||
deadline = time.monotonic() + 15
|
||||
while time.monotonic() < deadline:
|
||||
if (
|
||||
subprocess.run(
|
||||
["docker", "exec", container_name, "nginx", "-t"], capture_output=True
|
||||
).returncode
|
||||
== 0
|
||||
):
|
||||
break
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
console.print(" [red]✗ nginx failed to start[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
console.print(" [green]✓[/green] nginx ready")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
subprocess.run(["docker", "kill", container_name], capture_output=True)
|
||||
|
||||
|
||||
def cmd_bench(args: argparse.Namespace) -> None:
|
||||
instances = args.instances
|
||||
mode = "1 instance" if instances == 1 else f"{instances} instances, nginx LB"
|
||||
creds = (args.auth_username, args.auth_password) if args.auth else None
|
||||
|
||||
if args.url:
|
||||
console.print(
|
||||
Panel.fit(
|
||||
f"[bold]Gateway Benchmark[/bold] ({mode})\n"
|
||||
f"URL: [cyan]{args.url}[/cyan]\n"
|
||||
f"Auth: {'basic-auth as ' + args.auth_username if creds else 'disabled'}\n"
|
||||
f"Requests: {args.requests} · Concurrency: {args.max_concurrent}"
|
||||
f" · Runs: {args.runs}",
|
||||
border_style="cyan",
|
||||
)
|
||||
)
|
||||
console.print("\n[bold]Running benchmark[/bold]")
|
||||
_run_benchmark(
|
||||
args.url,
|
||||
args.requests,
|
||||
args.max_concurrent,
|
||||
args.runs,
|
||||
args.min_rps,
|
||||
args.max_p50_ms,
|
||||
args.max_p99_ms,
|
||||
args.output,
|
||||
creds,
|
||||
)
|
||||
return
|
||||
|
||||
needs_docker = instances > 1 or args.database == "postgres"
|
||||
if needs_docker:
|
||||
_check_docker()
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="mlflow-bench-") as work_dir:
|
||||
port = args.port
|
||||
fake_port = args.fake_server_port
|
||||
instance_ports = [args.base_port + i for i in range(instances)]
|
||||
|
||||
auth_line = f"basic-auth as {args.auth_username}" if creds else "disabled"
|
||||
if instances == 1:
|
||||
panel = (
|
||||
f"[bold]Gateway Benchmark[/bold] ({mode})\n"
|
||||
f"Workers: {args.workers} · DB: {args.database.upper()} · "
|
||||
f"Usage tracking: {args.usage_tracking} · Auth: {auth_line}\n"
|
||||
f"Requests: {args.requests} · Concurrency: {args.max_concurrent} · "
|
||||
f"Runs: {args.runs} · Fake delay: {args.fake_delay_ms}ms\n"
|
||||
f"Ports: MLflow :{port} · Fake server :{fake_port}"
|
||||
)
|
||||
else:
|
||||
panel = (
|
||||
f"[bold]Gateway Benchmark[/bold] ({mode})\n"
|
||||
f"Workers/instance: {args.workers} · "
|
||||
f"Total workers: {instances * args.workers} · "
|
||||
f"Usage tracking: {args.usage_tracking} · Auth: {auth_line}\n"
|
||||
f"Requests: {args.requests} · Concurrency: {args.max_concurrent} · "
|
||||
f"Runs: {args.runs} · Fake delay: {args.fake_delay_ms}ms\n"
|
||||
f"Ports: instances {instance_ports[0]}–{instance_ports[-1]}"
|
||||
f" · LB :{port} · Fake server :{fake_port}"
|
||||
)
|
||||
console.print(Panel.fit(panel, border_style="cyan"))
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
stack.callback(lambda: console.print("\n[dim]Cleaning up...[/dim]"))
|
||||
|
||||
# Backend
|
||||
if instances > 1 or args.database == "postgres":
|
||||
console.print("\n[bold]PostgreSQL[/bold]")
|
||||
backend_uri = stack.enter_context(_start_postgres())
|
||||
else:
|
||||
db_path = Path(work_dir) / "mlflow.db"
|
||||
backend_uri = f"sqlite:///{db_path}"
|
||||
console.print(f"\n[dim]Using SQLite: {db_path}[/dim]")
|
||||
|
||||
# Servers
|
||||
console.print("\n[bold]Starting servers[/bold]")
|
||||
stack.enter_context(
|
||||
_start_fake_server(work_dir, port=fake_port, workers=FAKE_SERVER_WORKERS)
|
||||
)
|
||||
|
||||
if instances == 1:
|
||||
stack.enter_context(
|
||||
_start_mlflow(work_dir, port, args.workers, backend_uri, auth=args.auth)
|
||||
)
|
||||
|
||||
console.print("\n[bold]Setting up gateway endpoint[/bold]")
|
||||
invoke_url = _setup_endpoint(
|
||||
f"http://127.0.0.1:{port}",
|
||||
f"http://127.0.0.1:{fake_port}/v1",
|
||||
ENDPOINT_NAME,
|
||||
usage_tracking=args.usage_tracking,
|
||||
creds=creds,
|
||||
)
|
||||
_sanity_check(invoke_url, creds)
|
||||
else:
|
||||
# Start instance 0 first — it initializes the DB schema.
|
||||
# All instances share the same PostgreSQL DB, so starting concurrently
|
||||
# can cause CREATE TABLE race conditions.
|
||||
stack.enter_context(
|
||||
_start_mlflow(
|
||||
work_dir,
|
||||
instance_ports[0],
|
||||
args.workers,
|
||||
backend_uri,
|
||||
"MLflow instance 0",
|
||||
host="0.0.0.0",
|
||||
auth=args.auth,
|
||||
)
|
||||
)
|
||||
for i, p in enumerate(instance_ports[1:], start=1):
|
||||
stack.enter_context(
|
||||
_start_mlflow(
|
||||
work_dir,
|
||||
p,
|
||||
args.workers,
|
||||
backend_uri,
|
||||
f"MLflow instance {i}",
|
||||
host="0.0.0.0",
|
||||
auth=args.auth,
|
||||
)
|
||||
)
|
||||
|
||||
console.print("\n[bold]Setting up gateway endpoint[/bold]")
|
||||
_setup_endpoint(
|
||||
f"http://127.0.0.1:{instance_ports[0]}",
|
||||
f"http://127.0.0.1:{fake_port}/v1",
|
||||
ENDPOINT_NAME,
|
||||
usage_tracking=args.usage_tracking,
|
||||
creds=creds,
|
||||
)
|
||||
|
||||
console.print("\n[bold]Starting nginx load balancer[/bold]")
|
||||
nginx_container = "benchmark-nginx"
|
||||
stack.enter_context(
|
||||
_start_nginx(
|
||||
work_dir, instance_ports, port=port, container_name=nginx_container
|
||||
)
|
||||
)
|
||||
subprocess.run(
|
||||
["docker", "exec", nginx_container, "nginx", "-s", "reload"],
|
||||
capture_output=True,
|
||||
)
|
||||
time.sleep(1)
|
||||
|
||||
invoke_url = f"http://127.0.0.1:{port}/gateway/{ENDPOINT_NAME}/mlflow/invocations"
|
||||
_sanity_check(invoke_url, creds)
|
||||
|
||||
console.print("\n[bold]Running benchmark[/bold]")
|
||||
_run_benchmark(
|
||||
invoke_url,
|
||||
args.requests,
|
||||
args.max_concurrent,
|
||||
args.runs,
|
||||
args.min_rps,
|
||||
args.max_p50_ms,
|
||||
args.max_p99_ms,
|
||||
args.output,
|
||||
creds,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="MLflow AI Gateway benchmark",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--url",
|
||||
metavar="URL",
|
||||
help="Benchmark this endpoint URL directly, skipping server setup entirely",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instances",
|
||||
type=int,
|
||||
default=int(os.environ.get("INSTANCES", "4")),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Number of MLflow instances to run (default: 4). "
|
||||
"Values >1 require Docker (postgres + nginx). "
|
||||
"Use --instances 1 for a single instance with optional SQLite."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--workers",
|
||||
type=int,
|
||||
default=int(os.environ.get("WORKERS_PER_INSTANCE", "4")),
|
||||
metavar="N",
|
||||
help="Gunicorn/uvicorn worker processes per MLflow instance (default: 4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--database",
|
||||
choices=["sqlite", "postgres"],
|
||||
default="sqlite",
|
||||
help=(
|
||||
"Database to use — only applies when --instances 1. "
|
||||
"'postgres' auto-starts a Docker container. (default: sqlite)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-usage-tracking",
|
||||
dest="usage_tracking",
|
||||
action="store_false",
|
||||
default=True,
|
||||
help="Disable usage tracking (tracing) on the benchmark endpoint",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=int(os.environ.get("MLFLOW_PORT", str(MLFLOW_PORT))),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Port the benchmark client sends requests to. "
|
||||
"For --instances 1 this is the MLflow port; "
|
||||
"for --instances >1 this is the nginx load balancer port. (default: 5731)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--base-port",
|
||||
type=int,
|
||||
default=int(os.environ.get("BASE_PORT", str(INSTANCE_BASE_PORT))),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Starting port for MLflow instances in multi mode. "
|
||||
"Instances listen on base-port, base-port+1, … (default: 5800)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fake-server-port",
|
||||
type=int,
|
||||
metavar="N",
|
||||
default=int(os.environ.get("FAKE_SERVER_PORT", str(FAKE_SERVER_PORT))),
|
||||
help="Port for the fake OpenAI server that simulates provider latency (default: 9137)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--requests",
|
||||
type=int,
|
||||
default=int(os.environ.get("REQUESTS", "2000")),
|
||||
metavar="N",
|
||||
help="Total requests to send per benchmark run (default: 2000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-concurrent",
|
||||
type=int,
|
||||
default=int(os.environ.get("MAX_CONCURRENT", "50")),
|
||||
metavar="N",
|
||||
help="Maximum number of in-flight requests at any time (default: 50)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--runs",
|
||||
type=int,
|
||||
default=int(os.environ.get("RUNS", "3")),
|
||||
metavar="N",
|
||||
help="Number of timed runs; results are reported per-run and averaged (default: 3)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fake-delay-ms",
|
||||
type=int,
|
||||
default=int(os.environ.get("FAKE_RESPONSE_DELAY_MS", "50")),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Simulated provider latency in ms. Set to 0 to measure pure MLflow overhead "
|
||||
"with no provider delay. (default: 50)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
default=None,
|
||||
metavar="FILE",
|
||||
help="Write benchmark results as JSON to FILE (useful for CI artifact upload)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-rps",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Exit 1 if average throughput across runs falls below N req/s (CI threshold)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p50-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Exit 1 if average P50 latency across runs exceeds N ms (CI threshold)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p99-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Exit 1 if average P99 latency across runs exceeds N ms (CI threshold)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth",
|
||||
action="store_true",
|
||||
default=os.environ.get("AUTH", "").lower() in ("1", "true"),
|
||||
help=(
|
||||
"Start MLflow with --app-name=basic-auth and authenticate every setup + "
|
||||
"benchmark request using --auth-username/--auth-password."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-username",
|
||||
default=os.environ.get("AUTH_USERNAME", "admin"),
|
||||
help="Basic auth username (default: admin, from basic_auth.ini)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-password",
|
||||
default=os.environ.get("AUTH_PASSWORD", "password1234"),
|
||||
help="Basic auth password (default: password1234, from basic_auth.ini)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
os.environ["FAKE_RESPONSE_DELAY_MS"] = str(args.fake_delay_ms)
|
||||
cmd_bench(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,11 @@
|
||||
# Tracing Benchmark
|
||||
|
||||
Per-commit tracing perf check. Runs on push to `master`; trend at https://mlflow.github.io/mlflow/dev/benchmarks/tracing/.
|
||||
|
||||
```bash
|
||||
uv run pytest dev/benchmarks/tracing/ --benchmark-only
|
||||
```
|
||||
|
||||
Add a scenario by writing a `test_*` function in `test_trace_perf.py` — it appears in the chart on the next master push. Renaming a test starts a new trend line.
|
||||
|
||||
Setup is modeled after [`opentelemetry-python`'s benchmark workflow](https://github.com/open-telemetry/opentelemetry-python/blob/main/.github/workflows/benchmarks.yml).
|
||||
@@ -0,0 +1,116 @@
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from opentelemetry import trace as trace_api
|
||||
from opentelemetry.sdk.resources import Resource as _OTelResource
|
||||
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
|
||||
from opentelemetry.trace import SpanContext
|
||||
|
||||
from mlflow.entities.span import Span, SpanType, create_mlflow_span
|
||||
from mlflow.entities.trace_info import TraceInfo
|
||||
from mlflow.entities.trace_location import TraceLocation
|
||||
from mlflow.entities.trace_state import TraceState
|
||||
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
||||
from mlflow.tracing.constant import SpanAttributeKey, TraceTagKey
|
||||
from mlflow.tracing.utils import TraceJSONEncoder
|
||||
|
||||
ENV_CHOICES = ["prod", "staging", "dev"]
|
||||
NAME_PREFIXES = ["agent_run", "qa_chain", "rag_pipeline", "summarizer"]
|
||||
WEEK_MS = 7 * 24 * 60 * 60 * 1000
|
||||
|
||||
SEED_TRACES = 1000
|
||||
SEED_SPANS_PER_TRACE = 10
|
||||
|
||||
|
||||
def generate_trace_data(
|
||||
experiment_id: str,
|
||||
num_spans: int,
|
||||
rng: random.Random,
|
||||
) -> tuple[TraceInfo, list[Span]]:
|
||||
trace_id = f"tr-{uuid.uuid4().hex}"
|
||||
request_time = int(time.time() * 1000) - rng.randint(0, WEEK_MS)
|
||||
name_prefix = rng.choice(NAME_PREFIXES)
|
||||
trace_info = TraceInfo(
|
||||
trace_id=trace_id,
|
||||
trace_location=TraceLocation.from_experiment_id(experiment_id),
|
||||
request_time=request_time,
|
||||
state=rng.choice([TraceState.OK, TraceState.OK, TraceState.OK, TraceState.ERROR]),
|
||||
execution_duration=rng.randint(100, 5000),
|
||||
tags={
|
||||
TraceTagKey.TRACE_NAME: f"{name_prefix}_{trace_id[-4:]}",
|
||||
"env": rng.choice(ENV_CHOICES),
|
||||
},
|
||||
)
|
||||
|
||||
span_types = [SpanType.LLM, SpanType.RETRIEVER, SpanType.TOOL, SpanType.CHAIN]
|
||||
base_ns = 1_000_000_000_000
|
||||
spans: list[Span] = []
|
||||
|
||||
for i in range(num_spans):
|
||||
is_root = i == 0
|
||||
span_type = SpanType.AGENT if is_root else rng.choice(span_types)
|
||||
parent_id = None if is_root else rng.choice(range(max(0, i - 3), i))
|
||||
|
||||
trace_num = rng.randint(1, 2**63)
|
||||
ctx = SpanContext(
|
||||
trace_id=trace_num,
|
||||
span_id=i + 1,
|
||||
is_remote=False,
|
||||
trace_flags=trace_api.TraceFlags(1),
|
||||
trace_state=trace_api.TraceState(),
|
||||
)
|
||||
|
||||
parent_ctx = None
|
||||
if parent_id is not None:
|
||||
parent_ctx = SpanContext(
|
||||
trace_id=trace_num,
|
||||
span_id=parent_id + 1,
|
||||
is_remote=False,
|
||||
trace_flags=trace_api.TraceFlags(1),
|
||||
trace_state=trace_api.TraceState(),
|
||||
)
|
||||
|
||||
attrs: dict[str, object] = {}
|
||||
if is_root:
|
||||
attrs[SpanAttributeKey.INPUTS] = json.dumps(
|
||||
{"query": "What is ML?"}, cls=TraceJSONEncoder
|
||||
)
|
||||
attrs[SpanAttributeKey.OUTPUTS] = json.dumps(
|
||||
{"response": "ML is..."}, cls=TraceJSONEncoder
|
||||
)
|
||||
|
||||
otel_span = OTelReadableSpan(
|
||||
name=f"{span_type.lower()}_{i}" if not is_root else "agent_run",
|
||||
context=ctx,
|
||||
parent=parent_ctx,
|
||||
attributes={
|
||||
"mlflow.traceRequestId": json.dumps(trace_id),
|
||||
"mlflow.spanType": json.dumps(span_type, cls=TraceJSONEncoder),
|
||||
**attrs,
|
||||
},
|
||||
start_time=base_ns + i * 10_000_000,
|
||||
end_time=base_ns + i * 10_000_000 + rng.randint(5_000_000, 50_000_000),
|
||||
status=trace_api.Status(trace_api.StatusCode.OK),
|
||||
resource=_OTelResource.get_empty(),
|
||||
)
|
||||
spans.append(create_mlflow_span(otel_span, trace_id, span_type))
|
||||
|
||||
return trace_info, spans
|
||||
|
||||
|
||||
def seed_traces(
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
count: int,
|
||||
spans_per_trace: int,
|
||||
) -> list[str]:
|
||||
rng = random.Random(123)
|
||||
trace_ids: list[str] = []
|
||||
for _ in range(count):
|
||||
ti, sp = generate_trace_data(experiment_id, spans_per_trace, rng)
|
||||
store.start_trace(ti)
|
||||
store.log_spans(experiment_id, sp)
|
||||
trace_ids.append(ti.trace_id)
|
||||
return trace_ids
|
||||
@@ -0,0 +1,39 @@
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from _data import SEED_SPANS_PER_TRACE, SEED_TRACES, seed_traces
|
||||
|
||||
import mlflow
|
||||
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def bench_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
||||
return tmp_path_factory.mktemp("bench")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def store(bench_dir: Path) -> SqlAlchemyStore:
|
||||
db_uri = f"sqlite:///{bench_dir / 'mlflow.db'}"
|
||||
(bench_dir / "artifacts").mkdir()
|
||||
artifact_root = (bench_dir / "artifacts").as_uri()
|
||||
return SqlAlchemyStore(db_uri, artifact_root)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def experiment_id(store: SqlAlchemyStore) -> str:
|
||||
return str(store.create_experiment("bench"))
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def seeded(store: SqlAlchemyStore, experiment_id: str) -> list[str]:
|
||||
return seed_traces(store, experiment_id, SEED_TRACES, SEED_SPANS_PER_TRACE)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def e2e_setup(bench_dir: Path) -> Iterator[None]:
|
||||
mlflow.set_tracking_uri(f"sqlite:///{bench_dir / 'e2e.db'}")
|
||||
mlflow.set_experiment("bench_e2e")
|
||||
yield
|
||||
mlflow.flush_trace_async_logging(terminate=True)
|
||||
@@ -0,0 +1,131 @@
|
||||
import random
|
||||
|
||||
from _data import generate_trace_data
|
||||
from pytest_benchmark.fixture import BenchmarkFixture
|
||||
|
||||
import mlflow
|
||||
from mlflow.entities.span import SpanType
|
||||
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
||||
|
||||
DEFAULT_SPANS = 100
|
||||
INGEST_ROUNDS = 20
|
||||
INGEST_WARMUP = 3
|
||||
|
||||
|
||||
def test_ingest(benchmark: BenchmarkFixture, store: SqlAlchemyStore, experiment_id: str) -> None:
|
||||
rng = random.Random(42)
|
||||
|
||||
def setup():
|
||||
ti, sp = generate_trace_data(experiment_id, DEFAULT_SPANS, rng)
|
||||
return (ti, sp), {}
|
||||
|
||||
def do(ti, sp):
|
||||
store.start_trace(ti)
|
||||
store.log_spans(experiment_id, sp)
|
||||
|
||||
benchmark.pedantic(
|
||||
do, setup=setup, iterations=1, rounds=INGEST_ROUNDS, warmup_rounds=INGEST_WARMUP
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_tag(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="tag.env = 'prod'",
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_state(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="status = 'ERROR'",
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_name_like(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="name LIKE 'rag_pipeline%'",
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_timestamp(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="timestamp > 0",
|
||||
order_by=["timestamp DESC"],
|
||||
)
|
||||
|
||||
|
||||
def _run_agent_workflow(num_tools: int, num_docs: int, query: str) -> None:
|
||||
with mlflow.start_span(name="agent_run", span_type=SpanType.AGENT) as root:
|
||||
root.set_inputs({"query": query})
|
||||
|
||||
with mlflow.start_span(name="retrieve", span_type=SpanType.RETRIEVER) as retr:
|
||||
retr.set_inputs({"query": query})
|
||||
docs = [
|
||||
{"id": f"doc_{i}", "score": 0.9 - i * 0.01, "text": f"doc text {i} " * 10}
|
||||
for i in range(num_docs)
|
||||
]
|
||||
retr.set_outputs({"documents": docs})
|
||||
|
||||
with mlflow.start_span(name="plan", span_type=SpanType.CHAIN) as planner:
|
||||
planner.set_inputs({"query": query, "num_docs": len(docs)})
|
||||
steps = [f"step_{i}" for i in range(num_tools)]
|
||||
planner.set_outputs({"steps": steps})
|
||||
|
||||
tool_results = []
|
||||
for step in steps:
|
||||
with mlflow.start_span(name=f"tool:{step}", span_type=SpanType.TOOL) as tool:
|
||||
tool.set_inputs({"step": step})
|
||||
result = {"step": step, "status": "ok", "value": len(step)}
|
||||
tool.set_outputs(result)
|
||||
tool_results.append(result)
|
||||
|
||||
with mlflow.start_span(name="summarize", span_type=SpanType.LLM) as summ:
|
||||
summ.set_inputs({"query": query, "tool_results": tool_results})
|
||||
response = f"Answer to {query!r} using {num_docs} docs and {num_tools} tool calls."
|
||||
summ.set_outputs({"response": response})
|
||||
summ.set_attribute("model", "gpt-test")
|
||||
summ.set_attribute("usage.input_tokens", 1234)
|
||||
summ.set_attribute("usage.output_tokens", 567)
|
||||
|
||||
root.set_outputs({"response": response})
|
||||
|
||||
|
||||
def test_e2e_agent(benchmark: BenchmarkFixture, e2e_setup: None) -> None:
|
||||
counter = [0]
|
||||
|
||||
def do():
|
||||
_run_agent_workflow(num_tools=20, num_docs=20, query=f"q-{counter[0]}")
|
||||
counter[0] += 1
|
||||
|
||||
benchmark(do)
|
||||
Executable
+229
@@ -0,0 +1,229 @@
|
||||
#!/usr/bin/env bash
|
||||
# =============================================================================
|
||||
# Static Site Build Script
|
||||
#
|
||||
# This script performs the following tasks:
|
||||
# 1. Checks that NodeJS (>=18.0) is installed; if not, it prints instructions
|
||||
# for installing via nvm.
|
||||
# 2. Changes directory into the docs folder so that all commands run from there.
|
||||
# (Note: This ensures that the DOCS_BASE_URL value is interpreted relative
|
||||
# to the docs folder. For example, if running from the project root, the
|
||||
# effective path will be: <root>/docs/docs/latest.)
|
||||
# 3. Installs dependencies via npm.
|
||||
# 4. (Optional) Builds the API docs:
|
||||
# - If --build-api-docs is provided, then the API docs are built.
|
||||
# - If --with-r-docs is also provided, the build includes R docs;
|
||||
# otherwise, R docs are skipped.
|
||||
# 5. Converts notebooks to MDX.
|
||||
# 6. Exports the DOCS_BASE_URL environment variable (default: /docs/latest) so
|
||||
# that Docusaurus uses the proper base URL.
|
||||
# 7. Builds the static site.
|
||||
#
|
||||
# Once complete, the script instructs the user to navigate into the docs folder
|
||||
# and run:
|
||||
#
|
||||
# npm run serve -- --port <your_port_number>
|
||||
#
|
||||
# Options:
|
||||
# --build-api-docs Opt in to build the API docs (default: do not build)
|
||||
# --with-r-docs When building API docs, include R documentation
|
||||
# (default: skip R docs)
|
||||
# --docs-base-url URL Override the default DOCS_BASE_URL (default: /docs/latest)
|
||||
# -h, --help Display this help message and exit
|
||||
#
|
||||
# Example:
|
||||
# ./dev/build-docs.sh --build-api-docs --with-r-docs --docs-base-url /docs/latest
|
||||
# =============================================================================
|
||||
|
||||
# Exit immediately if a command exits with a non-zero status,
|
||||
# treat unset variables as an error, and fail on pipeline errors.
|
||||
set -euo pipefail
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Define color and style variables for styled output.
|
||||
# -----------------------------------------------------------------------------
|
||||
BOLD='\033[1m'
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[0;33m'
|
||||
BLUE='\033[0;34m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Logging functions for consistent output styling.
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info() { echo -e "${BOLD}${BLUE}[INFO]${NC} $1"; }
|
||||
log_success() { echo -e "${BOLD}${GREEN}[SUCCESS]${NC} $1"; }
|
||||
log_warning() { echo -e "${BOLD}${YELLOW}[WARNING]${NC} $1"; }
|
||||
log_error() { echo -e "${BOLD}${RED}[ERROR]${NC} $1"; }
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Default configuration values.
|
||||
# -----------------------------------------------------------------------------
|
||||
BUILD_API_DOCS=false
|
||||
WITH_R_DOCS=false
|
||||
# Default DOCS_BASE_URL is set as expected when running from within the docs folder.
|
||||
# Note: When running from the project root, since we change into docs/,
|
||||
# the effective reference becomes: <root>/docs/docs/latest.
|
||||
DOCS_BASE_URL="/docs/latest"
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Display usage information.
|
||||
# -----------------------------------------------------------------------------
|
||||
usage() {
|
||||
cat <<EOF
|
||||
Usage: $0 [options]
|
||||
|
||||
Options:
|
||||
--build-api-docs Opt in to build the API docs (default: do not build)
|
||||
--with-r-docs When building API docs, include R documentation
|
||||
(default: skip R docs)
|
||||
--docs-base-url URL Override the default DOCS_BASE_URL (default: /docs/latest)
|
||||
-h, --help Display this help and exit
|
||||
|
||||
Example:
|
||||
$0 --build-api-docs --with-r-docs --docs-base-url /docs/latest
|
||||
EOF
|
||||
}
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Parse command-line arguments manually.
|
||||
# -----------------------------------------------------------------------------
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--build-api-docs)
|
||||
BUILD_API_DOCS=true
|
||||
shift
|
||||
;;
|
||||
--with-r-docs)
|
||||
WITH_R_DOCS=true
|
||||
shift
|
||||
;;
|
||||
--docs-base-url)
|
||||
if [[ -n "${2:-}" ]]; then
|
||||
DOCS_BASE_URL="$2"
|
||||
shift 2
|
||||
else
|
||||
log_error "--docs-base-url requires an argument."
|
||||
usage
|
||||
exit 1
|
||||
fi
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
log_error "Unknown option: $1"
|
||||
usage
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Function to compare semantic version numbers.
|
||||
# Returns 0 (true) if version $1 is greater than or equal to version $2.
|
||||
# -----------------------------------------------------------------------------
|
||||
version_ge() {
|
||||
# Usage: version_ge "18.15.0" "18.0.0"
|
||||
[ "$(printf '%s\n' "$2" "$1" | sort -V | head -n1)" = "$2" ]
|
||||
}
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Check that NodeJS is installed and meets the version requirement (>= 18.0).
|
||||
# -----------------------------------------------------------------------------
|
||||
if ! command -v node >/dev/null 2>&1; then
|
||||
log_error "NodeJS is not installed. Please install NodeJS (>= 18.0) from https://nodejs.org/ or via nvm."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
NODE_VERSION=$(node --version | sed 's/v//')
|
||||
REQUIRED_VERSION="18.0.0"
|
||||
if ! version_ge "$NODE_VERSION" "$REQUIRED_VERSION"; then
|
||||
log_error "Detected NodeJS version $NODE_VERSION. Please install NodeJS >= $REQUIRED_VERSION."
|
||||
log_info "If you have nvm installed, you can run:"
|
||||
echo -e "${BOLD}nvm install node && nvm use node${NC}"
|
||||
exit 1
|
||||
fi
|
||||
log_success "NodeJS version $NODE_VERSION is valid."
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Check that npm is installed.
|
||||
# -----------------------------------------------------------------------------
|
||||
if ! command -v npm >/dev/null 2>&1; then
|
||||
log_error "npm is not installed. Please install npm from https://nodejs.org/."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Change directory into the docs folder so that all commands run from there.
|
||||
# This ensures that DOCS_BASE_URL is interpreted correctly.
|
||||
# -----------------------------------------------------------------------------
|
||||
if [ ! -d "docs" ]; then
|
||||
log_error "The docs directory was not found. Make sure you're running this script from the project root."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
log_info "Changing directory to docs/ ..."
|
||||
cd docs
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Install dependencies via npm.
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "Installing dependencies with npm..."
|
||||
npm install
|
||||
log_success "Dependencies installed."
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Optionally build the API documentation.
|
||||
# This step is opt-in via the --build-api-docs flag.
|
||||
# If building API docs, the --with-r-docs flag controls whether R docs are included.
|
||||
# -----------------------------------------------------------------------------
|
||||
if [ "$BUILD_API_DOCS" = true ]; then
|
||||
if [ "$WITH_R_DOCS" = true ]; then
|
||||
log_info "Building API docs including R documentation..."
|
||||
npm run build-api-docs
|
||||
else
|
||||
log_info "Building API docs without R documentation..."
|
||||
npm run build-api-docs:no-r
|
||||
fi
|
||||
log_success "API docs built successfully."
|
||||
else
|
||||
log_info "Skipping API docs build phase."
|
||||
fi
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Update the API module references for link functionality
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "Updating API module links..."
|
||||
npm run update-api-modules
|
||||
log_success "Updated API module links."
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Convert notebooks to MDX format.
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "Converting notebooks to MDX..."
|
||||
npm run convert-notebooks
|
||||
log_success "Notebooks converted to MDX."
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Export DOCS_BASE_URL and build the static site.
|
||||
#
|
||||
# Since we're in the docs folder, exporting DOCS_BASE_URL as "/docs/latest"
|
||||
# means that when served from the project root, the built site will be available
|
||||
# at <root>/docs/docs/latest.
|
||||
# -----------------------------------------------------------------------------
|
||||
export DOCS_BASE_URL
|
||||
log_info "DOCS_BASE_URL set to '${DOCS_BASE_URL}'."
|
||||
log_info "Building static site files with npm..."
|
||||
npm run build
|
||||
log_success "Static site built successfully."
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Final instructions for the user.
|
||||
# -----------------------------------------------------------------------------
|
||||
log_info "To run the site locally, please navigate to the 'docs' folder and execute:"
|
||||
echo -e "${BOLD}npm run serve -- --port <your_port_number>${NC}"
|
||||
log_info "For example: ${BOLD}npm run serve -- --port 3000${NC}"
|
||||
log_success "Static site build process completed."
|
||||
+140
@@ -0,0 +1,140 @@
|
||||
import argparse
|
||||
import contextlib
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import zipfile
|
||||
from collections.abc import Generator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Package:
|
||||
# name of the package on PyPI.
|
||||
pypi_name: str
|
||||
# type of the package, one of "dev", "skinny", "tracing", "release"
|
||||
type: str
|
||||
# path to the package relative to the root of the repository
|
||||
build_path: str
|
||||
|
||||
|
||||
DEV = Package("mlflow", "dev", ".")
|
||||
RELEASE = Package("mlflow", "release", ".")
|
||||
SKINNY = Package("mlflow-skinny", "skinny", "libs/skinny")
|
||||
TRACING = Package("mlflow-tracing", "tracing", "libs/tracing")
|
||||
|
||||
PACKAGES = [
|
||||
DEV,
|
||||
SKINNY,
|
||||
RELEASE,
|
||||
TRACING,
|
||||
]
|
||||
|
||||
JS_BUILD_DIR = Path("mlflow/server/js/build")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Build MLflow package.")
|
||||
parser.add_argument(
|
||||
"--package-type",
|
||||
help="Package type to build. Default is 'dev'.",
|
||||
choices=[p.type for p in PACKAGES],
|
||||
default="dev",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sha",
|
||||
help="If specified, include the SHA in the wheel name as a build tag.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def restore_changes() -> Generator[None, None, None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
subprocess.check_call([
|
||||
"git",
|
||||
"restore",
|
||||
"README.md",
|
||||
"pyproject.toml",
|
||||
])
|
||||
|
||||
|
||||
def validate_ui_assets_pre_build(package: Package) -> None:
|
||||
if package != RELEASE:
|
||||
return
|
||||
if not JS_BUILD_DIR.exists() or not any(JS_BUILD_DIR.iterdir()):
|
||||
raise RuntimeError("Build the UI first before building the release package.")
|
||||
|
||||
|
||||
def validate_ui_assets_post_build(wheel_path: Path, package: Package) -> None:
|
||||
ui_asset_prefix = f"{JS_BUILD_DIR.as_posix()}/"
|
||||
with zipfile.ZipFile(wheel_path) as zf:
|
||||
has_ui_assets = any(name.startswith(ui_asset_prefix) for name in zf.namelist())
|
||||
if package == RELEASE:
|
||||
if not has_ui_assets:
|
||||
raise RuntimeError(
|
||||
f"UI assets are missing from the release wheel: {wheel_path}. "
|
||||
"Build the UI first before building the release package."
|
||||
)
|
||||
elif package in (SKINNY, TRACING):
|
||||
if has_ui_assets:
|
||||
raise RuntimeError(
|
||||
f"UI assets should not be included in the {package.type} wheel: {wheel_path}."
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
# Initialize submodules (e.g., mlflow/assistant/skills)
|
||||
subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"])
|
||||
|
||||
# Clean up build artifacts generated by previous builds
|
||||
paths_to_clean_up = ["build"]
|
||||
for pkg in PACKAGES:
|
||||
paths_to_clean_up += [
|
||||
f"{pkg.build_path}/dist",
|
||||
f"{pkg.build_path}/{pkg.pypi_name}.egg_info",
|
||||
]
|
||||
for path in map(Path, paths_to_clean_up):
|
||||
if not path.exists():
|
||||
continue
|
||||
if path.is_file():
|
||||
path.unlink()
|
||||
else:
|
||||
shutil.rmtree(path)
|
||||
|
||||
package = next(p for p in PACKAGES if p.type == args.package_type)
|
||||
|
||||
validate_ui_assets_pre_build(package)
|
||||
|
||||
with restore_changes():
|
||||
pyproject = Path("pyproject.toml")
|
||||
if package == RELEASE:
|
||||
pyproject.write_text(Path("pyproject.release.toml").read_text())
|
||||
|
||||
DIST_DIR = Path("dist").resolve()
|
||||
DIST_DIR.mkdir(exist_ok=True)
|
||||
subprocess.check_call([
|
||||
sys.executable,
|
||||
"-m",
|
||||
"build",
|
||||
package.build_path,
|
||||
"--outdir",
|
||||
DIST_DIR,
|
||||
])
|
||||
|
||||
wheel = next(DIST_DIR.glob("mlflow*.whl"))
|
||||
validate_ui_assets_post_build(wheel, package)
|
||||
|
||||
if args.sha:
|
||||
name, version, rest = wheel.name.split("-", 2)
|
||||
build_tag = f"0.sha.{args.sha}" # build tag must start with a digit
|
||||
wheel.rename(wheel.with_name(f"{name}-{version}-{build_tag}-{rest}"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,257 @@
|
||||
"""Validate GitHub Actions workflow and action files.
|
||||
|
||||
Complements `.github/policy.rego` with checks that need cross-file or remote
|
||||
context.
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
# Matches a `uses:` line that references a remote action (not a local `./` path).
|
||||
# Captures: owner/repo[/subpath] @ ref [ # comment ]
|
||||
_USES_RE = re.compile(
|
||||
r"""
|
||||
^\s*-?\s*uses:\s+ # leading `- uses:` or `uses:`
|
||||
(?P<action>[^@\s]+) # owner/repo[/subpath]
|
||||
@
|
||||
(?P<ref>[^\s#]+) # ref (SHA, tag, or branch)
|
||||
(?:\s+\#\s*(?P<comment>\S+))? # optional # comment
|
||||
""",
|
||||
re.VERBOSE,
|
||||
)
|
||||
|
||||
# A full 40-character hexadecimal SHA.
|
||||
_SHA_RE = re.compile(r"^[0-9a-f]{40}$")
|
||||
|
||||
# Requires at least vMAJOR.MINOR.PATCH to avoid ambiguous moving tags like v4.
|
||||
_VERSION_COMMENT_RE = re.compile(r"^v\d+\.\d+\.\d+(?:\.\d+)*$")
|
||||
|
||||
_CACHE_PATH = Path(".cache/action-pins.json")
|
||||
|
||||
|
||||
def _load_cache() -> dict[str, bool]:
|
||||
if _CACHE_PATH.exists():
|
||||
try:
|
||||
return json.loads(_CACHE_PATH.read_text()) # type: ignore[no-any-return]
|
||||
except (json.JSONDecodeError, OSError):
|
||||
pass
|
||||
return {}
|
||||
|
||||
|
||||
def _save_cache(cache: dict[str, bool]) -> None:
|
||||
try:
|
||||
_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
_CACHE_PATH.write_text(json.dumps(cache, indent=2, sort_keys=True))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _repo_from_action(action: str) -> str:
|
||||
match action.split("/"):
|
||||
case [owner, repo, *_]:
|
||||
return f"{owner}/{repo}"
|
||||
case _:
|
||||
raise ValueError(f"Invalid action format: {action!r}")
|
||||
|
||||
|
||||
def _verify_sha_tag(action: str, sha: str, tag: str, cache: dict[str, bool]) -> bool | None:
|
||||
|
||||
cache_key = f"{action}@{sha}#{tag}"
|
||||
if cache_key in cache:
|
||||
return cache[cache_key]
|
||||
|
||||
repo = _repo_from_action(action)
|
||||
try:
|
||||
result = _resolve_tag(repo, sha, tag)
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
||||
return None
|
||||
|
||||
cache[cache_key] = result
|
||||
return result
|
||||
|
||||
|
||||
def _resolve_tag(repo: str, sha: str, tag: str) -> bool:
|
||||
output = subprocess.check_output(
|
||||
["git", "ls-remote", "--tags", f"https://github.com/{repo}.git", tag],
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
return any(line.split()[0] == sha for line in output.splitlines() if line)
|
||||
|
||||
|
||||
def _iter_files() -> Iterator[Path]:
|
||||
root = Path(".github")
|
||||
for pattern in (
|
||||
"workflows/*.yml",
|
||||
"workflows/*.yaml",
|
||||
"actions/**/*.yml",
|
||||
"actions/**/*.yaml",
|
||||
):
|
||||
yield from root.glob(pattern)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ActionRef:
|
||||
prefix: str
|
||||
action: str
|
||||
ref: str
|
||||
comment: str | None
|
||||
|
||||
|
||||
def _iter_actions(path: Path) -> Iterator[ActionRef]:
|
||||
with path.open(encoding="utf-8") as f:
|
||||
for lineno, line in enumerate(f, start=1):
|
||||
if m := _USES_RE.match(line):
|
||||
action = m.group("action")
|
||||
if not action.startswith("./"):
|
||||
prefix = f"{path}:{lineno}: {line.strip()!r}"
|
||||
yield ActionRef(prefix, action, m.group("ref"), m.group("comment"))
|
||||
|
||||
|
||||
def _check_action(a: ActionRef, cache: dict[str, bool]) -> str | None:
|
||||
if not _SHA_RE.match(a.ref):
|
||||
return f"{a.prefix}\n error: ref '{a.ref}' is not a 40-character SHA"
|
||||
|
||||
if not a.comment or not _VERSION_COMMENT_RE.match(a.comment):
|
||||
return (
|
||||
f"{a.prefix}\n error: missing or invalid version comment"
|
||||
f" (expected '# vX.Y.Z', got {a.comment!r})"
|
||||
)
|
||||
|
||||
verified = _verify_sha_tag(a.action, a.ref, a.comment, cache)
|
||||
if verified is None:
|
||||
return (
|
||||
f"{a.prefix}\n error: could not verify SHA against tag '{a.comment}'"
|
||||
f" for {_repo_from_action(a.action)} (GitHub API unavailable)"
|
||||
)
|
||||
if not verified:
|
||||
return (
|
||||
f"{a.prefix}\n error: SHA '{a.ref}' does not match tag '{a.comment}'"
|
||||
f" for {_repo_from_action(a.action)}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
_LITERAL_CHARS = re.compile(r"[A-Za-z0-9._/\-]")
|
||||
|
||||
|
||||
def _glob_to_regex(pattern: str) -> re.Pattern[str]:
|
||||
# GitHub uses minimatch-style globs: `**` crosses `/`, `*` does not.
|
||||
i = 0
|
||||
n = len(pattern)
|
||||
parts = ["^"]
|
||||
while i < n:
|
||||
c = pattern[i]
|
||||
if c == "*" and i + 1 < n and pattern[i + 1] == "*":
|
||||
i += 2
|
||||
if i < n and pattern[i] == "/":
|
||||
# `**/` matches zero or more path segments
|
||||
parts.append("(?:.*/)?")
|
||||
i += 1
|
||||
else:
|
||||
parts.append(".*")
|
||||
elif c == "*":
|
||||
parts.append("[^/]*")
|
||||
i += 1
|
||||
elif _LITERAL_CHARS.match(c):
|
||||
parts.append(re.escape(c))
|
||||
i += 1
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported character {c!r} at position {i} in pattern {pattern!r}."
|
||||
" Extend _glob_to_regex if this is a valid GitHub path filter character."
|
||||
)
|
||||
parts.append("$")
|
||||
return re.compile("".join(parts))
|
||||
|
||||
|
||||
def _list_tracked_files() -> list[str]:
|
||||
return subprocess.check_output(["git", "ls-files"], text=True).splitlines()
|
||||
|
||||
|
||||
def _pattern_matches(pattern: str, files: list[str]) -> bool:
|
||||
regex = _glob_to_regex(pattern.removeprefix("!"))
|
||||
return any(regex.match(f) for f in files)
|
||||
|
||||
|
||||
def _iter_path_patterns(path: Path) -> Iterator[tuple[str, str, str]]:
|
||||
data = yaml.load(path.read_text(encoding="utf-8"), Loader=yaml.CSafeLoader)
|
||||
if not isinstance(data, dict):
|
||||
return
|
||||
# PyYAML parses the literal `on:` key as the boolean True (YAML 1.1).
|
||||
on = data.get("on", data.get(True))
|
||||
if not isinstance(on, dict):
|
||||
return
|
||||
for event, cfg in on.items():
|
||||
if not isinstance(cfg, dict):
|
||||
continue
|
||||
for key in ("paths", "paths-ignore"):
|
||||
for pattern in cfg.get(key) or []:
|
||||
yield str(event), key, pattern
|
||||
|
||||
|
||||
def _iter_workflow_files() -> Iterator[Path]:
|
||||
root = Path(".github/workflows")
|
||||
for ext in ("*.yml", "*.yaml"):
|
||||
yield from root.glob(ext)
|
||||
|
||||
|
||||
def _check_paths() -> Iterator[str]:
|
||||
files = _list_tracked_files()
|
||||
for path in sorted(_iter_workflow_files()):
|
||||
for event, key, pattern in _iter_path_patterns(path):
|
||||
if not _pattern_matches(pattern, files):
|
||||
yield (
|
||||
f"{path}: [on.{event}.{key}] pattern {pattern!r} does not"
|
||||
" match any tracked file"
|
||||
)
|
||||
|
||||
|
||||
def _check_version_consistency(all_action_refs: list[ActionRef]) -> Iterator[str]:
|
||||
by_action: dict[str, list[ActionRef]] = defaultdict(list)
|
||||
for action_ref in all_action_refs:
|
||||
by_action[action_ref.action].append(action_ref)
|
||||
|
||||
for action, refs in sorted(by_action.items()):
|
||||
versions = {(ref.ref, ref.comment) for ref in refs}
|
||||
if len(versions) > 1:
|
||||
lines = "\n".join(f" {ref.prefix}" for ref in sorted(refs, key=lambda r: r.prefix))
|
||||
yield f"{action} is pinned to multiple versions:\n{lines}"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
cache = _load_cache()
|
||||
all_errors: list[str] = []
|
||||
all_action_refs: list[ActionRef] = []
|
||||
try:
|
||||
for path in _iter_files():
|
||||
for action_ref in _iter_actions(path):
|
||||
if error := _check_action(action_ref, cache):
|
||||
all_errors.append(error)
|
||||
else:
|
||||
all_action_refs.append(action_ref)
|
||||
finally:
|
||||
_save_cache(cache)
|
||||
all_errors.extend(_check_version_consistency(all_action_refs))
|
||||
all_errors.extend(_check_paths())
|
||||
|
||||
if all_errors:
|
||||
print("check-actions: the following violations were found:\n", file=sys.stderr)
|
||||
for err in all_errors:
|
||||
print(err, file=sys.stderr)
|
||||
print(f"\n{len(all_errors)} violation(s) found.", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Executable
+382
@@ -0,0 +1,382 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def is_github_actions() -> bool:
|
||||
return os.environ.get("GITHUB_ACTIONS") == "true"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Error:
|
||||
file_path: Path
|
||||
line: int
|
||||
column: int
|
||||
lines: list[str]
|
||||
|
||||
def format(self, github: bool = False) -> str:
|
||||
message = " ".join(self.lines)
|
||||
if github:
|
||||
return f"::warning file={self.file_path},line={self.line},col={self.column}::{message}"
|
||||
else:
|
||||
return f"{self.file_path}:{self.line}:{self.column}: {message}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Parameter:
|
||||
name: str
|
||||
position: int | None # None for keyword-only
|
||||
is_required: bool
|
||||
is_positional_only: bool
|
||||
is_keyword_only: bool
|
||||
lineno: int
|
||||
col_offset: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class Signature:
|
||||
positional: list[Parameter] # Includes positional-only and regular positional
|
||||
keyword_only: list[Parameter]
|
||||
has_var_positional: bool # *args
|
||||
has_var_keyword: bool # **kwargs
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParameterError:
|
||||
message: str
|
||||
param_name: str
|
||||
lineno: int
|
||||
col_offset: int
|
||||
|
||||
|
||||
def parse_signature(args: ast.arguments) -> Signature:
|
||||
"""Convert ast.arguments to a Signature dataclass for easier processing."""
|
||||
parameters_positional: list[Parameter] = []
|
||||
parameters_keyword_only: list[Parameter] = []
|
||||
|
||||
# Process positional-only parameters
|
||||
for i, arg in enumerate(args.posonlyargs):
|
||||
parameters_positional.append(
|
||||
Parameter(
|
||||
name=arg.arg,
|
||||
position=i,
|
||||
is_required=True, # All positional-only are required
|
||||
is_positional_only=True,
|
||||
is_keyword_only=False,
|
||||
lineno=arg.lineno,
|
||||
col_offset=arg.col_offset,
|
||||
)
|
||||
)
|
||||
|
||||
# Process regular positional parameters
|
||||
offset = len(args.posonlyargs)
|
||||
first_optional_idx = len(args.posonlyargs + args.args) - len(args.defaults)
|
||||
|
||||
for i, arg in enumerate(args.args):
|
||||
pos = offset + i
|
||||
parameters_positional.append(
|
||||
Parameter(
|
||||
name=arg.arg,
|
||||
position=pos,
|
||||
is_required=pos < first_optional_idx,
|
||||
is_positional_only=False,
|
||||
is_keyword_only=False,
|
||||
lineno=arg.lineno,
|
||||
col_offset=arg.col_offset,
|
||||
)
|
||||
)
|
||||
|
||||
# Process keyword-only parameters
|
||||
for arg, default in zip(args.kwonlyargs, args.kw_defaults):
|
||||
parameters_keyword_only.append(
|
||||
Parameter(
|
||||
name=arg.arg,
|
||||
position=None,
|
||||
is_required=default is None,
|
||||
is_positional_only=False,
|
||||
is_keyword_only=True,
|
||||
lineno=arg.lineno,
|
||||
col_offset=arg.col_offset,
|
||||
)
|
||||
)
|
||||
|
||||
return Signature(
|
||||
positional=parameters_positional,
|
||||
keyword_only=parameters_keyword_only,
|
||||
has_var_positional=args.vararg is not None,
|
||||
has_var_keyword=args.kwarg is not None,
|
||||
)
|
||||
|
||||
|
||||
def check_signature_compatibility(
|
||||
old_fn: ast.FunctionDef | ast.AsyncFunctionDef,
|
||||
new_fn: ast.FunctionDef | ast.AsyncFunctionDef,
|
||||
) -> list[ParameterError]:
|
||||
"""
|
||||
Return list of error messages when *new_fn* is not backward-compatible with *old_fn*,
|
||||
or None if compatible.
|
||||
|
||||
Compatibility rules
|
||||
-------------------
|
||||
• Positional / positional-only parameters
|
||||
- Cannot be reordered, renamed, or removed.
|
||||
- Adding **required** ones is breaking.
|
||||
- Adding **optional** ones is allowed only at the end.
|
||||
- Making an optional parameter required is breaking.
|
||||
|
||||
• Keyword-only parameters (order does not matter)
|
||||
- Cannot be renamed or removed.
|
||||
- Making an optional parameter required is breaking.
|
||||
- Adding a required parameter is breaking; adding an optional parameter is fine.
|
||||
"""
|
||||
old_sig = parse_signature(old_fn.args)
|
||||
new_sig = parse_signature(new_fn.args)
|
||||
errors: list[ParameterError] = []
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 1. Positional / pos-only parameters
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
# (a) existing parameters must line up
|
||||
for idx, old_param in enumerate(old_sig.positional):
|
||||
if idx >= len(new_sig.positional):
|
||||
errors.append(
|
||||
ParameterError(
|
||||
message=f"Positional param '{old_param.name}' was removed.",
|
||||
param_name=old_param.name,
|
||||
lineno=old_param.lineno,
|
||||
col_offset=old_param.col_offset,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
new_param = new_sig.positional[idx]
|
||||
if old_param.name != new_param.name:
|
||||
errors.append(
|
||||
ParameterError(
|
||||
message=(
|
||||
f"Positional param order/name changed: "
|
||||
f"'{old_param.name}' -> '{new_param.name}'."
|
||||
),
|
||||
param_name=new_param.name,
|
||||
lineno=new_param.lineno,
|
||||
col_offset=new_param.col_offset,
|
||||
)
|
||||
)
|
||||
# Stop checking further positional params after first order/name mismatch
|
||||
break
|
||||
|
||||
if (not old_param.is_required) and new_param.is_required:
|
||||
errors.append(
|
||||
ParameterError(
|
||||
message=f"Optional positional param '{old_param.name}' became required.",
|
||||
param_name=new_param.name,
|
||||
lineno=new_param.lineno,
|
||||
col_offset=new_param.col_offset,
|
||||
)
|
||||
)
|
||||
|
||||
# (b) any extra new positional params must be optional and appended
|
||||
if len(new_sig.positional) > len(old_sig.positional):
|
||||
for idx in range(len(old_sig.positional), len(new_sig.positional)):
|
||||
new_param = new_sig.positional[idx]
|
||||
if new_param.is_required:
|
||||
errors.append(
|
||||
ParameterError(
|
||||
message=f"New required positional param '{new_param.name}' added.",
|
||||
param_name=new_param.name,
|
||||
lineno=new_param.lineno,
|
||||
col_offset=new_param.col_offset,
|
||||
)
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 2. Keyword-only parameters (order-agnostic)
|
||||
# ------------------------------------------------------------------ #
|
||||
old_kw_names = {p.name for p in old_sig.keyword_only}
|
||||
new_kw_names = {p.name for p in new_sig.keyword_only}
|
||||
|
||||
# Build mappings for easier lookup
|
||||
old_kw_by_name = {p.name: p for p in old_sig.keyword_only}
|
||||
new_kw_by_name = {p.name: p for p in new_sig.keyword_only}
|
||||
|
||||
# removed or renamed
|
||||
for name in old_kw_names - new_kw_names:
|
||||
old_param = old_kw_by_name[name]
|
||||
errors.append(
|
||||
ParameterError(
|
||||
message=f"Keyword-only param '{name}' was removed.",
|
||||
param_name=name,
|
||||
lineno=old_param.lineno,
|
||||
col_offset=old_param.col_offset,
|
||||
)
|
||||
)
|
||||
|
||||
# optional -> required upgrades
|
||||
for name in old_kw_names & new_kw_names:
|
||||
if not old_kw_by_name[name].is_required and new_kw_by_name[name].is_required:
|
||||
new_param = new_kw_by_name[name]
|
||||
errors.append(
|
||||
ParameterError(
|
||||
message=f"Keyword-only param '{name}' became required.",
|
||||
param_name=name,
|
||||
lineno=new_param.lineno,
|
||||
col_offset=new_param.col_offset,
|
||||
)
|
||||
)
|
||||
|
||||
# new required keyword-only params
|
||||
errors.extend(
|
||||
ParameterError(
|
||||
message=f"New required keyword-only param '{param.name}' added.",
|
||||
param_name=param.name,
|
||||
lineno=param.lineno,
|
||||
col_offset=param.col_offset,
|
||||
)
|
||||
for param in new_sig.keyword_only
|
||||
if param.is_required and param.name not in old_kw_names
|
||||
)
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
def _is_private(n: str) -> bool:
|
||||
return n.startswith("_") and not n.startswith("__") and not n.endswith("__")
|
||||
|
||||
|
||||
class FunctionSignatureExtractor(ast.NodeVisitor):
|
||||
def __init__(self) -> None:
|
||||
self.functions: dict[str, ast.FunctionDef | ast.AsyncFunctionDef] = {}
|
||||
self.stack: list[ast.ClassDef] = []
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef) -> None:
|
||||
self.stack.append(node)
|
||||
self.generic_visit(node)
|
||||
self.stack.pop()
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
|
||||
# Is this a private function or a function in a private class?
|
||||
# If so, skip it.
|
||||
if _is_private(node.name) or (self.stack and _is_private(self.stack[-1].name)):
|
||||
return
|
||||
|
||||
names = [*(c.name for c in self.stack), node.name]
|
||||
self.functions[".".join(names)] = node
|
||||
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
|
||||
if _is_private(node.name) or (self.stack and _is_private(self.stack[-1].name)):
|
||||
return
|
||||
|
||||
names = [*(c.name for c in self.stack), node.name]
|
||||
self.functions[".".join(names)] = node
|
||||
|
||||
|
||||
def get_changed_python_files(base_branch: str = "master") -> list[Path]:
|
||||
# In GitHub Actions PR context, we need to fetch the base branch first
|
||||
if is_github_actions():
|
||||
# Fetch the base branch to ensure we have it locally
|
||||
subprocess.check_call(
|
||||
["git", "fetch", "origin", f"{base_branch}:{base_branch}"],
|
||||
)
|
||||
|
||||
result = subprocess.check_output(
|
||||
["git", "diff", "--name-only", f"{base_branch}...HEAD"], text=True
|
||||
)
|
||||
files = [s.strip() for s in result.splitlines()]
|
||||
return [Path(f) for f in files if f]
|
||||
|
||||
|
||||
def parse_functions(content: str) -> dict[str, ast.FunctionDef | ast.AsyncFunctionDef]:
|
||||
tree = ast.parse(content)
|
||||
extractor = FunctionSignatureExtractor()
|
||||
extractor.visit(tree)
|
||||
return extractor.functions
|
||||
|
||||
|
||||
def get_file_content_at_revision(file_path: Path, revision: str) -> str | None:
|
||||
try:
|
||||
return subprocess.check_output(["git", "show", f"{revision}:{file_path}"], text=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Warning: Failed to get file content at revision: {e}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
|
||||
def compare_signatures(base_branch: str = "master") -> list[Error]:
|
||||
errors: list[Error] = []
|
||||
for file_path in get_changed_python_files(base_branch):
|
||||
# Ignore non-Python files
|
||||
if not file_path.suffix == ".py":
|
||||
continue
|
||||
|
||||
# Ignore files not in the mlflow directory
|
||||
if file_path.parts[0] != "mlflow":
|
||||
continue
|
||||
|
||||
# Ignore private modules
|
||||
if any(part.startswith("_") and part != "__init__.py" for part in file_path.parts):
|
||||
continue
|
||||
|
||||
base_content = get_file_content_at_revision(file_path, base_branch)
|
||||
if base_content is None:
|
||||
# Find not found in the base branch, likely added in the current branch
|
||||
continue
|
||||
|
||||
if not file_path.exists():
|
||||
# File not found, likely deleted in the current branch
|
||||
continue
|
||||
|
||||
current_content = file_path.read_text()
|
||||
base_functions = parse_functions(base_content)
|
||||
current_functions = parse_functions(current_content)
|
||||
for func_name in set(base_functions.keys()) & set(current_functions.keys()):
|
||||
base_func = base_functions[func_name]
|
||||
current_func = current_functions[func_name]
|
||||
if param_errors := check_signature_compatibility(base_func, current_func):
|
||||
# Create individual errors for each problematic parameter
|
||||
errors.extend(
|
||||
Error(
|
||||
file_path=file_path,
|
||||
line=param_error.lineno,
|
||||
column=param_error.col_offset + 1,
|
||||
lines=[
|
||||
"[Non-blocking | Ignore if not public API]",
|
||||
param_error.message,
|
||||
f"This change will break existing `{func_name}` calls.",
|
||||
"If this is not intended, please fix it.",
|
||||
],
|
||||
)
|
||||
for param_error in param_errors
|
||||
)
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
@dataclass
|
||||
class Args:
|
||||
base_branch: str
|
||||
|
||||
|
||||
def parse_args() -> Args:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Check for breaking changes in Python function signatures"
|
||||
)
|
||||
parser.add_argument("--base-branch", default=os.environ.get("GITHUB_BASE_REF", "master"))
|
||||
args = parser.parse_args()
|
||||
return Args(base_branch=args.base_branch)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
errors = compare_signatures(args.base_branch)
|
||||
for error in errors:
|
||||
print(error.format(github=is_github_actions()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
Pre-commit hook to check for missing `__init__.py` files in mlflow and tests directories.
|
||||
|
||||
This script ensures that all directories under the mlflow package and tests directory that contain
|
||||
Python files also have an `__init__.py` file. This prevents `setuptools` from excluding these
|
||||
directories during package build and ensures test modules are properly structured.
|
||||
|
||||
Usage:
|
||||
uv run dev/check_init_py.py
|
||||
|
||||
Requirements:
|
||||
- If `mlflow/foo/bar.py` exists, `mlflow/foo/__init__.py` must exist.
|
||||
- If `tests/foo/test_bar.py` exists, `tests/foo/__init__.py` must exist.
|
||||
- Only test files (starting with `test_`) in the tests directory are checked.
|
||||
- All parent directories of Python files are checked recursively for `__init__.py`.
|
||||
- Ignore directories that do not contain any Python files (e.g., `mlflow/server/js`).
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def get_tracked_python_files() -> list[Path]:
|
||||
try:
|
||||
result = subprocess.check_output(
|
||||
["git", "ls-files", "mlflow/**/*.py", "tests/**/*.py"],
|
||||
text=True,
|
||||
)
|
||||
paths = (Path(f) for f in result.splitlines() if f)
|
||||
return [p for p in paths if (not p.is_relative_to("tests") or p.name.startswith("test_"))]
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error running git ls-files: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
python_files = get_tracked_python_files()
|
||||
if not python_files:
|
||||
return 0
|
||||
|
||||
python_dirs = {p for f in python_files for p in f.parents if p != Path(".")}
|
||||
if missing_init_files := [d for d in python_dirs if not (d / "__init__.py").exists()]:
|
||||
print("Error: The following directories contain Python files but lack __init__.py:")
|
||||
for d in sorted(missing_init_files):
|
||||
print(f" {d.as_posix()}/")
|
||||
print("Please add __init__.py files to the directories listed above.")
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,290 @@
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import itertools
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from packaging.version import Version
|
||||
|
||||
MAX_COMMITS_PER_SCRIPT = 90
|
||||
|
||||
|
||||
def chunk_list(lst: list[str], size: int) -> list[list[str]]:
|
||||
return [lst[i : i + size] for i in range(0, len(lst), size)]
|
||||
|
||||
|
||||
def get_token() -> str | None:
|
||||
if token := os.environ.get("GH_TOKEN"):
|
||||
return token
|
||||
try:
|
||||
token = subprocess.check_output(
|
||||
["gh", "auth", "token"], text=True, stderr=subprocess.DEVNULL
|
||||
).strip()
|
||||
if token:
|
||||
return token
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def get_headers() -> dict[str, str]:
|
||||
if token := get_token():
|
||||
return {"Authorization": f"token {token}"}
|
||||
return {}
|
||||
|
||||
|
||||
def validate_version(version: str) -> None:
|
||||
"""
|
||||
Validate that the version has a micro version component.
|
||||
Raises ValueError if the version is invalid.
|
||||
"""
|
||||
parsed_version = Version(version)
|
||||
if len(parsed_version.release) != 3:
|
||||
raise ValueError(
|
||||
f"Invalid version: '{version}'. "
|
||||
"Version must be in the format <major>.<minor>.<micro> (e.g., '2.10.0')"
|
||||
)
|
||||
|
||||
|
||||
def get_release_branch(version: str) -> str:
|
||||
major_minor_version = ".".join(version.split(".")[:2])
|
||||
return f"branch-{major_minor_version}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Commit:
|
||||
sha: str
|
||||
pr_num: int
|
||||
date: str
|
||||
|
||||
|
||||
def get_commit_count(branch: str, since: str) -> int:
|
||||
"""
|
||||
Get the total count of commits in the branch since the given date using GraphQL API.
|
||||
"""
|
||||
query = """
|
||||
query($branch: String!, $since: GitTimestamp!) {
|
||||
repository(owner: "mlflow", name: "mlflow") {
|
||||
ref(qualifiedName: $branch) {
|
||||
target {
|
||||
... on Commit {
|
||||
history(since: $since) {
|
||||
totalCount
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
response = requests.post(
|
||||
"https://api.github.com/graphql",
|
||||
json={"query": query, "variables": {"branch": branch, "since": since}},
|
||||
headers=get_headers(),
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
ref = data["data"]["repository"]["ref"]
|
||||
if ref is None:
|
||||
raise ValueError(f"Branch '{branch}' not found")
|
||||
total_count: int = ref["target"]["history"]["totalCount"]
|
||||
return total_count
|
||||
|
||||
|
||||
def get_commits(branch: str) -> list[Commit]:
|
||||
"""
|
||||
Get the commits in the release branch via GitHub API (last 90 days).
|
||||
Returns commits sorted by date (oldest first).
|
||||
"""
|
||||
per_page = 100
|
||||
pr_rgx = re.compile(r".+\s+\(#(\d+)\)$")
|
||||
since = (datetime.now(timezone.utc) - timedelta(days=90)).isoformat()
|
||||
|
||||
# Get total commit count first
|
||||
total_count = get_commit_count(branch, since)
|
||||
if total_count == 0:
|
||||
print(f"No commits found in {branch} since {since}")
|
||||
return []
|
||||
|
||||
total_pages = (total_count + per_page - 1) // per_page
|
||||
print(f"Total commits: {total_count}, fetching {total_pages} page(s)...")
|
||||
|
||||
def fetch_page(page: int) -> list[Commit]:
|
||||
print(f"Fetching page {page}/{total_pages}...")
|
||||
params: dict[str, str | int] = {
|
||||
"sha": branch,
|
||||
"per_page": per_page,
|
||||
"page": page,
|
||||
"since": since,
|
||||
}
|
||||
response = requests.get(
|
||||
"https://api.github.com/repos/mlflow/mlflow/commits",
|
||||
params=params,
|
||||
headers=get_headers(),
|
||||
)
|
||||
response.raise_for_status()
|
||||
commits = []
|
||||
for item in response.json():
|
||||
msg = item["commit"]["message"].split("\n")[0]
|
||||
if m := pr_rgx.search(msg):
|
||||
# Use committer date (not author date) because cherry-picked commits
|
||||
# retain the original author date but get a new committer date.
|
||||
date = item["commit"]["committer"]["date"]
|
||||
commits.append(Commit(sha=item["sha"], pr_num=int(m.group(1)), date=date))
|
||||
return commits
|
||||
|
||||
# Fetch all pages in parallel. executor.map preserves order.
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
results = executor.map(fetch_page, range(1, total_pages + 1))
|
||||
|
||||
return sorted(itertools.chain.from_iterable(results), key=lambda c: c.date)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PR:
|
||||
pr_num: int
|
||||
merged: bool
|
||||
|
||||
|
||||
def is_closed(pr: dict[str, Any]) -> bool:
|
||||
return pr["state"] == "closed" and pr["pull_request"]["merged_at"] is None
|
||||
|
||||
|
||||
def fetch_patch_prs(version: str) -> dict[int, bool]:
|
||||
"""
|
||||
Fetch PRs labeled with `v{version}` from the MLflow repository.
|
||||
"""
|
||||
label = f"v{version}"
|
||||
per_page = 100
|
||||
page = 1
|
||||
pulls: list[dict[str, Any]] = []
|
||||
while True:
|
||||
response = requests.get(
|
||||
f'https://api.github.com/search/issues?q=is:pr+repo:mlflow/mlflow+label:"{label}"&per_page={per_page}&page={page}',
|
||||
headers=get_headers(),
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
# Exclude closed PRs that are not merged
|
||||
pulls.extend(pr for pr in data["items"] if not is_closed(pr))
|
||||
if len(data["items"]) < per_page:
|
||||
break
|
||||
page += 1
|
||||
|
||||
return {pr["number"]: pr["pull_request"].get("merged_at") is not None for pr in pulls}
|
||||
|
||||
|
||||
def main(version: str, dry_run: bool) -> None:
|
||||
validate_version(version)
|
||||
release_branch = get_release_branch(version)
|
||||
commits = get_commits(release_branch)
|
||||
patch_prs = fetch_patch_prs(version)
|
||||
if not_cherry_picked := set(patch_prs) - {c.pr_num for c in commits}:
|
||||
print(f"The following patch PRs are not cherry-picked to {release_branch}:")
|
||||
for idx, pr_num in enumerate(sorted(not_cherry_picked)):
|
||||
merged = patch_prs[pr_num]
|
||||
url = f"https://github.com/mlflow/mlflow/pull/{pr_num} (merged: {merged})"
|
||||
line = f" {idx + 1}. {url}"
|
||||
if not merged:
|
||||
line = f"\033[91m{line}\033[0m" # Red color using ANSI escape codes
|
||||
print(line)
|
||||
|
||||
master_commits = get_commits("master")
|
||||
cherry_picks = [c.sha for c in master_commits if c.pr_num in not_cherry_picked]
|
||||
|
||||
# Split into chunks if needed
|
||||
chunks = chunk_list(cherry_picks, MAX_COMMITS_PER_SCRIPT)
|
||||
|
||||
# Print warning if splitting
|
||||
if len(chunks) > 1:
|
||||
print(
|
||||
f"\n⚠️ WARNING: {len(cherry_picks)} commits will be split into "
|
||||
f"{len(chunks)} scripts."
|
||||
)
|
||||
print("Create one PR per script and merge them sequentially:")
|
||||
print(" file PR 1 → merge PR 1 → pull release branch → file PR 2 → merge PR 2 → ...")
|
||||
print("This is required to stay under GitHub's 100-commit rebase merge limit.\n")
|
||||
|
||||
print("\n# Steps to cherry-pick the patch PRs:")
|
||||
print(
|
||||
f"1. Make sure your local master and {release_branch} branches are synced with "
|
||||
"upstream."
|
||||
)
|
||||
print(f"2. Cut a new branch from {release_branch} (e.g. {release_branch}-cherry-picks).")
|
||||
|
||||
# Generate script(s)
|
||||
tmp_dir = Path(tempfile.gettempdir())
|
||||
script_paths: list[tuple[Path, int]] = []
|
||||
for i, chunk in enumerate(chunks, 1):
|
||||
if len(chunks) == 1:
|
||||
script_path = tmp_dir / "cherry-pick.sh"
|
||||
else:
|
||||
script_path = tmp_dir / f"cherry-pick-{i}.sh"
|
||||
|
||||
script_content = f"""\
|
||||
#!/usr/bin/env bash
|
||||
# Cherry-picks for v{version} -> {release_branch}
|
||||
# Generated: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")}
|
||||
# Commits: {len(chunk)} of {len(cherry_picks)} total
|
||||
#
|
||||
# If conflicts occur, resolve them and run:
|
||||
# git cherry-pick --continue
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Guard: Prevent running on master branch
|
||||
current_branch=$(git rev-parse --abbrev-ref HEAD)
|
||||
if [[ "$current_branch" == "master" ]]; then
|
||||
echo "ERROR: This script must not be run on the master branch."
|
||||
echo "Please checkout a release branch (e.g., {release_branch}) or a branch derived from it."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
git cherry-pick {" ".join(chunk)}
|
||||
"""
|
||||
script_path.write_text(script_content)
|
||||
script_path.chmod(0o755)
|
||||
script_paths.append((script_path, len(chunk)))
|
||||
|
||||
if len(chunks) == 1:
|
||||
print("3. Run the cherry-pick script on the new branch:\n")
|
||||
print(f"Cherry-pick script written to: {script_paths[0][0]}")
|
||||
print(f"\n4. File a PR against {release_branch}.")
|
||||
else:
|
||||
print("3. For each script (in order):")
|
||||
print(f" a. Create a new branch from {release_branch}")
|
||||
print(" b. Run the script")
|
||||
print(f" c. File a PR against {release_branch}")
|
||||
print(f" d. After merge, pull {release_branch} from remote before the next script\n")
|
||||
print(" Scripts:")
|
||||
for script_path, commit_count in script_paths:
|
||||
print(f" {script_path} ({commit_count} commits)")
|
||||
|
||||
sys.exit(0 if dry_run else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--version", required=True, help="The version to release")
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
default=os.environ.get("DRY_RUN", "true").lower() == "true",
|
||||
help="Dry run mode (default: True, can be set via DRY_RUN env var)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-dry-run",
|
||||
action="store_false",
|
||||
dest="dry_run",
|
||||
help="Disable dry run mode",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
main(args.version, args.dry_run)
|
||||
@@ -0,0 +1,65 @@
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
# https://agentskills.io/specification#frontmatter
|
||||
NAME_RE = re.compile(r"^[a-z0-9]+(-[a-z0-9]+)*$")
|
||||
NAME_MAX = 64
|
||||
DESCRIPTION_MAX = 1024
|
||||
|
||||
|
||||
def parse_frontmatter(text: str) -> dict[str, Any] | None:
|
||||
if not text.startswith("---\n"):
|
||||
return None
|
||||
end = text.find("\n---\n", 4)
|
||||
if end == -1:
|
||||
return None
|
||||
data = yaml.safe_load(text[4:end]) or {}
|
||||
return data if isinstance(data, dict) else None
|
||||
|
||||
|
||||
def validate_name(name: Any, parent: str) -> list[str]:
|
||||
if not name:
|
||||
return ["missing or empty `name`"]
|
||||
if not isinstance(name, str) or not NAME_RE.fullmatch(name) or len(name) > NAME_MAX:
|
||||
return [
|
||||
f"invalid `name`: must be 1-{NAME_MAX} lowercase alphanumeric/hyphen chars, "
|
||||
"no leading/trailing or consecutive hyphens"
|
||||
]
|
||||
if name != parent:
|
||||
return [f"`name` {name!r} does not match parent directory {parent!r}"]
|
||||
return []
|
||||
|
||||
|
||||
def validate_description(description: Any) -> list[str]:
|
||||
if not description:
|
||||
return ["missing or empty `description`"]
|
||||
if not isinstance(description, str) or len(description) > DESCRIPTION_MAX:
|
||||
return [f"invalid `description`: must be 1-{DESCRIPTION_MAX} characters"]
|
||||
return []
|
||||
|
||||
|
||||
def check(path: Path) -> list[str]:
|
||||
fm = parse_frontmatter(path.read_text(encoding="utf-8"))
|
||||
if fm is None:
|
||||
return ["missing frontmatter"]
|
||||
return [
|
||||
*validate_name(fm.get("name"), path.parent.name),
|
||||
*validate_description(fm.get("description")),
|
||||
]
|
||||
|
||||
|
||||
def main(argv: list[str]) -> bool:
|
||||
failed = False
|
||||
for arg in argv:
|
||||
for err in check(Path(arg)):
|
||||
print(f"{arg}: {err}", file=sys.stderr)
|
||||
failed = True
|
||||
return failed
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main(sys.argv[1:]))
|
||||
Executable
+148
@@ -0,0 +1,148 @@
|
||||
"""
|
||||
Detect files where all changes are whitespace-only.
|
||||
|
||||
This helps avoid unnecessary commit history noise from whitespace-only changes.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from typing import cast
|
||||
|
||||
BYPASS_LABEL = "allow-whitespace-only"
|
||||
|
||||
_MAX_ATTEMPTS = 3
|
||||
|
||||
|
||||
def _is_retryable(exc: Exception) -> bool:
|
||||
match exc:
|
||||
case urllib.error.HTTPError(code=code):
|
||||
return code >= 500
|
||||
case urllib.error.URLError():
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def _retry_urlopen(request: urllib.request.Request, timeout: int = 30) -> str:
|
||||
for i in range(_MAX_ATTEMPTS):
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
return cast(str, response.read().decode("utf-8"))
|
||||
except Exception as exc:
|
||||
if not _is_retryable(exc) or i == _MAX_ATTEMPTS - 1:
|
||||
raise
|
||||
time.sleep(2**i)
|
||||
raise RuntimeError("unreachable")
|
||||
|
||||
|
||||
def github_api_request(url: str, accept: str) -> str:
|
||||
headers = {
|
||||
"Accept": accept,
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
}
|
||||
|
||||
if github_token := os.environ.get("GH_TOKEN"):
|
||||
headers["Authorization"] = f"Bearer {github_token}"
|
||||
|
||||
request = urllib.request.Request(url, headers=headers)
|
||||
return _retry_urlopen(request)
|
||||
|
||||
|
||||
def get_pr_diff(owner: str, repo: str, pull_number: int) -> str:
|
||||
url = f"https://github.com/{owner}/{repo}/pull/{pull_number}.diff"
|
||||
request = urllib.request.Request(url)
|
||||
return _retry_urlopen(request)
|
||||
|
||||
|
||||
def get_pr_labels(owner: str, repo: str, pull_number: int) -> list[str]:
|
||||
url = f"https://api.github.com/repos/{owner}/{repo}/pulls/{pull_number}"
|
||||
data = json.loads(github_api_request(url, "application/vnd.github.v3+json"))
|
||||
return [label_obj["name"] for label_obj in data.get("labels", [])]
|
||||
|
||||
|
||||
def parse_diff(diff_text: str | None) -> list[str]:
|
||||
if not diff_text:
|
||||
return []
|
||||
|
||||
files: list[str] = []
|
||||
current_file: str | None = None
|
||||
changes: list[str] = []
|
||||
|
||||
for line in diff_text.split("\n"):
|
||||
if line.startswith("diff --git"):
|
||||
if current_file and changes and all(c.strip() == "" for c in changes):
|
||||
files.append(current_file)
|
||||
|
||||
current_file = None
|
||||
changes = []
|
||||
|
||||
elif line.startswith("--- a/"):
|
||||
current_file = None if line == "--- /dev/null" else line[6:]
|
||||
|
||||
elif line.startswith("+++ b/"):
|
||||
current_file = None if line == "+++ /dev/null" else line[6:]
|
||||
|
||||
elif line.startswith("+") or line.startswith("-"):
|
||||
content = line[1:]
|
||||
changes.append(content)
|
||||
|
||||
if current_file and changes and all(c.strip() == "" for c in changes):
|
||||
files.append(current_file)
|
||||
|
||||
return files
|
||||
|
||||
|
||||
def parse_args() -> tuple[str, str, int]:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Check for unnecessary whitespace-only changes in the diff"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--repo",
|
||||
required=True,
|
||||
help='Repository in the format "owner/repo" (e.g., "mlflow/mlflow")',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pr",
|
||||
type=int,
|
||||
required=True,
|
||||
help="Pull request number",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
owner, repo = args.repo.split("/")
|
||||
return owner, repo, args.pr
|
||||
|
||||
|
||||
def main() -> None:
|
||||
owner, repo, pull_number = parse_args()
|
||||
diff_text = get_pr_diff(owner, repo, pull_number)
|
||||
if files := parse_diff(diff_text):
|
||||
pr_labels = get_pr_labels(owner, repo, pull_number)
|
||||
has_bypass_label = BYPASS_LABEL in pr_labels
|
||||
|
||||
level = "warning" if has_bypass_label else "error"
|
||||
message = (
|
||||
f"This file only has whitespace changes (bypassed with '{BYPASS_LABEL}' label)."
|
||||
if has_bypass_label
|
||||
else (
|
||||
f"This file only has whitespace changes. "
|
||||
f"Please revert them or apply the '{BYPASS_LABEL}' label to bypass this check "
|
||||
f"if they are necessary."
|
||||
)
|
||||
)
|
||||
|
||||
for file_path in files:
|
||||
# https://docs.github.com/en/actions/reference/workflow-commands-for-github-actions
|
||||
print(f"::{level} file={file_path},line=1,col=1::{message}")
|
||||
|
||||
if not has_bypass_label:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,86 @@
|
||||
# Clint
|
||||
|
||||
A custom linter for mlflow to enforce rules that ruff doesn't cover.
|
||||
|
||||
## Installation
|
||||
|
||||
```
|
||||
pip install -e dev/clint
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
clint file.py ...
|
||||
```
|
||||
|
||||
## Integrating with Visual Studio Code
|
||||
|
||||
1. Install [the Pylint extension](https://marketplace.visualstudio.com/items?itemName=ms-python.pylint)
|
||||
2. Add the following setting in your `settings.json` file:
|
||||
|
||||
```json
|
||||
{
|
||||
"pylint.path": ["${interpreter}", "-m", "clint"]
|
||||
}
|
||||
```
|
||||
|
||||
## Ignoring Rules for Specific Files or Lines
|
||||
|
||||
**To ignore a rule on a specific line (recommended):**
|
||||
|
||||
```python
|
||||
foo() # clint: disable=<rule_name>
|
||||
```
|
||||
|
||||
Replace `<rule_name>` with the actual rule you want to disable.
|
||||
|
||||
To disable multiple rules on the same line, use comma-separated rule names:
|
||||
|
||||
```python
|
||||
foo() # clint: disable=rule-a,rule-b
|
||||
```
|
||||
|
||||
The rule name is shown in the error message. For example:
|
||||
|
||||
```
|
||||
test_file.py:4:2: pytest-mark-repeat: @pytest.mark.repeat decorator...
|
||||
```
|
||||
|
||||
Use the rule name (`pytest-mark-repeat`) in the disable comment:
|
||||
|
||||
```python
|
||||
@pytest.mark.repeat(3) # clint: disable=pytest-mark-repeat
|
||||
def test_something():
|
||||
pass
|
||||
```
|
||||
|
||||
**For multi-line constructs (docstrings, etc.), place the disable comment on the closing line:**
|
||||
|
||||
```python
|
||||
def func():
|
||||
"""
|
||||
Docstring with [markdown link](url).
|
||||
""" # clint: disable=markdown-link
|
||||
pass
|
||||
```
|
||||
|
||||
This works because the linter checks both the start and end lines of the violation range.
|
||||
|
||||
**To ignore a rule for an entire file:**
|
||||
|
||||
Add the file path to the `exclude` list in your `pyproject.toml`:
|
||||
|
||||
```toml
|
||||
[tool.clint]
|
||||
exclude = [
|
||||
# ...existing entries...
|
||||
"path/to/file.py",
|
||||
]
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
pytest dev/clint
|
||||
```
|
||||
@@ -0,0 +1,11 @@
|
||||
[project]
|
||||
name = "clint"
|
||||
version = "0.1.0"
|
||||
description = "A custom linter for mlflow to enforce rules that ruff doesn't cover."
|
||||
readme = "README.md"
|
||||
authors = [{ name = "mlflow", email = "mlflow@mlflow.com" }]
|
||||
dependencies = ["tomli", "typing_extensions"]
|
||||
|
||||
|
||||
[project.scripts]
|
||||
clint = "clint:main"
|
||||
@@ -0,0 +1,110 @@
|
||||
import argparse
|
||||
import itertools
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Violation, lint_file
|
||||
from clint.utils import get_repo_root, resolve_paths
|
||||
|
||||
_WORKER_INDEX: SymbolIndex | None = None
|
||||
_WORKER_CONFIG: Config | None = None
|
||||
|
||||
|
||||
def _init_worker(index_path: Path, config: Config) -> None:
|
||||
global _WORKER_INDEX, _WORKER_CONFIG
|
||||
_WORKER_INDEX = SymbolIndex.load(index_path)
|
||||
_WORKER_CONFIG = config
|
||||
|
||||
|
||||
def _worker_lint(path: Path, code: str) -> list[Violation]:
|
||||
if _WORKER_INDEX is None or _WORKER_CONFIG is None:
|
||||
raise RuntimeError(
|
||||
"Worker not initialized; _init_worker must be called before _worker_lint"
|
||||
)
|
||||
return lint_file(path, code, _WORKER_CONFIG, _WORKER_INDEX)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Args:
|
||||
files: list[str]
|
||||
output_format: Literal["text", "json"]
|
||||
|
||||
@classmethod
|
||||
def parse(cls) -> Self:
|
||||
parser = argparse.ArgumentParser(description="Custom linter for mlflow.")
|
||||
parser.add_argument(
|
||||
"files",
|
||||
nargs="*",
|
||||
help="Files to lint. If not specified, lints all files in the current directory.",
|
||||
)
|
||||
parser.add_argument("--output-format", default="text")
|
||||
args, _ = parser.parse_known_args()
|
||||
return cls(files=args.files, output_format=args.output_format)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
config = Config.load()
|
||||
args = Args.parse()
|
||||
|
||||
input_paths = [Path(f) for f in args.files]
|
||||
|
||||
resolved_files = resolve_paths(input_paths)
|
||||
|
||||
# Apply exclude filtering
|
||||
files: list[Path] = []
|
||||
if config.exclude:
|
||||
repo_root = get_repo_root()
|
||||
cwd = Path.cwd()
|
||||
regex = re.compile("|".join(map(re.escape, config.exclude)))
|
||||
for f in resolved_files:
|
||||
# Convert file path to be relative to repo root for exclude pattern matching
|
||||
repo_relative_path = (cwd / f).resolve().relative_to(repo_root)
|
||||
if not regex.match(repo_relative_path.as_posix()):
|
||||
files.append(f)
|
||||
else:
|
||||
files = resolved_files
|
||||
|
||||
# Exit early if no files to lint
|
||||
if not files:
|
||||
return
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
# Pickle `SymbolIndex` to avoid expensive serialization overhead when passing
|
||||
# the large index object to multiple worker processes
|
||||
index_path = Path(tmp_dir) / "symbol_index.pkl"
|
||||
SymbolIndex.build().save(index_path)
|
||||
with ProcessPoolExecutor(initializer=_init_worker, initargs=(index_path, config)) as pool:
|
||||
futures = [pool.submit(_worker_lint, f, f.read_text()) for f in files]
|
||||
violations_iter = itertools.chain.from_iterable(
|
||||
f.result() for f in as_completed(futures)
|
||||
)
|
||||
if violations := list(violations_iter):
|
||||
if args.output_format == "json":
|
||||
sys.stdout.write(json.dumps([v.json() for v in violations]))
|
||||
elif args.output_format == "text":
|
||||
sys.stderr.write("\n".join(map(str, violations)) + "\n")
|
||||
count = len(violations)
|
||||
label = "error" if count == 1 else "errors"
|
||||
rule_label = "this rule" if count == 1 else "these rules"
|
||||
print(
|
||||
f"Found {count} {label}\n"
|
||||
f"See dev/clint/README.md for instructions on ignoring {rule_label}.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("No errors found!", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,3 @@
|
||||
from clint import main
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,216 @@
|
||||
# https://github.com/PyCQA/isort/blob/b818cec889657cb786beafe94a6641f8fc0f0e64/isort/stdlibs/py311.py
|
||||
BUILTIN_MODULES = {
|
||||
"_ast",
|
||||
"_thread",
|
||||
"abc",
|
||||
"aifc",
|
||||
"argparse",
|
||||
"array",
|
||||
"ast",
|
||||
"asynchat",
|
||||
"asyncio",
|
||||
"asyncore",
|
||||
"atexit",
|
||||
"audioop",
|
||||
"base64",
|
||||
"bdb",
|
||||
"binascii",
|
||||
"bisect",
|
||||
"builtins",
|
||||
"bz2",
|
||||
"cProfile",
|
||||
"calendar",
|
||||
"cgi",
|
||||
"cgitb",
|
||||
"chunk",
|
||||
"cmath",
|
||||
"cmd",
|
||||
"code",
|
||||
"codecs",
|
||||
"codeop",
|
||||
"collections",
|
||||
"colorsys",
|
||||
"compileall",
|
||||
"concurrent",
|
||||
"configparser",
|
||||
"contextlib",
|
||||
"contextvars",
|
||||
"copy",
|
||||
"copyreg",
|
||||
"crypt",
|
||||
"csv",
|
||||
"ctypes",
|
||||
"curses",
|
||||
"dataclasses",
|
||||
"datetime",
|
||||
"dbm",
|
||||
"decimal",
|
||||
"difflib",
|
||||
"dis",
|
||||
"distutils",
|
||||
"doctest",
|
||||
"email",
|
||||
"encodings",
|
||||
"ensurepip",
|
||||
"enum",
|
||||
"errno",
|
||||
"faulthandler",
|
||||
"fcntl",
|
||||
"filecmp",
|
||||
"fileinput",
|
||||
"fnmatch",
|
||||
"fractions",
|
||||
"ftplib",
|
||||
"functools",
|
||||
"gc",
|
||||
"getopt",
|
||||
"getpass",
|
||||
"gettext",
|
||||
"glob",
|
||||
"graphlib",
|
||||
"grp",
|
||||
"gzip",
|
||||
"hashlib",
|
||||
"heapq",
|
||||
"hmac",
|
||||
"html",
|
||||
"http",
|
||||
"idlelib",
|
||||
"imaplib",
|
||||
"imghdr",
|
||||
"imp",
|
||||
"importlib",
|
||||
"inspect",
|
||||
"io",
|
||||
"ipaddress",
|
||||
"itertools",
|
||||
"json",
|
||||
"keyword",
|
||||
"lib2to3",
|
||||
"linecache",
|
||||
"locale",
|
||||
"logging",
|
||||
"lzma",
|
||||
"mailbox",
|
||||
"mailcap",
|
||||
"marshal",
|
||||
"math",
|
||||
"mimetypes",
|
||||
"mmap",
|
||||
"modulefinder",
|
||||
"msilib",
|
||||
"msvcrt",
|
||||
"multiprocessing",
|
||||
"netrc",
|
||||
"nis",
|
||||
"nntplib",
|
||||
"ntpath",
|
||||
"numbers",
|
||||
"operator",
|
||||
"optparse",
|
||||
"os",
|
||||
"ossaudiodev",
|
||||
"pathlib",
|
||||
"pdb",
|
||||
"pickle",
|
||||
"pickletools",
|
||||
"pipes",
|
||||
"pkgutil",
|
||||
"platform",
|
||||
"plistlib",
|
||||
"poplib",
|
||||
"posix",
|
||||
"posixpath",
|
||||
"pprint",
|
||||
"profile",
|
||||
"pstats",
|
||||
"pty",
|
||||
"pwd",
|
||||
"py_compile",
|
||||
"pyclbr",
|
||||
"pydoc",
|
||||
"queue",
|
||||
"quopri",
|
||||
"random",
|
||||
"re",
|
||||
"readline",
|
||||
"reprlib",
|
||||
"resource",
|
||||
"rlcompleter",
|
||||
"runpy",
|
||||
"sched",
|
||||
"secrets",
|
||||
"select",
|
||||
"selectors",
|
||||
"shelve",
|
||||
"shlex",
|
||||
"shutil",
|
||||
"signal",
|
||||
"site",
|
||||
"smtpd",
|
||||
"smtplib",
|
||||
"sndhdr",
|
||||
"socket",
|
||||
"socketserver",
|
||||
"spwd",
|
||||
"sqlite3",
|
||||
"sre",
|
||||
"sre_compile",
|
||||
"sre_constants",
|
||||
"sre_parse",
|
||||
"ssl",
|
||||
"stat",
|
||||
"statistics",
|
||||
"string",
|
||||
"stringprep",
|
||||
"struct",
|
||||
"subprocess",
|
||||
"sunau",
|
||||
"symtable",
|
||||
"sys",
|
||||
"sysconfig",
|
||||
"syslog",
|
||||
"tabnanny",
|
||||
"tarfile",
|
||||
"telnetlib",
|
||||
"tempfile",
|
||||
"termios",
|
||||
"test",
|
||||
"textwrap",
|
||||
"threading",
|
||||
"time",
|
||||
"timeit",
|
||||
"tkinter",
|
||||
"token",
|
||||
"tokenize",
|
||||
"tomllib",
|
||||
"trace",
|
||||
"traceback",
|
||||
"tracemalloc",
|
||||
"tty",
|
||||
"turtle",
|
||||
"turtledemo",
|
||||
"types",
|
||||
"typing",
|
||||
"unicodedata",
|
||||
"unittest",
|
||||
"urllib",
|
||||
"uu",
|
||||
"uuid",
|
||||
"venv",
|
||||
"warnings",
|
||||
"wave",
|
||||
"weakref",
|
||||
"webbrowser",
|
||||
"winreg",
|
||||
"winsound",
|
||||
"wsgiref",
|
||||
"xdrlib",
|
||||
"xml",
|
||||
"xmlrpc",
|
||||
"zipapp",
|
||||
"zipfile",
|
||||
"zipimport",
|
||||
"zlib",
|
||||
"zoneinfo",
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
import io
|
||||
import re
|
||||
import tokenize
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Iterator
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from clint.linter import Position
|
||||
|
||||
NOQA_REGEX = re.compile(r"#\s*noqa\s*:\s*([A-Z]\d+(?:\s*,\s*[A-Z]\d+)*)", re.IGNORECASE)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Noqa:
|
||||
start: "Position"
|
||||
end: "Position"
|
||||
rules: set[str]
|
||||
|
||||
@classmethod
|
||||
def from_token(cls, token: tokenize.TokenInfo) -> Self | None:
|
||||
from clint.linter import Position
|
||||
|
||||
if match := NOQA_REGEX.match(token.string):
|
||||
rules = {r.strip() for r in match.group(1).upper().split(",")}
|
||||
start = Position(token.start[0], token.start[1])
|
||||
end = Position(token.end[0], token.end[1])
|
||||
return cls(start=start, end=end, rules=rules)
|
||||
return None
|
||||
|
||||
|
||||
def iter_comments(code: str) -> Iterator[tokenize.TokenInfo]:
|
||||
readline = io.StringIO(code).readline
|
||||
try:
|
||||
for token in tokenize.generate_tokens(readline):
|
||||
if token.type == tokenize.COMMENT:
|
||||
yield token
|
||||
except tokenize.TokenError:
|
||||
# Handle incomplete tokens at end of file
|
||||
pass
|
||||
@@ -0,0 +1,110 @@
|
||||
import re
|
||||
import typing
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import tomli
|
||||
from typing_extensions import Self
|
||||
|
||||
from clint.rules import ALL_RULES
|
||||
from clint.utils import get_repo_root
|
||||
|
||||
|
||||
def _validate_exclude_paths(exclude_paths: list[str]) -> None:
|
||||
"""Validate that all paths in the exclude list exist.
|
||||
|
||||
Args:
|
||||
exclude_paths: List of file/directory paths to validate (relative to repo root)
|
||||
|
||||
Raises:
|
||||
ValueError: If any path in the exclude list does not exist
|
||||
"""
|
||||
if not exclude_paths:
|
||||
return
|
||||
|
||||
repo_root = get_repo_root()
|
||||
if non_existing_paths := [path for path in exclude_paths if not (repo_root / path).exists()]:
|
||||
raise ValueError(
|
||||
f"Non-existing paths found in exclude field: {non_existing_paths}. "
|
||||
f"All paths in the exclude list must exist."
|
||||
)
|
||||
|
||||
|
||||
def _validate_typing_extensions_allowlist(allowlist: list[str]) -> None:
|
||||
"""Validate that the typing-extensions-allowlist doesn't contain stdlib items.
|
||||
|
||||
Args:
|
||||
allowlist: List of typing_extensions items to validate
|
||||
|
||||
Raises:
|
||||
ValueError: If any item in the allowlist is available in stdlib typing
|
||||
"""
|
||||
if not allowlist:
|
||||
return
|
||||
|
||||
# Extract the item name from full paths like "typing_extensions.overload"
|
||||
# and check if it's available in stdlib typing
|
||||
stdlib_items = []
|
||||
for item in allowlist:
|
||||
name = item.rsplit(".", 1)[-1]
|
||||
if hasattr(typing, name):
|
||||
stdlib_items.append(item)
|
||||
|
||||
if stdlib_items:
|
||||
raise ValueError(
|
||||
f"Items in typing-extensions-allowlist are available in stdlib `typing`: "
|
||||
f"{stdlib_items}. Use `from typing import ...` instead."
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
select: set[str] = field(default_factory=set)
|
||||
exclude: list[str] = field(default_factory=list)
|
||||
# Path -> List of modules that should not be imported globally under that path
|
||||
forbidden_top_level_imports: dict[str, list[str]] = field(default_factory=dict)
|
||||
typing_extensions_allowlist: list[str] = field(default_factory=list)
|
||||
example_rules: list[str] = field(default_factory=list)
|
||||
# Compiled regex pattern -> Set of rule names to ignore for files matching the pattern
|
||||
per_file_ignores: dict[re.Pattern[str], set[str]] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def load(cls) -> Self:
|
||||
repo_root = get_repo_root()
|
||||
pyproject = repo_root / "pyproject.toml"
|
||||
if not pyproject.exists():
|
||||
return cls()
|
||||
|
||||
with pyproject.open("rb") as f:
|
||||
data = tomli.load(f)
|
||||
|
||||
clint = data.get("tool", {}).get("clint", {})
|
||||
if not clint:
|
||||
return cls()
|
||||
|
||||
per_file_ignores_raw = clint.get("per-file-ignores", {})
|
||||
per_file_ignores: dict[re.Pattern[str], set[str]] = {}
|
||||
for pattern, rules in per_file_ignores_raw.items():
|
||||
per_file_ignores[re.compile(pattern)] = set(rules)
|
||||
|
||||
select = clint.get("select")
|
||||
if select is None:
|
||||
select = ALL_RULES
|
||||
else:
|
||||
if unknown_rules := set(select) - ALL_RULES:
|
||||
raise ValueError(f"Unknown rules in 'select': {unknown_rules}")
|
||||
select = set(select)
|
||||
|
||||
exclude_paths = clint.get("exclude", [])
|
||||
_validate_exclude_paths(exclude_paths)
|
||||
|
||||
typing_extensions_allowlist = clint.get("typing-extensions-allowlist", [])
|
||||
_validate_typing_extensions_allowlist(typing_extensions_allowlist)
|
||||
|
||||
return cls(
|
||||
select=select,
|
||||
exclude=exclude_paths,
|
||||
forbidden_top_level_imports=clint.get("forbidden-top-level-imports", {}),
|
||||
typing_extensions_allowlist=typing_extensions_allowlist,
|
||||
example_rules=clint.get("example-rules", []),
|
||||
per_file_ignores=per_file_ignores,
|
||||
)
|
||||
@@ -0,0 +1,221 @@
|
||||
"""Symbol indexing for MLflow codebase.
|
||||
|
||||
This module provides efficient indexing and lookup of Python symbols (functions, classes)
|
||||
across the MLflow codebase using AST parsing and parallel processing.
|
||||
|
||||
Key components:
|
||||
- FunctionInfo: Lightweight function signature information
|
||||
- ModuleSymbolExtractor: AST visitor for extracting symbols from modules
|
||||
- SymbolIndex: Main index class for symbol resolution and lookup
|
||||
|
||||
Example usage:
|
||||
|
||||
```python
|
||||
# Build an index of all MLflow symbols
|
||||
index = SymbolIndex.build()
|
||||
|
||||
# Look up function signature information
|
||||
func_info = index.resolve("mlflow.log_metric")
|
||||
print(f"Arguments: {func_info.args}") # -> ['key, 'value', 'step', ...]
|
||||
```
|
||||
"""
|
||||
|
||||
import ast
|
||||
import multiprocessing
|
||||
import pickle
|
||||
import subprocess
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from clint.utils import get_repo_root
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionInfo:
|
||||
"""Lightweight function signature information for efficient serialization."""
|
||||
|
||||
has_vararg: bool # *args
|
||||
has_kwarg: bool # **kwargs
|
||||
args: list[str] = field(default_factory=list) # Regular arguments
|
||||
kwonlyargs: list[str] = field(default_factory=list) # Keyword-only arguments
|
||||
posonlyargs: list[str] = field(default_factory=list) # Positional-only arguments
|
||||
|
||||
@classmethod
|
||||
def from_func_def(
|
||||
cls, node: ast.FunctionDef | ast.AsyncFunctionDef, skip_self: bool = False
|
||||
) -> Self:
|
||||
"""Create FunctionInfo from an AST function definition node."""
|
||||
args = node.args.args
|
||||
if skip_self and args:
|
||||
args = args[1:] # Skip 'self' for methods
|
||||
|
||||
return cls(
|
||||
has_vararg=node.args.vararg is not None,
|
||||
has_kwarg=node.args.kwarg is not None,
|
||||
args=[arg.arg for arg in args],
|
||||
kwonlyargs=[arg.arg for arg in node.args.kwonlyargs],
|
||||
posonlyargs=[arg.arg for arg in node.args.posonlyargs],
|
||||
)
|
||||
|
||||
@property
|
||||
def all_args(self) -> list[str]:
|
||||
return self.posonlyargs + self.args + self.kwonlyargs
|
||||
|
||||
|
||||
class ModuleSymbolExtractor(ast.NodeVisitor):
|
||||
"""Extracts function definitions and import mappings from a Python module."""
|
||||
|
||||
def __init__(self, mod: str) -> None:
|
||||
self.mod = mod
|
||||
self.import_mapping: dict[str, str] = {}
|
||||
self.func_mapping: dict[str, FunctionInfo] = {}
|
||||
|
||||
def visit_Import(self, node: ast.Import) -> None:
|
||||
for alias in node.names:
|
||||
if not alias.name.startswith("mlflow."):
|
||||
continue
|
||||
if alias.asname:
|
||||
self.import_mapping[f"{self.mod}.{alias.asname}"] = alias.name
|
||||
|
||||
def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
|
||||
if node.module is None or not node.module.startswith("mlflow."):
|
||||
return
|
||||
for alias in node.names:
|
||||
if alias.name.startswith("_"):
|
||||
continue
|
||||
if alias.asname:
|
||||
self.import_mapping[f"{self.mod}.{alias.asname}"] = f"{node.module}.{alias.name}"
|
||||
else:
|
||||
self.import_mapping[f"{self.mod}.{alias.name}"] = f"{node.module}.{alias.name}"
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
|
||||
if node.name.startswith("_"):
|
||||
return
|
||||
self.func_mapping[f"{self.mod}.{node.name}"] = FunctionInfo.from_func_def(node)
|
||||
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
|
||||
if node.name.startswith("_"):
|
||||
return
|
||||
self.func_mapping[f"{self.mod}.{node.name}"] = FunctionInfo.from_func_def(node)
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef) -> None:
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, ast.FunctionDef):
|
||||
if stmt.name == "__init__":
|
||||
info = FunctionInfo.from_func_def(stmt, skip_self=True)
|
||||
self.func_mapping[f"{self.mod}.{node.name}"] = info
|
||||
elif any(
|
||||
isinstance(deco, ast.Name) and deco.id in ("classmethod", "staticmethod")
|
||||
for deco in stmt.decorator_list
|
||||
):
|
||||
info = FunctionInfo.from_func_def(stmt, skip_self=True)
|
||||
self.func_mapping[f"{self.mod}.{node.name}.{stmt.name}"] = info
|
||||
else:
|
||||
# If no __init__ found, still add the class with *args and **kwargs
|
||||
self.func_mapping[f"{self.mod}.{node.name}"] = FunctionInfo(
|
||||
has_vararg=True, has_kwarg=True
|
||||
)
|
||||
|
||||
|
||||
def extract_symbols_from_file(
|
||||
rel_path: str, content: str
|
||||
) -> tuple[dict[str, str], dict[str, FunctionInfo]] | None:
|
||||
"""Extract function definitions and import mappings from a Python file."""
|
||||
p = Path(rel_path)
|
||||
if not p.parts or p.parts[0] != "mlflow":
|
||||
return None
|
||||
|
||||
try:
|
||||
tree = ast.parse(content)
|
||||
except (SyntaxError, UnicodeDecodeError):
|
||||
return None
|
||||
|
||||
mod_name = (
|
||||
".".join(p.parts[:-1]) if p.name == "__init__.py" else ".".join([*p.parts[:-1], p.stem])
|
||||
)
|
||||
|
||||
extractor = ModuleSymbolExtractor(mod_name)
|
||||
extractor.visit(tree)
|
||||
return extractor.import_mapping, extractor.func_mapping
|
||||
|
||||
|
||||
class SymbolIndex:
|
||||
"""Index of all symbols (functions, classes) in the MLflow codebase."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
import_mapping: dict[str, str],
|
||||
func_mapping: dict[str, FunctionInfo],
|
||||
) -> None:
|
||||
self.import_mapping = import_mapping
|
||||
self.func_mapping = func_mapping
|
||||
|
||||
def save(self, path: Path) -> None:
|
||||
with path.open("wb") as f:
|
||||
pickle.dump((self.import_mapping, self.func_mapping), f)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: Path) -> Self:
|
||||
with path.open("rb") as f:
|
||||
import_mapping, func_mapping = pickle.load(f)
|
||||
return cls(import_mapping, func_mapping)
|
||||
|
||||
@classmethod
|
||||
def build(cls) -> Self:
|
||||
repo_root = get_repo_root()
|
||||
py_files = subprocess.check_output(
|
||||
["git", "-C", repo_root, "ls-files", "mlflow/*.py"], text=True
|
||||
).splitlines()
|
||||
|
||||
mapping: dict[str, str] = {}
|
||||
func_mapping: dict[str, FunctionInfo] = {}
|
||||
|
||||
# Ensure at least 1 worker to avoid ProcessPoolExecutor ValueError
|
||||
max_workers = max(1, min(multiprocessing.cpu_count(), len(py_files)))
|
||||
with ProcessPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = {}
|
||||
for py_file in py_files:
|
||||
abs_file_path = repo_root / py_file
|
||||
if not abs_file_path.exists():
|
||||
continue
|
||||
content = abs_file_path.read_text()
|
||||
f = executor.submit(extract_symbols_from_file, py_file, content)
|
||||
futures[f] = py_file
|
||||
|
||||
for future in as_completed(futures):
|
||||
if result := future.result():
|
||||
file_imports, file_functions = result
|
||||
mapping.update(file_imports)
|
||||
func_mapping.update(file_functions)
|
||||
|
||||
return cls(mapping, func_mapping)
|
||||
|
||||
def _resolve_import(self, target: str) -> str:
|
||||
resolved = target
|
||||
seen = {resolved}
|
||||
while v := self.import_mapping.get(resolved):
|
||||
if v in seen:
|
||||
# Circular import detected, break to avoid infinite loop
|
||||
break
|
||||
seen.add(v)
|
||||
resolved = v
|
||||
return resolved
|
||||
|
||||
def resolve(self, target: str) -> FunctionInfo | None:
|
||||
"""Resolve a symbol to its actual definition, following import chains."""
|
||||
if f := self.func_mapping.get(target):
|
||||
return f
|
||||
|
||||
resolved = self._resolve_import(target)
|
||||
if f := self.func_mapping.get(resolved):
|
||||
return f
|
||||
|
||||
target, tail = target.rsplit(".", 1)
|
||||
resolved = self._resolve_import(target)
|
||||
if f := self.func_mapping.get(f"{resolved}.{tail}"):
|
||||
return f
|
||||
|
||||
return None
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,89 @@
|
||||
import ast
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
class Resolver:
|
||||
def __init__(self) -> None:
|
||||
self.name_map: dict[str, list[str]] = {}
|
||||
self._scope_stack: list[dict[str, list[str]]] = []
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all name mappings. Useful when starting to process a new file."""
|
||||
self.name_map.clear()
|
||||
self._scope_stack.clear()
|
||||
|
||||
def enter_scope(self) -> None:
|
||||
"""Enter a new scope by taking a snapshot of current mappings."""
|
||||
self._scope_stack.append(self.name_map.copy())
|
||||
|
||||
def exit_scope(self) -> None:
|
||||
"""Exit current scope by restoring the previous snapshot."""
|
||||
if self._scope_stack:
|
||||
self.name_map = self._scope_stack.pop()
|
||||
|
||||
@contextmanager
|
||||
def scope(self) -> Iterator[None]:
|
||||
"""Context manager for automatic scope management."""
|
||||
self.enter_scope()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.exit_scope()
|
||||
|
||||
def add_import(self, node: ast.Import) -> None:
|
||||
for alias in node.names:
|
||||
if alias.asname:
|
||||
self.name_map[alias.asname] = alias.name.split(".")
|
||||
else:
|
||||
toplevel = alias.name.split(".", 1)[0]
|
||||
self.name_map[toplevel] = [toplevel]
|
||||
|
||||
def add_import_from(self, node: ast.ImportFrom) -> None:
|
||||
if node.module is None:
|
||||
return
|
||||
|
||||
for alias in node.names:
|
||||
name = alias.asname or alias.name
|
||||
module_parts = node.module.split(".")
|
||||
self.name_map[name] = module_parts + [alias.name]
|
||||
|
||||
def resolve(self, node: ast.expr) -> list[str] | None:
|
||||
"""
|
||||
Resolve a node to its fully qualified name parts.
|
||||
|
||||
Args:
|
||||
node: AST node to resolve, typically a Call, Name, or Attribute.
|
||||
|
||||
Returns:
|
||||
List of name parts (e.g., ["threading", "Thread"]) or None if unresolvable
|
||||
"""
|
||||
if isinstance(node, ast.Call):
|
||||
parts = self._extract_call_parts(node.func)
|
||||
elif isinstance(node, ast.Name):
|
||||
parts = [node.id]
|
||||
elif isinstance(node, ast.Attribute):
|
||||
parts = self._extract_call_parts(node)
|
||||
else:
|
||||
return None
|
||||
|
||||
return self._resolve_parts(parts) if parts else None
|
||||
|
||||
def _extract_call_parts(self, node: ast.expr) -> list[str]:
|
||||
if isinstance(node, ast.Name):
|
||||
return [node.id]
|
||||
elif isinstance(node, ast.Attribute) and (
|
||||
base_parts := self._extract_call_parts(node.value)
|
||||
):
|
||||
return base_parts + [node.attr]
|
||||
return []
|
||||
|
||||
def _resolve_parts(self, parts: list[str]) -> list[str] | None:
|
||||
if not parts:
|
||||
return None
|
||||
|
||||
# Check if the first part is in our name mapping
|
||||
if resolved_base := self.name_map.get(parts[0]):
|
||||
return resolved_base + parts[1:]
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,122 @@
|
||||
from clint.rules.assign_before_append import AssignBeforeAppend
|
||||
from clint.rules.base import Rule
|
||||
from clint.rules.do_not_disable import DoNotDisable
|
||||
from clint.rules.docstring_param_order import DocstringParamOrder
|
||||
from clint.rules.empty_notebook_cell import EmptyNotebookCell
|
||||
from clint.rules.example_syntax_error import ExampleSyntaxError
|
||||
from clint.rules.except_bool_op import ExceptBoolOp
|
||||
from clint.rules.extraneous_docstring_param import ExtraneousDocstringParam
|
||||
from clint.rules.forbidden_deprecation_warning import ForbiddenDeprecationWarning
|
||||
from clint.rules.forbidden_make_judge_in_builtin_scorers import (
|
||||
ForbiddenMakeJudgeInBuiltinScorers,
|
||||
)
|
||||
from clint.rules.forbidden_set_active_model_usage import ForbiddenSetActiveModelUsage
|
||||
from clint.rules.forbidden_top_level_import import ForbiddenTopLevelImport
|
||||
from clint.rules.forbidden_trace_ui_in_notebook import ForbiddenTraceUIInNotebook
|
||||
from clint.rules.get_artifact_uri import GetArtifactUri
|
||||
from clint.rules.implicit_optional import ImplicitOptional
|
||||
from clint.rules.incorrect_type_annotation import IncorrectTypeAnnotation
|
||||
from clint.rules.invalid_abstract_method import InvalidAbstractMethod
|
||||
from clint.rules.invalid_experimental_decorator import InvalidExperimentalDecorator
|
||||
from clint.rules.isinstance_union_syntax import IsinstanceUnionSyntax
|
||||
from clint.rules.lazy_import import LazyImport
|
||||
from clint.rules.lazy_module import LazyModule
|
||||
from clint.rules.log_model_artifact_path import LogModelArtifactPath
|
||||
from clint.rules.markdown_link import MarkdownLink
|
||||
from clint.rules.missing_docstring_param import MissingDocstringParam
|
||||
from clint.rules.missing_notebook_h1_header import MissingNotebookH1Header
|
||||
from clint.rules.mlflow_class_name import MlflowClassName
|
||||
from clint.rules.mock_patch_as_decorator import MockPatchAsDecorator
|
||||
from clint.rules.mock_patch_dict_environ import MockPatchDictEnviron
|
||||
from clint.rules.multi_assign import MultiAssign
|
||||
from clint.rules.nested_mock_patch import NestedMockPatch
|
||||
from clint.rules.no_class_based_tests import NoClassBasedTests
|
||||
from clint.rules.no_rst import NoRst
|
||||
from clint.rules.no_shebang import NoShebang
|
||||
from clint.rules.os_chdir_in_test import OsChdirInTest
|
||||
from clint.rules.os_environ_delete_in_test import OsEnvironDeleteInTest
|
||||
from clint.rules.os_environ_set_in_test import OsEnvironSetInTest
|
||||
from clint.rules.prefer_dict_union import PreferDictUnion
|
||||
from clint.rules.prefer_next import PreferNext
|
||||
from clint.rules.prefer_os_environ import PreferOsEnviron
|
||||
from clint.rules.pytest_mark_repeat import PytestMarkRepeat
|
||||
from clint.rules.redundant_mock_return_value import RedundantMockReturnValue
|
||||
from clint.rules.redundant_test_docstring import RedundantTestDocstring
|
||||
from clint.rules.subprocess_check_call import SubprocessCheckCall
|
||||
from clint.rules.tempfile_in_test import TempfileInTest
|
||||
from clint.rules.test_name_typo import TestNameTypo
|
||||
from clint.rules.typing_extensions import TypingExtensions
|
||||
from clint.rules.unknown_mlflow_arguments import UnknownMlflowArguments
|
||||
from clint.rules.unknown_mlflow_function import UnknownMlflowFunction
|
||||
from clint.rules.unnamed_thread import UnnamedThread
|
||||
from clint.rules.unnamed_thread_pool import UnnamedThreadPool
|
||||
from clint.rules.unparameterized_generic_type import UnparameterizedGenericType
|
||||
from clint.rules.unused_disable_comment import UnusedDisableComment
|
||||
from clint.rules.use_gh_token import UseGhToken
|
||||
from clint.rules.use_sys_executable import UseSysExecutable
|
||||
from clint.rules.use_walrus_operator import UseWalrusOperator, WalrusOperatorVisitor
|
||||
from clint.rules.version_major_check import MajorVersionCheck
|
||||
|
||||
ALL_RULES = {rule.name for rule in Rule.__subclasses__()}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ALL_RULES",
|
||||
"Rule",
|
||||
"DoNotDisable",
|
||||
"DocstringParamOrder",
|
||||
"EmptyNotebookCell",
|
||||
"ExampleSyntaxError",
|
||||
"ExceptBoolOp",
|
||||
"ExtraneousDocstringParam",
|
||||
"ForbiddenDeprecationWarning",
|
||||
"ForbiddenMakeJudgeInBuiltinScorers",
|
||||
"ForbiddenSetActiveModelUsage",
|
||||
"ForbiddenTopLevelImport",
|
||||
"GetArtifactUri",
|
||||
"ForbiddenTraceUIInNotebook",
|
||||
"ImplicitOptional",
|
||||
"IncorrectTypeAnnotation",
|
||||
"IsinstanceUnionSyntax",
|
||||
"InvalidAbstractMethod",
|
||||
"InvalidExperimentalDecorator",
|
||||
"LazyImport",
|
||||
"LazyModule",
|
||||
"LogModelArtifactPath",
|
||||
"MarkdownLink",
|
||||
"MissingDocstringParam",
|
||||
"MissingNotebookH1Header",
|
||||
"MlflowClassName",
|
||||
"MockPatchDictEnviron",
|
||||
"MockPatchAsDecorator",
|
||||
"NestedMockPatch",
|
||||
"NoClassBasedTests",
|
||||
"NoRst",
|
||||
"NoShebang",
|
||||
"OsChdirInTest",
|
||||
"OsEnvironDeleteInTest",
|
||||
"OsEnvironSetInTest",
|
||||
"PreferDictUnion",
|
||||
"PreferNext",
|
||||
"PreferOsEnviron",
|
||||
"PytestMarkRepeat",
|
||||
"RedundantMockReturnValue",
|
||||
"RedundantTestDocstring",
|
||||
"SubprocessCheckCall",
|
||||
"TempfileInTest",
|
||||
"TestNameTypo",
|
||||
"UnnamedThreadPool",
|
||||
"TypingExtensions",
|
||||
"UnknownMlflowArguments",
|
||||
"UnknownMlflowFunction",
|
||||
"MultiAssign",
|
||||
"UnnamedThread",
|
||||
"UnparameterizedGenericType",
|
||||
"UnusedDisableComment",
|
||||
"AssignBeforeAppend",
|
||||
"UseGhToken",
|
||||
"UseSysExecutable",
|
||||
"UseWalrusOperator",
|
||||
"WalrusOperatorVisitor",
|
||||
"MajorVersionCheck",
|
||||
]
|
||||
@@ -0,0 +1,64 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class AssignBeforeAppend(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Avoid unnecessary assignment before appending to a list. "
|
||||
"Use a list comprehension instead."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.For, prev_stmt: ast.stmt | None) -> bool:
|
||||
"""
|
||||
Returns True if the for loop contains exactly two statements:
|
||||
an assignment followed by appending that variable to a list, AND
|
||||
the loop is immediately preceded by an empty list initialization.
|
||||
|
||||
Examples that should be flagged:
|
||||
---
|
||||
items = []
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
items.append(item)
|
||||
---
|
||||
"""
|
||||
# Match: for loop with exactly 2 statements in body
|
||||
match node:
|
||||
case ast.For(body=[stmt1, stmt2]):
|
||||
pass
|
||||
case _:
|
||||
return False
|
||||
|
||||
# Match stmt1: simple assignment (item = x)
|
||||
match stmt1:
|
||||
case ast.Assign(targets=[ast.Name(id=assigned_var)]):
|
||||
pass
|
||||
case _:
|
||||
return False
|
||||
|
||||
# Match stmt2: list.append(item)
|
||||
match stmt2:
|
||||
case ast.Expr(
|
||||
value=ast.Call(
|
||||
func=ast.Attribute(value=ast.Name(id=list_name), attr="append"),
|
||||
args=[ast.Name(id=appended_var)],
|
||||
)
|
||||
):
|
||||
# Check if the appended variable is the same as the assigned variable
|
||||
if appended_var != assigned_var:
|
||||
return False
|
||||
case _:
|
||||
return False
|
||||
|
||||
# Only flag if prev_stmt is empty list initialization for the same list
|
||||
match prev_stmt:
|
||||
case ast.Assign(
|
||||
targets=[ast.Name(id=prev_list_name)],
|
||||
value=ast.List(elts=[]),
|
||||
) if prev_list_name == list_name:
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
@@ -0,0 +1,31 @@
|
||||
import inspect
|
||||
import itertools
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
_id_counter = itertools.count(start=1)
|
||||
_CLASS_NAME_TO_RULE_NAME_REGEX = re.compile(r"(?<!^)(?=[A-Z])")
|
||||
|
||||
|
||||
class Rule(ABC):
|
||||
id: str
|
||||
name: str
|
||||
|
||||
def __init_subclass__(cls, **kwargs: Any) -> None:
|
||||
super().__init_subclass__(**kwargs)
|
||||
# Only generate ID for concrete classes
|
||||
if not inspect.isabstract(cls):
|
||||
id_ = next(_id_counter)
|
||||
cls.id = f"MLF{id_:04d}"
|
||||
cls.name = _CLASS_NAME_TO_RULE_NAME_REGEX.sub("-", cls.__name__).lower()
|
||||
|
||||
@abstractmethod
|
||||
def _message(self) -> str:
|
||||
"""
|
||||
Return a message that explains this rule.
|
||||
"""
|
||||
|
||||
@property
|
||||
def message(self) -> str:
|
||||
return self._message()
|
||||
@@ -0,0 +1,29 @@
|
||||
from typing_extensions import Self
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class DoNotDisable(Rule):
|
||||
RULES = {
|
||||
"B006": "Use None as default and set value in function body instead of mutable defaults",
|
||||
"F821": "Use typing.TYPE_CHECKING for forward references to optional dependencies",
|
||||
}
|
||||
|
||||
def __init__(self, rules: set[str]) -> None:
|
||||
self.rules = rules
|
||||
|
||||
@classmethod
|
||||
def check(cls, rules: set[str]) -> Self | None:
|
||||
if s := rules.intersection(DoNotDisable.RULES.keys()):
|
||||
return cls(s)
|
||||
return None
|
||||
|
||||
def _message(self) -> str:
|
||||
# Build message for all rules (works for single and multiple rules)
|
||||
hints = []
|
||||
for rule in sorted(self.rules):
|
||||
if hint := DoNotDisable.RULES.get(rule):
|
||||
hints.append(f"{rule}: {hint}")
|
||||
else:
|
||||
hints.append(rule)
|
||||
return f"DO NOT DISABLE {', '.join(hints)}"
|
||||
@@ -0,0 +1,9 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class DocstringParamOrder(Rule):
|
||||
def __init__(self, params: list[str]) -> None:
|
||||
self.params = params
|
||||
|
||||
def _message(self) -> str:
|
||||
return f"Unordered parameters in docstring: {self.params}"
|
||||
@@ -0,0 +1,6 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class EmptyNotebookCell(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Empty notebook cell. Remove it or add some content."
|
||||
@@ -0,0 +1,6 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ExampleSyntaxError(Rule):
|
||||
def _message(self) -> str:
|
||||
return "This example has a syntax error."
|
||||
@@ -0,0 +1,14 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ExceptBoolOp(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Did you mean `except (X, Y):`? Using or/and in an except handler is likely a mistake."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.ExceptHandler) -> bool:
|
||||
return isinstance(node.type, ast.BoolOp)
|
||||
@@ -0,0 +1,9 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ExtraneousDocstringParam(Rule):
|
||||
def __init__(self, params: set[str]) -> None:
|
||||
self.params = params
|
||||
|
||||
def _message(self) -> str:
|
||||
return f"Extraneous parameters in docstring: {self.params}"
|
||||
@@ -0,0 +1,34 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
def _is_deprecation_warning(expr: ast.expr) -> bool:
|
||||
return isinstance(expr, ast.Name) and expr.id == "DeprecationWarning"
|
||||
|
||||
|
||||
class ForbiddenDeprecationWarning(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Do not use `DeprecationWarning` with `warnings.warn()`. "
|
||||
"Use `FutureWarning` instead since Python does not show `DeprecationWarning` "
|
||||
"by default."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> ast.expr | None:
|
||||
"""
|
||||
Checks if the given node is a call to `warnings.warn` with `DeprecationWarning`.
|
||||
"""
|
||||
# Check if this is a call to `warnings.warn`
|
||||
if (resolved := resolver.resolve(node.func)) and resolved == ["warnings", "warn"]:
|
||||
# Check if there's a `category` positional argument with `DeprecationWarning`
|
||||
if len(node.args) >= 2 and _is_deprecation_warning(node.args[1]):
|
||||
return node.args[1]
|
||||
# Check if there's a `category` keyword argument with `DeprecationWarning`
|
||||
elif kw := next((kw.value for kw in node.keywords if kw.arg == "category"), None):
|
||||
if _is_deprecation_warning(kw):
|
||||
return kw
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,46 @@
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ForbiddenMakeJudgeInBuiltinScorers(Rule):
|
||||
"""Ensure make_judge is not used in builtin_scorers.py.
|
||||
|
||||
After switching to InstructionsJudge in builtin_scorers.py, this rule
|
||||
prevents future regressions by detecting any usage of make_judge in that file.
|
||||
"""
|
||||
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Usage of `make_judge` is forbidden in builtin_scorers.py. "
|
||||
"Use `InstructionsJudge` directly instead."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver, path: Path) -> bool:
|
||||
"""Check if this is a call to make_judge in builtin_scorers.py.
|
||||
|
||||
Args:
|
||||
node: The AST Call node to check
|
||||
resolver: Resolver instance to resolve fully qualified names
|
||||
path: Path to the file being linted
|
||||
|
||||
Returns:
|
||||
True if this is a forbidden make_judge call, False otherwise
|
||||
"""
|
||||
if path.name != "builtin_scorers.py":
|
||||
return False
|
||||
|
||||
if names := resolver.resolve(node):
|
||||
match names:
|
||||
case ["mlflow", "genai", "judges", "make_judge", *_]:
|
||||
return True
|
||||
case ["mlflow", "genai", "make_judge", *_]:
|
||||
return True
|
||||
case ["make_judge", *_]:
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
return False
|
||||
@@ -0,0 +1,22 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ForbiddenSetActiveModelUsage(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Usage of `set_active_model` is not allowed in mlflow, use `_set_active_model` instead."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""Check if this is a call to set_active_model function."""
|
||||
if names := resolver.resolve(node):
|
||||
match names:
|
||||
case ["mlflow", *_, "set_active_model"]:
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
return False
|
||||
@@ -0,0 +1,12 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ForbiddenTopLevelImport(Rule):
|
||||
def __init__(self, module: str) -> None:
|
||||
self.module = module
|
||||
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
f"Importing module `{self.module}` at the top level is not allowed "
|
||||
"in this file. Use lazy import instead."
|
||||
)
|
||||
@@ -0,0 +1,11 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ForbiddenTraceUIInNotebook(Rule):
|
||||
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."
|
||||
)
|
||||
@@ -0,0 +1,9 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class GetArtifactUri(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"`mlflow.get_artifact_uri` should not be used in examples. "
|
||||
"Use the return value of `log_model` instead."
|
||||
)
|
||||
@@ -0,0 +1,60 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class ImplicitOptional(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Use `Optional` if default value is `None`"
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.AnnAssign) -> bool:
|
||||
"""
|
||||
Returns True if the value to assign is `None` but the type annotation is
|
||||
not `Optional[...]` or `... | None`. For example: `a: int = None`.
|
||||
"""
|
||||
if not ImplicitOptional._is_none(node.value):
|
||||
return False
|
||||
|
||||
# Parse stringified annotations
|
||||
if isinstance(node.annotation, ast.Constant) and isinstance(node.annotation.value, str):
|
||||
try:
|
||||
parsed = ast.parse(node.annotation.value, mode="eval")
|
||||
ann = parsed.body
|
||||
except (SyntaxError, ValueError):
|
||||
# If parsing fails, the annotation is invalid and we trigger the rule
|
||||
# since we cannot verify it contains Optional or | None
|
||||
return True
|
||||
else:
|
||||
ann = node.annotation
|
||||
|
||||
return not (ImplicitOptional._is_optional(ann) or ImplicitOptional._is_bitor_none(ann))
|
||||
|
||||
@staticmethod
|
||||
def _is_optional(ann: ast.expr) -> bool:
|
||||
"""
|
||||
Returns True if `ann` looks like `Optional[...]`.
|
||||
"""
|
||||
return (
|
||||
isinstance(ann, ast.Subscript)
|
||||
and isinstance(ann.value, ast.Name)
|
||||
and ann.value.id == "Optional"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_bitor_none(ann: ast.expr) -> bool:
|
||||
"""
|
||||
Returns True if `ann` looks like `... | None`.
|
||||
"""
|
||||
return (
|
||||
isinstance(ann, ast.BinOp)
|
||||
and isinstance(ann.op, ast.BitOr)
|
||||
and (isinstance(ann.right, ast.Constant) and ann.right.value is None)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_none(value: ast.expr | None) -> bool:
|
||||
"""
|
||||
Returns True if `value` represents `None`.
|
||||
"""
|
||||
return isinstance(value, ast.Constant) and value.value is None
|
||||
@@ -0,0 +1,25 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class IncorrectTypeAnnotation(Rule):
|
||||
MAPPING = {
|
||||
"callable": "Callable",
|
||||
"any": "Any",
|
||||
}
|
||||
|
||||
def __init__(self, type_hint: str) -> None:
|
||||
self.type_hint = type_hint
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Name) -> bool:
|
||||
return node.id in IncorrectTypeAnnotation.MAPPING
|
||||
|
||||
def _message(self) -> str:
|
||||
if correct_hint := self.MAPPING.get(self.type_hint):
|
||||
return f"Did you mean `{correct_hint}` instead of `{self.type_hint}`?"
|
||||
|
||||
raise ValueError(
|
||||
f"Unexpected type: {self.type_hint}. It must be one of {list(self.MAPPING)}."
|
||||
)
|
||||
@@ -0,0 +1,50 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class InvalidAbstractMethod(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Abstract method should only contain a single statement/expression, "
|
||||
"and it must be `pass`, `...`, or a docstring."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_abstract_method(
|
||||
node: ast.FunctionDef | ast.AsyncFunctionDef, resolver: Resolver
|
||||
) -> bool:
|
||||
return any(
|
||||
(resolved := resolver.resolve(d)) and resolved == ["abc", "abstractmethod"]
|
||||
for d in node.decorator_list
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _has_invalid_body(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
|
||||
# Does this abstract method have multiple statements/expressions?
|
||||
if len(node.body) > 1:
|
||||
return True
|
||||
|
||||
# This abstract method has a single statement/expression.
|
||||
# Check if it's `pass`, `...`, or a docstring. If not, it's invalid.
|
||||
stmt = node.body[0]
|
||||
|
||||
# Check for `pass`
|
||||
if isinstance(stmt, ast.Pass):
|
||||
return False
|
||||
|
||||
# Check for `...` or docstring
|
||||
if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant):
|
||||
value = stmt.value.value
|
||||
# `...` literal or docstring
|
||||
return not (value is ... or isinstance(value, str))
|
||||
|
||||
# Any other statement is invalid
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.FunctionDef | ast.AsyncFunctionDef, resolver: Resolver) -> bool:
|
||||
return InvalidAbstractMethod._is_abstract_method(
|
||||
node, resolver
|
||||
) and InvalidAbstractMethod._has_invalid_body(node)
|
||||
@@ -0,0 +1,53 @@
|
||||
import ast
|
||||
|
||||
from packaging.version import InvalidVersion, Version
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
def _is_valid_version(version: str) -> bool:
|
||||
try:
|
||||
v = Version(version)
|
||||
return not (v.is_devrelease or v.is_prerelease or v.is_postrelease)
|
||||
except InvalidVersion:
|
||||
return False
|
||||
|
||||
|
||||
class InvalidExperimentalDecorator(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Invalid usage of `@experimental` decorator. It must be used with a `version` "
|
||||
"argument that is a valid semantic version string."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.expr, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the `@experimental` decorator from mlflow.utils.annotations is used
|
||||
incorrectly.
|
||||
"""
|
||||
resolved = resolver.resolve(node)
|
||||
if not resolved:
|
||||
return False
|
||||
|
||||
if resolved != ["mlflow", "utils", "annotations", "experimental"]:
|
||||
return False
|
||||
|
||||
if not isinstance(node, ast.Call):
|
||||
return True
|
||||
|
||||
version = next((k.value for k in node.keywords if k.arg == "version"), None)
|
||||
if version is None:
|
||||
# No `version` argument, invalid usage
|
||||
return True
|
||||
|
||||
if not isinstance(version, ast.Constant) or not isinstance(version.value, str):
|
||||
# `version` is not a string literal, invalid usage
|
||||
return True
|
||||
|
||||
if not _is_valid_version(version.value):
|
||||
# `version` is not a valid semantic version, # invalid usage
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,52 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class IsinstanceUnionSyntax(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Use `isinstance(obj, (X, Y))` instead of `isinstance(obj, X | Y)`. "
|
||||
"The union syntax with `|` is slower than using a tuple of types."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call) -> bool:
|
||||
"""
|
||||
Returns True if the call is isinstance with union syntax (X | Y) in the second argument.
|
||||
|
||||
Examples that should be flagged:
|
||||
- isinstance(obj, str | int)
|
||||
- isinstance(obj, int | str | float)
|
||||
- isinstance(value, (dict | list))
|
||||
|
||||
Examples that should NOT be flagged:
|
||||
- isinstance(obj, (str, int))
|
||||
- isinstance(obj, str)
|
||||
- other_func(obj, str | int)
|
||||
"""
|
||||
# Check if this is an isinstance call
|
||||
if not (isinstance(node.func, ast.Name) and node.func.id == "isinstance"):
|
||||
return False
|
||||
|
||||
# Check if the second argument uses union syntax (X | Y)
|
||||
match node.args:
|
||||
case [_, type_arg]:
|
||||
return IsinstanceUnionSyntax._has_union_syntax(type_arg)
|
||||
case _:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _has_union_syntax(node: ast.expr) -> bool:
|
||||
"""
|
||||
Returns True if the node contains union syntax with BitOr operator.
|
||||
This handles nested cases like (A | B) | C.
|
||||
"""
|
||||
match node:
|
||||
case ast.BinOp(op=ast.BitOr()):
|
||||
return True
|
||||
case ast.Tuple(elts=elts):
|
||||
# Check if any element in the tuple has union syntax
|
||||
return any(map(IsinstanceUnionSyntax._has_union_syntax, elts))
|
||||
case _:
|
||||
return False
|
||||
@@ -0,0 +1,26 @@
|
||||
from clint.builtin import BUILTIN_MODULES
|
||||
from clint.rules.base import Rule
|
||||
|
||||
# Third-party packages that are always available as core dependencies of mlflow-tracing
|
||||
# (the smallest installable unit of MLflow). Lazy imports of these packages are flagged
|
||||
# the same way as stdlib lazy imports.
|
||||
_ALWAYS_AVAILABLE_MODULES = {
|
||||
"cachetools",
|
||||
"packaging",
|
||||
"pydantic",
|
||||
}
|
||||
|
||||
_LAZY_IMPORT_MODULES = BUILTIN_MODULES | _ALWAYS_AVAILABLE_MODULES
|
||||
|
||||
|
||||
class LazyImport(Rule):
|
||||
def _message(self) -> str:
|
||||
return "This module must be imported at the top level."
|
||||
|
||||
@staticmethod
|
||||
def check(module: str | None) -> bool:
|
||||
"""Check if importing the given module lazily should be flagged."""
|
||||
if module is None:
|
||||
return False
|
||||
root_module = module.split(".", 1)[0]
|
||||
return root_module in _LAZY_IMPORT_MODULES
|
||||
@@ -0,0 +1,6 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class LazyModule(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Module loaded by `LazyLoader` must be imported in `TYPE_CHECKING` block."
|
||||
@@ -0,0 +1,54 @@
|
||||
import ast
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from clint.rules.base import Rule
|
||||
from clint.utils import resolve_expr
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from clint.index import SymbolIndex
|
||||
|
||||
|
||||
class LogModelArtifactPath(Rule):
|
||||
def _message(self) -> str:
|
||||
return "`artifact_path` parameter of `log_model` is deprecated. Use `name` instead."
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, index: "SymbolIndex") -> bool:
|
||||
"""
|
||||
Returns True if the call looks like `mlflow.<flavor>.log_model(...)` and
|
||||
the `artifact_path` argument is specified.
|
||||
"""
|
||||
parts = resolve_expr(node.func)
|
||||
if not parts or len(parts) != 3:
|
||||
return False
|
||||
|
||||
first, second, third = parts
|
||||
if not (first == "mlflow" and third == "log_model"):
|
||||
return False
|
||||
|
||||
# TODO: Remove this once spark flavor supports logging models as logged model artifacts
|
||||
if second == "spark":
|
||||
return False
|
||||
|
||||
function_name = f"{first}.{second}.log_model"
|
||||
artifact_path_idx = LogModelArtifactPath._find_artifact_path_index(index, function_name)
|
||||
if artifact_path_idx is None:
|
||||
return False
|
||||
|
||||
if len(node.args) > artifact_path_idx:
|
||||
return True
|
||||
else:
|
||||
return any(kw.arg and kw.arg == "artifact_path" for kw in node.keywords)
|
||||
|
||||
@staticmethod
|
||||
def _find_artifact_path_index(index: "SymbolIndex", function_name: str) -> int | None:
|
||||
"""
|
||||
Finds the index of the `artifact_path` argument in the function signature of `log_model`
|
||||
using the SymbolIndex.
|
||||
"""
|
||||
if f := index.resolve(function_name):
|
||||
try:
|
||||
return f.all_args.index("artifact_path")
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
@@ -0,0 +1,9 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class MarkdownLink(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Markdown link is not supported in docstring. "
|
||||
"Use reST link instead (e.g., `Link text <link URL>`_)."
|
||||
)
|
||||
@@ -0,0 +1,9 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class MissingDocstringParam(Rule):
|
||||
def __init__(self, params: set[str]) -> None:
|
||||
self.params = params
|
||||
|
||||
def _message(self) -> str:
|
||||
return f"Missing parameters in docstring: {self.params}"
|
||||
@@ -0,0 +1,6 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class MissingNotebookH1Header(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Notebook should have at least one H1 header for the title."
|
||||
@@ -0,0 +1,6 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class MlflowClassName(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Should use `Mlflow` in class name, not `MLflow` or `MLFlow`."
|
||||
@@ -0,0 +1,27 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class MockPatchAsDecorator(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Do not use `unittest.mock.patch` as a decorator. "
|
||||
"Use it as a context manager to avoid patches being active longer than needed "
|
||||
"and to make it clear which code depends on them."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(decorator_list: list[ast.expr], resolver: Resolver) -> ast.expr | None:
|
||||
"""
|
||||
Returns the decorator node if it is a `@mock.patch` or `@patch` decorator.
|
||||
"""
|
||||
for deco in decorator_list:
|
||||
if res := resolver.resolve(deco):
|
||||
match res:
|
||||
# Resolver returns ["unittest", "mock", "patch", ...]
|
||||
# The *_ captures variants like "object", "dict", etc.
|
||||
case ["unittest", "mock", "patch", *_]:
|
||||
return deco
|
||||
return None
|
||||
@@ -0,0 +1,46 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class MockPatchDictEnviron(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Do not use `mock.patch.dict` to modify `os.environ` in tests; "
|
||||
"use pytest's monkeypatch fixture (monkeypatch.setenv / monkeypatch.delenv) instead."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the call is to mock.patch.dict with "os.environ" or os.environ as first arg.
|
||||
Handles:
|
||||
- mock.patch.dict("os.environ", {...})
|
||||
- mock.patch.dict(os.environ, {...})
|
||||
- @mock.patch.dict("os.environ", {...})
|
||||
"""
|
||||
if not isinstance(node, ast.Call):
|
||||
return False
|
||||
|
||||
# Check if this is mock.patch.dict
|
||||
resolved = resolver.resolve(node.func)
|
||||
if resolved != ["unittest", "mock", "patch", "dict"]:
|
||||
return False
|
||||
|
||||
# Check if the first argument is "os.environ" (string) or os.environ (expression)
|
||||
if not node.args:
|
||||
return False
|
||||
|
||||
first_arg = node.args[0]
|
||||
|
||||
# Check for string literal "os.environ"
|
||||
if isinstance(first_arg, ast.Constant) and first_arg.value == "os.environ":
|
||||
return True
|
||||
|
||||
# Check for os.environ as an expression
|
||||
resolved_arg = resolver.resolve(first_arg)
|
||||
if resolved_arg == ["os", "environ"]:
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,51 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class MultiAssign(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Avoid multiple assignment (e.g., `x, y = func()`). Use separate assignments "
|
||||
"instead for better readability and easier debugging."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Assign) -> bool:
|
||||
"""
|
||||
Returns True if the assignment is a tuple assignment where the number of
|
||||
targets matches the number of values, unless all values are constants.
|
||||
|
||||
Examples that should be flagged:
|
||||
- x, y = func1(), func2()
|
||||
- a, b = get_value(), other_value
|
||||
|
||||
Examples that should NOT be flagged:
|
||||
- x, y = 1, 1 (all constants)
|
||||
- a, b, c = 0, 0, 0 (all constants)
|
||||
- x, y = z (unpacking from single value)
|
||||
- a, b = func() (unpacking from function return)
|
||||
- x, y = get_coordinates() (unpacking from function return)
|
||||
"""
|
||||
# Check if we have exactly one target and it's a Tuple
|
||||
if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Tuple):
|
||||
return False
|
||||
|
||||
# Check if the value is also a Tuple
|
||||
if not isinstance(node.value, ast.Tuple):
|
||||
return False
|
||||
|
||||
# Get the number of targets and values
|
||||
num_targets = len(node.targets[0].elts)
|
||||
num_values = len(node.value.elts)
|
||||
|
||||
# Only flag when we have matching number of targets and values (at least 2)
|
||||
if not (num_targets == num_values and num_targets >= 2):
|
||||
return False
|
||||
|
||||
# Allow if all values are constants (e.g., x, y = 1, 1)
|
||||
all_constants = all(isinstance(elt, ast.Constant) for elt in node.value.elts)
|
||||
if all_constants:
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,54 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class NestedMockPatch(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Do not nest `unittest.mock.patch` context managers. "
|
||||
"Use multiple context managers in a single `with` statement instead: "
|
||||
"`with mock.patch(...), mock.patch(...): ...`"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.With, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the with statement uses mock.patch and contains only a single
|
||||
nested with statement that also uses mock.patch.
|
||||
"""
|
||||
# Check if the outer with statement uses mock.patch
|
||||
outer_has_mock_patch = any(
|
||||
NestedMockPatch._is_mock_patch(item.context_expr, resolver) for item in node.items
|
||||
)
|
||||
|
||||
if not outer_has_mock_patch:
|
||||
return False
|
||||
|
||||
# Check if the body has exactly one statement and it's a with statement
|
||||
if len(node.body) == 1 and isinstance(node.body[0], ast.With):
|
||||
# Check if the nested with statement also uses mock.patch
|
||||
inner_has_mock_patch = any(
|
||||
NestedMockPatch._is_mock_patch(item.context_expr, resolver)
|
||||
for item in node.body[0].items
|
||||
)
|
||||
|
||||
if inner_has_mock_patch:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _is_mock_patch(node: ast.expr, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the node is a call to mock.patch or any of its variants.
|
||||
"""
|
||||
# Handle direct calls: mock.patch(...), mock.patch.object(...), etc.
|
||||
if isinstance(node, ast.Call):
|
||||
if res := resolver.resolve(node.func):
|
||||
match res:
|
||||
# Matches unittest.mock.patch, unittest.mock.patch.object, etc.
|
||||
case ["unittest", "mock", "patch", *_]:
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,35 @@
|
||||
import ast
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class NoClassBasedTests(Rule):
|
||||
def __init__(self, class_name: str) -> None:
|
||||
self.class_name = class_name
|
||||
|
||||
@classmethod
|
||||
def check(cls, node: ast.ClassDef, path_name: str) -> Self | None:
|
||||
# Only check in test files
|
||||
if not path_name.startswith("test_"):
|
||||
return None
|
||||
|
||||
if not node.name.startswith("Test"):
|
||||
return None
|
||||
|
||||
# Check if the class has any test methods
|
||||
if any(
|
||||
isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
and stmt.name.startswith("test_")
|
||||
for stmt in node.body
|
||||
):
|
||||
return cls(node.name)
|
||||
|
||||
return None
|
||||
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
f"Class-based tests are not allowed. "
|
||||
f"Convert class '{self.class_name}' to function-based tests."
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class NoRst(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Do not use RST style. Use Google style instead."
|
||||
@@ -0,0 +1,15 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class NoShebang(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Python scripts should not contain shebang lines"
|
||||
|
||||
@staticmethod
|
||||
def check(file_content: str) -> bool:
|
||||
"""
|
||||
Returns True if the file contains a shebang line at the beginning.
|
||||
|
||||
A shebang line is a line that starts with '#!' (typically #!/usr/bin/env python).
|
||||
"""
|
||||
return file_content.startswith("#!")
|
||||
@@ -0,0 +1,16 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class OsChdirInTest(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Do not use `os.chdir` in test directly. Use `monkeypatch.chdir` (https://docs.pytest.org/en/stable/reference/reference.html#pytest.MonkeyPatch.chdir)."
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the call is to os.chdir().
|
||||
"""
|
||||
return resolver.resolve(node) == ["os", "chdir"]
|
||||
@@ -0,0 +1,30 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class OsEnvironDeleteInTest(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Do not delete `os.environ` in test directly (del os.environ[...] or "
|
||||
"os.environ.pop(...)). Use `monkeypatch.delenv` "
|
||||
"(https://docs.pytest.org/en/stable/reference/reference.html#pytest.MonkeyPatch.delenv)."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Delete | ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the operation is deletion from os.environ[...] or
|
||||
a call to os.environ.pop().
|
||||
"""
|
||||
if isinstance(node, ast.Delete):
|
||||
# Handle: del os.environ["KEY"]
|
||||
if len(node.targets) == 1 and isinstance(node.targets[0], ast.Subscript):
|
||||
resolved = resolver.resolve(node.targets[0].value)
|
||||
return resolved == ["os", "environ"]
|
||||
elif isinstance(node, ast.Call):
|
||||
# Handle: os.environ.pop("KEY")
|
||||
resolved = resolver.resolve(node)
|
||||
return resolved == ["os", "environ", "pop"]
|
||||
return False
|
||||
@@ -0,0 +1,19 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class OsEnvironSetInTest(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Do not set `os.environ` in test directly. Use `monkeypatch.setenv` (https://docs.pytest.org/en/stable/reference/reference.html#pytest.MonkeyPatch.setenv)."
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Assign, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the assignment is to os.environ[...].
|
||||
"""
|
||||
if len(node.targets) == 1 and isinstance(node.targets[0], ast.Subscript):
|
||||
resolved = resolver.resolve(node.targets[0].value)
|
||||
return resolved == ["os", "environ"]
|
||||
return False
|
||||
@@ -0,0 +1,56 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
def _is_simple_name_or_attribute(node: ast.expr) -> bool:
|
||||
"""
|
||||
Check if a node is a simple name (e.g., `a`) or a chain of attribute
|
||||
accesses on a simple name (e.g., `obj.attr` or `a.b.c`).
|
||||
"""
|
||||
if isinstance(node, ast.Name):
|
||||
return True
|
||||
if isinstance(node, ast.Attribute):
|
||||
return _is_simple_name_or_attribute(node.value)
|
||||
return False
|
||||
|
||||
|
||||
class PreferDictUnion(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Use `|` operator for dictionary merging (e.g., `a | b`) "
|
||||
"instead of `{**a, **b}` for better readability."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Dict) -> bool:
|
||||
"""
|
||||
Returns True if the dictionary is composed entirely of 2+ dictionary unpacking
|
||||
expressions that can be replaced with the `|` operator.
|
||||
|
||||
Examples that should be flagged:
|
||||
- {**a, **b}
|
||||
- {**a, **b, **c}
|
||||
- {**obj.attr, **b}
|
||||
- {**a.b.c, **d}
|
||||
|
||||
Examples that should NOT be flagged:
|
||||
- {**a} # Single unpack
|
||||
- {**a, "key": value} # Mixed with literal keys
|
||||
- {**data[0], **b}, {**func(), **b} # Complex expressions
|
||||
- {**a,\n**b} # Multi-line dicts
|
||||
"""
|
||||
# Need at least 2 elements for a merge
|
||||
if len(node.keys) < 2:
|
||||
return False
|
||||
|
||||
# Skip multi-line dicts
|
||||
if node.end_lineno and node.end_lineno > node.lineno:
|
||||
return False
|
||||
|
||||
# All keys must be None (indicating dictionary unpacking with **)
|
||||
if not all(key is None for key in node.keys):
|
||||
return False
|
||||
|
||||
# All values must be simple names or attribute access on a name
|
||||
return all(_is_simple_name_or_attribute(value) for value in node.values)
|
||||
@@ -0,0 +1,36 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class PreferNext(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Use `next(x for x in items if condition)` instead of "
|
||||
"`[x for x in items if condition][0]` for finding the first matching element."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Subscript) -> bool:
|
||||
"""
|
||||
Returns True if the node is a list comprehension with an `if` clause
|
||||
subscripted with `[0]`.
|
||||
|
||||
Examples that should be flagged:
|
||||
- [x for x in items if f(x)][0]
|
||||
- [x.name for x in items if x.active][0]
|
||||
|
||||
Examples that should NOT be flagged:
|
||||
- [x for x in items][0] (no if clause)
|
||||
- [x for x in items if f(x)][1] (not [0])
|
||||
- [x for x in items if f(x)][-1] (not [0])
|
||||
- (x for x in items if f(x)) (already a generator)
|
||||
"""
|
||||
match node:
|
||||
case ast.Subscript(
|
||||
value=ast.ListComp(generators=generators),
|
||||
slice=ast.Constant(value=0),
|
||||
) if any(gen.ifs for gen in generators):
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
@@ -0,0 +1,25 @@
|
||||
import ast
|
||||
from typing import Literal
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
# See https://github.com/astral-sh/ruff/issues/3608
|
||||
class PreferOsEnviron(Rule):
|
||||
def __init__(self, func: Literal["getenv", "putenv"]) -> None:
|
||||
self.func = func
|
||||
|
||||
def _message(self) -> str:
|
||||
if self.func == "putenv":
|
||||
return "Use `os.environ[key] = value` instead of `os.putenv()`."
|
||||
return "Use `os.environ.get()` instead of `os.getenv()`."
|
||||
|
||||
@classmethod
|
||||
def check(cls, node: ast.Call, resolver: Resolver) -> Self | None:
|
||||
match resolver.resolve(node.func):
|
||||
case ["os", ("getenv" | "putenv") as func]:
|
||||
return cls(func)
|
||||
return None
|
||||
@@ -0,0 +1,22 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class PytestMarkRepeat(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"@pytest.mark.repeat decorator should not be committed. "
|
||||
"This decorator is meant for local testing only to check for flaky tests."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(decorator_list: list[ast.expr], resolver: Resolver) -> ast.expr | None:
|
||||
"""
|
||||
Returns the decorator node if it is a `@pytest.mark.repeat` decorator.
|
||||
"""
|
||||
for deco in decorator_list:
|
||||
if (res := resolver.resolve(deco)) and res == ["pytest", "mark", "repeat"]:
|
||||
return deco
|
||||
return None
|
||||
@@ -0,0 +1,36 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class RedundantMockReturnValue(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Do not pass `return_value=MagicMock()` or `return_value=Mock()` to `patch()`. "
|
||||
"The default return value of a mock is already a new mock."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the call is a patch() with a redundant return_value=Mock() or
|
||||
return_value=MagicMock() (no arguments) keyword argument.
|
||||
"""
|
||||
match resolver.resolve(node.func):
|
||||
case ["unittest", "mock", "patch", *_]:
|
||||
pass
|
||||
case _:
|
||||
return False
|
||||
|
||||
for keyword in node.keywords:
|
||||
match keyword:
|
||||
case ast.keyword(
|
||||
arg="return_value",
|
||||
value=ast.Call(args=[], keywords=[]) as value,
|
||||
):
|
||||
match resolver.resolve(value.func):
|
||||
case ["unittest", "mock", "Mock" | "MagicMock"]:
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Rule to detect redundant docstrings in test files.
|
||||
|
||||
This rule flags:
|
||||
- ALL single-line docstrings in test functions and classes (multi-line
|
||||
function/class docstrings are allowed since they generally provide
|
||||
meaningful context).
|
||||
- ALL module-level docstrings in test files (single- or multi-line).
|
||||
"""
|
||||
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class RedundantTestDocstring(Rule):
|
||||
@staticmethod
|
||||
def check(
|
||||
node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef, path_name: str
|
||||
) -> ast.Constant | None:
|
||||
if not (path_name.startswith("test_") or path_name.endswith("_test.py")):
|
||||
return None
|
||||
|
||||
is_class = isinstance(node, ast.ClassDef)
|
||||
|
||||
if is_class and not node.name.startswith("Test"):
|
||||
return None
|
||||
if not is_class and not node.name.startswith("test_"):
|
||||
return None
|
||||
|
||||
# Check if docstring exists and get the raw docstring for multiline detection
|
||||
if (
|
||||
node.body
|
||||
and isinstance(node.body[0], ast.Expr)
|
||||
and isinstance(node.body[0].value, ast.Constant)
|
||||
and isinstance(node.body[0].value.value, str)
|
||||
):
|
||||
raw_docstring = node.body[0].value.value
|
||||
|
||||
# If raw docstring has newlines, it's multiline - always allow
|
||||
if "\n" in raw_docstring:
|
||||
return None
|
||||
|
||||
# Return the docstring node to flag
|
||||
return node.body[0].value
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def check_module(module: ast.Module, path_name: str) -> ast.Constant | None:
|
||||
if not (path_name.startswith("test_") or path_name.endswith("_test.py")):
|
||||
return None
|
||||
|
||||
if (
|
||||
module.body
|
||||
and isinstance(module.body[0], ast.Expr)
|
||||
and isinstance(module.body[0].value, ast.Constant)
|
||||
and isinstance(module.body[0].value.value, str)
|
||||
):
|
||||
return module.body[0].value
|
||||
|
||||
return None
|
||||
|
||||
def _message(self) -> str:
|
||||
return "Docstrings in test files rarely provide meaningful context. Consider removing it."
|
||||
@@ -0,0 +1,43 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class SubprocessCheckCall(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Use `subprocess.check_call(...)` instead of `subprocess.run(..., check=True)` "
|
||||
"for better readability. Only applies when check=True is the only keyword argument."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if `node` is `subprocess.run(..., check=True)` with no other keyword arguments.
|
||||
"""
|
||||
resolved = resolver.resolve(node)
|
||||
|
||||
# Check if this is subprocess.run
|
||||
if resolved != ["subprocess", "run"]:
|
||||
return False
|
||||
|
||||
# Check if there are any keyword arguments
|
||||
if not node.keywords:
|
||||
return False
|
||||
|
||||
# Check if the only keyword argument is check=True
|
||||
if len(node.keywords) != 1:
|
||||
return False
|
||||
|
||||
keyword = node.keywords[0]
|
||||
|
||||
# Check if the keyword is 'check' (not **kwargs)
|
||||
if keyword.arg != "check":
|
||||
return False
|
||||
|
||||
# Check if the value is True
|
||||
if not isinstance(keyword.value, ast.Constant):
|
||||
return False
|
||||
|
||||
return keyword.value.value is True
|
||||
@@ -0,0 +1,23 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class TempfileInTest(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Do not use `tempfile` in tests. Use the `tmp_path` fixture instead "
|
||||
"(https://docs.pytest.org/en/stable/reference/reference.html#tmp-path)."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
match resolver.resolve(node):
|
||||
case [
|
||||
"tempfile",
|
||||
"TemporaryDirectory" | "NamedTemporaryFile" | "TemporaryFile" | "mkdtemp",
|
||||
]:
|
||||
return True
|
||||
case _:
|
||||
return False
|
||||
@@ -0,0 +1,6 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class TestNameTypo(Rule):
|
||||
def _message(self) -> str:
|
||||
return "This function looks like a test, but its name does not start with 'test_'."
|
||||
@@ -0,0 +1,15 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class TypingExtensions(Rule):
|
||||
def __init__(self, *, full_name: str, allowlist: list[str]) -> None:
|
||||
self.full_name = full_name
|
||||
self.allowlist = allowlist
|
||||
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
f"`{self.full_name}` is not allowed to use. Only {self.allowlist} are allowed. "
|
||||
"You can extend `tool.clint.typing-extensions-allowlist` in `pyproject.toml` if needed "
|
||||
"but make sure that the version requirement for `typing-extensions` is compatible with "
|
||||
"the added types."
|
||||
)
|
||||
@@ -0,0 +1,14 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UnknownMlflowArguments(Rule):
|
||||
def __init__(self, function_name: str, unknown_args: set[str]) -> None:
|
||||
self.function_name = function_name
|
||||
self.unknown_args = unknown_args
|
||||
|
||||
def _message(self) -> str:
|
||||
args_str = ", ".join(f"`{arg}`" for arg in sorted(self.unknown_args))
|
||||
return (
|
||||
f"Unknown arguments {args_str} passed to `{self.function_name}`. "
|
||||
"Check the function signature for valid parameter names."
|
||||
)
|
||||
@@ -0,0 +1,12 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UnknownMlflowFunction(Rule):
|
||||
def __init__(self, function_name: str) -> None:
|
||||
self.function_name = function_name
|
||||
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
f"Unknown MLflow function: `{self.function_name}`. "
|
||||
"This function may not exist or could be misspelled."
|
||||
)
|
||||
@@ -0,0 +1,23 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UnnamedThread(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"`threading.Thread()` must be called with a `name` argument to improve debugging "
|
||||
"and traceability of thread-related issues."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the call is threading.Thread() without a name parameter.
|
||||
"""
|
||||
if names := resolver.resolve(node):
|
||||
return names == ["threading", "Thread"] and not any(
|
||||
keyword.arg == "name" for keyword in node.keywords
|
||||
)
|
||||
return False
|
||||
@@ -0,0 +1,23 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UnnamedThreadPool(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"`ThreadPoolExecutor()` must be called with a `thread_name_prefix` argument to improve "
|
||||
"debugging and traceability of thread-related issues."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the call is ThreadPoolExecutor() without a thread_name_prefix parameter.
|
||||
"""
|
||||
if names := resolver.resolve(node):
|
||||
return names == ["concurrent", "futures", "ThreadPoolExecutor"] and not any(
|
||||
kw.arg == "thread_name_prefix" for kw in node.keywords
|
||||
)
|
||||
return False
|
||||
@@ -0,0 +1,32 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UnparameterizedGenericType(Rule):
|
||||
def __init__(self, type_hint: str) -> None:
|
||||
self.type_hint = type_hint
|
||||
|
||||
@staticmethod
|
||||
def is_generic_type(node: ast.Name | ast.Attribute, resolver: Resolver) -> bool:
|
||||
if names := resolver.resolve(node):
|
||||
return tuple(names) in {
|
||||
("typing", "Callable"),
|
||||
("typing", "Sequence"),
|
||||
}
|
||||
elif isinstance(node, ast.Name):
|
||||
return node.id in {
|
||||
"dict",
|
||||
"list",
|
||||
"set",
|
||||
"tuple",
|
||||
"frozenset",
|
||||
}
|
||||
return False
|
||||
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
f"Generic type `{self.type_hint}` must be parameterized "
|
||||
"(e.g., `list[str]` rather than `list`)."
|
||||
)
|
||||
@@ -0,0 +1,9 @@
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UnusedDisableComment(Rule):
|
||||
def __init__(self, rule_name: str) -> None:
|
||||
self.rule_name = rule_name
|
||||
|
||||
def _message(self) -> str:
|
||||
return f"Unused disable comment for rule `{self.rule_name}`"
|
||||
@@ -0,0 +1,24 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UseGhToken(Rule):
|
||||
def _message(self) -> str:
|
||||
return "Use GH_TOKEN instead of GITHUB_TOKEN for the environment variable name."
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if the call reads the GITHUB_TOKEN environment variable.
|
||||
Handles:
|
||||
- os.getenv("GITHUB_TOKEN")
|
||||
- os.environ.get("GITHUB_TOKEN")
|
||||
"""
|
||||
match node:
|
||||
case ast.Call(args=[ast.Constant(value="GITHUB_TOKEN"), *_]):
|
||||
match resolver.resolve(node.func):
|
||||
case ["os", "getenv"] | ["os", "environ", "get"]:
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,34 @@
|
||||
import ast
|
||||
|
||||
from clint.resolver import Resolver
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UseSysExecutable(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Use `[sys.executable, '-m', 'mlflow', ...]` when running mlflow CLI in a subprocess."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Call, resolver: Resolver) -> bool:
|
||||
"""
|
||||
Returns True if `node` looks like `subprocess.Popen(["mlflow", ...])`.
|
||||
"""
|
||||
resolved = resolver.resolve(node)
|
||||
if (
|
||||
resolved
|
||||
and len(resolved) == 2
|
||||
and resolved[0] == "subprocess"
|
||||
and resolved[1] in ["Popen", "run", "check_output", "check_call"]
|
||||
and node.args
|
||||
):
|
||||
first_arg = node.args[0]
|
||||
if isinstance(first_arg, ast.List) and first_arg.elts:
|
||||
first_elem = first_arg.elts[0]
|
||||
return (
|
||||
isinstance(first_elem, ast.Constant)
|
||||
and isinstance(first_elem.value, str)
|
||||
and first_elem.value == "mlflow"
|
||||
)
|
||||
return False
|
||||
@@ -0,0 +1,189 @@
|
||||
import ast
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
|
||||
class UseWalrusOperator(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Use the walrus operator `:=` when a variable is assigned and only used "
|
||||
"within an `if` block that tests its truthiness. "
|
||||
"For example, replace `a = ...; if a: use_a(a)` with `if a := ...: use_a(a)`."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(
|
||||
if_node: ast.If,
|
||||
prev_stmt: ast.stmt,
|
||||
following_stmts: list[ast.stmt],
|
||||
) -> bool:
|
||||
"""
|
||||
Flags::
|
||||
|
||||
a = func()
|
||||
if a:
|
||||
use(a)
|
||||
|
||||
Ignores: comparisons, tuple unpacking, multi-line, used in elif/else,
|
||||
used after if, line > 100 chars
|
||||
"""
|
||||
# Check if previous statement is a simple assignment (not augmented, not annotated)
|
||||
if not isinstance(prev_stmt, ast.Assign):
|
||||
return False
|
||||
|
||||
# Skip if the assignment statement spans multiple lines
|
||||
if (
|
||||
prev_stmt.end_lineno is not None
|
||||
and prev_stmt.lineno is not None
|
||||
and prev_stmt.end_lineno > prev_stmt.lineno
|
||||
):
|
||||
return False
|
||||
|
||||
# Must be a single target assignment to a Name
|
||||
if len(prev_stmt.targets) != 1:
|
||||
return False
|
||||
|
||||
target = prev_stmt.targets[0]
|
||||
if not isinstance(target, ast.Name):
|
||||
return False
|
||||
|
||||
var_name = target.id
|
||||
|
||||
# The if condition must be just the variable name (truthiness test)
|
||||
if not isinstance(if_node.test, ast.Name):
|
||||
return False
|
||||
|
||||
if if_node.test.id != var_name:
|
||||
return False
|
||||
|
||||
# Check that the variable is used in the if body
|
||||
if not _name_used_in_stmts(var_name, if_node.body):
|
||||
return False
|
||||
|
||||
# Check that the variable is NOT used in elif/else branches
|
||||
if if_node.orelse and _name_used_in_stmts(var_name, if_node.orelse):
|
||||
return False
|
||||
|
||||
# Check that the variable is NOT used after the if statement
|
||||
if following_stmts and _name_used_in_stmts(var_name, following_stmts):
|
||||
return False
|
||||
|
||||
# Skip if the fixed code would exceed 100 characters
|
||||
# Original: "if var:" -> Fixed: "if var := value:"
|
||||
value = prev_stmt.value
|
||||
if (
|
||||
value.end_col_offset is None
|
||||
or value.col_offset is None
|
||||
or if_node.test.end_col_offset is None
|
||||
):
|
||||
return False
|
||||
value_width = value.end_col_offset - value.col_offset
|
||||
fixed_line_length = (
|
||||
if_node.test.end_col_offset
|
||||
+ 4 # len(" := ")
|
||||
+ value_width
|
||||
+ 1 # len(":")
|
||||
)
|
||||
if fixed_line_length > 100:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _name_used_in_stmts(name: str, stmts: list[ast.stmt]) -> bool:
|
||||
"""Check if a name is used (loaded) in a list of statements.
|
||||
|
||||
Skips nested function/class definitions to avoid false positives from
|
||||
inner scopes that shadow or independently use the same name.
|
||||
"""
|
||||
return any(_name_used_in_node(name, stmt) for stmt in stmts)
|
||||
|
||||
|
||||
def _name_used_in_node(name: str, node: ast.AST) -> bool:
|
||||
"""Recursively check if a name is used."""
|
||||
match node:
|
||||
case ast.Name(id=id, ctx=ast.Load()) if id == name:
|
||||
return True
|
||||
case _:
|
||||
return any(_name_used_in_node(name, child) for child in ast.iter_child_nodes(node))
|
||||
|
||||
|
||||
class WalrusOperatorVisitor(ast.NodeVisitor):
|
||||
"""Visits all statement blocks to check for walrus operator opportunities."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.violations: list[ast.stmt] = []
|
||||
|
||||
def _check_stmts(self, stmts: list[ast.stmt]) -> None:
|
||||
for idx, stmt in enumerate(stmts[1:], start=1):
|
||||
if isinstance(stmt, ast.If):
|
||||
prev_stmt = stmts[idx - 1]
|
||||
following_stmts = stmts[idx + 1 :]
|
||||
if UseWalrusOperator.check(stmt, prev_stmt, following_stmts):
|
||||
self.violations.append(prev_stmt)
|
||||
|
||||
def _visit_stmts(self, stmts: list[ast.stmt]) -> None:
|
||||
for stmt in stmts:
|
||||
self.visit(stmt)
|
||||
|
||||
# Walrus opportunities only exist inside statement blocks, so skip descent into
|
||||
# expressions, decorators, args, etc. This trims the AST walk dramatically on
|
||||
# expression-heavy functions.
|
||||
def generic_visit(self, node: ast.AST) -> None:
|
||||
return
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._visit_stmts(node.body)
|
||||
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._visit_stmts(node.body)
|
||||
|
||||
def visit_If(self, node: ast.If) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._check_stmts(node.orelse)
|
||||
self._visit_stmts(node.body)
|
||||
self._visit_stmts(node.orelse)
|
||||
|
||||
def visit_For(self, node: ast.For) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._check_stmts(node.orelse)
|
||||
self._visit_stmts(node.body)
|
||||
self._visit_stmts(node.orelse)
|
||||
|
||||
def visit_AsyncFor(self, node: ast.AsyncFor) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._check_stmts(node.orelse)
|
||||
self._visit_stmts(node.body)
|
||||
self._visit_stmts(node.orelse)
|
||||
|
||||
def visit_While(self, node: ast.While) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._check_stmts(node.orelse)
|
||||
self._visit_stmts(node.body)
|
||||
self._visit_stmts(node.orelse)
|
||||
|
||||
def visit_With(self, node: ast.With) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._visit_stmts(node.body)
|
||||
|
||||
def visit_AsyncWith(self, node: ast.AsyncWith) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._visit_stmts(node.body)
|
||||
|
||||
def visit_Try(self, node: ast.Try) -> None:
|
||||
self._check_stmts(node.body)
|
||||
self._check_stmts(node.orelse)
|
||||
self._check_stmts(node.finalbody)
|
||||
self._visit_stmts(node.body)
|
||||
for handler in node.handlers:
|
||||
self._check_stmts(handler.body)
|
||||
self._visit_stmts(handler.body)
|
||||
self._visit_stmts(node.orelse)
|
||||
self._visit_stmts(node.finalbody)
|
||||
|
||||
def visit_Match(self, node: ast.Match) -> None:
|
||||
for case in node.cases:
|
||||
self._check_stmts(case.body)
|
||||
self._visit_stmts(case.body)
|
||||
@@ -0,0 +1,57 @@
|
||||
import ast
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from clint.rules.base import Rule
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from clint.resolver import Resolver
|
||||
|
||||
|
||||
class MajorVersionCheck(Rule):
|
||||
def _message(self) -> str:
|
||||
return (
|
||||
"Use `.major` field for major version comparisons instead of full version strings. "
|
||||
"This is more explicit, and efficient (avoids creating a second Version object). "
|
||||
"For example, use `Version(__version__).major >= 1` instead of "
|
||||
'`Version(__version__) >= Version("1.0.0")`.'
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def check(node: ast.Compare, resolver: "Resolver") -> bool:
|
||||
if len(node.ops) != 1 or len(node.comparators) != 1:
|
||||
return False
|
||||
|
||||
if not isinstance(node.ops[0], (ast.GtE, ast.LtE, ast.Gt, ast.Lt, ast.Eq, ast.NotEq)):
|
||||
return False
|
||||
|
||||
if not (
|
||||
isinstance(node.left, ast.Call)
|
||||
and MajorVersionCheck._is_version_call(node.left, resolver)
|
||||
):
|
||||
return False
|
||||
|
||||
comparator = node.comparators[0]
|
||||
if not (
|
||||
isinstance(comparator, ast.Call)
|
||||
and MajorVersionCheck._is_version_call(comparator, resolver)
|
||||
):
|
||||
return False
|
||||
|
||||
match comparator.args:
|
||||
case [arg] if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
|
||||
version_str = arg.value
|
||||
return MajorVersionCheck._is_major_only_version(version_str)
|
||||
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _is_version_call(node: ast.Call, resolver: "Resolver") -> bool:
|
||||
if resolved := resolver.resolve(node.func):
|
||||
return resolved == ["packaging", "version", "Version"]
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _is_major_only_version(version_str: str) -> bool:
|
||||
pattern = r"^(\d+)\.0\.0$"
|
||||
return re.match(pattern, version_str) is not None
|
||||
@@ -0,0 +1,90 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import re
|
||||
import subprocess
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_repo_root() -> Path:
|
||||
"""Find the git repository root directory with caching."""
|
||||
try:
|
||||
result = subprocess.check_output(["git", "rev-parse", "--show-toplevel"], text=True).strip()
|
||||
return Path(result)
|
||||
except (OSError, subprocess.CalledProcessError) as e:
|
||||
raise RuntimeError("Failed to find git repository root") from e
|
||||
|
||||
|
||||
def resolve_expr(expr: ast.expr) -> list[str] | None:
|
||||
"""
|
||||
Resolves `expr` to a list of attribute names. For example, given `expr` like
|
||||
`some.module.attribute`, ['some', 'module', 'attribute'] is returned.
|
||||
If `expr` is not resolvable, `None` is returned.
|
||||
"""
|
||||
if isinstance(expr, ast.Attribute):
|
||||
base = resolve_expr(expr.value)
|
||||
if base is None:
|
||||
return None
|
||||
return base + [expr.attr]
|
||||
elif isinstance(expr, ast.Name):
|
||||
return [expr.id]
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_ignored_rules_for_file(
|
||||
file_path: Path, per_file_ignores: dict[re.Pattern[str], set[str]]
|
||||
) -> set[str]:
|
||||
"""
|
||||
Returns a set of rule names that should be ignored for the given file path.
|
||||
|
||||
Args:
|
||||
file_path: The file path to check
|
||||
per_file_ignores: Dict mapping compiled regex patterns to lists of rule names to ignore
|
||||
|
||||
Returns:
|
||||
Set of rule names to ignore for this file
|
||||
"""
|
||||
ignored_rules: set[str] = set()
|
||||
for pattern, rules in per_file_ignores.items():
|
||||
if pattern.fullmatch(file_path.as_posix()):
|
||||
ignored_rules |= rules
|
||||
return ignored_rules
|
||||
|
||||
|
||||
ALLOWED_EXTS = {".md", ".mdx", ".rst", ".py", ".ipynb"}
|
||||
|
||||
|
||||
def _git_ls_files(pathspecs: list[Path]) -> list[Path]:
|
||||
"""
|
||||
Return git-tracked and untracked (but not ignored) files matching the given pathspecs.
|
||||
Git does not filter by extension; filtering happens in Python.
|
||||
"""
|
||||
try:
|
||||
out = subprocess.check_output(
|
||||
["git", "ls-files", "--cached", "--others", "--exclude-standard", "--", *pathspecs],
|
||||
text=True,
|
||||
)
|
||||
except (OSError, subprocess.CalledProcessError) as e:
|
||||
raise RuntimeError("Failed to list git files") from e
|
||||
|
||||
return [Path(line) for line in out.splitlines() if line]
|
||||
|
||||
|
||||
def resolve_paths(paths: list[Path]) -> list[Path]:
|
||||
"""
|
||||
Resolve CLI arguments into a list of tracked and untracked files to lint.
|
||||
|
||||
- Includes git-tracked files and untracked files (but not ignored files)
|
||||
- Only includes: .md, .mdx, .rst, .py, .ipynb
|
||||
"""
|
||||
if not paths:
|
||||
paths = [Path(".")]
|
||||
|
||||
all_files = _git_ls_files(paths)
|
||||
|
||||
filtered = {p for p in all_files if p.suffix.lower() in ALLOWED_EXTS and p.exists()}
|
||||
|
||||
return sorted(filtered)
|
||||
@@ -0,0 +1,7 @@
|
||||
import pytest
|
||||
from clint.index import SymbolIndex
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def index() -> SymbolIndex:
|
||||
return SymbolIndex.build()
|
||||
@@ -0,0 +1,148 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules import AssignBeforeAppend
|
||||
|
||||
|
||||
def test_assign_before_append_basic(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
items.append(item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 1
|
||||
assert all(isinstance(r.rule, AssignBeforeAppend) for r in results)
|
||||
assert results[0].range == Range(Position(2, 0))
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_different_variable(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
items.append(other_var)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_no_empty_list_init(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
items.append(item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_different_list(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
other_list.append(item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_three_statements(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
print(item)
|
||||
items.append(item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_one_statement(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
items.append(transform(x))
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_list_with_initial_values(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = [1, 2, 3]
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
items.append(item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_multiple_violations(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
items.append(item)
|
||||
|
||||
results = []
|
||||
for y in other_data:
|
||||
result = process(y)
|
||||
results.append(result)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 2
|
||||
assert all(isinstance(r.rule, AssignBeforeAppend) for r in results)
|
||||
assert results[0].range == Range(Position(2, 0))
|
||||
assert results[1].range == Range(Position(7, 0))
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_complex_assignment(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
item, other = transform(x)
|
||||
items.append(item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_no_flag_attribute_assignment(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
for x in data:
|
||||
self.item = transform(x)
|
||||
items.append(self.item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_assign_before_append_separated_statements(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
items = []
|
||||
other_statement()
|
||||
for x in data:
|
||||
item = transform(x)
|
||||
items.append(item)
|
||||
"""
|
||||
config = Config(select={AssignBeforeAppend.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 0
|
||||
@@ -0,0 +1,49 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules.do_not_disable import DoNotDisable
|
||||
|
||||
|
||||
def test_do_not_disable(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
# Bad B006
|
||||
# noqa: B006
|
||||
|
||||
# Bad F821
|
||||
# noqa: F821
|
||||
|
||||
# Good
|
||||
# noqa: B004
|
||||
"""
|
||||
config = Config(select={DoNotDisable.name})
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(violations) == 2
|
||||
assert all(isinstance(v.rule, DoNotDisable) for v in violations)
|
||||
assert violations[0].range == Range(Position(2, 0))
|
||||
assert violations[1].range == Range(Position(5, 0))
|
||||
|
||||
|
||||
def test_do_not_disable_comma_separated(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
# Bad: B006 and F821 both should be caught
|
||||
# noqa: B006, F821
|
||||
|
||||
# Bad: B006 and F821 both should be caught (no space after comma)
|
||||
# noqa: B006,F821
|
||||
|
||||
# Good: B004 is allowed
|
||||
# noqa: B004, B005
|
||||
"""
|
||||
config = Config(select={DoNotDisable.name})
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(violations) == 2
|
||||
assert all(isinstance(v.rule, DoNotDisable) for v in violations)
|
||||
# Both violations should have both rules B006 and F821
|
||||
assert isinstance(violations[0].rule, DoNotDisable)
|
||||
assert isinstance(violations[1].rule, DoNotDisable)
|
||||
assert violations[0].rule.rules == {"B006", "F821"}
|
||||
assert violations[1].rule.rules == {"B006", "F821"}
|
||||
assert violations[0].range == Range(Position(2, 0))
|
||||
assert violations[1].range == Range(Position(5, 0))
|
||||
@@ -0,0 +1,31 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules.docstring_param_order import DocstringParamOrder
|
||||
|
||||
|
||||
def test_docstring_param_order(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
# Bad
|
||||
def f(x: int, y: str) -> None:
|
||||
'''
|
||||
Args:
|
||||
y: Second param.
|
||||
x: First param.
|
||||
'''
|
||||
|
||||
# Good
|
||||
def f(a: int, b: str) -> None:
|
||||
'''
|
||||
Args:
|
||||
a: First param.
|
||||
b: Second param.
|
||||
'''
|
||||
"""
|
||||
config = Config(select={DocstringParamOrder.name})
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(violations) == 1
|
||||
assert all(isinstance(v.rule, DocstringParamOrder) for v in violations)
|
||||
assert violations[0].range == Range(Position(2, 0))
|
||||
@@ -0,0 +1,47 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import lint_file
|
||||
from clint.rules.empty_notebook_cell import EmptyNotebookCell
|
||||
|
||||
|
||||
def test_empty_notebook_cell(index: SymbolIndex) -> None:
|
||||
notebook_content = {
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [], # Empty cell
|
||||
"metadata": {},
|
||||
"execution_count": None,
|
||||
"outputs": [],
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": ["x = 5"],
|
||||
"metadata": {},
|
||||
"execution_count": None,
|
||||
"outputs": [],
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [], # Another empty cell
|
||||
"metadata": {},
|
||||
"execution_count": None,
|
||||
"outputs": [],
|
||||
},
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4,
|
||||
}
|
||||
code = json.dumps(notebook_content)
|
||||
config = Config(select={EmptyNotebookCell.name})
|
||||
violations = lint_file(Path("test_notebook.ipynb"), code, config, index)
|
||||
assert len(violations) == 2
|
||||
assert all(isinstance(v.rule, EmptyNotebookCell) for v in violations)
|
||||
assert violations[0].cell == 1
|
||||
assert violations[1].cell == 3
|
||||
@@ -0,0 +1,46 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules.example_syntax_error import ExampleSyntaxError
|
||||
|
||||
|
||||
def test_example_syntax_error(index: SymbolIndex) -> None:
|
||||
code = '''
|
||||
def bad():
|
||||
"""
|
||||
.. code-block:: python
|
||||
|
||||
def f():
|
||||
|
||||
"""
|
||||
|
||||
def good():
|
||||
"""
|
||||
.. code-block:: python
|
||||
|
||||
def f():
|
||||
return "This is a good example"
|
||||
"""
|
||||
'''
|
||||
config = Config(select={ExampleSyntaxError.name})
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(violations) == 1
|
||||
assert all(isinstance(v.rule, ExampleSyntaxError) for v in violations)
|
||||
assert violations[0].range == Range(Position(5, 8))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("suffix", [".md", ".mdx"])
|
||||
def test_example_syntax_error_markdown(index: SymbolIndex, suffix: str) -> None:
|
||||
code = """
|
||||
```python
|
||||
def g():
|
||||
```
|
||||
"""
|
||||
config = Config(select={ExampleSyntaxError.name})
|
||||
violations = lint_file(Path("test").with_suffix(suffix), code, config, index)
|
||||
assert len(violations) == 1
|
||||
assert all(isinstance(v.rule, ExampleSyntaxError) for v in violations)
|
||||
assert violations[0].range == Range(Position(2, 0))
|
||||
@@ -0,0 +1,54 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules import ExceptBoolOp
|
||||
|
||||
|
||||
def test_except_bool_op(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
# Bad - or in except
|
||||
try:
|
||||
pass
|
||||
except ValueError or KeyError:
|
||||
pass
|
||||
|
||||
# Bad - and in except
|
||||
try:
|
||||
pass
|
||||
except ValueError and KeyError:
|
||||
pass
|
||||
|
||||
# Bad - chained or
|
||||
try:
|
||||
pass
|
||||
except ValueError or KeyError or TypeError:
|
||||
pass
|
||||
|
||||
# Good - tuple syntax
|
||||
try:
|
||||
pass
|
||||
except (ValueError, KeyError):
|
||||
pass
|
||||
|
||||
# Good - single exception
|
||||
try:
|
||||
pass
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Good - bare except
|
||||
try:
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
"""
|
||||
config = Config(select={ExceptBoolOp.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert all(isinstance(r.rule, ExceptBoolOp) for r in results)
|
||||
assert [r.range for r in results] == [
|
||||
Range(Position(4, 0)),
|
||||
Range(Position(10, 0)),
|
||||
Range(Position(16, 0)),
|
||||
]
|
||||
@@ -0,0 +1,34 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules.extraneous_docstring_param import ExtraneousDocstringParam
|
||||
|
||||
|
||||
def test_extraneous_docstring_param(index: SymbolIndex) -> None:
|
||||
code = '''
|
||||
def bad_function(param1: str) -> None:
|
||||
"""
|
||||
Example function docstring.
|
||||
|
||||
Args:
|
||||
param1: First parameter
|
||||
param2: This parameter doesn't exist in function signature
|
||||
param3: Another non-existent parameter
|
||||
"""
|
||||
|
||||
def good_function(param1: str, param2: int) -> None:
|
||||
"""
|
||||
Good function with matching parameters.
|
||||
|
||||
Args:
|
||||
param1: First parameter
|
||||
param2: Second parameter
|
||||
"""
|
||||
'''
|
||||
config = Config(select={ExtraneousDocstringParam.name})
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(violations) == 1
|
||||
assert all(isinstance(v.rule, ExtraneousDocstringParam) for v in violations)
|
||||
assert violations[0].range == Range(Position(1, 0))
|
||||
@@ -0,0 +1,82 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules import ForbiddenDeprecationWarning
|
||||
|
||||
|
||||
def test_forbidden_deprecation_warning(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
import warnings
|
||||
|
||||
# Bad - should be flagged
|
||||
warnings.warn("message", category=DeprecationWarning)
|
||||
warnings.warn(
|
||||
"multiline message",
|
||||
category=DeprecationWarning,
|
||||
stacklevel=2
|
||||
)
|
||||
|
||||
# Good - should not be flagged
|
||||
warnings.warn("message", category=FutureWarning)
|
||||
warnings.warn("message", category=UserWarning)
|
||||
warnings.warn("message") # no category specified
|
||||
warnings.warn("message", stacklevel=2) # no category specified
|
||||
other_function("message", category=DeprecationWarning) # not warnings.warn
|
||||
"""
|
||||
config = Config(select={ForbiddenDeprecationWarning.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 2
|
||||
assert all(isinstance(r.rule, ForbiddenDeprecationWarning) for r in results)
|
||||
assert results[0].range == Range(Position(4, 34)) # First warnings.warn call
|
||||
assert results[1].range == Range(Position(7, 13)) # Second warnings.warn call
|
||||
|
||||
|
||||
def test_forbidden_deprecation_warning_import_variants(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
import warnings
|
||||
from warnings import warn
|
||||
import warnings as w
|
||||
|
||||
# All of these should be flagged
|
||||
warnings.warn("message", category=DeprecationWarning)
|
||||
warn("message", category=DeprecationWarning)
|
||||
w.warn("message", category=DeprecationWarning)
|
||||
"""
|
||||
config = Config(select={ForbiddenDeprecationWarning.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 3
|
||||
assert all(isinstance(r.rule, ForbiddenDeprecationWarning) for r in results)
|
||||
|
||||
|
||||
def test_forbidden_deprecation_warning_parameter_order(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
import warnings
|
||||
|
||||
# Different parameter orders - should be flagged
|
||||
warnings.warn("message", category=DeprecationWarning)
|
||||
warnings.warn(category=DeprecationWarning, message="test")
|
||||
"""
|
||||
config = Config(select={ForbiddenDeprecationWarning.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 2
|
||||
assert all(isinstance(r.rule, ForbiddenDeprecationWarning) for r in results)
|
||||
|
||||
|
||||
def test_forbidden_deprecation_warning_positional_args(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
import warnings
|
||||
|
||||
# Positional arguments - should be flagged
|
||||
warnings.warn("message", DeprecationWarning)
|
||||
warnings.warn("message", DeprecationWarning, 2)
|
||||
|
||||
# Good - should not be flagged
|
||||
warnings.warn("message", FutureWarning)
|
||||
warnings.warn("message") # no category specified
|
||||
"""
|
||||
config = Config(select={ForbiddenDeprecationWarning.name})
|
||||
results = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(results) == 2
|
||||
assert all(isinstance(r.rule, ForbiddenDeprecationWarning) for r in results)
|
||||
@@ -0,0 +1,86 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import lint_file
|
||||
from clint.rules.forbidden_make_judge_in_builtin_scorers import (
|
||||
ForbiddenMakeJudgeInBuiltinScorers,
|
||||
)
|
||||
|
||||
|
||||
def test_forbidden_make_judge_in_builtin_scorers(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
from mlflow.genai.judges.make_judge import make_judge
|
||||
from mlflow.genai.judges import InstructionsJudge
|
||||
|
||||
# BAD - direct call after import
|
||||
judge1 = make_judge(name="test", instructions="test")
|
||||
|
||||
# BAD - module qualified call
|
||||
from mlflow.genai import judges
|
||||
judge2 = judges.make_judge(name="test", instructions="test")
|
||||
|
||||
# GOOD - using InstructionsJudge instead
|
||||
judge3 = InstructionsJudge(name="test", instructions="test")
|
||||
"""
|
||||
config = Config(select={ForbiddenMakeJudgeInBuiltinScorers.name})
|
||||
violations = lint_file(Path("builtin_scorers.py"), code, config, index)
|
||||
|
||||
# Should detect: 1 import + 2 calls = 3 violations
|
||||
assert len(violations) == 3
|
||||
assert all(isinstance(v.rule, ForbiddenMakeJudgeInBuiltinScorers) for v in violations)
|
||||
|
||||
|
||||
def test_make_judge_allowed_in_other_files(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
from mlflow.genai.judges.make_judge import make_judge
|
||||
|
||||
# GOOD - allowed in other files
|
||||
judge = make_judge(name="test", instructions="test")
|
||||
"""
|
||||
config = Config(select={ForbiddenMakeJudgeInBuiltinScorers.name})
|
||||
violations = lint_file(Path("some_other_file.py"), code, config, index)
|
||||
|
||||
# Should NOT trigger in other files
|
||||
assert len(violations) == 0
|
||||
|
||||
|
||||
def test_instructions_judge_not_flagged(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
from mlflow.genai.judges import InstructionsJudge
|
||||
|
||||
# GOOD - InstructionsJudge is the correct approach
|
||||
judge = InstructionsJudge(name="test", instructions="test")
|
||||
"""
|
||||
config = Config(select={ForbiddenMakeJudgeInBuiltinScorers.name})
|
||||
violations = lint_file(Path("builtin_scorers.py"), code, config, index)
|
||||
|
||||
assert len(violations) == 0
|
||||
|
||||
|
||||
def test_nested_make_judge_call(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
from mlflow.genai.judges.make_judge import make_judge
|
||||
|
||||
# BAD - nested call
|
||||
result = some_function(make_judge(name="test", instructions="test"))
|
||||
"""
|
||||
config = Config(select={ForbiddenMakeJudgeInBuiltinScorers.name})
|
||||
violations = lint_file(Path("builtin_scorers.py"), code, config, index)
|
||||
|
||||
# Should detect: 1 import + 1 call = 2 violations
|
||||
assert len(violations) == 2
|
||||
assert all(isinstance(v.rule, ForbiddenMakeJudgeInBuiltinScorers) for v in violations)
|
||||
|
||||
|
||||
def test_make_judge_in_comment_not_flagged(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
from mlflow.genai.judges import InstructionsJudge
|
||||
|
||||
# This comment mentions make_judge but should not trigger
|
||||
judge = InstructionsJudge(name="test", instructions="test")
|
||||
"""
|
||||
config = Config(select={ForbiddenMakeJudgeInBuiltinScorers.name})
|
||||
violations = lint_file(Path("builtin_scorers.py"), code, config, index)
|
||||
|
||||
assert len(violations) == 0
|
||||
@@ -0,0 +1,33 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules.forbidden_set_active_model_usage import ForbiddenSetActiveModelUsage
|
||||
|
||||
|
||||
def test_forbidden_set_active_model_usage(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
import mlflow
|
||||
|
||||
# Bad
|
||||
mlflow.set_active_model("model_name")
|
||||
|
||||
# Good
|
||||
mlflow._set_active_model("model_name")
|
||||
|
||||
# Bad - with aliasing
|
||||
from mlflow import set_active_model
|
||||
set_active_model("model_name")
|
||||
|
||||
# Good - with aliasing
|
||||
from mlflow import _set_active_model
|
||||
_set_active_model("model_name")
|
||||
"""
|
||||
config = Config(select={ForbiddenSetActiveModelUsage.name})
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(violations) == 3
|
||||
assert all(isinstance(v.rule, ForbiddenSetActiveModelUsage) for v in violations)
|
||||
assert violations[0].range == Range(Position(4, 0)) # mlflow.set_active_model call
|
||||
assert violations[1].range == Range(Position(10, 0)) # from mlflow import set_active_model
|
||||
assert violations[2].range == Range(Position(11, 0)) # direct set_active_model call
|
||||
@@ -0,0 +1,45 @@
|
||||
from pathlib import Path
|
||||
|
||||
from clint.config import Config
|
||||
from clint.index import SymbolIndex
|
||||
from clint.linter import Position, Range, lint_file
|
||||
from clint.rules.forbidden_top_level_import import ForbiddenTopLevelImport
|
||||
|
||||
|
||||
def test_forbidden_top_level_import(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
# Bad
|
||||
import foo
|
||||
from foo import bar
|
||||
|
||||
# Good
|
||||
import baz
|
||||
"""
|
||||
config = Config(
|
||||
select={ForbiddenTopLevelImport.name},
|
||||
forbidden_top_level_imports={"*": ["foo"]},
|
||||
)
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
assert len(violations) == 2
|
||||
assert all(isinstance(v.rule, ForbiddenTopLevelImport) for v in violations)
|
||||
assert violations[0].range == Range(Position(2, 0))
|
||||
assert violations[1].range == Range(Position(3, 0))
|
||||
|
||||
|
||||
def test_nested_if_in_type_checking_block(index: SymbolIndex) -> None:
|
||||
code = """
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
if True:
|
||||
pass
|
||||
import databricks # Should NOT be flagged
|
||||
from databricks import foo # Should NOT be flagged
|
||||
"""
|
||||
config = Config(
|
||||
select={ForbiddenTopLevelImport.name},
|
||||
forbidden_top_level_imports={"*": ["databricks"]},
|
||||
)
|
||||
violations = lint_file(Path("test.py"), code, config, index)
|
||||
# Should have no violations since imports are inside TYPE_CHECKING
|
||||
assert len(violations) == 0
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user