Feat playground gateway (#20)

* feat: playground in gateway

* fix: playground in gateway

* fix: do not import executor

* fix: make local work

* fix: fix typo

src/apis/gpt.py

@@ -78,7 +78,7 @@ If you have updated it already, please restart your terminal.
         print('\n')
         money_prompt = self._calculate_money_spent(self.chars_prompt_so_far, self.pricing_prompt)
         money_generation = self._calculate_money_spent(self.chars_generation_so_far, self.pricing_generation)
-        print('Total money spent so far on openai.com:', f'${money_prompt + money_generation}')
+        print('Total money spent so far on openai.com:', f'${money_prompt + money_generation:.3f}')
         print('\n')
 
     @staticmethod

src/apis/jina_cloud.py

@@ -33,8 +33,8 @@ def wait_until_app_is_ready(url):
         time.sleep(0.5)
 
 
-def open_streamlit_app():
+def open_streamlit_app(host: str):
-    url = "http://localhost:8081/playground"
+    url = f"{host}/playground"
     wait_until_app_is_ready(url)
     webbrowser.open(url, new=2)
 
@@ -121,7 +121,7 @@ def _deploy_on_jcloud(flow_yaml):
 
 def deploy_on_jcloud(executor_name, microservice_path):
     print('Deploy a jina flow')
-    full_flow_path = create_flow_yaml(microservice_path, executor_name, use_docker=True)
+    full_flow_path = create_flow_yaml(microservice_path, executor_name, use_docker=True, use_custom_gateway=True)
 
     for i in range(3):
         try:
@@ -145,18 +145,15 @@ Please try again later.
                         )
 
     print(f'''
-Your Microservice is deployed.
+Your Microservice is deployed at {host} and the playground is available at {host}/playground
-Run the following command to start the playground:
+We open now the playground in your browser.
-
+''')
-streamlit run {os.path.join(microservice_path, "app.py")} --server.port 8081 --server.address 0.0.0.0 -- --host {host}
+    open_streamlit_app(host)
-'''
-          )
     return host
 
 
 def run_streamlit_app(app_path):
-    subprocess.run(['streamlit', 'run', app_path, 'server.address', '0.0.0.0', '--server.port', '8081', '--', '--host',
+    subprocess.run(['streamlit', 'run', app_path, 'server.address', '0.0.0.0', '--server.port', '8081'])
-                    'http://localhost:8080'])
 
 
 def run_locally(executor_name, microservice_version_path):
@@ -178,7 +175,7 @@ c) try to run your microservice locally without docker. It is worth a try but mi
             exit(1)
         use_docker = False
     print('Run a jina flow locally')
-    full_flow_path = create_flow_yaml(microservice_version_path, executor_name, use_docker)
+    full_flow_path = create_flow_yaml(microservice_version_path, executor_name, use_docker, False)
     flow = Flow.load_config(full_flow_path)
     with flow:
         print(f'''
@@ -186,27 +183,25 @@ Your microservice started locally.
 We now start the playground for you.
 ''')
 
-        app_path = os.path.join(microservice_version_path, "app.py")
+        app_path = os.path.join(microservice_version_path, 'gateway', "app.py")
 
         # Run the Streamlit app in a separate thread
         streamlit_thread = threading.Thread(target=run_streamlit_app, args=(app_path,))
         streamlit_thread.start()
 
         # Open the Streamlit app in the user's default web browser
-        open_streamlit_app()
+        open_streamlit_app(host='http://localhost:8081')
 
         flow.block()
 
 
-def create_flow_yaml(dest_folder, executor_name, use_docker):
+def create_flow_yaml(dest_folder, executor_name, use_docker, use_custom_gateway):
     if use_docker:
         prefix = 'jinaai+docker'
     else:
         prefix = 'jinaai'
-    flow = f'''
+    flow = f'''jtype: Flow
-jtype: Flow
 with:
-  name: nowapi
   port: 8080
   protocol: http
 jcloud:
@@ -214,7 +209,9 @@ jcloud:
   labels:
     creator: microchain
   name: gptdeploy
-
+gateway:
+    {f"uses: {prefix}://{get_user_name(DEMO_TOKEN)}/Gateway{executor_name}:latest" if use_custom_gateway else ""}
+    {"" if use_docker else "install-requirements: True"}
 executors:
   - name: {executor_name.lower()}
     uses: {prefix}://{get_user_name(DEMO_TOKEN)}/{executor_name}:latest

src/options/__init__.py

@@ -31,7 +31,7 @@ def validate_folder_is_correct(microservice_path):
         raise ValueError(f'Path {microservice_path} needs to contain only one folder. Please make sure that you only have one microservice in this folder.')
     latest_version_path = get_latest_version_path(microservice_path)
     required_files = [
-        'app.py',
+        'gateway/app.py',
         'requirements.txt',
         'Dockerfile',
         'config.yml',

src/options/generate/generator.py

@@ -1,6 +1,7 @@
 import os
 import random
 import re
+import shutil
 
 from src.apis import gpt
 from src.apis.jina_cloud import process_error_message, push_executor
@@ -35,12 +36,12 @@ class Generator:
                 return single_code_block_match[0].strip()
         return ''
 
-    def write_config_yml(self, microservice_name, dest_folder):
+    def write_config_yml(self, class_name, dest_folder, python_file='microservice.py'):
-        config_content = f'''jtype: {microservice_name}
+        config_content = f'''jtype: {class_name}
 py_modules:
-  - microservice.py
+  - {python_file}
 metas:
-  name: {microservice_name}
+  name: {class_name}
 '''
         with open(os.path.join(dest_folder, 'config.yml'), 'w') as f:
             f.write(config_content)
