Port over the latest batch score changes from the ADO repository

assets/batch_score/components/driver/dev/README.md

@@ -36,7 +36,7 @@ This folder does not verify any yaml configurations.
             },
             "args": [
                 "--debug_mode", "true",
-                "--online_endpoint_url", "https://real-dv3-stable.centralus.inference.ml.azure.com/v1/engines/davinci/completions"
+                "--online_endpoint_url", "https://real-dv3-stable.centralus.inference.ml.azure.com/v1/engines/davinci/completions",
             ]
         }
     ]

assets/batch_score/components/driver/dev/batch_score_simulator.py

@@ -45,7 +45,9 @@ class Simulator:
     """PRS Simulator."""
 
     def __init__(self, data_input_folder_path):
-        """Init function."""
+        """Initialize PRS Simulator."""
+        os.getcwd()
+
         self.__mltable_data: mltable = mltable.load(data_input_folder_path)
         self.__df_data = self.__mltable_data.to_pandas_dataframe()
         self.__minibatch_size = 500  # lines

assets/batch_score/components/driver/dev/quota_simulator.py

@@ -24,7 +24,7 @@ class QuotaSimulator(ServiceSimulator):
     """Quota simulator."""
 
     def __init__(self):
-        """Init function."""
+        """Initialize QuotaSimulator."""
         super().__init__(handler=partial(self.RequestHandler, simulator=self))
 
         self._leases = {}
@@ -48,7 +48,7 @@ class RequestHandler(ServiceSimulator.RequestHandler):
         """Simulator request handler."""
 
         def __init__(self, *args, simulator: 'QuotaSimulator', **kwargs):
-            """Init function."""
+            """Initialize RequestHandler."""
             self.__simulator = simulator
             super().__init__(*args, **kwargs)
 

assets/batch_score/components/driver/dev/routing_simulator.py

@@ -19,7 +19,7 @@ class RoutingSimulator(ServiceSimulator):
     """Routing simulator."""
 
     def __init__(self):
-        """Init function."""
+        """Initialize RoutingSimulator."""
         super().__init__(handler=self.RequestHandler)
 
     @classmethod

assets/batch_score/components/driver/dev/service_simulator.py

@@ -15,7 +15,7 @@ class ServiceSimulator:
     """Service simulator."""
 
     def __init__(self, handler):
-        """Init function."""
+        """Initialize ServiceSimulator."""
         self.__server = ThreadingHTTPServer(("localhost", 0), handler)
         self.__thread = Thread(target=self.__server.serve_forever, daemon=True)
 

assets/batch_score/components/driver/dev/token_estimation/train_token_model.py

@@ -159,7 +159,7 @@ def print_stats(model: LinearRegression, test_x, test_y, *, name=''):
     print(f'Model {name}')
     print('=====')
     print()
-    print('Coefficients: ', model.coef_, sep='\n')
+    print('Coefficients:', model.coef_, sep='\n')
     print('---')
     print(f'Min: {test_y.min()}')
     print(f'Avg: {int(test_y.mean())}')

assets/batch_score/components/driver/src/__init__.py

@@ -1,3 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+
+# flake8: noqa: F401
+
 """__init__."""

assets/batch_score/components/driver/src/batch_score/__init__.py

@@ -1,9 +1,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+
+# flake8: noqa: F401
+
 """__init__."""
 
 import subprocess
 import sys
 
 # TODO: remove this hack after dedicated environment is published.
-subprocess.check_call([sys.executable, "-m", "pip", "install", 'tiktoken==0.5.2'])
+subprocess.check_call([sys.executable,"-m", "pip", "install",
+                       'PyDispatcher==2.0.7',
+                       'StrEnum==0.4.15',
+                       'tiktoken==0.5.2'])

assets/batch_score/components/driver/src/batch_score/aoai/__init__.py

@@ -1,3 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+
+# flake8: noqa: F401
+
 """__init__."""

assets/batch_score/components/driver/src/batch_score/aoai/scoring/__init__.py

@@ -1,3 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+
+# flake8: noqa: F401
+
 """__init__."""

