codecache: pull out some Graph serialization code into common helpers #141502
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/141502
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (2 Unrelated Failures)As of commit 530614a with merge base 4959784 ( UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
aorenste
changed the title
aorenste
changed the title
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
…mon helpers" Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
| # models to disk. | ||
| self.current_callable = None | ||
|
|
||
| def after_deserialization(self, gm: Optional[torch.fx.GraphModule]) -> str: |
There was a problem hiding this comment.
@ezyang , this is basically just pulling out the code you wanted to refactor into post_compile, right?
There was a problem hiding this comment.
Also qq (maybe for @aorenste) do we do this serialization even when the outputcode is generated via AOTI (not CompiledFxGraph)? I.e. does it belong in OutputCode, or does it belong in CompiledFxGraph?
There was a problem hiding this comment.
I'm not sure why we want this separate from post_compile (there's probably a good reason, I'm just not seeing it). I am not sure why this isn't a shared method in the top level protocol.
There was a problem hiding this comment.
There may not be a good reason. I did all this code before post_compile/OutputCode existed - so it may make sense to merge it all together better.
There was a problem hiding this comment.
The actual purpose of the PR seems good, but we may want to change/bikeshed some names around to figure out the order these functions should be called.
I think it's confusing to both ahve a post_compile step and a after_deserialization step, but I also think after_deserialization is a better name. I guess concretely I think, assuming this after_deserialization always runs on all types of OutputCodes:
- Define
OutputCode.post_serializeto be thisafter_deserializationfunction - Rename
CompiledFxGraph.post_compiletoCompiledFxGraph.post_serialize, and haveCompiledFxGraphinherit fromOutputCode's definition by callingOutputCodefirst.
That said, today CompiledFxGraph isn't even an actual child class of CompiledFxGraph. Not sure what's stopping us from making that happen though. Will wait for @ezyang 's thoughts.
| except OSError: | ||
| # Not expected, but in case the PyCodeCache entry is removed from | ||
| # underneath us, treat it as a cache miss and recompile. | ||
| log.error("Failed to load cached artifact: %s", artifact_path) |
There was a problem hiding this comment.
Removing the log line seems worse
There was a problem hiding this comment.
The log still happens - just inside after_deserialization which will raise an exception (because it doesn't know if ignoring the exception is the right thing or not).
| get_metrics_context().increment("num_triton_bundles", 1) | ||
|
|
||
| inductor_meta = autotune_cache.inductor_meta_from_config() | ||
| AutotuneCacheBundler.begin_compile(inductor_meta, code=code) |
There was a problem hiding this comment.
Moving this after the OSError catch looks like a bug fix
There was a problem hiding this comment.
It's not really. The problem is that in order to look up the autotune cache we need the artifact code - but we don't have the artifact code until after after_deserialization() is called. AutotuneCacheBundler.begin_compile() just needs to get called sometime after you have the artifact code and before the first autotune occurs.
There was a problem hiding this comment.
The bug fix is having begun the compile and then immediately erroring out (without ending compile)
There was a problem hiding this comment.
I see what you're saying. Yeah - the AutotuneCacheBundler is a weird thing because our compile doesn't really have a good place to "wrap" (it would have to be somewhere in the dynamo call bytecode handler). begin_compile really just means "go fetch the compiled artifacts and spit them out on disk as local artifacts". end_compile really just means "collect all the local artifacts and bundle them into a single remote artifact". So not having an end doesn't really mean anything bad - it's supposed to handle that case properly (since it can happen in lots of ways).
| # so we serialize their PyCodeCache disk cache location instead. | ||
| # TODO: This could be better if we're ever able to serialize compiled | ||
| # models to disk. | ||
| self.current_callable = None |
There was a problem hiding this comment.
This is OK for refactoring but I don't want a "mutate the object so I can serialize it" API to be the final way of doing this. Definitely want a better API.
| from .graph import GraphLowering | ||
|
|
||
| # This is used by tests to check the output for specific details. | ||
| GraphLowering.save_output_code(code) |
- Turn fx_codegen_and_compile() into a class (FxCompile) so we can override the implementation. - Pull the current body into an implementation (_InProcessFxCompile) which just performs the existing behavior. - Add an async interface. (See below) The intended future behavior of Async Compile will be to allow dynamo functions to start compiling in the background (and on a separate machine) while we continue to run eager in the foreground. As such we'll need to put the compilation behind some kind of Future implementation - it makes sense to reuse the existing python futures for that. An async function is just a syntactic way to return an asyncio.Future. Because asyncio.run() adds confusion to the stack traces when the called function isn't actually being used in an asynchronous way we also provide a synchronous interface which can be directly called. Pull Request resolved: #141505 Approved by: https://github.com/ezyang ghstack dependencies: #141502
…pytorch#141502) Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. Pull Request resolved: pytorch#141502 Approved by: https://github.com/ezyang
- Turn fx_codegen_and_compile() into a class (FxCompile) so we can override the implementation. - Pull the current body into an implementation (_InProcessFxCompile) which just performs the existing behavior. - Add an async interface. (See below) The intended future behavior of Async Compile will be to allow dynamo functions to start compiling in the background (and on a separate machine) while we continue to run eager in the foreground. As such we'll need to put the compilation behind some kind of Future implementation - it makes sense to reuse the existing python futures for that. An async function is just a syntactic way to return an asyncio.Future. Because asyncio.run() adds confusion to the stack traces when the called function isn't actually being used in an asynchronous way we also provide a synchronous interface which can be directly called. Pull Request resolved: pytorch#141505 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#141502
…pytorch#141502) Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile. Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache. Pull Request resolved: pytorch#141502 Approved by: https://github.com/ezyang
- Turn fx_codegen_and_compile() into a class (FxCompile) so we can override the implementation. - Pull the current body into an implementation (_InProcessFxCompile) which just performs the existing behavior. - Add an async interface. (See below) The intended future behavior of Async Compile will be to allow dynamo functions to start compiling in the background (and on a separate machine) while we continue to run eager in the foreground. As such we'll need to put the compilation behind some kind of Future implementation - it makes sense to reuse the existing python futures for that. An async function is just a syntactic way to return an asyncio.Future. Because asyncio.run() adds confusion to the stack traces when the called function isn't actually being used in an asynchronous way we also provide a synchronous interface which can be directly called. Pull Request resolved: pytorch#141505 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#141502
4 participants
Moved some code from FxGraphCache.lookup_graph() which dealt with serializing and deserializing CompiledFxGraph into CompiledFxGraph itself so it can be reused later by Async Compile.
Async Compile will need to serialize the compiled CompiledFxGraph from one process and deserialize it in another - so it's very similar to the cache.
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov