From cd021d048e4645d1bf9f193587ca0eda377a2721 Mon Sep 17 00:00:00 2001 From: Oguz Ulgen Date: Tue, 22 Apr 2025 13:52:51 -0700 Subject: [PATCH] Fix circular imports (#151939) Pull Request resolved: https://github.com/pytorch/pytorch/pull/151939 Approved by: https://github.com/jamesjwu --- torch/_dynamo/pgo.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/torch/_dynamo/pgo.py b/torch/_dynamo/pgo.py index 686bdd0536f..5f6967428cc 100644 --- a/torch/_dynamo/pgo.py +++ b/torch/_dynamo/pgo.py @@ -35,7 +35,6 @@ from torch._dynamo.utils import ( ) from torch._environment import is_fbcode from torch._logging._internal import trace_structured_artifact -from torch.compiler._cache import CacheArtifactManager, CacheArtifactType if TYPE_CHECKING: @@ -591,6 +590,8 @@ def get_code_state() -> defaultdict[CodeId, CodeState]: _INIT_CODE_STATE = copy.deepcopy(_CODE_STATE) return _CODE_STATE + from torch.compiler._cache import CacheArtifactManager, CacheArtifactType + # Attempt local path = code_state_path(cache_key) if path is not None and os.path.exists(path): @@ -712,6 +713,8 @@ def put_local_code_state(cache_key: str) -> None: pickled_code = pickle.dumps(_CODE_STATE) + from torch.compiler._cache import CacheArtifactManager, CacheArtifactType + CacheArtifactManager.record_artifact( CacheArtifactType.PGO, cache_key, pickled_code )