[Pre py3.12 Upgrade] Update black pylint

GitOrigin-RevId: fa61c32

Author: matthew-e-hopkins

.pylintrc

@@ -111,6 +111,7 @@ disable=abstract-method,
         no-member, # Config objects use dynamic members generated by attrs.
         no-name-in-module,
         no-self-use,
+        not-callable,
         nonzero-method,
         oct-method,
         old-division,
@@ -163,6 +164,8 @@ disable=abstract-method,
         xrange-builtin,
         zip-builtin-not-iterating,
 
+# Increase default from 5 to 9
+max-positional-arguments=9
 
 [REPORTS]
 

axlearn/audio/decoder_asr.py

@@ -1013,7 +1013,7 @@ def tokens_to_scores(
 
         return tokens_to_scores
 
-    def beam_search_decode(
+    def beam_search_decode(  # pytype: disable=signature-mismatch
         self,
         input_batch: Nested[Tensor],
         num_decodes: int,
@@ -1225,7 +1225,7 @@ def _validate_decode_batch(self, input_batch: Nested[Tensor]):
         if "prefix" not in input_batch:
             raise ValueError("Input batch is expected to contain `prefix`.")
 
-    def beam_search_decode(
+    def beam_search_decode(  # pytype: disable=signature-mismatch
         self,
         input_batch: Nested[Tensor],
         num_decodes: int,

axlearn/audio/evaler_asr_test.py

@@ -302,13 +302,19 @@ def test_brevity_penalty(self):
 
         # Without brevity penalty, outputs match exactly.
         outputs = _compute_metrics(hypotheses, references, vocab_file=_1K_VOCAB_FILE)
-        self.assertNestedAllClose(outputs["word_errors/wer"].mean, 0.0)
-        self.assertNestedAllClose(outputs["word_errors/sentence_accuracy"].mean, 1.0)
+        self.assertNestedAllClose(
+            outputs["word_errors/wer"].mean, 0.0  # pytype: disable=attribute-error
+        )
+        self.assertNestedAllClose(
+            outputs["word_errors/sentence_accuracy"].mean, 1.0  # pytype: disable=attribute-error
+        )
 
         # With brevity penalty, DummyModel should emit fewer tokens.
         brevity_penalty = brevity_penalty_fn(alpha=0.9, bp_type="t5")
         outputs = _compute_metrics(
             hypotheses, references, brevity_penalty=brevity_penalty, vocab_file=_1K_VOCAB_FILE
         )
-        self.assertGreater(outputs["word_errors/wer"].mean, 0.0)
-        self.assertNestedAllClose(outputs["word_errors/sentence_accuracy"].mean, 0.0)
+        self.assertGreater(outputs["word_errors/wer"].mean, 0.0)  # pytype: disable=attribute-error
+        self.assertNestedAllClose(
+            outputs["word_errors/sentence_accuracy"].mean, 0.0  # pytype: disable=attribute-error
+        )

axlearn/audio/model_asr.py

@@ -59,7 +59,7 @@ def predict(self, input_batch: Nested[Tensor]) -> Nested[Tensor]:
         )
         return dict(logits=logits)
 
-    def forward(
+    def forward(  # pytype: disable=signature-mismatch
         self, input_batch: Nested[Tensor], *, return_aux: bool = False
     ) -> tuple[Tensor, Nested[Tensor]]:
         """Computes loss and predictions (such as logits) in auxiliary outputs.
@@ -95,7 +95,7 @@ def forward(
         )
         return loss, aux_outputs if return_aux else {}
 
-    def beam_search_decode(
+    def beam_search_decode(  # pytype: disable=signature-mismatch
         self, input_batch: Nested[Tensor], num_decodes: int, **kwargs
     ) -> DecodeOutputs:
         """Performs beam search decoding.

axlearn/cloud/common/scheduler_test.py

@@ -210,7 +210,9 @@ def test_validate_tiers(self):
             sched.schedule(resource_limits=[], **common_kwargs)
 
         with self.assertRaisesRegex(ValueError, "sequence"):
-            sched.schedule(resource_limits={"v4": 10, "v3": 3}, **common_kwargs)
+            sched.schedule(
+                resource_limits={"v4": 10, "v3": 3}, **common_kwargs
+            )  # pytype: disable=wrong-arg-types
 
         sched.schedule(resource_limits=[{"v4": 10, "v3": 3}], **common_kwargs)
 
@@ -582,7 +584,7 @@ def test_schedule(
         {"unused_limits": [{"v4": 4}, {"v4": 8}]},
     )
     def test_schedule_result(self, unused_limits: Optional[Sequence[ResourceMap[int]]]):
-        schedule_result = BaseScheduler.ScheduleResults(
+        schedule_result = BaseScheduler.ScheduleResults(  # pytype: disable=wrong-arg-types
             project_limits={"a": {"v4": 5}, "b": {"unknown": 0, "v4": 1}},
             project_usages={"a": {"v4": 5}, "b": {"unknown": 0, "v4": 1}},
             job_verdicts={"a1": True, "a2": False, "b1": False, "b2": True},

axlearn/cloud/gcp/bundler_test.py

@@ -142,7 +142,7 @@ def test_wait_until_finished_is_no_op_if_async_false(self):
         # Should be a no-op if is_async=False.
         cfg = self._get_test_cloud_build_bundler()
 
-        with self._mock_status(None) as mock_status:
+        with self._mock_status(None) as mock_status:  # pytype: disable=wrong-arg-types
             b = cfg.set(is_async=False).instantiate()
             b.wait_until_finished("test-name")
             self.assertFalse(mock_status.called)
@@ -151,7 +151,7 @@ def test_wait_until_finished_is_called_with_happy_status_path(self):
         # Tests happy path: transitions from no status -> pending -> success.
         cfg = self._get_test_cloud_build_bundler()
 
-        with self._mock_status(
+        with self._mock_status(  # pytype: disable=wrong-arg-types
             None, CloudBuildStatus.PENDING, CloudBuildStatus.SUCCESS
         ) as mock_status:
             b = cfg.set(is_async=True).instantiate()
@@ -162,7 +162,7 @@ def test_wait_until_finished_raises_runtime_error_with_cloud_build_status_failur
         # Tests that we raise a runtime error if CloudBuildStatus.FAILURE status is returned.
         cfg = self._get_test_cloud_build_bundler()
 
-        with self._mock_status(
+        with self._mock_status(  # pytype: disable=wrong-arg-types
             None, CloudBuildStatus.PENDING, CloudBuildStatus.FAILURE
         ) as mock_status:
             b = cfg.set(is_async=True).instantiate()
@@ -177,7 +177,9 @@ def test_wait_until_finished_retries_with_runtime_error(self):
         # Tests that the query is retried if retrieving status fails with a RuntimeError.
         cfg = self._get_test_cloud_build_bundler()
 
-        with self._mock_status(RuntimeError("fake error"), CloudBuildStatus.SUCCESS) as mock_status:
+        with self._mock_status(
+            RuntimeError("fake error"), CloudBuildStatus.SUCCESS
+        ) as mock_status:  # pytype: disable=wrong-arg-types
             b = cfg.set(is_async=True).instantiate()
             b.wait_until_finished("test-name")
             self.assertEqual(2, mock_status.call_count)
@@ -189,7 +191,7 @@ def test_wait_until_finished_triggers_timeout(self):
         with mock.patch("time.perf_counter") as mock_perf_counter:
             mock_perf_counter.side_effect = [0, 10, 500, 3601]
 
-            with self._mock_status(
+            with self._mock_status(  # pytype: disable=wrong-arg-types
                 None, CloudBuildStatus.PENDING, CloudBuildStatus.PENDING
             ) as mock_status:
                 b = cfg.set(is_async=True).instantiate()

axlearn/cloud/gcp/jobs/launch.py

@@ -781,7 +781,7 @@ def _wrapped_usage(
 if __name__ == "__main__":
     # remove FLAGS injected by tensorflow and array record package during import
     # example: import grain.python as grain
-    # TODO: remove this after we stop using the global FLAG for axlearn flags
+    # TODO: remove this after we stop using the global FLAG for ajax/axlearn flags
     flags_to_remove = []
     modules_to_exclude = ["tensorflow", "grain", "array_record", "orbax"]
     for module_name, flags_list in FLAGS.flags_by_module_dict().items():

axlearn/cloud/gcp/jobs/launch_test.py

@@ -393,7 +393,7 @@ def test_list(self):
 class TestBastionManagedGKEJob(TestWithTemporaryCWD):
     """Tests BastionManagedGKEJob."""
 
-    def run(self, **kwargs):
+    def run(self, **kwargs):  # pytype: disable=signature-mismatch
         # Run tests under mock settings.
         self._settings = default_mock_settings()
         patch_name = mock.patch(f"{launch.__name__}.generate_job_name", return_value="job-name")
@@ -439,6 +439,8 @@ def run(self, **kwargs):
         ],
         action=_ACTIONS,
     )
+    # TODO: Try to reduce positional arguments
+    # pylint: disable-next=too-many-positional-arguments
     def test_tpu_flags(
         self,
         action,

axlearn/cloud/gcp/jobset_utils_test.py

@@ -190,6 +190,8 @@ def test_validate_jobset_name(self):
         additional_node_networks=[None, "network-1:subnet-1,network-2:subnet-2"],
         image_id=[None, "my-image-id"],
     )
+    # TODO: Try to reduce positional arguments
+    # pylint: disable-next=too-many-positional-arguments
     def test_build_pod(
         self,
         bundler_cls: type[Bundler],

axlearn/cloud/gcp/lws_utils.py

@@ -96,7 +96,7 @@ class TPULeaderWorkerTemplate(TPUJobBuilder):
 
     def __call__(self) -> Sequence[Nested[Any]]:
         system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type]
-        return dict(
+        return dict(  # pytype: disable=bad-return-type
             size=system.vms_per_slice,
             workerTemplate=self._build_pod(),
         )