Merge pull request #91 from pinecone-io/app_config

Configure the service + CLI from config file

pyproject.toml

@@ -23,6 +23,8 @@ uvicorn = "^0.20.0"
 tenacity = "^8.2.1"
 sse-starlette = "^1.6.5"
 types-tqdm = "^4.61.0"
+gunicorn = "^21.2.0"
+types-pyyaml = "^6.0.12.12"
 
 
 [tool.poetry.group.dev.dependencies]

src/resin_cli/app.py

@@ -6,11 +6,13 @@
 
 import openai
 from multiprocessing import current_process
+
+import yaml
 from dotenv import load_dotenv
 
 from resin.llm import BaseLLM
 from resin.llm.models import UserMessage
-from resin.tokenizer import OpenAITokenizer, Tokenizer
+from resin.tokenizer import Tokenizer
 from resin.knowledge_base import KnowledgeBase
 from resin.context_engine import ContextEngine
 from resin.chat_engine import ChatEngine
@@ -28,6 +30,7 @@
     ContextUpsertRequest, HealthStatus, ContextDeleteRequest
 
 from resin.llm.openai import OpenAILLM
+from resin_cli.errors import ConfigError
 
 load_dotenv()  # load env vars before import of openai
 openai.api_key = os.getenv("OPENAI_API_KEY")
@@ -198,24 +201,66 @@ def _init_logging():
 
 
 def _init_engines():
-    global kb, context_engine, chat_engine, llm
-    Tokenizer.initialize(OpenAITokenizer, model_name='gpt-3.5-turbo-0613')
+    global kb, context_engine, chat_engine, llm, logger
 
-    INDEX_NAME = os.getenv("INDEX_NAME")
-    if not INDEX_NAME:
+    index_name = os.getenv("INDEX_NAME")
+    if not index_name:
         raise ValueError("INDEX_NAME environment variable must be set")
 
-    kb = KnowledgeBase(index_name=INDEX_NAME)
-    context_engine = ContextEngine(knowledge_base=kb)
-    llm = OpenAILLM()
-    chat_engine = ChatEngine(context_engine=context_engine, llm=llm)
+    config_file = os.getenv("RESIN_CONFIG_FILE")
+    if config_file:
+        _load_config(config_file)
+
+    else:
+        logger.info("Did not find config file. Initializing engines with default "
+                    "configuration")
+        Tokenizer.initialize()
+        kb = KnowledgeBase(index_name=index_name)
+        context_engine = ContextEngine(knowledge_base=kb)
+        llm = OpenAILLM()
+        chat_engine = ChatEngine(context_engine=context_engine, llm=llm)
 
     kb.connect()
 
 
-def start(host="0.0.0.0", port=8000, reload=False, workers=1):
-    uvicorn.run("resin_cli.app:app",
-                host=host, port=port, reload=reload, workers=workers)
+def _load_config(config_file):
+    global chat_engine, llm, context_engine, kb, logger
+    logger.info(f"Initializing engines with config file {config_file}")
+    try:
+        with open(config_file, "r") as f:
+            config = yaml.safe_load(f)
+    except Exception as e:
+        logger.exception(f"Failed to load config file {config_file}")
+        raise ConfigError(
+            f"Failed to load config file {config_file}. Error: {str(e)}"
+        )
+    tokenizer_config = config.get("tokenizer", {})
+    Tokenizer.initialize_from_config(tokenizer_config)
+    if "chat_engine" not in config:
+        raise ConfigError(
+            f"Config file {config_file} must contain a 'chat_engine' section"
+        )
+    chat_engine_config = config["chat_engine"]
+    try:
+        chat_engine = ChatEngine.from_config(chat_engine_config)
+    except Exception as e:
+        logger.exception(
+            f"Failed to initialize chat engine from config file {config_file}"
+        )
+        raise ConfigError(
+            f"Failed to initialize chat engine from config file {config_file}."
+            f" Error: {str(e)}"
+        )
+    llm = chat_engine.llm
+    context_engine = chat_engine.context_engine
+    kb = context_engine.knowledge_base
+
+
+def start(host="0.0.0.0", port=8000, reload=False, config_file=None):
+    if config_file:
+        os.environ["RESIN_CONFIG_FILE"] = config_file
+
+    uvicorn.run("resin_cli.app:app", host=host, port=port, reload=reload, workers=0)
 
 
 if __name__ == "__main__":

