Code Review & Testing Checklist for AI & Machine Learning

Run TensorFlow Data Validation or Great Expectations in CI to assert column presence, types, ranges, and categorical domains. Catch silent breaking changes in upstream tables before training or feature generation.

Use DVC, LakeFS, or Delta Lake versioned tables with content hashes so every training run is traceable to exact data commits. Review PRs for updated data pointers and ensure lineage is reflected in experiment metadata.

Add tests that recompute train, validation, and test splits using scikit-learn GroupKFold or StratifiedKFold to verify no entity or time leakage. Fail PRs if distributions shifted or overlap is detected between folds.

Unit test Feast or Tecton feature definitions ensuring transformations match between batch and online retrieval. Compare point-in-time correct joins with known golden rows to prevent temporal leakage.

Use Hypothesis or pytest parameterization to generate edge-case rows that include NaNs, extreme values, and low-frequency classes. Assert preprocessing pipelines handle them deterministically without silent casting or clipping.

Run Microsoft Presidio, BigQuery DLP, or homemade regex scans to detect PII in datasets and logs. Block merges unless the transformation pipeline masks or tokenizes sensitive fields and unit tests verify redaction.

Create synthetic fixtures for SQL or Spark transforms that cover tie-breaking and null handling in joins, windows, and group-bys. Assert row counts and key-level metrics match expected golden outputs across engines.

Set seeds for NumPy, PyTorch, and CUDA, enable torch.use_deterministic_algorithms and configure cuDNN deterministic flags. Add a smoke test that ensures two runs on the same commit produce identical metrics and weights within tolerance.

Use TorchMetrics or scikit-learn to compute precision, recall, F1, and AUC with explicit micro or macro averaging, then compare against hand-checked mini datasets. Lock thresholding behavior and class ordering to avoid accidental shifts.

Run bootstrap confidence intervals or McNemar’s test for classification, and fail the PR if the new model underperforms a pinned baseline beyond a small epsilon. Store baselines as versioned artifacts and refresh them intentionally.

Add gradcheck and gradgradcheck for differentiable modules in PyTorch, and verify shape invariants and dtype propagation. For data loaders, assert shuffling, batching, and collation behavior with synthetic datasets under multiple seeds.

Prefer safetensors for PyTorch or ONNX for portable inference, and disable pickle loading in production paths. Add static analysis or tests that fail when pickle or eval-based deserialization is introduced.

Adopt Hydra or OmegaConf with explicit default lists and compose-only rules, then unit test that config merges do not change critical hyperparameters. Require PRs to document any config deltas in experiment logs.

Use MLflow or Weights & Biases to log git commit, requirements lockfile, data version IDs, and environment snapshots. Block merges if new runs lack complete metadata for reproducibility or if artifacts are missing signatures.

Check small, deterministic samples into the repo or store under DVC with content hashes. Assert model forward pass logits, serialization round-trip, and pipeline outputs match golden values within a tolerance window.

Execute tests across key Python versions and PyTorch or TensorFlow minors, and include CPU-only and GPU paths. Use pytest markers to skip heavy training in quick checks while preserving a nightly job for full coverage.

Use pytest-benchmark or asv to track training step time, latency, and memory, then fail PRs exceeding defined budgets. For inference, include cold and warm-start timings with representative batch sizes to catch micro-regressions.

Spin up ephemeral dependencies like DuckDB, Postgres, or MinIO in CI, and mock third-party APIs with responses or localstack for AWS. Validate end-to-end tasks in Airflow, Prefect, or Dagster using a minimal DAG against fixtures.

Enable pre-commit hooks to strip notebook outputs and run nbval or papermill smoke tests on critical notebooks. Fail PRs if notebooks rely on unstated environment variables or non-reproducible paths.

Base Dockerfiles on specific CUDA and cuDNN tags, run as non-root, and enable nvidia-container-toolkit where GPUs are required. Validate at build time that driver-toolkit compatibility is met and libraries resolve to exact versions.

Cache pip or poetry dependencies, Hugging Face models, and Torch hubs between runs, and store them per hash. Use path-based filters to run only impacted tests, and pull artifacts from DVC remote instead of recomputing.

Use promptfoo, TruLens, or OpenAI Evals to run prompts against a fixed evaluation set with deterministic seeds where possible. Gate merges on pass rates for rubric-scored outputs and specific fixtures representing hard cases.

Validate LLM outputs against Pydantic models, JSON Schema, or Guardrails to ensure exact keys, types, and enumerations. Add retry and repair logic tests that simulate malformed outputs and confirm graceful recovery.

Integrate Perspective API, Detoxify, or Azure Content Safety in tests with seeded prompts known to trigger edge cases. Set clear threshold guards and fail PRs that degrade safety scores or bypass filters.

Measure recall@k, MRR, and nDCG using LangSmith, LlamaIndex evals, or custom scripts on a canonical corpus. Fail PRs if chunking, indexing, or embedding changes reduce retrieval performance beyond defined tolerances.

Store prompts as templates with variables in version control and keep API keys in secret managers, never in code. Add tests that enforce placeholder presence, fallbacks, and redaction of secrets in logs.

Estimate tokens with tiktoken or tokenizer APIs and run load tests to measure end-to-end latency. Fail PRs that exceed per-request or per-batch budgets, and verify caching strategies via LangChain or custom middlewares.

Integrate gitleaks or git-secrets in CI to detect tokens, keys, and passwords in commits, configs, and notebooks. Block merges on findings and require rotation via your secret manager before unblocking.

Run Snyk, Trivy, or Grype on lockfiles and container images, and fail builds on high severity CVEs with defined exceptions documented in policy. Alert reviewers when inference images pull new base layers.

Generate Software Bill of Materials with Syft and flag incompatible licenses for transitive dependencies and pre-trained models. Require PRs to include dataset source, license, and usage scope before merging.

Test S3 or GCS policies for least privilege, enforce server-side encryption, and require KMS keys for training and inference artifacts. Use policy-as-code tools like Open Policy Agent to validate access rules in CI.

Run mypy for type safety on tensors and arrays, ruff for style and bug-prone patterns, and bandit to catch eval or pickle in critical paths. Add rules specific to deserialization and command execution in data pipelines.

Generate or update Model Cards and Datasheets for Datasets via templates when metrics or data sources change. Fail PRs if documentation is missing required sections like evaluation conditions or known limitations.