[BE] Enable UFMT on torchci in CI (#4700)

Adds UFMT to be run in CI on python files in `torchci`.  

Given that it wasn't being run in CI before this, I think it's ok to
make the include patterns more restrictive.

.github/workflows/lint.yml

@@ -42,7 +42,7 @@ jobs:
       - name: Run lintrunner on all files - Linux
         run: |
           set +e
-          if ! lintrunner -v --force-color --all-files --tee-json=lint.json --take ACTIONLINT,MYPY,RUSTFMT,COPYRIGHT,LINTRUNNER_VERSION; then
+          if ! lintrunner -v --force-color --all-files --tee-json=lint.json --take ACTIONLINT,MYPY,RUSTFMT,COPYRIGHT,LINTRUNNER_VERSION,UFMT; then
               echo ""
               echo -e "\e[1m\e[36mYou can reproduce these results locally by using \`lintrunner -m main\`.\e[0m"
               exit 1

.lintrunner.toml

@@ -241,7 +241,7 @@ init_command = [
 [[linter]]
 code = 'UFMT'
 include_patterns = [
-    '**/*.py',
+    'torchci/**/*.py',
 ]
 command = [
     'python3',

torchci/scripts/calculate_file_test_rating.py

@@ -94,33 +94,32 @@ def evaluate(failing_tests, merge_bases, rev_mapping, get_test_name_fn):
     print(f"# of failing tests: {len(all_failing_tests)}")
     print()
 
+
 def extract_test_file_name(test_row):
     return f"{test_row['invoking_file']}"
 
+
 def extract_test_class_name(test_row):
     if test_row["classname"]:
         return f"{test_row['invoking_file']}::{test_row['classname']}"
     else:
         return f"{test_row['invoking_file']}"
 
+
 def calculate_test_file_ratings(tests, merge_bases):
     # Should return a mapping of changed file -> failing test files -> confidence score
 
     return calculate_generic_test_ratings(
-        tests,
-        merge_bases,
-        get_test_name_fn=extract_test_file_name
-        )
+        tests, merge_bases, get_test_name_fn=extract_test_file_name
+    )
 
 
 def calculate_test_class_ratings(tests, merge_bases):
     # Should return a mapping of changed file -> failing test classes -> confidence score
 
     return calculate_generic_test_ratings(
-        tests,
-        merge_bases,
-        get_test_name_fn=extract_test_class_name
-        )
+        tests, merge_bases, get_test_name_fn=extract_test_class_name
+    )
 
 
 def calculate_generic_test_ratings(tests, merge_bases, get_test_name_fn):
@@ -136,7 +135,7 @@ def calculate_generic_test_ratings(tests, merge_bases, get_test_name_fn):
     # Make mapping of failing test -> changed file -> confidence score
     failing_tests_to_causes = {}
     for failing_test in failing_tests_to_sha:
-        score_dict = defaultdict(int) # changed file -> confidence score
+        score_dict = defaultdict(int)  # changed file -> confidence score
         for sha in failing_tests_to_sha[failing_test]:
             changed_files = merge_bases[sha]["changed_files"]
             for changed_file in changed_files:
@@ -147,7 +146,9 @@ def calculate_generic_test_ratings(tests, merge_bases, get_test_name_fn):
     rev_mapping = defaultdict(lambda: defaultdict(float))
     for failing_test in failing_tests_to_causes:
         for changed_file in failing_tests_to_causes[failing_test]:
-            rev_mapping[changed_file][failing_test] = failing_tests_to_causes[failing_test][changed_file]
+            rev_mapping[changed_file][failing_test] = failing_tests_to_causes[
+                failing_test
+            ][changed_file]
     return rev_mapping
 
 

torchci/scripts/compute_correlation.py

@@ -2,14 +2,12 @@
 import math
 import os
 from collections import defaultdict
-
 from typing import Any, List
 
 import pandas as pd
 from rockset import RocksetClient
 from sklearn.metrics.pairwise import pairwise_distances
 
-
 ROCKSET_API_KEY = os.environ.get("ROCKSET_API_KEY")
 if ROCKSET_API_KEY is None:
     raise RuntimeError("ROCKSET_API_KEY not set")

torchci/scripts/get_merge_base_info.py

@@ -3,14 +3,8 @@
 from pathlib import Path
 from typing import List
 
-from pathlib import Path
-
-from utils_td_heuristics import (
-    list_past_year_shas,
-    run_command,
-)
-
 from rockset_utils import query_rockset, remove_from_rockset, upload_to_rockset
+from utils_td_heuristics import list_past_year_shas, run_command
 
 REPO_ROOT = Path(__file__).resolve().parent.parent.parent
 

torchci/scripts/queue_alert.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, List, NamedTuple
 
 import rockset  # type: ignore[import]
+
 from check_alerts import clear_alerts, create_issue, fetch_alerts, update_issue
 from setuptools import distutils  # type: ignore[import]
 

torchci/scripts/reverts.py

@@ -7,7 +7,6 @@
 from typing import Dict, List, Optional, Tuple
 
 from rockset import RocksetClient  # type: ignore[import]
-
 from torchci.scripts.github_analyze import GitCommit, GitRepo  # type: ignore[import]
 
 # Should match the contents produced by trymerge on revert

torchci/scripts/rockset_utils.py

@@ -1,5 +1,5 @@
-from functools import lru_cache
 import os
+from functools import lru_cache
 from typing import Any, Dict, List, Optional
 
 import rockset

torchci/scripts/td_heuristic_profiling.py

@@ -1,7 +1,12 @@
 import json
 
 import requests
-from utils_td_heuristics import evaluate, get_filtered_failed_tests, get_merge_bases_dict
+
+from utils_td_heuristics import (
+    evaluate,
+    get_filtered_failed_tests,
+    get_merge_bases_dict,
+)
 
 
 def get_profiling_dict():

torchci/scripts/td_rockset_analysis.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+
 from rockset_utils import query_rockset
 
 REPO_ROOT = Path(__file__).resolve().parent.parent.parent

torchci/scripts/test_calculate_file_test_rating.py

@@ -1,6 +1,10 @@
 from unittest import main, TestCase
 
-from calculate_file_test_rating import calculate_test_file_ratings, calculate_test_class_ratings, filter_tests
+from calculate_file_test_rating import (
+    calculate_test_class_ratings,
+    calculate_test_file_ratings,
+    filter_tests,
+)
 
 
 class TestCalculateFileTestRating(TestCase):
@@ -75,7 +79,10 @@ def test_calculate_test_class_ratings(self):
                 self.gen_merge_base("head_sha_2", ["a.txt"], "merge_base_2"),
             ]
         )
