OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Add cache size to pt2_compile_events (#139627)

Browse Source 
Summary:
I realized I wanted to check "are my cache entries/IO unreasonably large"
and there's no easy way to do it.  This lets me do it.

Test Plan: servicelab

Differential Revision: D65390363

Pull Request resolved: https://github.com/pytorch/pytorch/pull/139627
Approved by: https://github.com/c00w
This commit is contained in:
Edward Yang 2024-11-05 00:30:07 +00:00 committed by PyTorch MergeBot
parent 0058f71002
commit 639162f39a
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
torch/_dynamo/pgo.py
View File

@ -541,6 +541,7 @@ def get_code_state() -> DefaultDict[CodeId, CodeState]:
with open(path, "rb") as f:
try:
_CODE_STATE = pickle.load(f)
chromium_log.add_event_data(name, cache_size_bytes=f.tell())
except Exception:
log.warning(
"get_code_state failed while reading %s", path, exc_info=True
@ -569,6 +570,7 @@ def get_code_state() -> DefaultDict[CodeId, CodeState]:
data = cache_data["data"]
assert isinstance(data, str)
payload = base64.b64decode(data)
chromium_log.add_event_data(name, cache_size_bytes=len(payload))
_CODE_STATE = pickle.loads(payload)
except Exception:
log.warning(
@ -631,6 +633,7 @@ def put_local_code_state(cache_key: str) -> None:
with FileLock(lock_path, timeout=LOCK_TIMEOUT):
with open(tmp_path, "wb") as f:
pickle.dump(_CODE_STATE, f)
chromium_log.add_event_data(name, cache_size_bytes=f.tell())
os.rename(tmp_path, path)
log.info(
"put_code_state: wrote local %s, %d entries", path, len(_CODE_STATE)
@ -655,6 +658,7 @@ def put_remote_code_state(cache_key: str) -> None:
return
content = pickle.dumps(_CODE_STATE)
chromium_log.add_event_data(name, cache_size_bytes=len(content))
cache_data: JsonDataTy = {
"data": base64.b64encode(content).decode("ascii"),
}