Fix sequential bottleneck in parallel parsing#21291
Conversation
The file is now usually only read in the Rust extension. This improves parallel scaling, as `get_source()` was a sequential bottleneck. I measured ~5% improvement to parallel type checking times in some cases on macOS (though it was a bit noisy).
This comment has been minimized.
This comment has been minimized.
ilevkivskyi
left a comment
ilevkivskyi
left a comment
There was a problem hiding this comment.
I think I found a logical flaw that may limit the performance improvement from this. As mentioned in parse_all() docstring, it is a god idea to keep it in roughly 1:1 correspondence with parse_file(), see more details in review comments.
| state.needs_parse = False | ||
| # New parser reads source from file directly, we do this only for | ||
| # the side effect of parsing inline mypy configurations. | ||
| state.get_source() |
There was a problem hiding this comment.
I think you need to (conditionally) remove the same call in State.parse_file(), otherwise the worker will call it when loading the tree (look for state.parse_file(raw_data=raw_data) in worker.py).
| self.errors.set_file(state.xpath, state.id, state.options) | ||
| for lineno, error in config_errors: | ||
| self.error(lineno, error) | ||
| state.check_for_invalid_options() |
There was a problem hiding this comment.
This method is getting really long, maybe it is possible to factor out tree loading logic to a separate method? (Especially in the view of the comment above, since you will probably need this in parse_file() as well).
There was a problem hiding this comment.
Refactored the parallel part into a separate method.
| self.log(f"Using cached AST for {state.xpath} ({state.id})") | ||
| state.tree, state.early_errors = self.ast_cache[state.id] | ||
| state.tree, state.early_errors, source_hash = self.ast_cache[state.id] | ||
| if state.source_hash is None: |
There was a problem hiding this comment.
Why the is None check needed here?
| manager.log(f"Using cached AST for {self.xpath} ({self.id})") | ||
| self.tree, self.early_errors = manager.ast_cache[self.id] | ||
| self.tree, self.early_errors, source_hash = manager.ast_cache[self.id] | ||
| if self.source_hash is None: |
There was a problem hiding this comment.
Same here, not sure why is None is needed.
| self.is_partial_stub_package = is_partial_stub_package | ||
| self.uses_template_strings = uses_template_strings | ||
| self.source_hash = source_hash | ||
| self.mypy_comments = mypy_comments if mypy_comments is not None else [] |
There was a problem hiding this comment.
I think these two (or at least the second one) need to be sent to the worker, i.e. you will need to handle them in write() and read(). The worker needs to know the full options, since we don't send options over the socket for each module (it is a big object). I guess tests pass now, because the worker still calls get_source().
There was a problem hiding this comment.
Added serialiation back (I had it removed since I thought it's not needed).
|
According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅ |
ilevkivskyi
left a comment
ilevkivskyi
left a comment
There was a problem hiding this comment.
LG, thanks! We will be able to simplify this when we will have just one parser.
2 participants
Previously we always read the file, processed inline comments, and calculated sha1 for each parsed file sequentially in Python. Now these are mostly moved to the Rust extension, which allows better parallel scaling.
I measured ~5% improvement to parallel type checking times in some cases on macOS (though it was a bit noisy, and used an earlier version of this PR).
Related to #21215.