chore: switch from mypy to pyrefly (#675)

Signed-off-by: Terry Kong <terryk@nvidia.com>

Author: terrykong

.github/workflows/cicd-main.yml

@@ -139,31 +139,18 @@ jobs:
           uv venv
           uv run --group dev pre-commit install
           uv run --group dev pre-commit run --all-files --show-diff-on-failure --color=always
-      - name: Minimize uv cache
-        run: uv cache prune --ci
-
-  mypy-check:
-    name: Mypy check
-    needs: [pre-flight]
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
-      - name: Install uv
-        uses: astral-sh/setup-uv@v5
-        with:
-          version: "0.7.2"
-          enable-cache: true
-          prune-cache: false
-      # Faster than uv python install since it caches python alongside runner
-      - name: "Set up Python"
-        uses: actions/setup-python@v5
-        with:
-          python-version-file: ".python-version"
-      - name: Check mypy
+      # TODO: this is a temporary check and should be removed once we have 100% correctness
+      - name: Check if any files with zero errors not in whitelist
         run: |
-          uv venv
-          uv run --group test mypy nemo_rl examples
+          missing_count=0
+          for file in $(uv run --group dev pyrefly check $(git ls-files 'nemo_rl/**/*.py' 'examples/**/*.py' 'docs/*.py' 'tools/**/*.py') --output-format json | jq -r --slurpfile all_files <(git ls-files 'nemo_rl/**/*.py' 'examples/**/*.py' 'docs/*.py' 'tools/**/*.py' | jq -R -s 'split("\n")[:-1]') --arg pwd "$(pwd)/" '(.errors | group_by(.path) | map({(.[0].path | sub($pwd; "")): length}) | add // {}) as $error_counts | $all_files[0][] | . as $file | if ($error_counts[$file] // 0) == 0 then $file else empty end'); do
+            if ! fgrep -q "$file" pyrefly.toml; then
+              echo "File $file has zero errors but is not in pyrefly.toml in the 'project-includes' list. Please add it to this whitelist."
+              ((missing_count++))
+            fi
+          done
+          
+          exit $missing_count
       - name: Minimize uv cache
         run: uv cache prune --ci
 
@@ -319,8 +306,7 @@ jobs:
                 (
                   needs.pre-flight.outputs.test_level != 'none' &&
                   needs.sphinx-build.result == 'success' &&
-                  needs.tests.result == 'success' &&
-                  (needs.mypy-check.result == 'success' || true)
+                  needs.tests.result == 'success'
                 )
               )
             }}

.pre-commit-config.yaml

@@ -35,3 +35,9 @@ repos:
         files: '.*\/[^\/]*_[^\/]*\.md$'
         exclude: '^\.github/'
         types: [file]
+
+  - repo: https://github.com/facebook/pyrefly
+    rev: 0.24.2
+    hooks:
+      - id: pyrefly-typecheck
+        files: \.py$

examples/configs/dpo.yaml

@@ -156,6 +156,7 @@ logger:
   tensorboard_enabled: false
   mlflow_enabled: false  # Disable MLflow logging
   monitor_gpus: true  # If true, will monitor GPU usage and log to wandb and/or tensorboard
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: "dpo-dev"
     name: "dpo"

examples/configs/recipes/llm/dpo-llama3.1-8b-instruct-4n8g-fsdp2tp1.v2.yaml

@@ -80,6 +80,7 @@ logger:
   tensorboard_enabled: true
   mlflow_enabled: false
   monitor_gpus: true
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: nemo-rl
     name: dpo-llama3.1-8b-instruct-4n8g-fsdp2tp1

examples/configs/recipes/llm/dpo-llama3.1-8b-instruct-4n8g-fsdp2tp2-quick.v2.yaml

@@ -80,6 +80,7 @@ logger:
   tensorboard_enabled: true
   mlflow_enabled: false
   monitor_gpus: true
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: nemo-rl
     name: dpo-llama3.1-8b-instruct-4n8g-fsdp2tp1

examples/configs/recipes/llm/dpo-llama3.1-8b-instruct-4n8g-megatron.yaml

@@ -113,6 +113,7 @@ logger:
   tensorboard_enabled: true
   mlflow_enabled: false
   monitor_gpus: true
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: nemo-rl
     name: dpo-llama3.1-8b-instruct-4n8g-fsdp2tp1

examples/configs/recipes/llm/dpo-llama3.1-8b-instruct-4n8g-megatrontp2pp2-quick.yaml

@@ -113,6 +113,7 @@ logger:
   tensorboard_enabled: true
   mlflow_enabled: false
   monitor_gpus: true
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: nemo-rl
     name: dpo-llama3.1-8b-instruct-4n8g-fsdp2tp1

examples/configs/recipes/llm/dpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v2.yaml

@@ -80,6 +80,7 @@ logger:
   tensorboard_enabled: true
   mlflow_enabled: false
   monitor_gpus: true
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: nemo-rl
     name: dpo-llama3.2-1b-instruct-1n8g-fsdp2tp1

examples/configs/recipes/llm/sft-llama3.1-8b-instruct-1n8g-fsdp2tp1-long.v2.yaml

@@ -60,6 +60,7 @@ logger:
   tensorboard_enabled: true
   mlflow_enabled: false
   monitor_gpus: true
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: nemo-rl
     name: sft-llama3.1-8b-instruct-1n8g-fsdp2tp1-long

examples/configs/recipes/llm/sft-llama3.1-8b-instruct-1n8g-fsdp2tp2sp.v2.yaml

@@ -60,6 +60,7 @@ logger:
   tensorboard_enabled: true
   mlflow_enabled: false
   monitor_gpus: true
+  num_val_samples_to_print: 0  # Number of validation samples to pretty print on terminal
   wandb:
     project: nemo-rl
     name: sft-llama3.1-8b-instruct-1n8g-fsdp2tp2sp