assets/batch_score/components/driver/src/batch_score/aoai/scoring/aoai_header_provider.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""This file contains the definition for AOAI header provider."""
+
+import json
+import uuid
+
+from ...common.auth.auth_provider import AuthProvider
+from ...common.scoring.generic_scoring_client import HeaderProvider
+
+
+class AoaiHeaderProvider(HeaderProvider):
+    """Defines the AOAI header provider."""
+
+    def __init__(
+        self,
+        auth_provider: AuthProvider,
+        additional_headers: str = None,
+    ):
+        """Initialize AoaiHeaderProvider."""
+        self._auth_provider = auth_provider
+
+        if additional_headers is not None:
+            self._additional_headers = json.loads(additional_headers)
+        else:
+            self._additional_headers = {}
+
+    def get_headers(self) -> dict:
+        """Get the headers from the auth provider and additional headers."""
+        headers = {
+            'Content-Type': 'application/json',
+            'x-ms-client-request-id': str(uuid.uuid4()),
+        }
+        headers.update(self._auth_provider.get_auth_headers())
+        headers.update(self._additional_headers)
+
+        return headers

assets/batch_score/components/driver/src/batch_score/aoai/scoring/aoai_response_handler.py

@@ -0,0 +1,167 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""This file contains the definition for AOAI response handler."""
+
+import asyncio
+
+import aiohttp
+
+from ...common.scoring.generic_scoring_client import (
+    HttpResponseHandler,
+    HttpScoringResponse,
+)
+from ...common.scoring.scoring_request import ScoringRequest
+from ...common.scoring.scoring_result import (
+    RetriableException,
+    ScoringResult,
+    ScoringResultStatus,
+)
+from ...common.scoring.tally_failed_request_handler import TallyFailedRequestHandler
+from ...common.telemetry.events import event_utils
+from ...common.telemetry.events.batch_score_request_completed_event import BatchScoreRequestCompletedEvent
+
+
+class AoaiHttpResponseHandler(HttpResponseHandler):
+    """Defines the AOAI HTTP response handler."""
+
+    RETRIABLE_STATUS_CODES = [408, 429, 500, 502, 503, 504]
+
+    def __init__(self, tally_handler: TallyFailedRequestHandler):
+        """Initialize AoaiHttpResponseHandler."""
+        self.__tally_handler = tally_handler
+
+    def handle_response(
+            self,
+            http_response: HttpScoringResponse,
+            scoring_request: ScoringRequest,
+            x_ms_client_request_id: str,
+            start: float,
+            end: float,
+            worker_id: str) -> ScoringResult:
+        """Handle the response from the model for the provided scoring request."""
+        # Emit request completed event
+        self._emit_request_completed_event(
+            http_response=http_response,
+            scoring_request=scoring_request,
+            x_ms_client_request_id=x_ms_client_request_id,
+            start=start,
+            end=end,
+            worker_id=worker_id
+        )
+
+        if http_response.exception:
+            return self._handle_exception(
+                http_response=http_response,
+                scoring_request=scoring_request,
+                start=start,
+                end=end,
+            )
+
+        response_status = http_response.status
+        if response_status == 200:
+            return self._create_scoring_result(
+                status=ScoringResultStatus.SUCCESS,
+                scoring_request=scoring_request,
+                start=start,
+                end=end,
+                http_post_response=http_response,
+            )
+
+        if self.is_retriable(response_status):
+            raise RetriableException(
+                status_code=http_response.status,
+                response_payload=http_response.payload)
+
+        result = self._create_scoring_result(
+            status=ScoringResultStatus.FAILURE,
+            scoring_request=scoring_request,
+            start=start,
+            end=end,
+            http_post_response=http_response,
+        )
+
+        model_response_code = http_response.get_model_response_code()
+        if (
+            result.status == ScoringResultStatus.FAILURE
+            and self.__tally_handler.should_tally(
+                response_status=response_status,
+                model_response_status=model_response_code)
+        ):
+            result.omit = True
+
+        return result
+
+    def is_retriable(self, http_status: int) -> bool:
+        """Is the http status retriable."""
+        return http_status in self.RETRIABLE_STATUS_CODES
+
+    def _handle_exception(
+            self,
+            http_response: HttpScoringResponse,
+            scoring_request: ScoringRequest,
+            start: float,
+            end: float) -> ScoringResult:
+        """Handle exception by raising retriable exception or creating scoring result for non-retriable exception."""
+        try:
+            raise http_response.exception
+        except (
+            aiohttp.ClientConnectorError,
+            aiohttp.ServerConnectionError,
+            asyncio.TimeoutError,
+        ):
+            raise RetriableException(
+                status_code=http_response.status,
+                response_payload=http_response.payload,
+            )
+        except Exception:
+            return self._create_scoring_result(
+                status=ScoringResultStatus.FAILURE,
+                scoring_request=scoring_request,
+                start=start,
+                end=end,
+                http_post_response=http_response,
+            )
+
+    @event_utils.catch_and_log_all_exceptions
+    def _emit_request_completed_event(
+            self,
+            http_response: HttpScoringResponse,
+            scoring_request: ScoringRequest,
+            x_ms_client_request_id: str,
+            start: float,
+            end: float,
+            worker_id: str) -> None:
+
+        def get_prompt_tokens(response_body: any):
+            if not isinstance(response_body, dict):
+                return None
+
+            return response_body.get("usage", {}).get("prompt_tokens")
+
+        def get_completion_tokens(response_body: any):
+            if not isinstance(response_body, dict):
+                return None
+
+            return response_body.get("usage", {}).get("completion_tokens")
+
+        def get_mini_batch_id(mini_batch_context: any):
+            if mini_batch_context:
+                return mini_batch_context.mini_batch_id
+
+        request_completed_event = BatchScoreRequestCompletedEvent(
+            minibatch_id=get_mini_batch_id(scoring_request.mini_batch_context),
+            input_row_id=scoring_request.internal_id,
+            x_ms_client_request_id=x_ms_client_request_id,
+            worker_id=worker_id,
+            scoring_url=scoring_request.scoring_url,
+            is_successful=http_response.status == 200,
+            is_retriable=self.is_retriable(http_response.status),
+            response_code=http_response.status,
+            model_response_code=http_response.get_model_response_code(),
+            prompt_tokens=get_prompt_tokens(http_response.payload),
+            completion_tokens=get_completion_tokens(http_response.payload),
+            duration_ms=(end - start) * 1000,
+            segmented_request_id=scoring_request.segment_id
+        )
+        event_utils.emit_event(batch_score_event=request_completed_event)