Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 2 additions & 6 deletions nodes/src/nodes/ocr/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
7 changes: 7 additions & 0 deletions nodes/src/nodes/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
55 changes: 55 additions & 0 deletions packages/server/engine-lib/rocketlib-python/lib/depends.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@

from __future__ import annotations

import functools
import hashlib
import os
import platform
Expand Down Expand Up @@ -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'
Expand All @@ -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:
Expand Down Expand Up @@ -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():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 nit: the variant decision is invisible — when a user reports slow polars or an Illegal-instruction crash, there's no trace of which branch ran. Add one debug(f'polars variant: {"lts-cpu" if missing_avx2 else "standard"} (avx2={not missing_avx2})') (call _is_x86_64_missing_avx2() once into a local). Also worth updating the docstring above, which still documents only the uv/onnxruntime excludes. Optional: @functools.lru_cache on _is_x86_64_missing_avx2 — it currently re-runs sysctl/cpuinfo reads on every excludes write (twice per depends() call across all nodes at startup).

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'
Comment thread
coderabbitai[bot] marked this conversation as resolved.

with open(excludes_path, 'w', encoding='utf-8') as f:
f.write(excludes)
return excludes_path
Expand Down
118 changes: 118 additions & 0 deletions packages/server/engine-lib/rocketlib-python/tests/test_depends.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
import os

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 bug: nothing runs pytest against the repo-root test/ directory, so this suite never executes — not locally via ./builder test and not in CI. The engine-lib pytest task (packages/server/scripts/tasks.js:1062) targets packages/server/engine-lib/rocketlib-python/tests/; move the file there. The sys.path bootstrap then shrinks to one .. hop as well.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 risk: module-level sys.modules['engLib'] = MagicMock() persists for the entire pytest session — any later test importing a rocketlib module gets the mock instead of the real thing. Use the scoped save/patch/restore pattern the other suites use (e.g. nodes/test/guardrails/test_all.py), or @patch.dict(sys.modules, ...) at class level.

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()

Loading