From 11e62b75b0d59cd038aaaaed92056b90a5e720e8 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Sat, 7 Dec 2024 19:41:47 -0800
Subject: [PATCH] Set up pre-commit

---
 .github/workflows/pre-commit.yml      | 11 +++++++++++
 .pre-commit-config.yaml               | 23 +++++++++++++++++++++++
 .ruff.toml                            |  7 +++++++
 actions/load-base-env-vars/action.yml |  2 +-
 containers/docker_build.py            |  2 +-
 5 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/pre-commit.yml
 create mode 100644 .pre-commit-config.yaml
 create mode 100644 .ruff.toml

diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
new file mode 100644
index 0000000..82ed202
--- /dev/null
+++ b/.github/workflows/pre-commit.yml
@@ -0,0 +1,11 @@
+name: pre-commit
+
+on: [push, pull_request]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+      - uses: pre-commit/action@v3.0.1
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..4dd38c6
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,23 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v5.0.0
+  hooks:
+  - id: trailing-whitespace
+  - id: end-of-file-fixer
+  - id: check-yaml
+  - id: check-added-large-files
+    args: ["--maxkb=4000"]
+- repo: https://github.com/psf/black
+  rev: 24.10.0
+  hooks:
+  - id: black
+- repo: https://github.com/pycqa/isort
+  rev: 5.13.2
+  hooks:
+  - id: isort
+    args: ["--profile", "black", "--filter-files"]
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.8.2
+  hooks:
+    - id: ruff
+      args: ["--config", ".ruff.toml"]
diff --git a/.ruff.toml b/.ruff.toml
new file mode 100644
index 0000000..8580674
--- /dev/null
+++ b/.ruff.toml
@@ -0,0 +1,7 @@
+line-length = 120
+
+[lint]
+select = [
+    # numpy 2.0 deprecations/removals
+    "NPY201",
+]
diff --git a/actions/load-base-env-vars/action.yml b/actions/load-base-env-vars/action.yml
index 6974301..8a76c33 100644
--- a/actions/load-base-env-vars/action.yml
+++ b/actions/load-base-env-vars/action.yml
@@ -21,4 +21,4 @@ runs:
         path: ./xgboost-devops
         sparse-checkout: actions
     - name: Load base env vars
-      uses: ./xgboost-devops/actions/impls/load-base-env-vars
\ No newline at end of file
+      uses: ./xgboost-devops/actions/impls/load-base-env-vars
diff --git a/containers/docker_build.py b/containers/docker_build.py
index 907cdd5..91a5657 100644
--- a/containers/docker_build.py
+++ b/containers/docker_build.py
@@ -115,7 +115,7 @@ def main(*, args: argparse.Namespace) -> None:
         required=True,
         help=(
             "Tag to assign to the newly built container, e.g. "
-            "492475357299.dkr.ecr.us-west-2.amazonaws.com/xgb-ci.gpu:master"
+            "492475357299.dkr.ecr.us-west-2.amazonaws.com/xgb-ci.gpu:main"
         ),
     )
     parser.add_argument(