src/resin_cli/cli.py

@@ -1,10 +1,13 @@
 import os
-from typing import List, Optional
+import signal
+import subprocess
+from typing import Dict, Any, Optional, List
 
 import click
 import time
 
 import requests
+import yaml
 from dotenv import load_dotenv
 from tenacity import retry, stop_after_attempt, wait_fixed
 from tqdm import tqdm
@@ -19,9 +22,9 @@
 from resin.tokenizer import Tokenizer
 from resin_cli.data_loader import (
     load_from_path,
-    CLIError,
     IDsNotUniqueError,
     DocumentsValidationError)
+from resin_cli.errors import CLIError
 
 from resin import __version__
 
@@ -30,8 +33,7 @@
 from .api_models import ChatDebugInfo
 
 
-dotenv_path = os.path.join(os.path.dirname(__file__), ".env")
-load_dotenv(dotenv_path)
+load_dotenv()
 if os.getenv("OPENAI_API_KEY"):
     openai.api_key = os.getenv("OPENAI_API_KEY")
 
@@ -99,6 +101,34 @@ def _initialize_tokenizer():
         raise CLIError(msg)
 
 
+def _load_kb_config(config_file: Optional[str]) -> Dict[str, Any]:
+    if config_file is None:
+        return {}
+
+    try:
+        with open(os.path.join("config", config_file), 'r') as f:
+            config = yaml.safe_load(f)
+    except Exception as e:
+        msg = f"Failed to load config file {config_file}. Reason:\n{e}"
+        raise CLIError(msg)
+
+    if "knowledge_base" in config:
+        kb_config = config.get("knowledge_base", None)
+    elif "chat_engine" in config:
+        kb_config = config["chat_engine"]\
+            .get("context_engine", {})\
+            .get("knowledge_base", None)
+    else:
+        kb_config = None
+
+    if kb_config is None:
+        msg = (f"Did not find a `knowledge_base` configuration in {config_file}, "
+               "Would you like to use the default configuration?")
+        click.confirm(click.style(msg, fg="red"), abort=True)
+        kb_config = {}
+    return kb_config
+
+
 @click.group(invoke_without_command=True, context_settings=CONTEXT_SETTINGS)
 @click.version_option(__version__, "-v", "--version", prog_name="Resin")
 @click.pass_context
@@ -137,9 +167,13 @@ def health(url):
     )
 )
 @click.argument("index-name", nargs=1, envvar="INDEX_NAME", type=str, required=True)
-def new(index_name):
+@click.option("--config", "-c", default=None,
+              help="Path to a resin config file. Optional, otherwise configuration "
+                   "defaults will be used.")
+def new(index_name: str, config: Optional[str]):
     _initialize_tokenizer()
-    kb = KnowledgeBase(index_name=index_name)
+    kb_config = _load_kb_config(config)
+    kb = KnowledgeBase.from_config(kb_config, index_name=index_name)
     click.echo("Resin is going to create a new index: ", nl=False)
     click.echo(click.style(f"{kb.index_name}", fg="green"))
     click.confirm(click.style("Do you want to continue?", fg="red"), abort=True)
@@ -179,7 +213,14 @@ def new(index_name):
                    "be uploaded. "
                    "When set to True, the upsert process will continue on failure, as "
                    "long as less than 10% of the documents have failed to be uploaded.")
