test

Author: JenniferWang

.flake8

@@ -8,8 +8,9 @@ max-line-length = 120
 # N817 ignored because importing using acronyms is convention (DistributedDataParallel as DDP)
 # E731 allow usage of assigning lambda expressions
 # N803,N806 allow caps and mixed case in function params. This is to work with Triton kernel coding style.
+# E704 ignored to allow black's formatting of Protocol stub methods (def method(self) -> None: ...)
 ignore =
-    E203,E305,E402,E501,E721,E741,F405,F821,F841,F999,W503,W504,C408,E302,W291,E303,N812,N817,E731,N803,N806
+    E203,E305,E402,E501,E704,E721,E741,F405,F821,F841,F999,W503,W504,C408,E302,W291,E303,N812,N817,E731,N803,N806
     # shebang has extra meaning in fbcode lints, so I think it's not worth trying
     # to line this up with executable bit
     EXE001,

pyproject.toml

@@ -100,23 +100,9 @@ prerelease = "allow"
 environments = [
     "sys_platform == 'linux'",
 ]
-# override-dependencies = ["torch>2.7.1", "torchaudio>=2.7.1", "torchvision>=0.22.0"]
-
-[tool.ufmt]
-formatter = "ruff-api"
-sorter = "usort"
 
 [tool.black]
-target-version = ["py310"]
+target-version = ["py310"] # match the minium supported python version
 
 [tool.usort]
 first_party_detection = false
-
-[tool.ruff]
-target-version = "py310"
-
-[tool.ruff.lint.isort]
-case-sensitive = false
-combine-as-imports = true
-detect-same-package = false
-order-by-type = false

src/forge/observability/perf_tracker.py

@@ -276,9 +276,9 @@ def __init__(self, max_workers: int = 2) -> None:
         if not torch.cuda.is_available():
             raise RuntimeError("CUDA is not available for timing")
         self._executor = ThreadPoolExecutor(max_workers=max_workers)
-        self._futures: list[
-            tuple[str, Future[float], int]
-        ] = []  # (name, future, submission_index)
+        self._futures: list[tuple[str, Future[float], int]] = (
+            []
+        )  # (name, future, submission_index)
         self._durations: list[tuple[str, float]] = []
         self._chain_start: torch.cuda.Event | None = None
 

tests/unit_tests/datasets/test_hf.py

@@ -229,12 +229,12 @@ def test_shuffling_behavior(self, dataset_factory, small_dataset_file):
         ), f"Shuffled epochs should be shuffled differently, got {first_epoch_ids} and {second_epoch_ids}"
 
         # But should contain the same set of IDs
-        assert (
-            set(first_epoch_ids) == set(range(SMALL_DATASET_SIZE))
-        ), f"First epoch samples should be (0-{SMALL_DATASET_SIZE-1}), got {first_epoch_ids}"
-        assert (
-            set(second_epoch_ids) == set(range(SMALL_DATASET_SIZE))
-        ), f"Second epoch samples should be (0-{SMALL_DATASET_SIZE-1}), got {second_epoch_ids}"
+        assert set(first_epoch_ids) == set(
+            range(SMALL_DATASET_SIZE)
+        ), f"First epoch samples should be (0-{SMALL_DATASET_SIZE - 1}), got {first_epoch_ids}"
+        assert set(second_epoch_ids) == set(
+            range(SMALL_DATASET_SIZE)
+        ), f"Second epoch samples should be (0-{SMALL_DATASET_SIZE - 1}), got {second_epoch_ids}"
 
     def test_epoch_tracking(self, dataset_factory, small_dataset_file):
         """Test that epoch number is correctly tracked across dataset restarts."""

tests/unit_tests/observability/conftest.py

@@ -47,9 +47,10 @@ def mock_rank():
 @pytest.fixture
 def mock_actor_context():
     """Mock Monarch actor context for testing actor name generation."""
-    with patch("forge.observability.metrics.context") as mock_context, patch(
-        "forge.observability.metrics.current_rank"
-    ) as mock_rank:
+    with (
+        patch("forge.observability.metrics.context") as mock_context,
+        patch("forge.observability.metrics.current_rank") as mock_rank,
+    ):
         # Setup mock context
         ctx = MagicMock()
         actor_instance = MagicMock()

tests/unit_tests/observability/test_metrics.py

@@ -94,10 +94,11 @@ async def test_backend_role_usage(self):
         )
 
         # Mock all the WandB init methods to focus only on role validation
-        with patch.object(wandb_backend, "_init_global"), patch.object(
-            wandb_backend, "_init_shared_global"
-        ), patch.object(wandb_backend, "_init_shared_local"), patch.object(
-            wandb_backend, "_init_per_rank"
+        with (
+            patch.object(wandb_backend, "_init_global"),
+            patch.object(wandb_backend, "_init_shared_global"),
+            patch.object(wandb_backend, "_init_shared_local"),
+            patch.object(wandb_backend, "_init_per_rank"),
         ):
             # Should not raise error for valid roles (type system prevents invalid values)
             await wandb_backend.init(role=BackendRole.GLOBAL)
@@ -433,11 +434,19 @@ async def _test_fetcher_registration(self, env_var_value, should_register_fetche
 
         # Assert based on expected behavior
         if should_register_fetchers:
-            assert proc_has_fetcher, f"Expected process to have _local_fetcher when FORGE_DISABLE_METRICS={env_var_value}"
-            assert global_has_fetcher, f"Expected global logger to have fetcher registered when FORGE_DISABLE_METRICS={env_var_value}"
+            assert (
+                proc_has_fetcher
+            ), f"Expected process to have _local_fetcher when FORGE_DISABLE_METRICS={env_var_value}"
+            assert (
+                global_has_fetcher
+            ), f"Expected global logger to have fetcher registered when FORGE_DISABLE_METRICS={env_var_value}"
         else:
-            assert not proc_has_fetcher, f"Expected process to NOT have _local_fetcher when FORGE_DISABLE_METRICS={env_var_value}"
-            assert not global_has_fetcher, f"Expected global logger to NOT have fetcher registered when FORGE_DISABLE_METRICS={env_var_value}"
+            assert (
+                not proc_has_fetcher
+            ), f"Expected process to NOT have _local_fetcher when FORGE_DISABLE_METRICS={env_var_value}"
+            assert (
+                not global_has_fetcher
+            ), f"Expected global logger to NOT have fetcher registered when FORGE_DISABLE_METRICS={env_var_value}"
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(

tests/unit_tests/observability/test_perf_tracker.py

@@ -389,8 +389,9 @@ def test_metric_timer_uses_gpu_override(
         if env_value == "true" and not torch.cuda.is_available():
             pytest.skip("CUDA not available")
 
-        with patch("torch.cuda.is_available", return_value=True), patch(
-            "forge.observability.perf_tracker.record_metric"
+        with (
+            patch("torch.cuda.is_available", return_value=True),
+            patch("forge.observability.perf_tracker.record_metric"),
         ):
             monkeypatch.setenv(METRIC_TIMER_USES_GPU.name, env_value)