frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[MegaCache] Rename the PGO artifact when used between different jobs (#151482)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/151482
Approved by: https://github.com/bobrenjc93, https://github.com/jamesjwu
This commit is contained in:
Oguz Ulgen
2025-04-16 22:05:19 -07:00
committed by PyTorch MergeBot
parent fe90a5c140
commit 8404c09b15
3 changed files with 68 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/compiler/_cache.py
View File

@ -206,7 +206,7 @@ class CacheArtifactManager:
log.warning("Failed to un-pickle cache artifacts", exc_info=True)
return None
from torch._dynamo.pgo import write_local_impl
from torch._dynamo.pgo import rewrite_cache_key_for_mega_cache, write_local_impl
from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
from torch._inductor.codecache import FxGraphCache
from torch._inductor.runtime.autotune_cache import _LocalAutotuneCacheBackend
@ -226,7 +226,9 @@ class CacheArtifactManager:
elif artifact.type == CacheArtifactType.AOT_AUTOGRAD:
AOTAutogradCache._write_to_local_cache(artifact.key, artifact.content)
elif artifact.type == CacheArtifactType.PGO:
meta = write_local_impl(artifact.key, artifact.content)
meta = write_local_impl(
rewrite_cache_key_for_mega_cache(artifact.key), artifact.content
)
assert meta is not None
else:
log.warning(f"Unsupported artifact type {artifact.type}") # noqa: G004