[WIP] feat: executor to gcp custom container

jina/enums.py

@@ -269,6 +269,7 @@ class ProviderType(BetterEnum):
 
  NONE = 0 #: no provider
  SAGEMAKER = 1 #: AWS SageMaker
+ GCP = 2 #: GCP
 
 
 def replace_enum_to_str(obj):

jina/orchestrate/deployments/__init__.py

@@ -386,7 +386,7 @@ def __init__(
  :param port_monitoring: The port on which the prometheus server is exposed, default is a random port between [49152, 65535]
  :param prefer_platform: The preferred target Docker platform. (e.g. "linux/amd64", "linux/arm64")
  :param protocol: Communication protocol of the server exposed by the Executor. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param py_modules: The customized python modules need to be imported before loading the executor
 
  Note that the recommended way is to only import a single module - a simple python file, if your
@@ -476,21 +476,21 @@ def __init__(
  args = ArgNamespace.kwargs2namespace(kwargs, parser, True)
  self.args = args
  self._gateway_load_balancer = False
- if self.args.provider == ProviderType.SAGEMAKER:
+ if self.args.provider in (ProviderType.SAGEMAKER, ProviderType.GCP):
  if self._gateway_kwargs.get('port', 0) == 8080:
  raise ValueError(
- 'Port 8080 is reserved for Sagemaker deployment. '
+ 'Port 8080 is reserved for CSP deployment. '
  'Please use another port'
  )
  if self.args.port != [8080]:
  warnings.warn(
- 'Port is changed to 8080 for Sagemaker deployment. '
+ 'Port is changed to 8080 for CSP deployment. '
  f'Port {self.args.port} is ignored'
  )
  self.args.port = [8080]
  if self.args.protocol != [ProtocolType.HTTP]:
  warnings.warn(
- 'Protocol is changed to HTTP for Sagemaker deployment. '
+ 'Protocol is changed to HTTP for CSP deployment. '
  f'Protocol {self.args.protocol} is ignored'
  )
  self.args.protocol = [ProtocolType.HTTP]

jina/orchestrate/flow/base.py

@@ -273,7 +273,7 @@ def __init__(
 
  Used to control the speed of data input into a Flow. 0 disables prefetch (1000 requests is the default)
  :param protocol: Communication protocol of the server exposed by the Gateway. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param proxy: If set, respect the http_proxy and https_proxy environment variables. otherwise, it will unset these proxy variables before start. gRPC seems to prefer no proxy
  :param py_modules: The customized python modules need to be imported before loading the gateway
 
@@ -464,7 +464,7 @@ def __init__(
 
  Used to control the speed of data input into a Flow. 0 disables prefetch (1000 requests is the default)
  :param protocol: Communication protocol of the server exposed by the Gateway. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param proxy: If set, respect the http_proxy and https_proxy environment variables. otherwise, it will unset these proxy variables before start. gRPC seems to prefer no proxy
  :param py_modules: The customized python modules need to be imported before loading the gateway
 
@@ -969,7 +969,7 @@ def add(
  :param port_monitoring: The port on which the prometheus server is exposed, default is a random port between [49152, 65535]
  :param prefer_platform: The preferred target Docker platform. (e.g. "linux/amd64", "linux/arm64")
  :param protocol: Communication protocol of the server exposed by the Executor. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param py_modules: The customized python modules need to be imported before loading the executor
 
  Note that the recommended way is to only import a single module - a simple python file, if your
@@ -1132,7 +1132,7 @@ def add(
  :param port_monitoring: The port on which the prometheus server is exposed, default is a random port between [49152, 65535]
  :param prefer_platform: The preferred target Docker platform. (e.g. "linux/amd64", "linux/arm64")
  :param protocol: Communication protocol of the server exposed by the Executor. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param py_modules: The customized python modules need to be imported before loading the executor
 
  Note that the recommended way is to only import a single module - a simple python file, if your
@@ -1396,7 +1396,7 @@ def config_gateway(
 
  Used to control the speed of data input into a Flow. 0 disables prefetch (1000 requests is the default)
  :param protocol: Communication protocol of the server exposed by the Gateway. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param proxy: If set, respect the http_proxy and https_proxy environment variables. otherwise, it will unset these proxy variables before start. gRPC seems to prefer no proxy
  :param py_modules: The customized python modules need to be imported before loading the gateway
 
@@ -1496,7 +1496,7 @@ def config_gateway(
 
  Used to control the speed of data input into a Flow. 0 disables prefetch (1000 requests is the default)
  :param protocol: Communication protocol of the server exposed by the Gateway. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param proxy: If set, respect the http_proxy and https_proxy environment variables. otherwise, it will unset these proxy variables before start. gRPC seems to prefer no proxy
  :param py_modules: The customized python modules need to be imported before loading the gateway
 

jina/serve/executors/__init__.py

@@ -1084,7 +1084,7 @@ def serve(
  :param port_monitoring: The port on which the prometheus server is exposed, default is a random port between [49152, 65535]
  :param prefer_platform: The preferred target Docker platform. (e.g. "linux/amd64", "linux/arm64")
  :param protocol: Communication protocol of the server exposed by the Executor. This can be a single value or a list of protocols, depending on your chosen Gateway. Choose the convenient protocols from: ['GRPC', 'HTTP', 'WEBSOCKET'].
- :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER'].
+ :param provider: If set, Executor is translated to a custom container compatible with the chosen provider. Choose the convenient providers from: ['NONE', 'SAGEMAKER', 'GCP'].
  :param py_modules: The customized python modules need to be imported before loading the executor
 
  Note that the recommended way is to only import a single module - a simple python file, if your

jina/serve/runtimes/servers/http.py

@@ -297,3 +297,40 @@ def app(self):
  cors=self.cors,
  logger=self.logger,
  )
+
+
+class GCPHTTPServer(FastAPIBaseServer):
+ """
+ :class:`GCPHTTPServer` is a FastAPIBaseServer that uses a custom FastAPI app for GCP endpoints
+
+ """
+
+ @property
+ def port(self):
+ """Get the port for the GCP server
+ :return: Return the port for the GCP server, always 8080"""
+ return 8080
+
+ @property
+ def ports(self):
+ """Get the port for the GCP server
+ :return: Return the port for the GCP server, always 8080"""
+ return [8080]
+
+ @property
+ def app(self):
+ """Get the GCP fastapi app
+ :return: Return a FastAPI app for the GCP container
+ """
+ return self._request_handler._http_fastapi_gcp_app(
+ title=self.title,
+ description=self.description,
+ no_crud_endpoints=self.no_crud_endpoints,
+ no_debug_endpoints=self.no_debug_endpoints,
+ expose_endpoints=self.expose_endpoints,
+ expose_graphql_endpoint=self.expose_graphql_endpoint,
+ tracing=self.tracing,
+ tracer_provider=self.tracer_provider,
+ cors=self.cors,
+ logger=self.logger,
+ )

jina/serve/runtimes/worker/http_gcp_app.py

@@ -0,0 +1,192 @@
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union, Any
+
+from jina._docarray import docarray_v2
+from jina.importer import ImportExtensions
+from jina.types.request.data import DataRequest
+
+if TYPE_CHECKING:
+ from jina.logging.logger import JinaLogger
+
+if docarray_v2:
+ from docarray import BaseDoc, DocList
+
+
+def get_fastapi_app(
+ request_models_map: Dict,
+ caller: Callable,
+ logger: 'JinaLogger',
+ cors: bool = False,
+ **kwargs,
+):
+ """
+ Get the app from FastAPI as the REST interface.
+
+ :param request_models_map: Map describing the endpoints and its Pydantic models
+ :param caller: Callable to be handled by the endpoints of the returned FastAPI app
+ :param logger: Logger object
+ :param cors: If set, a CORS middleware is added to FastAPI frontend to allow cross-origin access.
+ :param kwargs: Extra kwargs to make it compatible with other methods
+ :return: fastapi app
+ """
+ with ImportExtensions(required=True):
+ import pydantic
+ from fastapi import FastAPI, HTTPException, Request
+ from fastapi.middleware.cors import CORSMiddleware
+ from pydantic import BaseModel, Field
+ from pydantic.config import BaseConfig, inherit_config
+
+ import os
+
+ from jina.proto import jina_pb2
+ from jina.serve.runtimes.gateway.models import _to_camel_case
+
+ if not docarray_v2:
+ logger.warning('Only docarray v2 is supported with Sagemaker. ')
+ return
+
+ class Header(BaseModel):
+ request_id: Optional[str] = Field(
+ description='Request ID', example=os.urandom(16).hex()
+ )
+
+ class Config(BaseConfig):
+ alias_generator = _to_camel_case
+ allow_population_by_field_name = True
+
+ class InnerConfig(BaseConfig):
+ alias_generator = _to_camel_case
+ allow_population_by_field_name = True
+
+ class VertexAIResponse(BaseModel):
+ predictions: Any = Field(
+ description='Prediction results',
+ )
+
+ app = FastAPI()
+
+ if cors:
+ app.add_middleware(
+ CORSMiddleware,
+ allow_origins=['*'],
+ allow_credentials=True,
+ allow_methods=['*'],
+ allow_headers=['*'],
+ )
+ logger.warning('CORS is enabled. This service is accessible from any website!')
+
+ def add_post_route(
+ endpoint_path,
+ input_model,
+ output_model,
+ input_doc_list_model=None,
+ ):
+ from docarray.base_doc.docarray_response import DocArrayResponse
+
+ app_kwargs = dict(
+ path=f'/{endpoint_path.strip("/")}',
+ methods=['POST'],
+ summary=f'Endpoint {endpoint_path}',
+ response_model=Union[output_model, List[output_model]],
+ response_class=DocArrayResponse,
+ )
+
+ def is_valid_csv(content: str) -> bool:
+ import csv
+ from io import StringIO
+
+ try:
+ f = StringIO(content)
+ reader = csv.DictReader(f)
+ for _ in reader:
+ pass
+
+ return True
+ except Exception:
+ return False
+
+ async def process(body) -> output_model:
+ req = DataRequest()
+ if body.header is not None:
+ req.header.request_id = body.header.request_id
+
+ if body.parameters is not None:
+ req.parameters = body.parameters
+ req.header.exec_endpoint = endpoint_path
+ req.document_array_cls = DocList[input_doc_model]
+
+ data = body.data
+ if isinstance(data, list):
+ req.data.docs = DocList[input_doc_list_model](data)
+ else:
+ req.data.docs = DocList[input_doc_list_model]([data])
+ if body.header is None:
+ req.header.request_id = req.docs[0].id
+
+ resp = await caller(req)
+ status = resp.header.status
+
+ if status.code == jina_pb2.StatusProto.ERROR:
+ raise HTTPException(status_code=499, detail=status.description)
+ else:
+ return {"predictions": resp.docs}
+ return output_model(predictions=resp.docs)
+
+ @app.api_route(**app_kwargs)
+ async def post(request: Request):
+ content_type = request.headers.get('content-type')
+ if content_type == 'application/json':
+ json_body = await request.json()
+ transformed_json_body = {"data": [{"text": instance} for instance in json_body["instances"]]}
+ return await process(input_model(**transformed_json_body))
+
+ elif content_type in ('text/csv', 'application/csv'):
+ # TODO: fix here
+ return await process(input_model(data=[]))
+ else:
+ raise HTTPException(
+ status_code=400,
+ detail=f'Invalid content-type: {content_type}. '
+ f'Please use either application/json or text/csv.',
+ )
+
+ for endpoint, input_output_map in request_models_map.items():
+ if endpoint != '_jina_dry_run_':
+ input_doc_model = input_output_map['input']['model']
+ parameters_model = input_output_map['parameters']['model'] or Optional[Dict]
+ default_parameters = (
+ ... if input_output_map['parameters']['model'] else None
+ )
+
+ _config = inherit_config(InnerConfig, BaseDoc.__config__)
+ endpoint_input_model = pydantic.create_model(
+ f'{endpoint.strip("/")}_input_model',
+ data=(Union[List[input_doc_model], input_doc_model], ...),
+ parameters=(parameters_model, default_parameters),
+ header=(Optional[Header], None),
+ __config__=_config,
+ )
+
+ add_post_route(
+ endpoint,
+ input_model=endpoint_input_model,
+ output_model=VertexAIResponse,
+ input_doc_list_model=input_doc_model,
+ )
+
+ from jina.serve.runtimes.gateway.health_model import JinaHealthModel
+
+ # `/ping` route is required by AWS Sagemaker
+ @app.get(
+ path='/ping',
+ summary='Get the health of Jina Executor service',
+ response_model=JinaHealthModel,
+ )
+ async def _executor_health():
+ """
+ Get the health of this Gateway service.
+ .. # noqa: DAR201
+
+ """
+ return {}
+
+ return app

jina/serve/runtimes/worker/request_handling.py

@@ -212,6 +212,28 @@ async def _shutdown():
 
  return app
 
+ def _http_fastapi_gcp_app(self, **kwargs):
+ from jina.serve.runtimes.worker.http_gcp_app import get_fastapi_app
+
+ request_models_map = self._executor._get_endpoint_models_dict()
+
+ def call_handle(request):
+ is_generator = request_models_map[request.header.exec_endpoint][
+ 'is_generator'
+ ]
+
+ return self.process_single_data(request, None, is_generator=is_generator)
+
+ app = get_fastapi_app(
+ request_models_map=request_models_map, caller=call_handle, **kwargs
+ )
+
+ @app.on_event('shutdown')
+ async def _shutdown():
+ await self.close()
+
+ return app
+
  async def _hot_reload(self):
  import inspect
 

tests/integration/docarray_v2/gcp/Dockerfile

@@ -0,0 +1,7 @@
+FROM jinaai/jina:test-pip
+
+COPY . /executor_root/
+
+WORKDIR /executor_root/SampleExecutor
+
+ENTRYPOINT ["jina", "executor", "--uses", "config.yml"]

tests/integration/docarray_v2/gcp/SampleExecutor/README.md

@@ -0,0 +1,2 @@
+# SampleExecutor
+

tests/integration/docarray_v2/gcp/SampleExecutor/config.yml

@@ -0,0 +1,8 @@
+jtype: SampleExecutor
+py_modules:
+ - executor.py
+metas:
+ name: SampleExecutor
+ description: 
+ url: 
+ keywords: []