-        expected = {"a.txt": {"invoking_file::classname": 1.5}, "b.txt": {"invoking_file::classname": 0.5}}
+        expected = {
+            "a.txt": {"invoking_file::classname": 1.5},
+            "b.txt": {"invoking_file::classname": 0.5},
+        }
         scores = calculate_test_class_ratings(tests, merge_bases)
         self.assertDictEqual(scores, expected)
 
@@ -89,8 +96,14 @@ def test_calculate_test_class_ratings(self):
             ]
         )
         expected = {
-            "a.txt": {"invoking_file_1::classname1": 0.5, "invoking_file_2::classname2": 0.5},
-            "b.txt": {"invoking_file_1::classname1": 0.5, "invoking_file_2::classname2": 0.5},
+            "a.txt": {
+                "invoking_file_1::classname1": 0.5,
+                "invoking_file_2::classname2": 0.5,
+            },
+            "b.txt": {
+                "invoking_file_1::classname1": 0.5,
+                "invoking_file_2::classname2": 0.5,
+            },
         }
         scores = calculate_test_class_ratings(tests, merge_bases)
         self.assertDictEqual(scores, expected)

torchci/scripts/test_check_alerts.py

@@ -13,7 +13,6 @@
     PYTORCH_ALERT_LABEL,
 )
 
