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

[PGO] distinguish sticky PGO put (#163799)

Browse Source 
Summary: put_remote_code_state vs. put_extra_remote_code_state

Test Plan: test_pgo

Differential Revision: D83195687

Pull Request resolved: https://github.com/pytorch/pytorch/pull/163799
Approved by: https://github.com/bobrenjc93
This commit is contained in:
Pian Pawakapan
2025-09-25 06:59:22 +00:00
committed by PyTorch MergeBot
parent 2b6a74abf1
commit 353991dd92
1 changed files with 10 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
torch/_dynamo/pgo.py
View File

@ -952,9 +952,14 @@ def put_local_code_state(cache_key: str) -> None:
) )
def put_remote_code_state(cache_key: str) -> None: def put_remote_code_state(cache_key: str, extra_code_state: bool = False) -> None:
event_name = (
"put_remote_code_state"
if not extra_code_state
else "put_extra_remote_code_state"
)
with dynamo_timed( with dynamo_timed(
name := "pgo.put_remote_code_state", name := f"pgo.{event_name}",
log_pt2_compile_event=True, log_pt2_compile_event=True,
dynamo_compile_column_us="pgo_put_remote_code_state_time_us", dynamo_compile_column_us="pgo_put_remote_code_state_time_us",
): ):
@ -964,7 +969,7 @@ def put_remote_code_state(cache_key: str) -> None:
remote_cache = get_remote_cache() remote_cache = get_remote_cache()
if remote_cache is None: if remote_cache is None:
log.info("put_code_state: remote cache disabled") log.info("%s: remote cache disabled", event_name)
return return
content = pickle.dumps(_CODE_STATE) content = pickle.dumps(_CODE_STATE)
@ -974,11 +979,11 @@ def put_remote_code_state(cache_key: str) -> None:
} }
remote_cache.put(cache_key, cache_data) remote_cache.put(cache_key, cache_data)
log.info( log.info(
"put_code_state: wrote remote %s, %d entries", cache_key, len(_CODE_STATE) "%s: wrote remote %s, %d entries", event_name, cache_key, len(_CODE_STATE)
) )
# TODO: don't log this multiple times # TODO: don't log this multiple times
trace_structured_artifact( trace_structured_artifact(
"put_remote_code_state", event_name,
"string", "string",
lambda: render_code_state(_CODE_STATE), lambda: render_code_state(_CODE_STATE),
) )