From 5115df3f0be4f277c058ff5023c62652365659b7 Mon Sep 17 00:00:00 2001
From: Logan Adams <114770087+loadams@users.noreply.github.com>
Date: Mon, 26 Feb 2024 12:37:32 -0800
Subject: [PATCH] Add script to check for `--extra-index-url` (#5184)

Co-authored-by: Michael Wyatt <michaelwyatt@microsoft.com>
---
 .pre-commit-config.yaml        |  9 ++++++++
 scripts/check-extraindexurl.py | 42 ++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)
 create mode 100755 scripts/check-extraindexurl.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 2432a7a24..11a142568 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -78,3 +78,12 @@ repos:
         language: python
         exclude: ^(.github/workflows/|scripts/check-torchcuda.py|docs/_tutorials/accelerator-abstraction-interface.md|accelerator/cuda_accelerator.py|deepspeed/inference/engine.py|deepspeed/model_implementations/transformers/clip_encoder.py|deepspeed/model_implementations/diffusers/vae.py|deepspeed/model_implementations/diffusers/unet.py|op_builder/spatial_inference.py|op_builder/transformer_inference.py|op_builder/builder.py|setup.py|tests/unit/ops/sparse_attention/test_sparse_attention.py)
         # Specific deepspeed/ files are excluded for now until we wrap ProcessGroup in deepspeed.comm
+
+-   repo: local
+    hooks:
+    -   id: check-extraindexurl
+        name: check-extraindexurl
+        entry: ./scripts/check-extraindexurl.py
+        language: python
+        files: \.(yml|yaml|sh|py)$
+        exclude: ^(scripts/check-extraindexurl.py)
diff --git a/scripts/check-extraindexurl.py b/scripts/check-extraindexurl.py
new file mode 100755
index 000000000..01b506dc9
--- /dev/null
+++ b/scripts/check-extraindexurl.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from __future__ import annotations
+'''Copyright The Microsoft DeepSpeed Team'''
+"""
+Checks each file in sys.argv for the string "--extra-index-url".
+Modified from https://github.com/jlebar/pre-commit-hooks/blob/master/check_do_not_submit.py
+"""
+
+import subprocess
+import sys
+
+
+def err(s: str) -> None:
+    print(s, file=sys.stderr)
+
+
+print(*sys.argv[1:])
+
+# There are many ways we could search for the string "--extra-index-url", but `git
+# grep --no-index` is nice because
+#  - it's very fast (as compared to iterating over the file in Python)
+#  - we can reasonably assume it's available on all machines
+#  - unlike plain grep, which is slower and has different flags on MacOS versus
+#    Linux, git grep is always the same.
+res = subprocess.run(
+    ["git", "grep", "-Hn", "--no-index", "-e", r"--extra-index-url", *sys.argv[1:]],
+    capture_output=True,
+)
+if res.returncode == 0:
+    err('Error: The string "--extra-index-url" was found.\nPlease replace all calls to --extra-index-url with "--index-url"'
+        )
+    err(res.stdout.decode("utf-8"))
+    sys.exit(1)
+elif res.returncode == 2:
+    err(f"Error invoking grep on {', '.join(sys.argv[1:])}:")
+    err(res.stderr.decode("utf-8"))
+    sys.exit(2)