Enable pyrefly on all libs (#2425)

Author: ravwojdyla

.pyrefly-baseline.json

@@ -571,7 +571,14 @@ class PrecommitConfig:
         ],
     ),
     PrecommitConfig(
-        patterns=["lib/marin/src/**/*.py", "lib/levanter/src/**/*.py"],
+        patterns=[
+            "lib/marin/src/**/*.py",
+            "lib/levanter/src/**/*.py",
+            "lib/haliax/src/**/*.py",
+            "lib/fray/src/**/*.py",
+            "lib/iris/src/**/*.py",
+            "lib/zephyr/src/**/*.py",
+        ],
         checks=[
             check_pyrefly,
         ],

infra/pre-commit.py

@@ -14,6 +14,8 @@
 
 """Ray-based cluster implementation."""
 
+from typing import Any, cast
+
 import asyncio
 import logging
 import os
@@ -307,8 +309,8 @@ def _get_runtime_env(self, request: JobRequest) -> dict | None:
         logger.info("Ray runtime env: %s", runtime_env)
         return runtime_env
 
-    def _get_entrypoint_params(self, request: JobRequest) -> dict:
-        params = {}
+    def _get_entrypoint_params(self, request: JobRequest) -> dict[str, Any]:
+        params: dict[str, Any] = {}
 
         if request.resources.cpu > 0:
             params["entrypoint_num_cpus"] = float(request.resources.cpu)
@@ -333,7 +335,7 @@ def monitor(self, job_id: JobId) -> JobInfo:
         job = self._jobs[job_id]
         if job.submission_id is None:
             logger.info("Job is a remote ref, monitoring is automatic, waiting.")
-            return self.wait(job_id)
+            return cast(JobInfo, self.wait(job_id))
 
         async def stream_logs():
             async for line in self._job_client().tail_job_logs(job_id):

lib/fray/src/fray/cluster/ray/cluster.py

@@ -21,7 +21,7 @@
 import logging
 import os
 import threading
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Iterator
 from concurrent.futures import Future, ThreadPoolExecutor, wait
 from contextlib import contextmanager
 from contextvars import ContextVar
@@ -126,7 +126,7 @@ class ActorHandle:
     Provides a unified interface for calling actor methods with .remote() and .call().
     """
 
-    def __getattr__(self, method_name: str):
+    def __getattr__(self, method_name: str) -> "ActorMethod":
         """Get a callable method wrapper for the actor."""
         raise NotImplementedError
 
@@ -150,7 +150,7 @@ def __init__(self, instance: Any, lock: threading.Lock, context):
         self._lock = lock  # Serializes all method calls
         self._context = context
 
-    def __getattr__(self, method_name: str):
+    def __getattr__(self, method_name: str) -> "ThreadActorMethod":
         method = getattr(self._instance, method_name)
         return ThreadActorMethod(method, self._lock, self._context)
 
@@ -183,7 +183,7 @@ class _ImmediateFuture:
 
     def __init__(self, result: Any):
         self._result = result
-        self._iterator = None
+        self._iterator: Iterator[Any] | None = None
 
     def result(self) -> Any:
         return self._result
@@ -206,7 +206,7 @@ class GeneratorFuture:
 
     def __init__(self, future: Future):
         self._future = future
-        self._iterator = None
+        self._iterator: Iterator[Any] | None = None
 
     def result(self) -> Any:
         """Get the underlying result from the future."""
@@ -422,7 +422,7 @@ def create_actor(
         num_cpus: float | None = None,
         **kwargs,
     ) -> ActorHandle:
-        options = {}
+        options: dict[str, Any] = {}
         if name is not None:
             options["name"] = name
         options["get_if_exists"] = get_if_exists

lib/fray/src/fray/job/context.py

@@ -77,7 +77,7 @@ def __init__(self, host: str = "0.0.0.0", port: int = 9999):
 
         config = uvicorn.Config(self.app, host=host, port=port, log_level="error", access_log=False)
         self.server = uvicorn.Server(config)
-        self.server_thread = None
+        self.server_thread: ServerThread | None = None
 
     def _create_app(self) -> FastAPI:
         """Create FastAPI app with namespaced queue endpoints."""

lib/fray/src/fray/queue/http.py

@@ -124,12 +124,11 @@ def __init__(
         self._config = config
 
         # Setup cache directory
+        self._temp_dir: tempfile.TemporaryDirectory[str] | None = None
         if cache_dir:
             self._cache_dir = cache_dir
-            self._temp_dir = None
         elif config.cache_dir:
             self._cache_dir = config.cache_dir
-            self._temp_dir = None
         else:
             # Create temporary cache
             self._temp_dir = tempfile.TemporaryDirectory(prefix="worker_cache_")

lib/iris/src/iris/cluster/worker/worker.py

@@ -20,7 +20,7 @@
 import re
 from collections.abc import Callable, Iterable, Iterator
 from dataclasses import dataclass, field
-from typing import Any, Generic, Literal, TypeVar
+from typing import Any, Generic, Literal, TypeVar, cast
 
 import fsspec
 from braceexpand import braceexpand
@@ -830,7 +830,7 @@ def reduce(
             4950
         """
         if global_reducer is None:
-            global_reducer = local_reducer
+            global_reducer = cast(Callable[[Iterator[R]], R], local_reducer)
 
         return Dataset(self.source, [*self.operations, ReduceOp(local_reducer, global_reducer)])
 

lib/zephyr/src/zephyr/dataset.py

@@ -51,6 +51,7 @@
     WindowOp,
     WriteOp,
 )
