diff --git a/nodes/src/nodes/ocr/requirements.txt b/nodes/src/nodes/ocr/requirements.txt index 0d430a193..7e372d093 100644 --- a/nodes/src/nodes/ocr/requirements.txt +++ b/nodes/src/nodes/ocr/requirements.txt @@ -2,11 +2,7 @@ # # Note: PyTorch and OCR model dependencies (EasyOCR, DocTR, etc.) are handled # by ai.common.models which uses ai.common.torch for torch installation. -# This node only needs img2table for table structure detection. +# This node only needs table structure detection. # -# numpy and pillow are base engine dependencies - don't reinstall them here +# numpy and the table detection library are base engine dependencies, and pillow arrives transitively via them - don't reinstall them here # as that can corrupt already-loaded modules in the server process. -# -img2table -pillow -numpy diff --git a/nodes/src/nodes/requirements.txt b/nodes/src/nodes/requirements.txt index 87c8fa14a..65217e79e 100644 --- a/nodes/src/nodes/requirements.txt +++ b/nodes/src/nodes/requirements.txt @@ -23,3 +23,10 @@ safetensors # Needed for dependencies compilations Cython versioneer + +# OCR base dependencies +img2table +# Explicit polars requirements +polars>=1.0.0 +# Included explicitly for older/emulated CPUs without AVX2 (see depends.py) +polars-lts-cpu>=1.0.0 diff --git a/packages/server/engine-lib/rocketlib-python/lib/depends.py b/packages/server/engine-lib/rocketlib-python/lib/depends.py index 2517959c8..d64ddd61f 100644 --- a/packages/server/engine-lib/rocketlib-python/lib/depends.py +++ b/packages/server/engine-lib/rocketlib-python/lib/depends.py @@ -32,6 +32,7 @@ from __future__ import annotations +import functools import hashlib import os import platform @@ -632,6 +633,11 @@ def _find_requirement_files() -> list[str]: def _compute_hash(file_paths: list[str]) -> str: """Compute a fast hash from file metadata (mtime + size).""" hasher = hashlib.md5() + + # Mix the AVX2 verdict into the hash so a copied cache invalidates if the CPU architecture changes + avx2_missing = '1' if _is_x86_64_missing_avx2() else '0' + hasher.update(f'avx2_missing:{avx2_missing}\n'.encode()) + for path in sorted(file_paths): stat = os.stat(path) entry = f'{path}:{stat.st_size}:{stat.st_mtime_ns}\n' @@ -654,6 +660,41 @@ def _save_hash(hash_file: str, hash_value: str): f.write(hash_value) +@functools.lru_cache() +def _is_x86_64_missing_avx2() -> bool: + """Check if CPU is x86_64 but lacks AVX2 (e.g. Rosetta 2 or old Intel).""" + machine = platform.machine().lower() + if machine not in ('x86_64', 'amd64'): + return False + + system = platform.system() + if system == 'Darwin': + # Check Rosetta or old Mac without AVX2 + try: + output = subprocess.check_output(['sysctl', '-n', 'hw.optional.avx2_0'], text=True) + return output.strip() != '1' + except Exception: + return True + elif system == 'Linux': + try: + with open('/proc/cpuinfo', 'r', encoding='utf-8') as f: + return 'avx2' not in f.read() + except Exception: + return True + elif system == 'Windows': + try: + import ctypes + kernel32 = ctypes.windll.kernel32 + kernel32.IsProcessorFeaturePresent.argtypes = [ctypes.c_uint32] + kernel32.IsProcessorFeaturePresent.restype = ctypes.c_int + # PF_AVX2_INSTRUCTIONS_AVAILABLE is 40 + return not bool(kernel32.IsProcessorFeaturePresent(40)) + except Exception: + return True + + return False + + def _combine_requirements(file_paths: list[str], output_path: str): """Concatenate all requirement files into one.""" with open(output_path, 'w', encoding='utf-8') as out: @@ -759,11 +800,25 @@ def _write_excludes_file() -> str: Excludes `uv` (bootstrapped by depends.py; pip-installing it crashes on Windows) and, on non-Darwin, plain `onnxruntime` (it clobbers onnxruntime-gpu in the same folder; the gpu build provides `import onnxruntime`). + + Also conditionally excludes the incorrect polars distribution for the current CPU architecture. + Since polars and polars-lts-cpu both provide `import polars` and can't be installed together, + we rely on `--excludes` to prevent uv from resolving both. """ excludes_path = os.path.join(engine_cache_dir(), 'excludes.txt') excludes = 'uv\n' if platform.system() != 'Darwin': excludes += 'onnxruntime\n' + + if _is_x86_64_missing_avx2(): + debug('AVX2 instructions not detected; using polars-lts-cpu fallback') + # Exclude standard polars so uv doesn't install it alongside polars-lts-cpu + excludes += 'polars\n' + else: + debug('AVX2 fallback not required; using standard polars') + # Exclude the cpu variant so it isn't installed on standard AVX2 hardware + excludes += 'polars-lts-cpu\n' + with open(excludes_path, 'w', encoding='utf-8') as f: f.write(excludes) return excludes_path diff --git a/packages/server/engine-lib/rocketlib-python/tests/test_depends.py b/packages/server/engine-lib/rocketlib-python/tests/test_depends.py new file mode 100644 index 000000000..a72e10266 --- /dev/null +++ b/packages/server/engine-lib/rocketlib-python/tests/test_depends.py @@ -0,0 +1,118 @@ +import os +import sys +import unittest +from unittest.mock import patch, mock_open, MagicMock + +# Add depends.py directory to path +DEPENDS_DIR = os.path.abspath( + os.path.join(os.path.dirname(__file__), '..', 'lib') +) +sys.path.append(DEPENDS_DIR) + +# Dynamically stub engLib using a scoped pattern to prevent session pollution +def _load_depends_module(): + """Load depends.py as a standalone module.""" + saved_englib = sys.modules.get('engLib') + try: + sys.modules['engLib'] = MagicMock() + import depends + return depends + finally: + if saved_englib is None: + sys.modules.pop('engLib', None) + else: + sys.modules['engLib'] = saved_englib + +depends = _load_depends_module() + + +class TestDepends(unittest.TestCase): + def setUp(self): + # Clear lru_cache before tests so we can evaluate multiple mock states per test + depends._is_x86_64_missing_avx2.cache_clear() + + @patch('depends.platform.machine') + @patch('depends.platform.system') + def test_missing_avx2_non_x86(self, mock_system, mock_machine): + # If architecture is ARM (e.g. apple silicon without rosetta), it should return False + # as it doesn't even need to check for AVX2 instructions + mock_machine.return_value = 'arm64' + self.assertFalse(depends._is_x86_64_missing_avx2()) + + mock_machine.return_value = 'aarch64' + self.assertFalse(depends._is_x86_64_missing_avx2()) + + @patch('depends.subprocess.check_output') + @patch('depends.platform.system') + @patch('depends.platform.machine') + def test_missing_avx2_darwin_rosetta(self, mock_machine, mock_system, mock_subprocess): + mock_machine.return_value = 'x86_64' + mock_system.return_value = 'Darwin' + + # Test case: sysctl returns 0 (missing AVX2) + mock_subprocess.return_value = '0\n' + depends._is_x86_64_missing_avx2.cache_clear() + self.assertTrue(depends._is_x86_64_missing_avx2()) + + # Test case: sysctl returns 1 (has AVX2) + mock_subprocess.return_value = '1\n' + depends._is_x86_64_missing_avx2.cache_clear() + self.assertFalse(depends._is_x86_64_missing_avx2()) + + # Test case: sysctl fails (failsafe to True) + mock_subprocess.side_effect = Exception('sysctl failed') + depends._is_x86_64_missing_avx2.cache_clear() + self.assertTrue(depends._is_x86_64_missing_avx2()) + + @patch('depends.open', new_callable=mock_open) + @patch('depends.platform.system') + @patch('depends.platform.machine') + def test_missing_avx2_linux(self, mock_machine, mock_system, mock_file): + mock_machine.return_value = 'x86_64' + mock_system.return_value = 'Linux' + + # Test case: cpuinfo has avx2 + mock_file.return_value.read.return_value = 'flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single pti intel_ppin ssbd mba ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts pku ospke md_clear pconfig stibp_always_on flush_l1d arch_capabilities' + depends._is_x86_64_missing_avx2.cache_clear() + self.assertFalse(depends._is_x86_64_missing_avx2()) + + # Test case: cpuinfo does not have avx2 + mock_file.return_value.read.return_value = 'flags: fpu vme de pse tsc' + depends._is_x86_64_missing_avx2.cache_clear() + self.assertTrue(depends._is_x86_64_missing_avx2()) + + # Test case: file open fails + mock_file.side_effect = Exception('Permission denied') + depends._is_x86_64_missing_avx2.cache_clear() + self.assertTrue(depends._is_x86_64_missing_avx2()) + + @patch.dict('sys.modules', {'ctypes': MagicMock()}) + @patch('depends.platform.system') + @patch('depends.platform.machine') + def test_missing_avx2_windows(self, mock_machine, mock_system): + mock_machine.return_value = 'AMD64' # Windows uses AMD64 + mock_system.return_value = 'Windows' + + # Test case: kernel32.IsProcessorFeaturePresent(40) returns True + mock_ctypes = sys.modules['ctypes'] + mock_kernel32 = MagicMock() + mock_kernel32.IsProcessorFeaturePresent.return_value = True + mock_ctypes.windll.kernel32 = mock_kernel32 + + depends._is_x86_64_missing_avx2.cache_clear() + self.assertFalse(depends._is_x86_64_missing_avx2()) + + # Test case: returns False + mock_kernel32.IsProcessorFeaturePresent.return_value = False + depends._is_x86_64_missing_avx2.cache_clear() + self.assertTrue(depends._is_x86_64_missing_avx2()) + + # Test case: ctypes fails + mock_ctypes.windll.kernel32.IsProcessorFeaturePresent.side_effect = Exception('Failed') + depends._is_x86_64_missing_avx2.cache_clear() + self.assertTrue(depends._is_x86_64_missing_avx2()) + + +if __name__ == '__main__': + unittest.main() +