Write serialized options to a file to broadcast to workers (#21306)

Fixes #21305

In theory we can also send options over the socket, but they are needed
very early, so it is easier to read them from a file with how the code
is organized now.

Author: ilevkivskyi

mypy/build.py

@@ -18,7 +18,6 @@
 import gc
 import json
 import os
-import pickle
 import platform
 import re
 import stat
@@ -44,7 +43,6 @@
     final,
 )
 
-from librt.base64 import b64encode
 from librt.internal import (
     cache_version,
     read_bool,
@@ -281,7 +279,7 @@ def __init__(self, status_file: str, options_data: str, env: Mapping[str, str])
             "-m",
             "mypy.build_worker",
             f"--status-file={status_file}",
-            f'--options-data="{options_data}"',
+            f"--options-data={options_data}",
         ]
         # Return early without waiting, caller must call connect() before using the client.
         self.proc = subprocess.Popen(command, env=env)
@@ -389,10 +387,12 @@ def default_flush_errors(
     connect_threads = []
     # A quasi-unique ID for this specific mypy invocation.
     build_id = os.urandom(4).hex()
+    options_data = None
     if options.num_workers > 0:
-        # TODO: switch to something more efficient than pickle (also in the daemon).
-        pickled_options = pickle.dumps(options.snapshot())
-        options_data = b64encode(pickled_options).decode()
+        os.makedirs(options.cache_dir, exist_ok=True)
+        options_data = os_path_join(options.cache_dir, f".worker_options.{build_id}.data")
+        with open(options_data, "wb") as f:
+            f.write(options.to_bytes())
         workers = [
             WorkerClient(
                 f".mypy_worker.{build_id}.{idx}.json", options_data, worker_env or os.environ
@@ -448,6 +448,8 @@ def connect(wc: WorkerClient, data: bytes) -> None:
         # shut them down cleanly. Otherwise, they will linger until connection timeout.
         for thread in connect_threads:
             thread.join()
+        if options_data is not None:
+            os.unlink(options_data)
         for worker in workers:
             if not worker.connected:
                 continue
@@ -560,10 +562,11 @@ def build_inner(
 def warn_unused_configs(
     options: Options, flush_errors: Callable[[str | None, list[str], bool], None]
 ) -> None:
-    if options.warn_unused_configs and options.unused_configs and not options.non_interactive:
+    unused_configs = options.get_unused_configs()
+    if options.warn_unused_configs and unused_configs and not options.non_interactive:
         unused = get_config_module_names(
             options.config_file,
-            [glob for glob in options.per_module_options.keys() if glob in options.unused_configs],
+            [glob for glob in options.per_module_options.keys() if glob in unused_configs],
         )
         flush_errors(
             None, ["{}: note: unused section(s): {}".format(options.config_file, unused)], False

mypy/build_worker/worker.py

@@ -11,7 +11,7 @@
 * In a loop:
   - Receive an SCC id from coordinator, and start processing it.
   - SCC is processed in two phases: interface and implementation, send a response after each.
-  - When prompted by coordinator (with a scc_id=None message), cleanup and shutdown.
+  - When prompted by coordinator (with a scc_ids=[] message), cleanup and shutdown.
 """
 
 from __future__ import annotations
@@ -20,13 +20,11 @@
 import gc
 import json
 import os
-import pickle
 import platform
 import sys
 import time
 from typing import NamedTuple
 
-from librt.base64 import b64decode
 from librt.internal import ReadBuffer, read_tag
 
 from mypy import util
@@ -48,7 +46,7 @@
     process_stale_scc_implementation,
     process_stale_scc_interface,
 )
-from mypy.cache import Tag, read_int_list
+from mypy.cache import Tag, read_int_list, read_json
 from mypy.defaults import RECURSION_LIMIT, WORKER_CONNECTION_TIMEOUT, WORKER_IDLE_TIMEOUT
 from mypy.errors import CompileError, ErrorInfo, Errors, report_internal_error
 from mypy.fscache import FileSystemCache
@@ -61,7 +59,7 @@
 
 parser = argparse.ArgumentParser(prog="mypy_worker", description="Mypy build worker")
 parser.add_argument("--status-file", help="status file to communicate worker details")
-parser.add_argument("--options-data", help="serialized mypy options")
+parser.add_argument("--options-data", help="file with serialized mypy options")
 
 CONNECTION_NAME = "build_worker"
 
@@ -85,11 +83,12 @@ def main(argv: list[str]) -> None:
     # This mimics how daemon receives the options. Note we need to postpone
     # processing error codes after plugins are loaded, because plugins can add
     # custom error codes.
-    options_dict = pickle.loads(b64decode(args.options_data))
-    options_obj = Options()
+    with open(args.options_data, "rb") as f:
+        buf = ReadBuffer(f.read())
+    options_dict = read_json(buf)
     disable_error_code = options_dict.pop("disable_error_code", [])
     enable_error_code = options_dict.pop("enable_error_code", [])
-    options = options_obj.apply_changes(options_dict)
+    options = Options().apply_changes(options_dict)
 
     status_file = args.status_file
     server = IPCServer(CONNECTION_NAME, WORKER_CONNECTION_TIMEOUT)

mypy/options.py

@@ -11,6 +11,7 @@
 from librt.internal import WriteBuffer, write_bool, write_str
 
 from mypy import defaults
+from mypy.cache import write_json
 from mypy.errorcodes import ErrorCode, error_codes
 from mypy.util import get_class_descriptors, replace_object_state
 
@@ -343,7 +344,7 @@ def __init__(self) -> None:
         # Per-module options (raw)
         self.per_module_options: dict[str, dict[str, object]] = {}
         self._glob_options: list[tuple[str, Pattern[str]]] = []
-        self.unused_configs: set[str] = set()
+        self._unused_configs: set[str] = set()
 
         # -- development options --
         self.verbosity = 0  # More verbose messages (for troubleshooting)
@@ -451,6 +452,17 @@ def snapshot(self) -> dict[str, object]:
         d = {k: v for k, v in d.items() if not k.startswith("_")}
         return d
 
+    def to_bytes(self) -> bytes:
+        """Serialize this options object to binary data."""
+        assert self.transform_source is None, "Source transform cannot be serialized"
+        snapshot = self.snapshot()
+        # Caller will need to use process_error_codes() to re-compute these.
+        del snapshot["disabled_error_codes"]
+        del snapshot["enabled_error_codes"]
+        buf = WriteBuffer()
+        write_json(buf, snapshot)
+        return buf.getvalue()
+
     def __repr__(self) -> str:
         return f"Options({pprint.pformat(self.snapshot())})"
 
@@ -560,7 +572,7 @@ def build_per_module_cache(self) -> None:
         # sections as used if any real modules use them or if any
         # concrete config sections use them. This means we need to
         # track which get used while constructing.
-        self.unused_configs = set(unstructured_glob_keys)
+        self._unused_configs = set(unstructured_glob_keys)
 
         for key in wildcards + concrete:
             # Find what the options for this key would be, just based
@@ -571,7 +583,7 @@ def build_per_module_cache(self) -> None:
 
         # Add the more structured sections into unused configs, since
         # they only count as used if actually used by a real module.
-        self.unused_configs.update(structured_keys)
+        self._unused_configs.update(structured_keys)
 
     def clone_for_module(self, module: str) -> Options:
         """Create an Options object that incorporates per-module options.
@@ -585,7 +597,7 @@ def clone_for_module(self, module: str) -> Options:
 
         # If the module just directly has a config entry, use it.
         if module in self._per_module_cache:
-            self.unused_configs.discard(module)
+            self._unused_configs.discard(module)
             return self._per_module_cache[module]
 
         # If not, search for glob paths at all the parents. So if we are looking for
@@ -598,7 +610,7 @@ def clone_for_module(self, module: str) -> Options:
         for i in range(len(path), 0, -1):
             key = ".".join(path[:i] + ["*"])
             if key in self._per_module_cache:
-                self.unused_configs.discard(key)
+                self._unused_configs.discard(key)
                 options = self._per_module_cache[key]
                 break
 
@@ -607,7 +619,7 @@ def clone_for_module(self, module: str) -> Options:
         if not module.endswith(".*"):
             for key, pattern in self._glob_options:
                 if pattern.match(module):
-                    self.unused_configs.discard(key)
+                    self._unused_configs.discard(key)
                     options = options.apply_changes(self.per_module_options[key])
 
         # We could update the cache to directly point to modules once
@@ -616,6 +628,9 @@ def clone_for_module(self, module: str) -> Options:
 
         return options
 
+    def get_unused_configs(self) -> set[str]:
+        return self._unused_configs.copy()
+
     def compile_glob(self, s: str) -> Pattern[str]:
         # Compile one of the glob patterns to a regex so that '.*' can
         # match *zero or more* module sections. This means we compile