Merge remote-tracking branch 'origin/next' into nme_gpu

Author: Jeronymous

celery_app/register.py

@@ -60,9 +60,10 @@ def register(is_heartbeat: bool = False) -> bool:
 def unregister() -> None:
     """Un-register the service"""
     try:
+        password = os.environ.get("BROKER_PASS", None)
         host, port = os.environ.get("SERVICES_BROKER").split("//")[1].split(":")
         r = redis.Redis(
-            host=host, port=int(port), db=SERVICE_DISCOVERY_DB, password="password"
+            host=host, port=int(port), db=SERVICE_DISCOVERY_DB, password=password
         )
         r.json().delete(f"service:{host_name}")
     except Exception as error:

docker-entrypoint.sh

@@ -50,7 +50,7 @@ run_celery_worker() {
     python -c "from celery_app.register import register; register()" || exit $?
     echo "Service registered"
     ## WORKER
-    celery --app=celery_app.celeryapp worker $OPT -Ofair -n diarization_worker@%h --queues=$QUEUE -c ${CONCURRENCY:-1} || exit $?
+    celery --app=celery_app.celeryapp worker $OPT -Ofair -n ${SERVICE_NAME}_worker@%h --queues=$QUEUE -c ${CONCURRENCY:-1} || exit $?
     ## UNREGISTERING
     python -c "from celery_app.register import unregister; unregister()" || exit $?
     echo "Service unregistered"

healthcheck.sh

@@ -6,26 +6,27 @@ if [ "$SERVICE_MODE" = "http" ]
 then
     curl --fail http://localhost:80/healthcheck || exit 1
 else
-    # Update last alive, using Heartbeat(register(True))
-    python -c "from celery_app.register import register; register(True)"
+    # Update last alive
+    python -c "from celery_app.register import register; register(False)"
 
     # Check if Celery worker process is running
     # warning : We might rework this when switching to non-root user. Which is required for security reasons.
     # Our docker images are currently running as root and this is bad :)
-    if ! pgrep -f "celeryapp worker"; then
+    PID=`pgrep -f "celeryapp worker"`
+    if [ -z "$PID" ]; then
         echo "HealtchCheck FAIL : Celery worker process not running"
         exit 1
     fi
 
-    # Check GPU utilization
-    if nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits | grep -v '^0$'; then
-        #"GPU is being utilized, assuming healthy"
-        exit 0
-    else
-        # If GPU is not being utilized, attempt a ping as a secondary check
-        if ! celery --app=celery_app.celeryapp inspect ping -d ${SERVICE_NAME}_worker@$HOSTNAME --timeout=20; then
-            echo "HealtchCheck FAIL : Celery worker not responding in time and GPU is not being utilized"
-            exit 1
+    # Attempt a ping
+    if ! celery --app=celery_app.celeryapp inspect ping -d ${SERVICE_NAME}_worker@$HOSTNAME --timeout=20; then
+        # Check GPU utilization
+        if nvidia-smi --query-compute-apps pid --format=csv,noheader | grep $PID; then
+            # GPU is being utilized, assuming healthy
+            echo "Celery worker not responding in time but GPU is being utilized (trying to ping again)"
+            exit 0
         fi
+        echo "HealtchCheck FAIL : Celery worker not responding in time and GPU is not being utilized"
+        exit 1
     fi
 fi

simple/.envdefault

@@ -11,6 +11,7 @@ CONCURRENCY=2
 
 # Enforce used device ("cpu" or "cuda")
 # DEVICE=cpu
+# DEVICE_CLUSTERING=cpu
 
 # Maximum number of threads on CPU
 NUM_THREADS=4

simple/Dockerfile

@@ -1,30 +1,4 @@
-# base stage target - common dependencies
-FROM python:3.10-slim as base
-
-# Set a common work directory
-WORKDIR /usr/src/app
-
-# Installing dependencies
-
-# Copy application code
-COPY celery_app /usr/src/app/celery_app
-COPY http_server /usr/src/app/http_server
-COPY document /usr/src/app/document
-COPY docker-entrypoint.sh wait-for-it.sh healthcheck.sh ./
-COPY simple/diarization /usr/src/app/diarization
-COPY simple/requirements.txt /usr/src/app/
-COPY simple/RELEASE.md ./
-
-############################################################################################################
-
-# Here we might build a specific image for the CPU target, use different dependencies, entrypoints, etc.
-# using  : FROM base as cpu
-# not relevant for this dockerfile
-
-############################################################################################################
-
-# gpu stage target - python app with nvidia/cuda base image
-FROM nvidia/cuda:12.3.2-base-ubuntu22.04 as gpu
+FROM python:3.10
 LABEL maintainer="[email protected], [email protected], [email protected]"
 
 # Re-Sets work directory (not sure if this is necessary)
@@ -37,13 +11,23 @@ RUN apt-get update && \
     rm -rf /var/lib/apt/lists/* && \
     ln -s /usr/bin/python3.10 /usr/bin/python
 
+# # Set a common work directory
+WORKDIR /usr/src/app
+
 # Install python packages
 # Reuse labels and base dependencies from the base stage
-COPY --from=base /usr/src/app/requirements.txt /usr/src/app/requirements.txt
-RUN pip install --no-cache-dir -r requirements.txt
+COPY simple/requirements.txt /usr/src/app/requirements.txt
+RUN pip install --upgrade pip && pip install --no-cache-dir -r requirements.txt
 
-# Reuse labels and base dependencies from the base stage
-COPY --from=base /usr/src/app /usr/src/app
+
+# Copy application code
+COPY celery_app /usr/src/app/celery_app
+COPY http_server /usr/src/app/http_server
+COPY document /usr/src/app/document
+COPY docker-entrypoint.sh wait-for-it.sh healthcheck.sh ./
+COPY simple/diarization /usr/src/app/diarization
+COPY simple/requirements.txt /usr/src/app/
+COPY simple/RELEASE.md ./
 
 # Set path
 ENV PYTHONPATH="${PYTHONPATH}:/usr/src/app/diarization"

simple/Dockerfile.stagewise

@@ -0,0 +1,56 @@
+# base stage target - common dependencies
+FROM python:3.10-slim as base
+
+# Set a common work directory
+WORKDIR /usr/src/app
+
+# Installing dependencies
+
+# Copy application code
+COPY celery_app /usr/src/app/celery_app
+COPY http_server /usr/src/app/http_server
+COPY document /usr/src/app/document
+COPY docker-entrypoint.sh wait-for-it.sh healthcheck.sh ./
+COPY simple/diarization /usr/src/app/diarization
+COPY simple/requirements.txt /usr/src/app/
+COPY simple/RELEASE.md ./
+
+############################################################################################################
+
+# Here we might build a specific image for the CPU target, use different dependencies, entrypoints, etc.
+# using  : FROM base as cpu
+# not relevant for this dockerfile
+
+############################################################################################################
+
+# gpu stage target - python app with nvidia/cuda base image
+FROM nvidia/cuda:12.3.2-base-ubuntu22.04 as gpu
+LABEL maintainer="[email protected], [email protected], [email protected]"
+
+# Re-Sets work directory (not sure if this is necessary)
+WORKDIR /usr/src/app
+
+# Install python and other dependencies
+RUN apt-get update && \
+    apt-get install -y python3.10 python3-pip ffmpeg git && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* && \
+    ln -s /usr/bin/python3.10 /usr/bin/python
+
+# Install python packages
+# Reuse labels and base dependencies from the base stage
+COPY --from=base /usr/src/app/requirements.txt /usr/src/app/requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Reuse labels and base dependencies from the base stage
+COPY --from=base /usr/src/app /usr/src/app
+
+# Set path
+ENV PYTHONPATH="${PYTHONPATH}:/usr/src/app/diarization"
+
+# Extract version information
+RUN export VERSION=$(awk -v RS='' '/#/ {print; exit}' RELEASE.md | head -1 | sed 's/#//' | sed 's/ //')
+
+# Health check and entrypoint for GPU target
+HEALTHCHECK CMD ./healthcheck.sh
+ENTRYPOINT ["./docker-entrypoint.sh"]

simple/README.md

@@ -59,6 +59,8 @@ An example of .env file is provided in [simple/.envdefault](https://github.com/l
 |:-|:-|:-|
 | `SERVING_MODE` | (Required) Specify launch mode | `http` |
 | `CONCURRENCY` | Number of worker(s) additional to the main worker | `0` \| `1` \| `2` \| ... |
+| `DEVICE` | Device to use for the embedding model (by default, GPU/CUDA is used if it is available, CPU otherwise) | `cpu` \| `cuda` |
+| `DEVICE_CLUSTERING` | Device to use for clustering (by default, GPU/CUDA is used if it is available, CPU otherwise) | `cpu` \| `cuda` |
 | `NUM_THREADS` | Number of threads (maximum) to use for things running on CPU | `1` \| `4` \| ... |
 | `DEVICE` | Device to use for the embeddings model (by default, GPU/CUDA is used if it is available, CPU otherwise) | `cpu` \| `cuda` \| `cuda:1` ... |
 | `DEVICE_VAD` | Device to use for the Voice Activity Detection (by default, CPU) | `cpu` \| `cuda` \| `cuda:1` ... |

simple/RELEASE.md

@@ -1,4 +1,6 @@
-# 1.0.2
+# 1.1.0
+- Use NEMO spectral clustering, with DEVICE_CLUSTERING environment variable to enforce using CPU ("cpu") or GPU ("cuda")
+- Switch from silero VAD version 4 to version 3.1 (less problems)
 - Fix healthcheck on GPU
 
 # 1.0.1

simple/diarization/processing/__init__.py

@@ -3,10 +3,16 @@
 
 
 device = os.environ.get("DEVICE")
-if device is None:
-   USE_GPU = torch.cuda.is_available()
+device_vad = "cpu" # Not implemented os.environ.get("DEVICE_VAD", "cpu")
+device_clustering = os.environ.get("DEVICE_CLUSTERING")
+if torch.cuda.is_available():
+   USE_GPU = (
+      device != "cpu"
+      or device_clustering  != "cpu"
+      or device_vad != "cpu"
+   )
 else:
-   USE_GPU = (device != "cpu")
+   USE_GPU = False
 
 device_vad = os.environ.get("DEVICE_VAD", "cpu")
 device_clustering = os.environ.get("DEVICE_CLUSTERING", "cpu")
@@ -21,6 +27,6 @@
 
 from .speakerdiarization import SpeakerDiarization
 
-diarizationworker = SpeakerDiarization(device=device, device_vad=device_vad, device_clustering=device_clustering, num_threads=NUM_THREADS)
+diarizationworker = SpeakerDiarization(device=device, device_clustering=device_clustering, device_vad=device_vad, num_threads=NUM_THREADS)
 
 __all__ = ["diarizationworker"]

simple/diarization/processing/simple_diarizer/simple_diarizer/cluster.py

@@ -6,7 +6,14 @@
 from sklearn.cluster import AgglomerativeClustering, KMeans, SpectralClustering
 from sklearn.metrics import pairwise_distances
 from .spectral_clustering import NME_SpectralClustering
-from sklearn.metrics.pairwise import cosine_similarity
+
+from .nmesc_clustering import (
+    NMESC,
+    getCosAffinityMatrix,
+    getAffinityGraphMat,
+    _SpectralClustering,
+)
+
 
 def cluster_AHC(embeds, n_clusters=None, threshold=None, metric="cosine", **kwargs):
     """
@@ -80,21 +87,42 @@ def cluster_SC(embeds, n_clusters=None, max_speakers= None, threshold=None, enha
 def cluster_NME_SC(embeds, n_clusters=None, max_speakers= None, threshold=None, enhance_sim=True, device=None, **kwargs):
     """
     Cluster embeds using NME-Spectral Clustering
+    """
 
-    if n_clusters is None:
-        assert threshold, "If num_clusters is not defined, threshold must be defined"
     """
     
-    #S = cos_similarity(embeds)
-    S =  getCosAffinityMatrix(embeds)
+    S =  getCosAffinityMatrix(embeds)  #cos_similarity(embeds)
 
     labels = NME_SpectralClustering(
             S,
             num_clusters=n_clusters,
             max_num_clusters=max_speakers,
             device=device,
         )
-            
+    
+    """ 
+    
+    mat = getCosAffinityMatrix(embeds)
+    nmesc = NMESC(
+        mat,
+        max_num_speaker=max_speakers,
+        max_rp_threshold=0.25,
+        sparse_search=True,
+        sparse_search_volume=30,
+        fixed_thres=None,
+        NME_mat_size=512,
+        device=device
+        
+    )
+    
+    
+    est_num_of_spk, p_hat_value = nmesc.NMEanalysis()    
+    affinity_mat = getAffinityGraphMat(mat, p_hat_value)
+    if n_clusters is not None:
+        est_num_of_spk = n_clusters
+    spectral_model = _SpectralClustering(n_clusters=est_num_of_spk, device=device)
+    labels = spectral_model.predict(affinity_mat)
+          
     return labels
 
 from sklearn.preprocessing import MinMaxScaler