Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,387 @@
|
||||
"""
|
||||
User-facing API for controlling DSL operation overrides.
|
||||
|
||||
The torch.backends.python_native module provides control over DSL (Domain Specific Language)
|
||||
operation overrides defined in torch._native. This allows users to selectively enable or disable
|
||||
high-performance implementations from various DSLs like Triton and CuteDSL.
|
||||
|
||||
The module supports both coarse-grained control (entire DSLs) and fine-grained control
|
||||
(individual operations or dispatch keys). All control operations support context managers
|
||||
for temporary state changes.
|
||||
|
||||
Example usage::
|
||||
|
||||
import torch.backends.python_native as pn
|
||||
|
||||
# DSL-level control
|
||||
pn.triton.enabled = False # Disable all triton ops
|
||||
pn.cutedsl.enabled = True # Enable all cutedsl ops
|
||||
|
||||
# Individual operation control
|
||||
pn.disable_operations("scaled_mm") # Disable specific op across all DSLs
|
||||
pn.enable_operations("scaled_mm") # Re-enable specific op
|
||||
|
||||
# Context manager support
|
||||
with pn.triton.disabled():
|
||||
result = some_computation() # Triton ops disabled here
|
||||
|
||||
# Query capabilities
|
||||
print(pn.available_dsls) # ['triton', 'cutedsl']
|
||||
print(pn.get_dsl_operations("triton")) # Operations for triton
|
||||
"""
|
||||
|
||||
import functools
|
||||
import sys
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
|
||||
from torch.backends import ContextProp, flags_frozen, PropModule
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _preserve_filter_state():
|
||||
"""Context manager to save and restore registry filter state."""
|
||||
filter_state = _get_filter_state()
|
||||
|
||||
# Save original state
|
||||
original_state = (
|
||||
set(filter_state._dsl_names),
|
||||
set(filter_state._op_symbols),
|
||||
set(filter_state._dispatch_keys),
|
||||
)
|
||||
|
||||
try:
|
||||
yield filter_state
|
||||
finally:
|
||||
# Restore original state
|
||||
filter_state._dsl_names.clear()
|
||||
filter_state._op_symbols.clear()
|
||||
filter_state._dispatch_keys.clear()
|
||||
|
||||
filter_state._dsl_names.update(original_state[0])
|
||||
filter_state._op_symbols.update(original_state[1])
|
||||
filter_state._dispatch_keys.update(original_state[2])
|
||||
|
||||
|
||||
def _get_dsl_registry():
|
||||
"""Lazy import to avoid circular imports."""
|
||||
from torch._native.dsl_registry import dsl_registry
|
||||
|
||||
return dsl_registry
|
||||
|
||||
|
||||
def _get_registry_functions():
|
||||
"""Lazy import of registry functions."""
|
||||
from torch._native.registry import (
|
||||
_filter_state,
|
||||
_graphs,
|
||||
deregister_op_overrides,
|
||||
reenable_op_overrides,
|
||||
)
|
||||
|
||||
return deregister_op_overrides, reenable_op_overrides, _graphs, _filter_state
|
||||
|
||||
|
||||
def _get_filter_state():
|
||||
"""Direct access to filter state."""
|
||||
return _get_registry_functions()[3]
|
||||
|
||||
|
||||
def _get_dsl_module(dsl_name: str):
|
||||
"""Get the registered DSL module for direct control.
|
||||
|
||||
Uses the DSL registry to dynamically look up DSL modules instead of
|
||||
hard-coding the mapping. This makes the function automatically extensible
|
||||
for new DSLs without code changes.
|
||||
|
||||
Args:
|
||||
dsl_name (str): Name of the DSL to retrieve.
|
||||
|
||||
Returns:
|
||||
DSLModuleProtocol: The registered DSL module.
|
||||
|
||||
Raises:
|
||||
ValueError: If the DSL is not registered.
|
||||
"""
|
||||
registry = _get_dsl_registry()
|
||||
|
||||
# Use the public API to get the DSL module
|
||||
dsl_module = registry.get_dsl_module(dsl_name)
|
||||
if dsl_module is not None:
|
||||
return dsl_module
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown DSL: {dsl_name}. Available DSLs: {registry.list_all_dsls()}"
|
||||
)
|
||||
|
||||
|
||||
class DSLController:
|
||||
"""Controller for a specific DSL."""
|
||||
|
||||
def __init__(self, dsl_name: str):
|
||||
self._dsl_name = dsl_name
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._dsl_name
|
||||
|
||||
@property
|
||||
def available(self) -> bool:
|
||||
"""Check if DSL runtime is available."""
|
||||
registry = _get_dsl_registry()
|
||||
return registry.is_dsl_available(self._dsl_name)
|
||||
|
||||
@property
|
||||
def version(self):
|
||||
"""Get DSL version."""
|
||||
registry = _get_dsl_registry()
|
||||
return registry.get_dsl_version(self._dsl_name)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
"""Check if DSL is currently enabled."""
|
||||
filter_state = _get_filter_state()
|
||||
return self._dsl_name not in filter_state._dsl_names
|
||||
|
||||
@enabled.setter
|
||||
def enabled(self, value: bool):
|
||||
"""Enable or disable the DSL."""
|
||||
if flags_frozen():
|
||||
raise RuntimeError(
|
||||
f"not allowed to set {self._dsl_name} DSL flags "
|
||||
"after disable_global_flags; please use flags() context manager instead"
|
||||
)
|
||||
if value:
|
||||
self.enable()
|
||||
else:
|
||||
self.disable()
|
||||
|
||||
def disable(self):
|
||||
"""Disable all operations for this DSL."""
|
||||
dsl_module = _get_dsl_module(self._dsl_name)
|
||||
dsl_module.deregister_op_overrides()
|
||||
|
||||
def enable(self):
|
||||
"""Re-enable all operations for this DSL."""
|
||||
reenable_op_overrides = _get_registry_functions()[1]
|
||||
reenable_op_overrides(enable_dsl_names=self._dsl_name)
|
||||
|
||||
@contextmanager
|
||||
def disabled(self):
|
||||
"""Context manager to temporarily disable DSL."""
|
||||
original_state = self.enabled
|
||||
try:
|
||||
self.disable()
|
||||
yield
|
||||
finally:
|
||||
if original_state:
|
||||
self.enable()
|
||||
|
||||
def __repr__(self):
|
||||
status = "available" if self.available else "unavailable"
|
||||
enabled_status = "enabled" if self.enabled else "disabled"
|
||||
return f"DSLController({self._dsl_name}, {status}, {enabled_status})"
|
||||
|
||||
|
||||
class PythonNativeModule(PropModule):
|
||||
"""Main module for python_native DSL control."""
|
||||
|
||||
def __init__(self, original_module):
|
||||
super().__init__(original_module, original_module.__name__)
|
||||
|
||||
@property
|
||||
def available_dsls(self) -> list[str]:
|
||||
"""Get list of available DSLs."""
|
||||
registry = _get_dsl_registry()
|
||||
result = registry.list_available_dsls()
|
||||
return list(result) if not isinstance(result, list) else result
|
||||
|
||||
@property
|
||||
def all_dsls(self) -> list[str]:
|
||||
"""Get list of all registered DSLs."""
|
||||
registry = _get_dsl_registry()
|
||||
result = registry.list_all_dsls()
|
||||
return list(result) if not isinstance(result, list) else result
|
||||
|
||||
def get_dsl_operations(self, dsl_name: str) -> list[str]:
|
||||
"""Get list of operations registered by a specific DSL.
|
||||
|
||||
Args:
|
||||
dsl_name (str): Name of the DSL to query (e.g., 'triton', 'cutedsl').
|
||||
|
||||
Returns:
|
||||
list[str]: Sorted list of operation names registered by the DSL.
|
||||
|
||||
Example::
|
||||
|
||||
ops = torch.backends.python_native.get_dsl_operations("triton")
|
||||
print(ops) # ['triton_to_mxfp8_dim0', ...]
|
||||
"""
|
||||
from torch._native.registry import get_dsl_operations
|
||||
|
||||
return get_dsl_operations(dsl_name)
|
||||
|
||||
def disable_operations(self, *op_symbols: str):
|
||||
"""Disable specific operations across all DSLs.
|
||||
|
||||
Args:
|
||||
*op_symbols (str): Names of operations to disable.
|
||||
|
||||
Example::
|
||||
|
||||
# Disable scaled matrix multiply across all DSLs
|
||||
torch.backends.python_native.disable_operations("scaled_mm")
|
||||
|
||||
# Disable multiple operations
|
||||
torch.backends.python_native.disable_operations(
|
||||
"scaled_mm", "flash_attention"
|
||||
)
|
||||
"""
|
||||
deregister_op_overrides = _get_registry_functions()[0]
|
||||
deregister_op_overrides(disable_op_symbols=list(op_symbols))
|
||||
|
||||
def enable_operations(self, *op_symbols: str):
|
||||
"""Re-enable specific operations across all DSLs.
|
||||
|
||||
Args:
|
||||
*op_symbols (str): Names of operations to re-enable.
|
||||
|
||||
Example::
|
||||
|
||||
# Re-enable previously disabled operations
|
||||
torch.backends.python_native.enable_operations(
|
||||
"scaled_mm", "flash_attention"
|
||||
)
|
||||
"""
|
||||
reenable_op_overrides = _get_registry_functions()[1]
|
||||
reenable_op_overrides(enable_op_symbols=list(op_symbols))
|
||||
|
||||
def disable_dispatch_keys(self, *dispatch_keys: str):
|
||||
"""Disable operations at specific dispatch keys.
|
||||
|
||||
Args:
|
||||
*dispatch_keys (str): Dispatch keys to disable (e.g., 'CUDA', 'CPU').
|
||||
|
||||
Example::
|
||||
|
||||
# Disable all native operations on CUDA
|
||||
torch.backends.python_native.disable_dispatch_keys("CUDA")
|
||||
"""
|
||||
deregister_op_overrides = _get_registry_functions()[0]
|
||||
deregister_op_overrides(disable_dispatch_keys=list(dispatch_keys))
|
||||
|
||||
def enable_dispatch_keys(self, *dispatch_keys: str):
|
||||
"""Re-enable operations at specific dispatch keys.
|
||||
|
||||
Args:
|
||||
*dispatch_keys (str): Dispatch keys to re-enable (e.g., 'CUDA', 'CPU').
|
||||
|
||||
Example::
|
||||
|
||||
# Re-enable native operations on CUDA
|
||||
torch.backends.python_native.enable_dispatch_keys("CUDA")
|
||||
"""
|
||||
reenable_op_overrides = _get_registry_functions()[1]
|
||||
reenable_op_overrides(enable_dispatch_keys=list(dispatch_keys))
|
||||
|
||||
@contextmanager
|
||||
def operations_disabled(self, *op_symbols: str):
|
||||
"""Context manager to temporarily disable operations.
|
||||
|
||||
Args:
|
||||
*op_symbols (str): Names of operations to temporarily disable.
|
||||
|
||||
Example::
|
||||
|
||||
with torch.backends.python_native.operations_disabled("scaled_mm"):
|
||||
# scaled_mm is disabled across all DSLs
|
||||
result = model(input)
|
||||
# scaled_mm is automatically re-enabled here
|
||||
"""
|
||||
filter_state = _get_filter_state()
|
||||
previously_disabled_ops = {
|
||||
op for op in op_symbols if op in filter_state._op_symbols
|
||||
}
|
||||
|
||||
self.disable_operations(*op_symbols)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Only re-enable operations that weren't already disabled
|
||||
ops_to_reenable = [
|
||||
op for op in op_symbols if op not in previously_disabled_ops
|
||||
]
|
||||
if ops_to_reenable:
|
||||
self.enable_operations(*ops_to_reenable)
|
||||
|
||||
@functools.lru_cache(maxsize=16) # noqa: B019
|
||||
def _get_dsl_controller(self, name: str) -> "DSLController":
|
||||
"""Get or create a DSL controller (cached)."""
|
||||
return DSLController(name)
|
||||
|
||||
def _get_registry_functions(self):
|
||||
"""Expose registry functions for testing."""
|
||||
return _get_registry_functions()
|
||||
|
||||
def is_operation_disabled(self, op_symbol: str) -> bool:
|
||||
"""Check if an operation is currently disabled."""
|
||||
filter_state = _get_filter_state()
|
||||
return op_symbol in filter_state._op_symbols
|
||||
|
||||
def is_dsl_disabled(self, dsl_name: str) -> bool:
|
||||
"""Check if a DSL is currently disabled."""
|
||||
filter_state = _get_filter_state()
|
||||
return dsl_name in filter_state._dsl_names
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
"""Dynamic attribute access for DSL controllers."""
|
||||
# Skip dunder attributes to avoid triggering DSL registry lookups
|
||||
# during torch initialization. inspect.getmodule() calls
|
||||
# hasattr(module, '__file__') which would otherwise cause a circular
|
||||
# import through _get_dsl_registry() while torch is still loading.
|
||||
if name.startswith("__") and name.endswith("__"):
|
||||
raise AttributeError(f"module '{self.__name__}' has no attribute '{name}'")
|
||||
|
||||
if name in self.all_dsls:
|
||||
return self._get_dsl_controller(name)
|
||||
|
||||
# Expose private functions for testing
|
||||
if name == "_get_dsl_module":
|
||||
return _get_dsl_module
|
||||
if name == "_get_registry_functions":
|
||||
return self._get_registry_functions
|
||||
if name == "_get_filter_state":
|
||||
return _get_filter_state
|
||||
|
||||
raise AttributeError(f"module '{self.__name__}' has no attribute '{name}'")
|
||||
|
||||
def __dir__(self):
|
||||
"""Return available attributes including DSL names."""
|
||||
attrs = set(super().__dir__())
|
||||
attrs.update(
|
||||
{
|
||||
"available_dsls",
|
||||
"all_dsls",
|
||||
"get_dsl_operations",
|
||||
"disable_operations",
|
||||
"enable_operations",
|
||||
"disable_dispatch_keys",
|
||||
"enable_dispatch_keys",
|
||||
"operations_disabled",
|
||||
"is_operation_disabled",
|
||||
"is_dsl_disabled",
|
||||
}
|
||||
)
|
||||
|
||||
# Add DSL names
|
||||
try:
|
||||
attrs.update(self.all_dsls)
|
||||
except Exception:
|
||||
# If registry not available yet, skip DSL names
|
||||
pass
|
||||
|
||||
return sorted(attrs)
|
||||
|
||||
|
||||
# Replace the current module with our enhanced version
|
||||
sys.modules[__name__] = PythonNativeModule(sys.modules[__name__])
|
||||
Reference in New Issue
Block a user