Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,610 @@
|
||||
import copy
|
||||
import logging
|
||||
import traceback
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import contextmanager
|
||||
from enum import Enum
|
||||
from typing import Any, Optional, ParamSpec, TypeVar, Union
|
||||
|
||||
from torch._utils_internal import signpost_event
|
||||
|
||||
from ._compatibility import compatibility
|
||||
from .graph import Graph
|
||||
from .graph_module import GraphModule
|
||||
from .node import Node
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"annotate",
|
||||
"annotate_fn",
|
||||
"preserve_node_meta",
|
||||
"has_preserved_node_meta",
|
||||
"set_stack_trace",
|
||||
"set_grad_fn_seq_nr",
|
||||
"reset_grad_fn_seq_nr",
|
||||
"format_stack",
|
||||
"set_current_meta",
|
||||
"get_current_meta",
|
||||
"NodeSource",
|
||||
"NodeSourceAction",
|
||||
"get_graph_provenance_json",
|
||||
"set_current_replay_node",
|
||||
"get_current_replay_node",
|
||||
]
|
||||
|
||||
current_meta: dict[str, Any] = {}
|
||||
current_replay_node: Node | None = None
|
||||
# Preserve the node meta fields in torch.fx.proxy._COPY_META_FIELDS
|
||||
should_preserve_node_meta = False
|
||||
# Preserve the "seq_nr" node meta field
|
||||
_should_preserve_node_meta = False
|
||||
|
||||
GRADIENT_ACC_SPECIAL_STACK = (
|
||||
"Gradient addition node due to multiple use of tensor around:"
|
||||
)
|
||||
# =============================================================================
|
||||
# FX Metadata Registry for Memory Profiler
|
||||
# =============================================================================
|
||||
# Global in-memory registry for FX metadata
|
||||
# Maps module_name -> metadata dict containing lineno_map and node_metadata
|
||||
_FX_METADATA_REGISTRY: dict[str, dict[str, Any]] = {}
|
||||
|
||||
|
||||
def _register_fx_metadata(module_name: str, metadata: dict[str, Any]) -> None:
|
||||
"""
|
||||
Register FX metadata in the global in-memory registry.
|
||||
|
||||
This is called automatically during graph module compilation to store metadata
|
||||
for later use by memory profiler augmentation.
|
||||
|
||||
Args:
|
||||
module_name: The module identifier (content-addressed filename)
|
||||
metadata: Metadata dict containing lineno_map, node_metadata, and source_code
|
||||
"""
|
||||
# TODO: add logging to tlparse
|
||||
_FX_METADATA_REGISTRY[module_name] = metadata
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
class NodeSourceAction(Enum):
|
||||
CREATE = "create"
|
||||
REPLACE = "replace"
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
class NodeSource:
|
||||
"""
|
||||
NodeSource is a data structure that contains the provenance information of a node.
|
||||
If node `a` is created from node `b`, then `a.meta["from_node"]` may contain NodeSource(b).
|
||||
"""
|
||||
|
||||
class NodeInfo:
|
||||
def __init__(self, name: str, target: str, graph_id: int) -> None:
|
||||
self.name = name
|
||||
self.target = target
|
||||
self.graph_id = graph_id
|
||||
|
||||
pass_name: str
|
||||
action: list["NodeSourceAction"]
|
||||
from_node: list["NodeSource"]
|
||||
node_info: Optional["NodeInfo"]
|
||||
_dict: dict[str, Any] | None
|
||||
_action_string: str | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
node: Node | None,
|
||||
pass_name: str = "",
|
||||
action: Union["NodeSourceAction", list["NodeSourceAction"]] | None = None,
|
||||
) -> None:
|
||||
self.pass_name = pass_name
|
||||
|
||||
if action is None:
|
||||
action = []
|
||||
elif not isinstance(action, list):
|
||||
action = [action]
|
||||
for a in action:
|
||||
if not isinstance(a, NodeSourceAction):
|
||||
raise AssertionError(f"Expected NodeSourceAction, got {type(a)}")
|
||||
self.action = action
|
||||
if node:
|
||||
self.node_info = self.NodeInfo(
|
||||
name=node.name, target=str(node.target), graph_id=id(node.graph)
|
||||
)
|
||||
self.from_node = (
|
||||
copy.deepcopy(node.meta["from_node"])
|
||||
if "from_node" in node.meta
|
||||
else []
|
||||
)
|
||||
else:
|
||||
self.node_info = None
|
||||
self.from_node = []
|
||||
|
||||
# cache the action string and dict representation for performance.
|
||||
self._action_string: str | None = None
|
||||
self._dict: dict[str, Any] | None = None
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self.node_info.name if self.node_info else ""
|
||||
|
||||
@property
|
||||
def target(self) -> str:
|
||||
return self.node_info.target if self.node_info else ""
|
||||
|
||||
@property
|
||||
def graph_id(self) -> int:
|
||||
return self.node_info.graph_id if self.node_info else -1
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.print_readable()
|
||||
|
||||
def _get_action_string(self) -> str:
|
||||
if self._action_string is None:
|
||||
self._action_string = "+".join([a.name.lower() for a in self.action])
|
||||
return self._action_string
|
||||
|
||||
def print_readable(self, indent: int = 0) -> str:
|
||||
if indent > 9:
|
||||
return ""
|
||||
result = ""
|
||||
action_string = self._get_action_string()
|
||||
result += (
|
||||
" " * indent * 4
|
||||
+ f"(name={self.name}, pass_name={self.pass_name}, action={action_string}, graph_id={self.graph_id})\n"
|
||||
)
|
||||
for item in self.from_node:
|
||||
result += item.print_readable(indent + 1)
|
||||
return result
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
if self._dict is None:
|
||||
# Convert the object to a dictionary
|
||||
action_string = self._get_action_string()
|
||||
self._dict = {
|
||||
"name": self.name,
|
||||
"target": self.target,
|
||||
"graph_id": self.graph_id,
|
||||
"pass_name": self.pass_name,
|
||||
"action": action_string,
|
||||
"from_node": [node.to_dict() for node in self.from_node],
|
||||
}
|
||||
|
||||
if self._dict is None:
|
||||
raise AssertionError("_dict is None after initialization")
|
||||
return self._dict
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, NodeSource):
|
||||
return False
|
||||
return self.to_dict() == other.to_dict()
|
||||
|
||||
def __hash__(self) -> int:
|
||||
# Create a hash based on the dictionary representation
|
||||
# We need to convert the dict to a hashable form
|
||||
def _make_hashable(obj: Any) -> Any:
|
||||
if isinstance(obj, dict):
|
||||
return tuple(sorted((k, _make_hashable(v)) for k, v in obj.items()))
|
||||
elif isinstance(obj, list):
|
||||
return tuple(_make_hashable(item) for item in obj)
|
||||
else:
|
||||
return obj
|
||||
|
||||
return hash(_make_hashable(self.to_dict()))
|
||||
|
||||
@classmethod
|
||||
def _from_dict(cls, d: dict[str, Any] | None) -> Optional["NodeSource"]:
|
||||
"""
|
||||
Recursively deserialize from_node metadata from dictionary data.
|
||||
It is used to deserialize the from_node field from serialized metadata.
|
||||
Please use constructor NodeSource(node, ...) to create a NodeSource object.
|
||||
"""
|
||||
if d is None:
|
||||
return None
|
||||
|
||||
if not isinstance(d, dict):
|
||||
raise AssertionError(f"Expected a dict, got {type(d)}")
|
||||
|
||||
# Create a NodeSource object directly without going through the constructor
|
||||
# to avoid issues with graph ID and node creation
|
||||
node_source = NodeSource.__new__(NodeSource)
|
||||
|
||||
# Reset the cached properties
|
||||
node_source._action_string = None
|
||||
node_source._dict = None
|
||||
|
||||
# Set the basic attributes
|
||||
node_source.pass_name = d.get("pass_name", "")
|
||||
|
||||
# Parse action string back to NodeSourceAction enum list
|
||||
action_str = d.get("action", "")
|
||||
actions = []
|
||||
if action_str:
|
||||
for action_name in action_str.split("+"):
|
||||
if action_name.upper() == "CREATE":
|
||||
actions.append(NodeSourceAction.CREATE)
|
||||
elif action_name.upper() == "REPLACE":
|
||||
actions.append(NodeSourceAction.REPLACE)
|
||||
node_source.action = actions
|
||||
|
||||
# Create the NodeInfo object directly
|
||||
if "name" in d and "target" in d and "graph_id" in d:
|
||||
node_info = NodeSource.NodeInfo(
|
||||
d.get("name", ""), d.get("target", ""), d.get("graph_id", -1)
|
||||
)
|
||||
node_source.node_info = node_info
|
||||
else:
|
||||
node_source.node_info = None
|
||||
|
||||
# Recursively deserialize nested from_node
|
||||
if d.get("from_node", None) is not None:
|
||||
node_source.from_node = [
|
||||
result
|
||||
for fn in d.get("from_node", [])
|
||||
if (result := cls._from_dict(fn)) is not None
|
||||
]
|
||||
else:
|
||||
node_source.from_node = []
|
||||
return node_source
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
@contextmanager
|
||||
def preserve_node_meta(enable: bool = True) -> Iterator[None]:
|
||||
global should_preserve_node_meta
|
||||
global current_meta
|
||||
saved_should_preserve_node_meta = should_preserve_node_meta
|
||||
# Shallow copy is OK since fields of current_meta are not mutated
|
||||
saved_current_meta = current_meta.copy()
|
||||
try:
|
||||
should_preserve_node_meta = enable
|
||||
yield
|
||||
finally:
|
||||
should_preserve_node_meta = saved_should_preserve_node_meta
|
||||
current_meta = saved_current_meta
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _preserve_node_seq_nr(preserve_seq_nr: bool = True) -> Iterator[None]:
|
||||
"""
|
||||
Temporarily enables or disables the preservation of node.meta["seq_nr"] in the
|
||||
tracing context.
|
||||
"""
|
||||
global _should_preserve_node_meta
|
||||
saved = _should_preserve_node_meta
|
||||
|
||||
try:
|
||||
_should_preserve_node_meta = preserve_seq_nr
|
||||
yield
|
||||
finally:
|
||||
_should_preserve_node_meta = saved
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def set_stack_trace(stack: list[str]) -> None:
|
||||
global current_meta
|
||||
|
||||
if should_preserve_node_meta:
|
||||
if stack:
|
||||
current_meta["stack_trace"] = "".join(stack)
|
||||
else:
|
||||
# when the stack is empty, we explicitly clear the stack_trace to avoid
|
||||
# propagating it to future node.˙
|
||||
current_meta.pop("stack_trace", None)
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
@contextmanager
|
||||
def annotate(annotation_dict: dict[str, Any]) -> Iterator[None]:
|
||||
"""
|
||||
Temporarily adds custom annotations to the current tracing context.
|
||||
The fx_node produced from this tracing context will have the
|
||||
custom annotations in node.metadata["custom"] field.
|
||||
|
||||
This context manager allows you to insert arbitrary metadata into the PT2
|
||||
tracing system by updating the global `current_meta["custom"]` dictionary.
|
||||
The annotations are automatically reverted after the context exits.
|
||||
|
||||
Gradient accumulation nodes will not be annotated.
|
||||
|
||||
This is intended for advanced users who need to attach additional metadata to the fx nodes
|
||||
(e.g., for debugging, analysis, or external tooling) during export tracing.
|
||||
|
||||
Note:
|
||||
This API is **not backward compatible** and may evolve in future releases.
|
||||
|
||||
Note:
|
||||
This API is not compatible with fx.symbolic_trace or jit.trace. It's intended
|
||||
to be used with PT2 family of tracers, e.g. torch.export and dynamo.
|
||||
|
||||
Args:
|
||||
annotation_dict (dict): A dictionary of custom key-value pairs to inject
|
||||
into the FX trace metadata.
|
||||
|
||||
Example:
|
||||
After exiting the context, custom annotations are removed.
|
||||
|
||||
>>> with annotate({"source": "custom_pass", "tag": 42}):
|
||||
... pass # Your computation here
|
||||
"""
|
||||
|
||||
global current_meta
|
||||
|
||||
has_custom = "custom" in current_meta
|
||||
old_custom = copy.copy(current_meta.get("custom", {}))
|
||||
|
||||
try:
|
||||
if not has_custom:
|
||||
current_meta["custom"] = dict[str, Any]()
|
||||
|
||||
# Update with all key-value pairs from the input dict
|
||||
current_meta["custom"].update(annotation_dict)
|
||||
yield
|
||||
finally:
|
||||
if has_custom:
|
||||
# Restore the original custom dict
|
||||
current_meta["custom"] = old_custom
|
||||
else:
|
||||
del current_meta["custom"]
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def annotate_fn(
|
||||
annotation_dict: dict[str, Any],
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
|
||||
"""
|
||||
A decorator that wraps a function with the annotate context manager.
|
||||
Use this when you want to annotate an entire function instead of a specific code block.
|
||||
|
||||
Note:
|
||||
This API is **not backward compatible** and may evolve in future releases.
|
||||
|
||||
Note:
|
||||
This API is not compatible with fx.symbolic_trace or jit.trace. It's intended
|
||||
to be used with PT2 family of tracers, e.g. torch.export and dynamo.
|
||||
|
||||
Args:
|
||||
annotation_dict (dict): A dictionary of custom key-value pairs to inject
|
||||
into the FX trace metadata for all operations in the function.
|
||||
|
||||
Example:
|
||||
All operations in my_function will have {"pp_stage": 1} in their metadata.
|
||||
|
||||
>>> @annotate_fn({"pp_stage": 1})
|
||||
... def my_function(x):
|
||||
... return x + 1
|
||||
"""
|
||||
from functools import wraps
|
||||
|
||||
def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
@wraps(func)
|
||||
# NB: Do not annotate with _P.args/_P.kwargs here. Dynamo guards on
|
||||
# the identity of ParamSpec annotation objects, causing guard failures.
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
with annotate(annotation_dict):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper # type: ignore[return-value]
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _set_autograd_backward(enable: bool = True) -> Iterator[None]:
|
||||
global current_meta
|
||||
|
||||
had_autograd_backward = "autograd_backward" in current_meta
|
||||
old_autograd_backward = current_meta.get("autograd_backward", False)
|
||||
|
||||
if enable:
|
||||
_mark_autograd_backward()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if had_autograd_backward:
|
||||
current_meta["autograd_backward"] = old_autograd_backward
|
||||
else:
|
||||
_reset_autograd_backward()
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def _mark_autograd_backward() -> None:
|
||||
global current_meta
|
||||
|
||||
current_meta["autograd_backward"] = True
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def _reset_autograd_backward() -> None:
|
||||
global current_meta
|
||||
|
||||
current_meta.pop("autograd_backward", None)
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def set_grad_fn_seq_nr(seq_nr: int) -> None:
|
||||
global current_meta
|
||||
|
||||
if should_preserve_node_meta:
|
||||
# The seq_nr is captured by eager mode in the grad_fn during forward
|
||||
current_meta["grad_fn_seq_nr"] = current_meta.get("grad_fn_seq_nr", []) + [
|
||||
seq_nr
|
||||
]
|
||||
current_meta["in_grad_fn"] = current_meta.get("in_grad_fn", 0) + 1
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def reset_grad_fn_seq_nr() -> None:
|
||||
# NB: reset state properly, this would be helpful towards supporting
|
||||
# reentrant autograd if we actually wanted to do that.
|
||||
global current_meta
|
||||
if should_preserve_node_meta:
|
||||
current_level = current_meta.get("in_grad_fn", 0)
|
||||
if current_level <= 0:
|
||||
raise AssertionError(f"Expected current_level > 0, got {current_level}")
|
||||
if current_level == 1:
|
||||
del current_meta["in_grad_fn"]
|
||||
del current_meta["grad_fn_seq_nr"]
|
||||
else:
|
||||
current_meta["in_grad_fn"] = current_level - 1
|
||||
current_meta["grad_fn_seq_nr"] = current_meta["grad_fn_seq_nr"][:-1]
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def format_stack() -> list[str]:
|
||||
if should_preserve_node_meta:
|
||||
return [current_meta.get("stack_trace", "")]
|
||||
else:
|
||||
# fallback to traceback.format_stack()
|
||||
return traceback.format_list(traceback.extract_stack()[:-1])
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def has_preserved_node_meta() -> bool:
|
||||
return should_preserve_node_meta
|
||||
|
||||
|
||||
def _is_preserving_node_seq_nr() -> bool:
|
||||
return _should_preserve_node_meta
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
@contextmanager
|
||||
def set_current_meta(node: Node, pass_name: str = "") -> Iterator[None]:
|
||||
global current_meta
|
||||
if should_preserve_node_meta and node.meta:
|
||||
saved_meta = current_meta
|
||||
try:
|
||||
current_meta = node.meta.copy()
|
||||
|
||||
# Update the "from_node" field in current_meta for provenance tracking.
|
||||
# Instead of appending, overwrite the "from_node" field because current_meta
|
||||
# will be assigned to the new node. The new NodeSource(node, ...) will
|
||||
# include the information from the previous current_meta["from_node"].
|
||||
current_meta["from_node"] = [
|
||||
NodeSource(node, pass_name, NodeSourceAction.CREATE)
|
||||
]
|
||||
yield
|
||||
finally:
|
||||
current_meta = saved_meta
|
||||
else:
|
||||
yield
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def get_current_meta() -> dict[str, Any]:
|
||||
return current_meta
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
@contextmanager
|
||||
def set_current_replay_node(node: Node | None) -> Iterator[None]:
|
||||
"""
|
||||
Set the currently replay node. If `current_replay_node` is not None,
|
||||
then we're re-generating the `current_replay_node` in FunctionalTensorMode.
|
||||
"""
|
||||
# See [Note] annotation for more details.
|
||||
global current_replay_node
|
||||
saved_current_replay_node = current_replay_node
|
||||
try:
|
||||
current_replay_node = node
|
||||
yield
|
||||
finally:
|
||||
current_replay_node = saved_current_replay_node
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def get_current_replay_node() -> Node | None:
|
||||
"""
|
||||
Get the currently replay node
|
||||
"""
|
||||
return current_replay_node
|
||||
|
||||
|
||||
@compatibility(is_backward_compatible=False)
|
||||
def get_graph_provenance_json(graph: Graph) -> dict[str, Any]:
|
||||
"""
|
||||
Given an fx.Graph, return a json that contains the provenance information of each node.
|
||||
"""
|
||||
try:
|
||||
provenance_tracking_json = {}
|
||||
for node in graph.nodes:
|
||||
if node.op == "call_function":
|
||||
provenance_tracking_json[node.name] = (
|
||||
[source.to_dict() for source in node.meta["from_node"]]
|
||||
if "from_node" in node.meta
|
||||
else []
|
||||
)
|
||||
return provenance_tracking_json
|
||||
except Exception as e:
|
||||
# Since this is just debugging, it should never interfere with regular
|
||||
# program execution, so we use this try-except to guard against any error
|
||||
signpost_event(
|
||||
"inductor",
|
||||
"provenance_tracking_error",
|
||||
{
|
||||
"function": "get_graph_provenance_json",
|
||||
"error_msg": str(e),
|
||||
"stack_trace": traceback.format_exc(),
|
||||
},
|
||||
)
|
||||
return {}
|
||||
|
||||
|
||||
def _get_custom_metadata(gm: GraphModule) -> str:
|
||||
if not isinstance(gm, GraphModule):
|
||||
raise AssertionError(f"Expected GraphModule, got {type(gm)}")
|
||||
|
||||
def helper(gm: GraphModule) -> list[Any]:
|
||||
custom_metadata = []
|
||||
for node in gm.graph.nodes:
|
||||
if hasattr(node, "meta") and node.meta.get("custom", None):
|
||||
custom_metadata.append((node.op, node.name, node.meta["custom"]))
|
||||
if node.op == "get_attr" and isinstance(
|
||||
getattr(gm, node.target), GraphModule
|
||||
):
|
||||
custom_metadata.append(
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
helper(getattr(gm, node.target))
|
||||
)
|
||||
return custom_metadata
|
||||
|
||||
return "\n".join(str(x) for x in helper(gm))
|
||||
|
||||
|
||||
def _get_ordered_seq_nr_groups(
|
||||
gm: GraphModule | list[GraphModule],
|
||||
) -> list[list[str]]:
|
||||
"""
|
||||
Group call_function nodes by seq_nr, order by seq_nr value,
|
||||
and return a list of lists of node names (sorted alphabetically).
|
||||
|
||||
Args:
|
||||
gm: A single GraphModule or a list of GraphModules to process.
|
||||
When a list is provided, nodes from all graphs are grouped together.
|
||||
|
||||
Returns:
|
||||
A list of lists, where each inner list contains node names that share the same seq_nr,
|
||||
sorted alphabetically. The outer list is ordered by seq_nr value.
|
||||
"""
|
||||
# Normalize input to a list
|
||||
if isinstance(gm, GraphModule):
|
||||
gms = [gm]
|
||||
else:
|
||||
gms = gm
|
||||
|
||||
seq_nr_dict: dict[int, list[str]] = defaultdict(list)
|
||||
for graph_module in gms:
|
||||
for node in graph_module.graph.nodes:
|
||||
if node.op == "call_function":
|
||||
seq_nr = node.meta.get("seq_nr")
|
||||
if seq_nr is not None:
|
||||
seq_nr_dict[seq_nr].append(node.name)
|
||||
# Sort by seq_nr and return list of sorted lists
|
||||
return [sorted(seq_nr_dict[k]) for k in sorted(seq_nr_dict.keys())]
|
||||
Reference in New Issue
Block a user