Supporting PyStein programming model in the Worker

azure_functions_worker/dispatcher.py

@@ -7,22 +7,20 @@
 
 import asyncio
 import concurrent.futures
+import json
 import logging
 import os
 import queue
 import sys
 import threading
+import uuid
+import grpc
+
 from asyncio import BaseEventLoop
 from logging import LogRecord
 from typing import List, Optional
 
-import grpc
-
-from . import bindings
-from . import constants
-from . import functions
-from . import loader
-from . import protos
+from . import bindings, constants, functions, loader, protos
 from .bindings.shared_memory_data_transfer import SharedMemoryManager
 from .constants import (PYTHON_THREADPOOL_THREAD_COUNT,
                         PYTHON_THREADPOOL_THREAD_COUNT_DEFAULT,
@@ -294,14 +292,58 @@ async def _handle__worker_init_request(self, req):
                 result=protos.StatusResult(
                     status=protos.StatusResult.Success)))
 
-    async def _handle__worker_status_request(self, req):
+    async def _handle__worker_status_request(self, request):
         # Logging is not necessary in this request since the response is used
         # for host to judge scale decisions of out-of-proc languages.
         # Having log here will reduce the responsiveness of the worker.
         return protos.StreamingMessage(
-            request_id=req.request_id,
+            request_id=request.request_id,
             worker_status_response=protos.WorkerStatusResponse())
 
+    async def _handle__functions_metadata_request(self, request):
+        metadata_request = request.functions_metadata_request
+        directory = metadata_request.function_app_directory
+        indexed_functions = loader.index_function_app(directory)
+
+        fx_metadata_results = []
+        for indexed_function in indexed_functions:
+            function_id = str(uuid.uuid4())
+            function_info = self._functions.add_indexed_function(
+                function_id,
+                function=indexed_function)
+
+            binding_protos = {}
+
+            for binding in indexed_function.get_bindings():
+
+                binding_protos[binding.name] = protos.BindingInfo(
+                    type=binding.type,
+                    data_type=binding.data_type,
+                    direction=binding.direction)
+
+            function_load_request = protos.RpcFunctionMetadata(
+                name=function_info.name,
+                function_id=function_id,
+                managed_dependency_enabled=False,
+                directory=function_info.directory,
+                script_file=indexed_function.function_script_file,
+                entry_point=function_info.name,
+                is_proxy=False,
+                language="python",
+                bindings=binding_protos,
+                raw_bindings=[json.dumps(i) for i in
+                              indexed_function.get_bindings_dict()[
+                                  "bindings"]])
+
+            fx_metadata_results.append(function_load_request)
+
+        return protos.StreamingMessage(
+            request_id=request.request_id,
+            function_metadata_response=protos.FunctionMetadataResponse(
+                function_metadata_results=fx_metadata_results,
+                result=protos.StatusResult(
+                    status=protos.StatusResult.Success)))
+
     async def _handle__function_load_request(self, req):
         func_request = req.function_load_request
         function_id = func_request.function_id
@@ -312,24 +354,25 @@ async def _handle__function_load_request(self, req):
                     f'function ID: {function_id}'
                     f'function Name: {function_name}')
         try:
-            func = loader.load_function(
-                func_request.metadata.name,
-                func_request.metadata.directory,
-                func_request.metadata.script_file,
-                func_request.metadata.entry_point)
-
-            self._functions.add_function(
-                function_id, func, func_request.metadata)
-
-            ExtensionManager.function_load_extension(
-                function_name,
-                func_request.metadata.directory
-            )
+            if not self._functions.get_function(function_id):
+                func = loader.load_function(
+                    func_request.metadata.name,
+                    func_request.metadata.directory,
+                    func_request.metadata.script_file,
+                    func_request.metadata.entry_point)
+
+                self._functions.add_function(
+                    function_id, func, func_request.metadata)
+
+                ExtensionManager.function_load_extension(
+                    function_name,
+                    func_request.metadata.directory
+                )
 
-            logger.info('Successfully processed FunctionLoadRequest, '
-                        f'request ID: {self.request_id}, '
-                        f'function ID: {function_id},'
-                        f'function Name: {function_name}')
+                logger.info('Successfully processed FunctionLoadRequest, '
+                            f'request ID: {self.request_id}, '
+                            f'function ID: {function_id},'
+                            f'function Name: {function_name}')
 
             return protos.StreamingMessage(
                 request_id=self.request_id,
@@ -361,6 +404,7 @@ async def _handle__invocation_request(self, req):
         try:
             fi: functions.FunctionInfo = self._functions.get_function(
                 function_id)
+            assert fi is not None
 
             function_invocation_logs: List[str] = [
                 'Received FunctionInvocationRequest',
@@ -751,7 +795,12 @@ def get_current_invocation_id() -> Optional[str]:
             if task_invocation_id is not None:
                 return task_invocation_id
 
-    return getattr(_invocation_id_local, 'v', None)
+    #return getattr(_invocation_id_local, 'v', None)
+
+    invocation_local_id = getattr(logging.getLogger(), 'invocation_local_id', None)
+    return getattr(invocation_local_id, 'v', None)
+
+
 
 
 _invocation_id_local = threading.local()