""" Memory Violation Detection Test Detects bad memory patterns in the LiteLLM codebase that can lead to memory leaks or OOMs. The detector uses a modular pattern-based system. To add detection for new memory patterns: 1. Create a Pattern subclass implementing get_pattern_name(), visit_assign(), and check_cleanup() - You can extend the Pattern class with additional methods as needed for your detection logic 2. Add the pattern to MemoryViolationDetector.DEFAULT_PATTERNS Currently detects: - queue.get() / queue.get_nowait() operations where variables aren't set to None - Class-level data structures that have add operations during runtime without size limits: * Built-in: list, dict, set * Collections: deque, defaultdict, Counter, OrderedDict, ChainMap * Queues: queue.Queue, asyncio.Queue (if unbounded, i.e., no maxsize parameter) * Heap operations: heapq.heappush(), heapq.heapreplace(), heapq.heappushpop() on class-level lists """ import ast import os from abc import ABC, abstractmethod from typing import List, Dict, Any, Optional, Sequence class Pattern(ABC): """Base class for memory violation detection patterns""" @abstractmethod def get_pattern_name(self) -> str: """Return unique identifier for this violation type""" pass @abstractmethod def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]: """Detect memory-sensitive operations in assignment. Returns list of {line, var_name, call} dicts.""" pass @abstractmethod def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt], context: Dict[str, Any]) -> List[Dict[str, Any]]: """Verify variables are set to None. Returns list of violation dicts.""" pass class QueueGetPattern(Pattern): """Detects queue.get()/get_nowait() operations that aren't cleared""" def get_pattern_name(self) -> str: return "queue_reference_not_cleared" def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]: """Detect queue.get() or queue.get_nowait() calls where object name contains 'queue'""" operations = [] if isinstance(node.value, ast.Call): func = node.value.func if isinstance(func, ast.Attribute) and func.attr in ("get", "get_nowait"): obj_name = context["get_attr_string"](func.value) if "queue" in obj_name.lower() and node.targets and isinstance(node.targets[0], ast.Name): operations.append({ "line": node.lineno, "var_name": node.targets[0].id, "call": context["get_call_string"](node.value), }) return operations def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt], context: Dict[str, Any]) -> List[Dict[str, Any]]: """Flag queue variables that aren't set to None""" violations = [] is_var_set_to_none = context["is_var_set_to_none"] current_function = context["current_function"] file_path = context["file_path"] queue_vars = {op["var_name"]: op["line"] for op in operations} for var_name, line_num in queue_vars.items(): if not is_var_set_to_none(var_name, function_body): violations.append({ "line": line_num, "type": self.get_pattern_name(), "var_name": var_name, "function": current_function, "file_path": file_path, "message": ( f"Queue variable '{var_name}' in function " f"'{current_function}' is not set to None after use. " f"If the runtime is overwhelmed, this can cause OOM (Out of Memory) errors." ), }) return violations class UnboundedDataStructurePattern(Pattern): """Detects class-level data structures (lists, dicts, sets) that can grow unbounded""" def get_pattern_name(self) -> str: return "unbounded_data_structure" def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]: """Detect list/dict/set creations that are at class level""" operations = [] # Check if this is a data structure creation is_data_structure = False structure_type = None if isinstance(node.value, (ast.List, ast.Dict, ast.Set)): is_data_structure = True if isinstance(node.value, ast.List): structure_type = "list" elif isinstance(node.value, ast.Dict): structure_type = "dict" elif isinstance(node.value, ast.Set): structure_type = "set" elif isinstance(node.value, ast.Call): # Check for list(), dict(), set() calls func = node.value.func if isinstance(func, ast.Name): if func.id in ("list", "dict", "set"): is_data_structure = True structure_type = func.id elif isinstance(func, ast.Attribute): # Handle cases like collections.defaultdict(list), collections.deque(), etc. obj_name = context["get_attr_string"](func.value) attr_name = func.attr # Check for collections module data structures if "collections" in obj_name.lower() or "collections" in str(func.value): if attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"): # For deque, we track it and let size checks determine if it's bounded # (deque with maxlen parameter is bounded, but we detect that via size checks) is_data_structure = True structure_type = attr_name elif attr_name in ("list", "dict", "set"): # collections.defaultdict(list) pattern is_data_structure = True structure_type = "defaultdict" if "defaultdict" in obj_name.lower() else attr_name # Check for queue.Queue, asyncio.Queue (if unbounded) elif "queue" in obj_name.lower() or "asyncio" in obj_name.lower(): if attr_name == "Queue": # Check if maxsize is set (bounded queue) has_maxsize = False for keyword in node.value.keywords: if keyword.arg == "maxsize": has_maxsize = True break if not has_maxsize: is_data_structure = True structure_type = "queue" # Direct attribute access like deque(), Counter(), etc. elif attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"): is_data_structure = True structure_type = attr_name if is_data_structure and node.targets and isinstance(node.targets[0], ast.Name): scope = context.get("current_scope", "function") # Only track if it's at class level (not module level) if scope == "class": operations.append({ "line": node.lineno, "var_name": node.targets[0].id, "structure_type": structure_type, "scope": scope, "call": context["get_call_string"](node.value) if isinstance(node.value, ast.Call) else f"{structure_type}()", }) return operations def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt], context: Dict[str, Any]) -> List[Dict[str, Any]]: """Flag persistent data structures that have add operations without size limits""" violations = [] current_function = context["current_function"] current_scope = context.get("current_scope", "function") file_path = context["file_path"] get_attr_string = context["get_attr_string"] # Skip if this is initialization code (module-level, class-level, or __init__ methods) # Only flag operations in regular methods/functions that can be called during runtime is_initialization = ( current_scope in ("module", "class") or current_function in ("__init__", "__new__", "__class_init__") or current_function is None # Module-level code ) if is_initialization: return violations # Don't flag initialization code # Track which variables have add operations and size checks var_add_operations = {} # var_name -> list of lines with add operations var_size_checks = {} # var_name -> has size limit check # Build a set of variable names to check tracked_vars = {op["var_name"]: op for op in operations} # Scan body for operations on these variables for stmt in function_body: for node in ast.walk(stmt): # Check for method calls that add items if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): attr_name = node.func.attr obj_name = get_attr_string(node.func.value) # Check if this is an add operation on one of our tracked variables for var_name, op in tracked_vars.items(): structure_type = op["structure_type"] # Match variable name (exact or as attribute) if obj_name == var_name or obj_name.endswith(f".{var_name}") or obj_name.endswith(f"['{var_name}']"): # Check for add operations add_ops = { "list": ["append", "extend", "insert"], "dict": ["update", "setdefault"], "set": ["add", "update"], "deque": ["append", "appendleft", "extend", "extendleft", "insert"], "defaultdict": ["update", "setdefault"], "Counter": ["update"], "OrderedDict": ["update", "setdefault"], "ChainMap": ["new_child"], "queue": ["put", "put_nowait"], } if attr_name in add_ops.get(structure_type, []): if var_name not in var_add_operations: var_add_operations[var_name] = [] var_add_operations[var_name].append(node.lineno) # Check for size limit checks (len() calls, maxsize/maxlen attributes) if (attr_name in ("__len__",) or "maxsize" in attr_name.lower() or "max_size" in attr_name.lower() or attr_name == "maxlen"): # For deque var_size_checks[var_name] = True # Check for heapq operations on tracked lists (heapq.heappush, heapq.heappop) if isinstance(node, ast.Call): func = node.func # Check for heapq.heappush(list_var, item) or heapq.heappop(list_var) if isinstance(func, ast.Attribute): func_obj = get_attr_string(func.value) func_name = func.attr # Check if it's a heapq operation if func_obj == "heapq" and func_name in ("heappush", "heapreplace", "heappushpop"): # First argument should be our tracked variable if len(node.args) > 0: arg_name = get_attr_string(node.args[0]) for var_name, op in tracked_vars.items(): if op["structure_type"] == "list" and ( arg_name == var_name or arg_name.endswith(f".{var_name}") ): if var_name not in var_add_operations: var_add_operations[var_name] = [] var_add_operations[var_name].append(node.lineno) # Check for dict item assignment: dict[key] = value if isinstance(node, ast.Assign): for target in node.targets: if isinstance(target, ast.Subscript): target_name = get_attr_string(target.value) for var_name in tracked_vars: if target_name == var_name or target_name.endswith(f".{var_name}"): if var_name not in var_add_operations: var_add_operations[var_name] = [] var_add_operations[var_name].append(node.lineno) # Check for augmented assignment: list += [...] if isinstance(node, ast.AugAssign): target_name = get_attr_string(node.target) for var_name in tracked_vars: if target_name == var_name or target_name.endswith(f".{var_name}"): if var_name not in var_add_operations: var_add_operations[var_name] = [] var_add_operations[var_name].append(node.lineno) # Check for size comparisons in conditionals if isinstance(node, (ast.If, ast.While, ast.Assert)): test = getattr(node, "test", None) if test: for comp_node in ast.walk(test): if isinstance(comp_node, ast.Compare): left_str = get_attr_string(comp_node.left) if hasattr(comp_node, "left") else "" # Check for len() calls if isinstance(comp_node.left, ast.Call): call_func = comp_node.left.func if isinstance(call_func, ast.Name) and call_func.id == "len": if len(comp_node.left.args) > 0: arg_name = get_attr_string(comp_node.left.args[0]) for var_name in tracked_vars: if arg_name == var_name or arg_name.endswith(f".{var_name}"): # Check if comparing to a limit for comparator in comp_node.comparators: if isinstance(comparator, ast.Constant): var_size_checks[var_name] = True elif isinstance(comparator, ast.Name): # Could be a constant like MAX_SIZE if "max" in comparator.id.lower() or "limit" in comparator.id.lower(): var_size_checks[var_name] = True # Handle deprecated ast.Num for Python < 3.8 try: Num = getattr(ast, "Num", None) if Num and isinstance(comparator, Num): var_size_checks[var_name] = True except (AttributeError, TypeError): pass # Check for direct variable comparisons for var_name in tracked_vars: if var_name in left_str: for comparator in comp_node.comparators: if isinstance(comparator, ast.Constant): var_size_checks[var_name] = True # Handle deprecated ast.Num for Python < 3.8 try: Num = getattr(ast, "Num", None) if Num and isinstance(comparator, Num): var_size_checks[var_name] = True except (AttributeError, TypeError): pass # Flag violations: persistent structures with add operations but no size checks for op in operations: var_name = op["var_name"] structure_type = op["structure_type"] if var_name in var_add_operations and var_name not in var_size_checks: violations.append({ "line": op["line"], "type": self.get_pattern_name(), "var_name": var_name, "function": current_function or "class-level", "file_path": file_path, "message": ( f"Class-level {structure_type} '{var_name}' " f"has add operations (lines {var_add_operations[var_name]}) but no size limit checks. " f"This can lead to unbounded memory growth and OOM errors during runtime." ), }) return violations class MemoryViolationDetector(ast.NodeVisitor): """AST visitor that detects memory violations using registered patterns""" DEFAULT_PATTERNS: List[Pattern] = [QueueGetPattern(), UnboundedDataStructurePattern()] def __init__(self, file_path: str, patterns: Optional[Sequence[Pattern]] = None): self.file_path = file_path self.violations: List[Dict[str, Any]] = [] self.current_function: Optional[str] = None self.current_scope: str = "module" # Track current scope: module, class, function self.patterns = self.DEFAULT_PATTERNS if patterns is None else patterns self.ast_tree: Optional[ast.Module] = None # Store full AST for module-level checks self.pattern_operations: Dict[str, List[Dict[str, Any]]] = { pattern.get_pattern_name(): [] for pattern in self.patterns } # Track class-level operations separately (for checking in functions) self.class_level_operations: Dict[str, List[Dict[str, Any]]] = { pattern.get_pattern_name(): [] for pattern in self.patterns } self._context = { "get_call_string": self._get_call_string, "get_attr_string": self._get_attr_string, "is_var_set_to_none": self._is_var_set_to_none, "current_function": None, "current_scope": "module", "file_path": file_path, } def visit_ClassDef(self, node): """Track class scope""" old_scope = self.current_scope self.current_scope = "class" self._context["current_scope"] = "class" self.generic_visit(node) self.current_scope = old_scope self._context["current_scope"] = old_scope def visit_FunctionDef(self, node): """Track function scope and check cleanup after visiting""" old_function = self.current_function old_scope = self.current_scope self.current_function = node.name self.current_scope = "function" self._context["current_function"] = node.name self._context["current_scope"] = "function" for pattern_name in self.pattern_operations: self.pattern_operations[pattern_name] = [] self.generic_visit(node) self._check_function_cleanup(node) self.current_function = old_function self.current_scope = old_scope self._context["current_function"] = old_function self._context["current_scope"] = old_scope def visit_AsyncFunctionDef(self, node): """Track async function scope and check cleanup after visiting""" old_function = self.current_function old_scope = self.current_scope self.current_function = node.name self.current_scope = "function" self._context["current_function"] = node.name self._context["current_scope"] = "function" for pattern_name in self.pattern_operations: self.pattern_operations[pattern_name] = [] self.generic_visit(node) self._check_function_cleanup(node) self.current_function = old_function self.current_scope = old_scope self._context["current_function"] = old_function self._context["current_scope"] = old_scope def visit_Assign(self, node): """Detect memory-sensitive operations in assignments""" for pattern in self.patterns: operations = pattern.visit_assign(node, self._context) # Track function-level operations self.pattern_operations[pattern.get_pattern_name()].extend(operations) # Track class-level operations separately (for checking in functions) for op in operations: if op.get("scope") == "class": self.class_level_operations[pattern.get_pattern_name()].append(op) self.generic_visit(node) def _check_function_cleanup(self, node): """Check cleanup for all detected operations""" for pattern in self.patterns: operations = self.pattern_operations[pattern.get_pattern_name()] if operations: violations = pattern.check_cleanup(operations, node.body, self._context) self.violations.extend(violations) # For UnboundedDataStructurePattern, also check if this function modifies class-level structures if isinstance(pattern, UnboundedDataStructurePattern): class_ops = self.class_level_operations[pattern.get_pattern_name()] if class_ops and self.current_function not in ("__init__", "__new__", "__class_init__", None): # Check if this regular function modifies class-level structures violations = pattern.check_cleanup(class_ops, node.body, self._context) self.violations.extend(violations) def _check_module_level_cleanup(self): """Check cleanup for module/class level operations""" # Module-level operations are now checked when visiting functions # This method is kept for potential future use but doesn't need to do anything # since we only want to flag runtime modifications in functions, not initialization code pass def _is_var_set_to_none(self, var_name: str, body: List[ast.stmt]) -> bool: """Check if variable is set to None after its initial assignment""" assignment_line = None for stmt in body: for node in ast.walk(stmt): if isinstance(node, ast.Assign): for target in node.targets: if isinstance(target, ast.Name) and target.id == var_name: assignment_line = node.lineno break if assignment_line: break if assignment_line: break if not assignment_line: return False for stmt in body: for node in ast.walk(stmt): if isinstance(node, ast.Assign): for target in node.targets: if isinstance(target, ast.Name) and target.id == var_name and node.lineno > assignment_line: if isinstance(node.value, ast.Constant) and node.value.value is None: return True try: NameConstant = getattr(ast, "NameConstant", None) if NameConstant and isinstance(node.value, NameConstant): if getattr(node.value, "value", None) is None: return True except (AttributeError, TypeError): pass return False def _get_call_string(self, node: ast.Call) -> str: """Get string representation of function call""" try: if hasattr(ast, "unparse"): return ast.unparse(node) elif isinstance(node.func, ast.Attribute): return f"{self._get_attr_string(node.func.value)}.{node.func.attr}()" return str(node) except Exception: return str(node) def _get_attr_string(self, node: ast.AST) -> str: """Get string representation of attribute access""" if isinstance(node, ast.Name): return node.id elif isinstance(node, ast.Attribute): return f"{self._get_attr_string(node.value)}.{node.attr}" return str(node) def check_file_for_memory_violations(file_path: str, patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]: """Check a single file for memory violations""" try: with open(file_path, "r", encoding="utf-8") as f: content = f.read() if "test" in file_path.lower() or "__pycache__" in file_path: return [] tree = ast.parse(content, filename=file_path) detector = MemoryViolationDetector(file_path, patterns) detector.ast_tree = tree # Store AST for potential future use detector.visit(tree) # Class-level operations are checked when visiting functions return detector.violations except Exception as e: print(f"Error parsing {file_path}: {e}") return [] def check_directory_for_memory_violations(directory_path: str, ignore_patterns: Optional[List[str]] = None, patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]: """Recursively scan directory for memory violations""" if ignore_patterns is None: ignore_patterns = ["__pycache__", ".pyc", "site-packages", "venv", ".venv", "env", ".env", "node_modules", "tests"] all_violations = [] for root, _dirs, files in os.walk(directory_path): if any(pattern in root for pattern in ignore_patterns): continue for file in files: if file.endswith(".py"): violations = check_file_for_memory_violations(os.path.join(root, file), patterns) all_violations.extend(violations) return all_violations def main(): """Run memory violation detection on codebase""" codebase_path = "./litellm" print("=" * 80) print("MEMORY VIOLATION DETECTION TEST") print("=" * 80) print(f"Scanning: {codebase_path}") print(f"Active patterns: {', '.join(p.get_pattern_name() for p in MemoryViolationDetector.DEFAULT_PATTERNS)}") print() violations = check_directory_for_memory_violations(codebase_path) if violations: by_type = {} for v in violations: vtype = v["type"] if vtype not in by_type: by_type[vtype] = [] by_type[vtype].append(v) print("MEMORY VIOLATIONS FOUND:") print("=" * 80) total = len(violations) for vtype, vlist in by_type.items(): print(f"\n{vtype.upper().replace('_', ' ')}: {len(vlist)} violation(s)") print("-" * 80) for v in vlist[:10]: print(f" [VIOLATION] {v['file_path'] if 'file_path' in v else 'unknown'}:{v['line']}") print(f" Function: {v['function']}") print(f" Variable: {v['var_name']}") print(f" {v['message']}") print() if len(vlist) > 10: print(f" ... and {len(vlist) - 10} more violations of this type") print("=" * 80) print(f"TOTAL VIOLATIONS: {total}") print() print("RECOMMENDATIONS:") print(" 1. Set queue variables to None after use: obj = queue.get(); ...; obj = None") print(" 2. Use bounded queues to prevent unbounded accumulation") print(" 3. Process items faster than they're added, or drain queues periodically") print(" 4. For class-level data structures (lists, dicts, sets) that are modified at runtime:") print(" - Add size limit checks: if len(data) >= MAX_SIZE: ...") print(" - Implement periodic cleanup or use bounded collections") print(" - Consider using collections.deque with maxlen for lists") print("=" * 80) first_v = violations[0] raise Exception( f"Found {total} memory violations! " f"First violation: {first_v.get('file_path', 'unknown')}:{first_v['line']} - " f"{first_v['message']}" ) else: print("OK No memory violations found!") if __name__ == "__main__": main()