diff --git a/src/nv_ingest/framework/orchestration/morpheus/util/pipeline/stage_builders.py b/src/nv_ingest/framework/orchestration/morpheus/util/pipeline/stage_builders.py index fa656c028..9e09f0318 100644 --- a/src/nv_ingest/framework/orchestration/morpheus/util/pipeline/stage_builders.py +++ b/src/nv_ingest/framework/orchestration/morpheus/util/pipeline/stage_builders.py @@ -408,29 +408,38 @@ def get_audio_retrieval_service(env_var_prefix): "", ) auth_token = os.environ.get( - "RIVA_NGC_API_KEY", + "NVIDIA_BUILD_API_KEY", + "", + ) or os.environ.get( + "NGC_API_KEY", "", ) infer_protocol = os.environ.get( "AUDIO_INFER_PROTOCOL", "http" if http_endpoint else "grpc" if grpc_endpoint else "", ) + function_id = os.environ.get( + "AUDIO_FUNCTION_ID", + "", + ) - logger.info(f"{prefix}_GRPC_TRITON: {grpc_endpoint}") - logger.info(f"{prefix}_HTTP_TRITON: {http_endpoint}") + logger.info(f"{prefix}_GRPC_ENDPOINT: {grpc_endpoint}") + logger.info(f"{prefix}_HTTP_ENDPOINT: {http_endpoint}") logger.info(f"{prefix}_INFER_PROTOCOL: {infer_protocol}") + logger.info(f"{prefix}_FUNCTION_ID: {function_id}") - return grpc_endpoint, http_endpoint, auth_token, infer_protocol + return grpc_endpoint, http_endpoint, auth_token, infer_protocol, function_id def add_audio_extractor_stage(pipe, morpheus_pipeline_config, ingest_config, default_cpu_count): - audio_grpc, audio_http, audio_auth, audio_infer_protocol = get_audio_retrieval_service("audio") + audio_grpc, audio_http, audio_auth, audio_infer_protocol, audio_function_id = get_audio_retrieval_service("audio") audio_extractor_config = ingest_config.get( "audio_extraction_module", { "audio_extraction_config": { "audio_endpoints": (audio_grpc, audio_http), "audio_infer_protocol": audio_infer_protocol, + "function_id": audio_function_id, "auth_token": audio_auth, # All auth tokens are the same for the moment } diff --git a/tests/integration/test_extract_audio.py b/tests/integration/test_extract_audio.py new file mode 100644 index 000000000..1e62d0eeb --- /dev/null +++ b/tests/integration/test_extract_audio.py @@ -0,0 +1,32 @@ +import subprocess +import sys +import time + +import pytest +from nv_ingest_api.util.message_brokers.simple_message_broker import SimpleClient +from nv_ingest_client.client import Ingestor +from nv_ingest_client.client import NvIngestClient + + +def test_images_extract_only( + pipeline_process, +): + client = NvIngestClient( + message_client_allocator=SimpleClient, + message_client_port=7671, + message_client_hostname="localhost", + ) + + ingestor = ( + Ingestor(client=client) + .files("./data/multimodal_test.wav") + .extract( + extract_method="audio", + ) + ) + + results = ingestor.ingest() + + assert len(results) == 1 + + print(results)