libs fixes

.github/workflows/ci.yml

@@ -14,6 +14,9 @@ jobs:
       - name: build image
         run: make build
 
+      - name: test-no-docker
+        run: make test-no-docker
+
       - name: test
         run: make test
 

Dockerfile

@@ -22,5 +22,10 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     pip install .[dev]
 
 COPY pitch_detectors /app/pitch_detectors
+
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install --no-deps .
+
 COPY tests /app/tests
+COPY scripts/ /app/scripts
 COPY data /app/data

Makefile

@@ -9,8 +9,12 @@ push:
 	docker push $(IMAGE)
 	docker push tandav/pitch-detectors:latest
 
+# python -m pitch_detectors.util ld_library_path
+
 .PHONY: test-no-docker
 test-no-docker:
+	PITCH_DETECTORS_PENN_CHECKPOINT_PATH=/home/tandav/docs/bhairava/libmv/data/fcnf0++.pt \
+	PITCH_DETECTORS_SPICE_MODEL_PATH=/home/tandav/docs/bhairava/libmv/data/spice_model \
 	pytest -x -v --cov pitch_detectors
 
 .PHONY: test

pitch_detectors/util.py

@@ -37,18 +37,30 @@ def source_hashes() -> dict[str, str]:
     return hashes
 
 
-def python_version() -> str:
-    return f'python{sys.version_info.major}.{sys.version_info.minor}'
+def ld_library_path() -> str:
+    site_packages = f'{sys.exec_prefix}/lib/python{sys.version_info.major}.{sys.version_info.minor}/site-packages'
+    libs = [
+        f'{site_packages}/nvidia/curand/lib',
+        f'{site_packages}/nvidia/cuda_runtime/lib',
+        f'{site_packages}/nvidia/cusparse/lib',
+        f'{site_packages}/nvidia/cudnn/lib',
+        f'{site_packages}/nvidia/cuda_nvrtc/lib',
+        f'{site_packages}/nvidia/cuda_cupti/lib',
+        f'{site_packages}/nvidia/nccl/lib',
+        f'{site_packages}/nvidia/cusolver/lib',
+        f'{site_packages}/nvidia/nvtx/lib',
+        f'{site_packages}/nvidia/cufft/lib',
+        f'{site_packages}/nvidia/cublas/lib',
+        f'{site_packages}/tensorrt',
+    ]
+    return ':'.join(libs)
 
 
-def ld_library_path() -> str:
-    cuda_home = f'{sys.exec_prefix}/lib/{python_version()}/site-packages/nvidia'
-    paths = []
-    for p in Path(cuda_home).iterdir():
-        if not p.is_dir() or p.name == '__pycache__':
-            continue
-        lib = p / 'lib'
-        if not lib.exists():
-            raise FileNotFoundError(f'lib not found: {lib}')
-        paths.append(str(lib))
-    return ':'.join(paths)
+if __name__ == '__main__':
+    supported_actions = {'ld_library_path'}
+    if len(sys.argv) != 2:  # noqa: PLR2004
+        raise ValueError('Pass action as argument. Supported_actions:', supported_actions)
+    if sys.argv[1] == 'ld_library_path':
+        print(ld_library_path())
+    else:
+        raise ValueError(f'Action {sys.argv[1]} not supported. Supported_actions:', supported_actions)

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "resampy",
     "scipy",
     "tensorflow<2.12.0",
+    # "tf-nightly", # trying this instead tensorflow. It support tensorrt8. (regular tensorflow only supports outdated tensorrt 7 which is python3.8 only)
     "tensorflow-hub",
     "torch",
     "torch-yin",

tests/algorithms_test.py

@@ -1,3 +1,4 @@
+import os
 import subprocess
 import sys
 
@@ -8,7 +9,7 @@
 
 @pytest.mark.order(3)
 @pytest.mark.parametrize('algorithm', ALGORITHMS)
-@pytest.mark.parametrize('gpu', ['true', 'false'])
+@pytest.mark.parametrize('gpu', ['false'] if os.environ.get('PITCH_DETECTORS_GPU') == 'false' else ['true', 'false'])
 def test_detection(algorithm, environ, gpu):
     env = environ | {
         'PITCH_DETECTORS_ALGORITHM': algorithm.name(),

tests/conftest.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 import pytest
@@ -14,9 +15,14 @@ def record():
 
 @pytest.fixture
 def environ():
-    return {
+    env = {
         'PITCH_DETECTORS_GPU_MEMORY_LIMIT': 'true',
-        'PITCH_DETECTORS_PENN_CHECKPOINT_PATH': '/home/tandav/docs/bhairava/libmv/data/fcnf0++.pt',
-        'PITCH_DETECTORS_SPICE_MODEL_PATH': '/home/tandav/docs/bhairava/libmv/data/spice_model',
+        # 'PITCH_DETECTORS_PENN_CHECKPOINT_PATH': os.environ.get('PITCH_DETECTORS_PENN_CHECKPOINT_PATH'),
+        # 'PITCH_DETECTORS_SPICE_MODEL_PATH': '/home/tandav/docs/bhairava/libmv/data/spice_model',
         'LD_LIBRARY_PATH': util.ld_library_path(),
     }
+    if 'PITCH_DETECTORS_PENN_CHECKPOINT_PATH' in os.environ:
+        env['PITCH_DETECTORS_PENN_CHECKPOINT_PATH'] = os.environ['PITCH_DETECTORS_PENN_CHECKPOINT_PATH']
+    if 'PITCH_DETECTORS_SPICE_MODEL_PATH' in os.environ:
+        env['PITCH_DETECTORS_SPICE_MODEL_PATH'] = os.environ['PITCH_DETECTORS_SPICE_MODEL_PATH']
+    return env

tests/ensemble_test.py

@@ -8,7 +8,11 @@
 
 
 def test_ensemble(record, environ):
-    with mock.patch.dict(os.environ, environ | {'PITCH_DETECTORS_GPU': 'true'}):
+    if os.environ.get('PITCH_DETECTORS_GPU') == 'false':
+        env = environ
+    else:
+        env = environ | {'PITCH_DETECTORS_GPU': 'true'}
+    with mock.patch.dict(os.environ, env):
         alg = algorithms.Ensemble(record.a, record.fs, algorithms=algorithms.ALGORITHMS)
         assert alg.f0.shape == alg.t.shape
         algorithms_cache = {k: F0(alg.t, alg.f0) for k, alg in alg._algorithms.items()}