feat: support passing metadata (#53)

* feat: support passing metadata * pass credentials via metadata call credentials * Update common proto submodule
nvidia-riva · Aug 23, 2023 · b367963 · b367963
1 parent 185e3ff
commit b367963
Show file tree

Hide file tree

Showing 15 changed files with 31 additions and 17 deletions.
diff --git a/common b/common
diff --git a/riva/client/argparse_utils.py b/riva/client/argparse_utils.py
@@ -57,4 +57,5 @@ def add_connection_argparse_parameters(parser: argparse.ArgumentParser) -> argpa
     parser.add_argument(
         "--use-ssl", action='store_true', help="Boolean to control if SSL/TLS encryption should be used."
     )
+    parser.add_argument("--metadata", action='append', nargs='+', help="Send HTTP Header(s) to server")
     return parser
diff --git a/riva/client/auth.py b/riva/client/auth.py
@@ -4,20 +4,26 @@
 import os
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
-
 import grpc
 
 
 def create_channel(
-    ssl_cert: Optional[Union[str, os.PathLike]] = None, use_ssl: bool = False, uri: str = "localhost:50051",
+    ssl_cert: Optional[Union[str, os.PathLike]] = None, use_ssl: bool = False, uri: str = "localhost:50051", metadata: Optional[List[Tuple[str, str]]] = None,
 ) -> grpc.Channel:
+
+    def metadata_callback(context, callback):
+        callback(metadata, None)
+
     if ssl_cert is not None or use_ssl:
         root_certificates = None
         if ssl_cert is not None:
             ssl_cert = Path(ssl_cert).expanduser()
             with open(ssl_cert, 'rb') as f:
                 root_certificates = f.read()
         creds = grpc.ssl_channel_credentials(root_certificates)
+        if metadata:
+            auth_creds = grpc.metadata_call_credentials(metadata_callback)
+            creds = grpc.composite_channel_credentials(creds, auth_creds)
         channel = grpc.secure_channel(uri, creds)
     else:
         channel = grpc.insecure_channel(uri)
@@ -30,6 +36,7 @@ def __init__(
         ssl_cert: Optional[Union[str, os.PathLike]] = None,
         use_ssl: bool = False,
         uri: str = "localhost:50051",
+        metadata_args: List[List[str]] = None,
     ) -> None:
         """
         A class responsible for establishing connection with a server and providing security metadata.
@@ -44,7 +51,13 @@ def __init__(
         self.ssl_cert: Optional[Path] = None if ssl_cert is None else Path(ssl_cert).expanduser()
         self.uri: str = uri
         self.use_ssl: bool = use_ssl
-        self.channel: grpc.Channel = create_channel(self.ssl_cert, self.use_ssl, self.uri)
+        self.metadata = []
+        if metadata_args:
+            for meta in metadata_args:
+                if len(meta) != 2:
+                    raise ValueError(f"Metadata should have 2 parameters in \"key\" \"value\" pair. Receieved {len(meta)} parameters.")
+                self.metadata.append(tuple(meta))
+        self.channel: grpc.Channel = create_channel(self.ssl_cert, self.use_ssl, self.uri, self.metadata)
 
     def get_auth_metadata(self) -> List[Tuple[str, str]]:
         """

diff --git a/scripts/asr/riva_streaming_asr_client.py b/scripts/asr/riva_streaming_asr_client.py
@@ -50,7 +50,7 @@ def streaming_transcription_worker(
 ) -> None:
     output_file = Path(output_file).expanduser()
     try:
-        auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+        auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
         asr_service = riva.client.ASRService(auth)
         config = riva.client.StreamingRecognitionConfig(
             config=riva.client.RecognitionConfig(

diff --git a/scripts/asr/transcribe_file.py b/scripts/asr/transcribe_file.py
@@ -66,7 +66,7 @@ def main() -> None:
     if args.list_devices:
         riva.client.audio_io.list_output_devices()
         return
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     asr_service = riva.client.ASRService(auth)
     config = riva.client.StreamingRecognitionConfig(
         config=riva.client.RecognitionConfig(

diff --git a/scripts/asr/transcribe_file_offline.py b/scripts/asr/transcribe_file_offline.py
@@ -26,7 +26,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     asr_service = riva.client.ASRService(auth)
     config = riva.client.RecognitionConfig(
         language_code=args.language_code,

diff --git a/scripts/asr/transcribe_mic.py b/scripts/asr/transcribe_mic.py
@@ -41,7 +41,7 @@ def main() -> None:
     if args.list_devices:
         riva.client.audio_io.list_input_devices()
         return
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     asr_service = riva.client.ASRService(auth)
     config = riva.client.StreamingRecognitionConfig(
         config=riva.client.RecognitionConfig(

diff --git a/scripts/nlp/eval_intent_slot.py b/scripts/nlp/eval_intent_slot.py
@@ -340,7 +340,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     intent_report, slot_report = intent_slots_classification_report(
         args.input_file,

diff --git a/scripts/nlp/intentslot_client.py b/scripts/nlp/intentslot_client.py
@@ -44,7 +44,7 @@ def pretty_print_result(
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     if args.interactive:
         while True:

diff --git a/scripts/nlp/ner_client.py b/scripts/nlp/ner_client.py
@@ -31,7 +31,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     tokens, slots, slot_confidences, starts, ends = riva.client.extract_most_probable_token_classification_predictions(
         service.classify_tokens(input_strings=args.query, model_name=args.model)

diff --git a/scripts/nlp/punctuation_client.py b/scripts/nlp/punctuation_client.py
@@ -39,7 +39,7 @@ def parse_args() -> argparse.Namespace:
 
 
 def run_punct_capit(args: argparse.Namespace) -> None:
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     nlp_service = riva.client.NLPService(auth)
     if args.interactive:
         while True:
@@ -134,7 +134,7 @@ def run_tests(args: argparse.Namespace) -> int:
         ],
     }
 
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     nlp_service = riva.client.NLPService(auth)
 
     fail_count = 0

diff --git a/scripts/nlp/qa_client.py b/scripts/nlp/qa_client.py
@@ -31,7 +31,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     resp = service.natural_query(args.query, args.context)
     print(resp)

diff --git a/scripts/nlp/text_classify_client.py b/scripts/nlp/text_classify_client.py
@@ -20,7 +20,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     print(riva.client.nlp.extract_most_probable_text_class_and_confidence(service.classify_text(args.query, args.model)))
 

diff --git a/scripts/nmt/nmt.py b/scripts/nmt/nmt.py
@@ -81,7 +81,7 @@ def request(inputs,args):
 
     args = parse_args()
 
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     nmt_client = riva.client.NeuralMachineTranslationClient(auth)
 
     if args.list_models:

diff --git a/scripts/tts/talk.py b/scripts/tts/talk.py
@@ -61,7 +61,7 @@ def main() -> None:
     if args.list_devices:
         riva.client.audio_io.list_output_devices()
         return
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.SpeechSynthesisService(auth)
     nchannels = 1
     sampwidth = 2
+12 −3		riva/proto/BUILD
+15 −0		riva/proto/riva_asr.proto
+16 −0		riva/proto/riva_common.proto
+35 −0		riva/proto/riva_nlp.proto
+30 −11		riva/proto/riva_nmt.proto
+8 −1		riva/proto/riva_tts.proto