@@ -148,7 +149,6 @@ metas:
             template_generate_playground.format(
                 code_files_wrapped=self.files_to_string(file_name_to_content, ['microservice.py', 'test_microservice.py']),
                 microservice_name=microservice_name,
-
             )
         )
         playground_content_raw = conversation.chat(
@@ -159,7 +159,33 @@ metas:
             )
         )
         playground_content = self.extract_content_from_result(playground_content_raw, 'app.py', match_single_block=True)
-        persist_file(playground_content, os.path.join(microservice_path, 'app.py'))
+        if playground_content == '':
+            content_raw = conversation.chat(f'You must add the app.py code. You most not output any other code')
+            playground_content = self.extract_content_from_result(
+                content_raw, 'app.py', match_single_block=True
+            )
+
+        gateway_path = os.path.join(microservice_path, 'gateway')
+        shutil.copytree(os.path.join(os.path.dirname(__file__), 'static_files', 'gateway'), gateway_path)
+        persist_file(playground_content, os.path.join(gateway_path, 'app.py'))
+
+        # fill-in name of microservice
+        gateway_name = f'Gateway{microservice_name}'
+        custom_gateway_path = os.path.join(gateway_path, 'custom_gateway.py')
+        with open(custom_gateway_path, 'r') as f:
+            custom_gateway_content = f.read()
+        custom_gateway_content = custom_gateway_content.replace(
+            'class CustomGateway(CompositeGateway):',
+            f'class {gateway_name}(CompositeGateway):'
+        )
+        with open(custom_gateway_path, 'w') as f:
+            f.write(custom_gateway_content)
+
+        # write config.yml
+        self.write_config_yml(gateway_name, gateway_path, 'custom_gateway.py')
+
+        # push the gateway
+        hubble_log = push_executor(gateway_path)
 
     def debug_microservice(self, path, microservice_name, num_approach, packages):
         for i in range(1, MAX_DEBUGGING_ITERATIONS):

src/options/generate/static_files/gateway/Dockerfile