-def upsert(index_name: str, data_path: str, batch_size: int, allow_failures: bool):
+@click.option("--config", "-c", default=None,
+              help="Path to a resin config file. Optional, otherwise configuration "
+                   "defaults will be used.")
+def upsert(index_name: str,
+           data_path: str,
+           batch_size: int,
+           allow_failures: bool,
+           config: Optional[str]):
     if index_name is None:
         msg = (
             "No index name provided. Please set --index-name or INDEX_NAME environment "
@@ -189,7 +230,8 @@ def upsert(index_name: str, data_path: str, batch_size: int, allow_failures: boo
 
     _initialize_tokenizer()
 
-    kb = KnowledgeBase(index_name=index_name)
+    kb_config = _load_kb_config(config)
+    kb = KnowledgeBase.from_config(kb_config, index_name=index_name)
     try:
         kb.connect()
     except RuntimeError as e:
@@ -427,10 +469,23 @@ def chat(chat_service_url, baseline, debug, stream):
               help="TCP port to bind the server to. Defaults to 8000")
 @click.option("--reload/--no-reload", default=False,
               help="Set the server to reload on code changes. Defaults to False")
-@click.option("--workers", default=1, help="Number of worker processes. Defaults to 1")
-def start(host, port, reload, workers):
+@click.option("--config", "-c", default=None,
+              help="Path to a resin config file. Optional, otherwise configuration "
+                   "defaults will be used.")
+def start(host: str, port: str, reload: bool, config: Optional[str]):
+    note_msg = (
+        "🚨 Note 🚨\n"
+        "For debugging only. To run the Resin service in production, run the command:\n"
+        "gunicorn resin_cli.app:app --worker-class uvicorn.workers.UvicornWorker "
+        f"--bind {host}:{port} --workers <num_workers>"
+    )
+    for c in note_msg:
+        click.echo(click.style(c, fg="red"), nl=False)
+        time.sleep(0.01)
+    click.echo()
+
     click.echo(f"Starting Resin service on {host}:{port}")
-    start_service(host, port=port, reload=reload, workers=workers)
+    start_service(host, port=port, reload=reload, config_file=config)
 
 
 @cli.command(
@@ -445,6 +500,28 @@ def start(host, port, reload, workers):
 @click.option("url", "--url", default="http://0.0.0.0:8000",
               help="URL of the Resin service to use. Defaults to http://0.0.0.0:8000")
 def stop(url):
+    # Check if the service was started using Gunicorn
+    res = subprocess.run(["pgrep", "-f", "gunicorn resin_cli.app:app"],
+                         capture_output=True)
+    output = res.stdout.decode("utf-8").split()
+
+    # If Gunicorn was used, kill all Gunicorn processes
+    if output:
+        msg = ("It seems that Resin service was launched using Gunicorn.\n"
+               "Do you want to kill all Gunicorn processes?")
+        click.confirm(click.style(msg, fg="red"), abort=True)
+        try:
+            subprocess.run(["pkill", "-f", "gunicorn resin_cli.app:app"], check=True)
+        except subprocess.CalledProcessError:
+            try:
+                [os.kill(int(pid), signal.SIGINT) for pid in output]
+            except OSError:
+                msg = (
+                    "Could not kill Gunicorn processes. Please kill them manually."
+                    f"Found process ids: {output}"
+                )
+                raise CLIError(msg)
+
     try:
         res = requests.get(urljoin(url, "/shutdown"))
         res.raise_for_status()

src/resin_cli/data_loader/__init__.py

@@ -1,6 +1,5 @@
 from .data_loader import (
     load_from_path,
-    CLIError,
     IDsNotUniqueError,
     DocumentsValidationError
 )

src/resin_cli/data_loader/data_loader.py

@@ -5,10 +5,8 @@
 from typing import List
 from textwrap import dedent
 
-import click
 import numpy as np
 import pandas as pd
-from click import ClickException
 
 from pydantic import ValidationError
 
@@ -27,11 +25,6 @@ def format_multiline(msg):
     return dedent(msg).strip()
 
 
-class CLIError(ClickException):
-    def format_message(self) -> str:
-        return click.style(format_multiline(self.message), fg='red')
-
-
 def _process_metadata(value):
     if pd.isna(value):
         return {}

src/resin_cli/errors.py

@@ -0,0 +1,13 @@
+import click
+from click import ClickException
+
+from resin_cli.data_loader.data_loader import format_multiline
+
+
+class CLIError(ClickException):
+    def format_message(self) -> str:
+        return click.style(format_multiline(self.message), fg='red')
+
+
+class ConfigError(RuntimeError):
+    pass