-
 JOB_NAME = "periodic / linux-xenial-cuda10.2-py3-gcc7-slow-gradcheck / test (default, 2, 2, linux.4xlarge.nvidia.gpu)"
 DISABLED_JOB_NAMES = [
     "linux-focal-rocm5.3-py3.8-slow / test (slow, 1, 1, linux.rocm.gpu, rerun_disabled_tests)",

torchci/scripts/test_update_test_times.py

@@ -1,7 +1,7 @@
 import json
 import unittest
 
-from update_test_times import gen_test_file_times, gen_test_class_times
+from update_test_times import gen_test_class_times, gen_test_file_times
 
 
 class TestUpdateTestTimesFile(unittest.TestCase):
@@ -94,8 +94,16 @@ def test_gen_test_file_times_old_values_still_present(self) -> None:
 
 
 class TestUpdateTestTimesClass(unittest.TestCase):
-    def make_rockset_row(self, job: str, config: str, file: str, classname: str, time: float):
-        return {"base_name": job, "test_config": config, "file": file, "classname": classname, "time": time}
+    def make_rockset_row(
+        self, job: str, config: str, file: str, classname: str, time: float
+    ):
+        return {
+            "base_name": job,
+            "test_config": config,
+            "file": file,
+            "classname": classname,
+            "time": time,
+        }
 
     def test_gen_test_class_times_create_default(self) -> None:
         data = [
@@ -141,7 +149,10 @@ def test_gen_test_class_times_override_default(self) -> None:
         ]
         res = gen_test_class_times(data, {})
         expected = {
-            "default": {"config": {"a": {"classa": 3.5}}, "default": {"a": {"classa": 4.0}}},
+            "default": {
+                "config": {"a": {"classa": 3.5}},
+                "default": {"a": {"classa": 4.0}},
+            },
             "job": {"config": {"a": {"classa": 6}}},
         }
         self.assertDictEqual(res, expected)
@@ -155,7 +166,10 @@ def test_gen_test_class_times_override_old_default(self) -> None:
         ]
         res = gen_test_class_times(data, {"default": {"config": {"a": {"classa": 57}}}})
         expected = {
-            "default": {"config": {"a": {"classa": 3.5}}, "default": {"a": {"classa": 4.0}}},
+            "default": {
+                "config": {"a": {"classa": 3.5}},
+                "default": {"a": {"classa": 4.0}},
+            },
             "job": {"config": {"a": {"classa": 6}}},
         }
         self.assertDictEqual(res, expected)
@@ -164,10 +178,19 @@ def test_gen_test_class_times_override_old_default(self) -> None:
             self.make_rockset_row("env", "config", "a", "classa", 1),
         ]
         res = gen_test_class_times(
-            data, {"default": {"config": {"a": {"classa": 57}}, "default": {"a": {"classa": 100}}}}
+            data,
+            {
+                "default": {
+                    "config": {"a": {"classa": 57}},
+                    "default": {"a": {"classa": 100}},
+                }
+            },
         )
         expected = {
-            "default": {"config": {"a": {"classa": 1.0}}, "default": {"a": {"classa": 1.0}}},
+            "default": {
+                "config": {"a": {"classa": 1.0}},
+                "default": {"a": {"classa": 1.0}},
+            },
             "env": {"config": {"a": {"classa": 1}}},
         }
         self.assertDictEqual(res, expected)
@@ -179,7 +202,10 @@ def test_gen_test_class_times_old_values_still_present(self) -> None:
         res = gen_test_class_times(data, {"env": {"config": {"b": {"classb": 57}}}})
         expected = {
             "env": {"config": {"b": {"classb": 57}, "a": {"classa": 5}}},
-            "default": {"config": {"a": {"classa": 5.0}}, "default": {"a": {"classa": 5.0}}},
+            "default": {
+                "config": {"a": {"classa": 5.0}},
+                "default": {"a": {"classa": 5.0}},
+            },
         }
         self.assertDictEqual(res, expected)