Put remote fx cache dynamo_timed definition in OSS location (#140016)

Summary: I'm refactoring dynamo_timed and updating the params. It will be much easier to do this refactor entirely in OSS. So this diff essentially provides a couple aliases in the OSS area that I can update without affecting the internal usage.

Test Plan: Ran locally and made sure I still got samples: https://fburl.com/scuba/dynamo_compile/sandbox/qub89lwj

Reviewed By: oulgen

Differential Revision: D65580302

Pull Request resolved: https://github.com/pytorch/pytorch/pull/140016
Approved by: https://github.com/oulgen
This commit is contained in:
Sam Larsen 2024-11-07 23:51:45 +00:00 committed by PyTorch MergeBot
parent ae01f2b61b
commit dbaa431dfb

View File

@ -3,6 +3,7 @@ from __future__ import annotations
import atexit
import collections
import dataclasses
import functools
import json
import logging
import os
@ -12,6 +13,7 @@ from abc import abstractmethod
from typing import Any, Callable, Dict, Generic, List, Optional, Type, TypeVar, Union
from typing_extensions import override, TypeAlias
from torch._dynamo.utils import dynamo_timed
from torch._inductor import config
@ -38,6 +40,22 @@ _T = TypeVar("_T")
_U = TypeVar("_U")
remote_fx_cache_get_timed = functools.partial(
dynamo_timed,
"FbRemoteFxGraphCache.get",
phase_name="remote_fx_graph_cache_get",
log_pt2_compile_event=False,
fwd_only=False,
)
remote_fx_cache_put_timed = functools.partial(
dynamo_timed,
"FbRemoteFxGraphCache.put",
phase_name="remote_fx_graph_cache_put",
log_pt2_compile_event=False,
fwd_only=False,
)
class RemoteCacheBackend(Generic[_T]):
"""
A backend implementation for accessing a remote/distributed cache. Only