@@ -0,0 +1,14 @@
+FROM jinaai/jina:3.14.1-py310-standard
+
+
+RUN apt-get update && apt-get install --no-install-recommends -y git pip nginx && rm -rf /var/lib/apt/lists/*
+
+## install requirements for the executor
+COPY requirements.txt .
+RUN pip install --compile -r requirements.txt
+
+# setup the workspace
+COPY . /workdir/
+WORKDIR /workdir
+
+ENTRYPOINT ["jina", "gateway", "--uses", "config.yml"]

src/options/generate/static_files/gateway/app_config.toml

@@ -0,0 +1,4 @@
+[server]
+
+baseUrlPath = "/playground"
+headless = true

src/options/generate/static_files/gateway/custom_gateway.py

@@ -0,0 +1,153 @@
+import os
+import shutil
+import subprocess
+from time import sleep
+from typing import List, Tuple
+
+import streamlit.web.bootstrap
+from jina import Gateway
+from jina.serve.runtimes.gateway.composite import CompositeGateway
+from streamlit.file_util import get_streamlit_file_path
+from streamlit.web.server import Server as StreamlitServer
+
+
+cur_dir = os.path.dirname(__file__)
+
+
+def cmd(command, std_output=False, wait=True):
+    if isinstance(command, str):
+        command = command.split()
+    if not std_output:
+        process = subprocess.Popen(
+            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+    else:
+        process = subprocess.Popen(command)
+    if wait:
+        output, error = process.communicate()
+        return output, error
+
+
+class PlaygroundGateway(Gateway):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.streamlit_script = 'app.py'
+        # copy playground/config.toml to streamlit config.toml
+        streamlit_config_toml_src = os.path.join(cur_dir, 'app_config.toml')
+        streamlit_config_toml_dest = get_streamlit_file_path("config.toml")
+        # create streamlit_config_toml_dest if it doesn't exist
+        os.makedirs(os.path.dirname(streamlit_config_toml_dest), exist_ok=True)
+        shutil.copyfile(streamlit_config_toml_src, streamlit_config_toml_dest)
+
+    async def setup_server(self):
+        streamlit.web.bootstrap._fix_sys_path(self.streamlit_script)
+        streamlit.web.bootstrap._fix_matplotlib_crash()
+        streamlit.web.bootstrap._fix_tornado_crash()
+        streamlit.web.bootstrap._fix_sys_argv(self.streamlit_script, ())
+        streamlit.web.bootstrap._fix_pydeck_mapbox_api_warning()
+        streamlit_cmd = f'streamlit run {self.streamlit_script}'
+
+        self.streamlit_server = StreamlitServer(
+            os.path.join(cur_dir, self.streamlit_script), streamlit_cmd
+        )
+
+    async def run_server(self):
+        await self.streamlit_server.start()
+        streamlit.web.bootstrap._on_server_start(self.streamlit_server)
+        streamlit.web.bootstrap._set_up_signal_handler(self.streamlit_server)
+
+    async def shutdown(self):
+        self.streamlit_server.stop()
+
+
+class CustomGateway(CompositeGateway):
+    """The CustomGateway assumes that the gateway has been started with http on port 8080.
+    This is the port on which the nginx process listens. After nginx has been started,
+    it will start the playground on port 8501 and the actual HTTP gateway will start on port 8082.
+
+    Nginx is configured to route the requests in the following way:
+    - /playground -> playground on port 8501
+    - / -> HTTP gateway on port 8082
+    """
+
+    def __init__(self, **kwargs):
+        # need to update port to 8082, as nginx will listen on 8080
+        http_idx = 0
+        http_port = kwargs['runtime_args']['port'][http_idx]
+        if kwargs['runtime_args']['port'][http_idx] != 8080:
+            raise ValueError(
+                f'Please, let http port ({http_port}) be 8080 for nginx to work'
+            )
+        kwargs['runtime_args']['port'][http_idx] = 8082
+        super().__init__(**kwargs)
+
+        # remove potential clashing arguments from kwargs
+        kwargs.pop("port", None)
+        kwargs.pop("protocol", None)
+
+        # note order is important
+        self._add_gateway(
+            PlaygroundGateway,
+            8501,
+            **kwargs,
+        )
+
+        self.setup_nginx()
+        self.nginx_was_shutdown = False
+
+    async def shutdown(self):
+        await super().shutdown()
+        if not self.nginx_was_shutdown:
+            self.shutdown_nginx()
+            self.nginx_was_shutdown = True
+
+    def setup_nginx(self):
+        command = [
+            'nginx',
+            '-c',
+            os.path.join(cur_dir, '', 'nginx.conf'),
+        ]
+        output, error = self._run_nginx_command(command)
+        self.logger.info('Nginx started')
+        self.logger.info(f'nginx output: {output}')
+        self.logger.info(f'nginx error: {error}')
+
+    def shutdown_nginx(self):
+        command = ['nginx', '-s', 'stop']
+        output, error = self._run_nginx_command(command)
+        self.logger.info('Nginx stopped')
+        self.logger.info(f'nginx output: {output}')
+        self.logger.info(f'nginx error: {error}')
+
+    def _run_nginx_command(self, command: List[str]) -> Tuple[bytes, bytes]:
+        self.logger.info(f'Running command: {command}')
+        output, error = cmd(command)
+        if error != b'':
+            # on CI we need to use sudo; using NOW_CI_RUN isn't good if running test locally
+            self.logger.info(f'nginx error: {error}')
+            command.insert(0, 'sudo')
+            self.logger.info(f'So running command: {command}')
+            output, error = cmd(command)
+        sleep(10)
+        return output, error
+
+    def _add_gateway(self, gateway_cls, port, protocol='http', **kwargs):
+        # ignore metrics_registry since it is not copyable
+        runtime_args = self._deepcopy_with_ignore_attrs(
+            self.runtime_args,
+            [
+                'metrics_registry',
+                'tracer_provider',
+                'grpc_tracing_server_interceptors',
+                'aio_tracing_client_interceptors',
+                'tracing_client_interceptor',
+                'monitoring',  # disable it for fastapi gateway
+            ],
+        )
+        runtime_args.port = [port]
+        runtime_args.protocol = [protocol]
+        gateway_kwargs = {k: v for k, v in kwargs.items() if k != 'runtime_args'}
+        gateway_kwargs['runtime_args'] = dict(vars(runtime_args))
+        gateway = gateway_cls(**gateway_kwargs)
+        gateway.streamer = self.streamer
+        self.gateways.insert(0, gateway)

src/options/generate/static_files/gateway/nginx.conf

@@ -0,0 +1,62 @@
+events {
+  worker_connections  4096;  ## Default: 1024
+}
+
+http  {
+    server {
+        listen 8080;
+        server_name localhost;
+
+
+        # from https://medium.com/@dasirra/using-streamlit-nginx-docker-to-build-and-put-in-production-dashboards-in-aws-lightsail-781dab8f2836
+        location ^~ /static {
+            proxy_pass http://localhost:8501/static/;
+        }
+        location ^~ /healthz {
+            proxy_pass http://localhost:8501/healthz;
+        }
+        location ^~ /vendor {
+            proxy_pass http://localhost:8501/vendor;
+        }
+        location ^~ /st-allowed-message-origins {
+            proxy_pass http://localhost:8501/st-allowed-message-origins;
+        }
+
+        # for jcloud deployment, very important; actually talks via websocket
+        location ^~ /stream {
+            # inspired from https://discuss.streamlit.io/t/how-to-use-streamlit-with-nginx/378/7
+            proxy_pass http://localhost:8501/stream;
+            proxy_http_version 1.1;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header Host $host;
+            proxy_set_header Upgrade $http_upgrade;
+            proxy_set_header Connection "upgrade";
+            proxy_read_timeout 86400;
+        }
+        location ^~ /favicon.png {
+            proxy_pass http://localhost:8501/favicon.png;
+        }
+        # to make extra components work
+        location ^~ /component {
+            proxy_pass http://localhost:8501/component;
+        }
+
+        location /playground {
+            # streamlit specific from https://discuss.streamlit.io/t/streamlit-docker-nginx-ssl-https/2195
+            proxy_http_version 1.1;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header Host $host;
+            proxy_set_header Upgrade $http_upgrade;
+            proxy_set_header Connection "upgrade";
+            proxy_read_timeout 86400;
+            proxy_pass http://localhost:8501;
+            client_max_body_size 50M;
+        }
+
+        location / {
+            proxy_pass http://localhost:8082;
+            client_max_body_size 50M;
+        }
+
+     }
+}

src/options/generate/static_files/gateway/requirements.txt

@@ -0,0 +1,3 @@
+streamlit==1.16.0
+extra-streamlit-components==0.1.55
+jina==3.14.1

src/options/generate/templates_user.py

@@ -299,14 +299,15 @@ The playground contains many emojis that fit the theme of the playground and has
 This is an example how you can connect to the executor assuming the document (d) is already defined:
 ```
 from jina import Client, Document, DocumentArray
-client = Client(host=host)
+client = Client(host='http://localhost:8080')
 response = client.post('/', inputs=DocumentArray([d])) # always use '/'
 print(response[0].text) # can also be blob in case of image/audio..., this should be visualized in the streamlit app
 ```
 Note that the response will always be in response[0].text
 You must provide the complete app.py file with the exact same syntax to wrap the code.
-The playground (app.py) must read the host from sys.argv because it will be started with a custom host: streamlit run app.py -- --host http(s)://...
+The playground (app.py) must always use the host on http://localhost:8080.
-The playground (app.py) must not let the user configure the host on the ui.
+The playground (app.py) must not let the user configure the host on the UI.
+The playground (app.py) must not import the executor.
 '''
 )