+from zephyr.expr import Expr
 from zephyr.readers import InputFileSpec
 
 if TYPE_CHECKING:
@@ -512,7 +513,7 @@ def _compute_file_pushdown(
     Returns:
         Tuple of (source_items, remaining_operations), where filter/select have been pushed down.
     """
-    filter_expr = None
+    filter_expr: Expr | None = None
     select_columns = load_op.columns
     ops_to_skip: set[int] = set()
 
@@ -882,7 +883,7 @@ class StageContext:
     total_shards: int
     chunk_size: int
     aux_shards: dict[int, list[Any]] = field(default_factory=dict)
-    execution_context: JobContext = None
+    execution_context: JobContext | None = None
 
     def get_right_shard(self, op_index: int) -> Any:
         """Get right shard for join at given op index.

lib/zephyr/src/zephyr/plan.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from dataclasses import asdict, dataclass, is_dataclass
+from dataclasses import asdict, is_dataclass
 import itertools
 import os
 from collections.abc import Iterable
@@ -129,7 +129,7 @@ def infer_parquet_type(value):
         return pa.string()
 
 
-def infer_parquet_schema(record: dict | dataclass):
+def infer_parquet_schema(record: dict[str, Any] | Any):
     """Infer PyArrow schema from a dictionary record."""
     import pyarrow as pa
 

lib/zephyr/src/zephyr/writers.py

@@ -81,7 +81,14 @@ exclude = ["marin/", "scripts/"]
 
 # Pyrefly type checker configuration
 [tool.pyrefly]
-project-includes = ["lib/marin/src/**/*.py", "lib/levanter/src/**/*.py"]
+project-includes = [
+    "lib/marin/src/**/*.py",
+    "lib/levanter/src/**/*.py",
+    "lib/haliax/src/**/*.py",
+    "lib/fray/src/**/*.py",
+    "lib/iris/src/**/*.py",
+    "lib/zephyr/src/**/*.py",
+]
 
 # Explicitly tell Pyrefly where our editable packages live so it resolves imports
 # against the library sources instead of the top-level `src` directory, which only
@@ -92,6 +99,7 @@ search-path = [
     "lib/haliax/src",
     "lib/zephyr/src",
     "lib/fray/src",
+    "lib/iris/src",
 ]
 disable-search-path-heuristics = true
 
@@ -107,6 +115,7 @@ project-excludes = [
     "examples/**",           # Example code doesn't need strict typing
     "lib/**/crawl/**",       # Crawl scripts have library typing issues with smart_open
     "lib/marin/src/marin/processing/classification/deduplication/vendor", # Exclude vendor stuff in dedupe
+    "lib/iris/src/iris/rpc/*_pb2*",  # Generated protobuf files
 ]
 
 # Disable specific error codes that are primarily noise from missing type stubs