Setup scripts for local dev

Setup scripts will assume it's in the root plugin directory of OSD and places the `scripts` folder into the OSD/scripts folder for run purposes with the following command: ``` yarn setup ``` This then will allow from the root OSD directory to run ``` yarn assistant:start ``` which will setup up the OpenSearch cluster as expected. Please note, at the time of writing this the feature/os-assistant and feature/2.11/os-assistant added the feature for installing OpenSearch assistant related plugins. Not main or 2.x. Signed-off-by: Kawika Avilla <[email protected]>
opensearch-project · Nov 30, 2023 · acc3701 · acc3701
1 parent a5c2025
commit acc3701
Show file tree

Hide file tree

Showing 6 changed files with 301 additions and 1 deletion.
diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md
@@ -24,6 +24,34 @@ Ultimately, your directory structure should look like this:
 │       └── dashboards-assistant
 ```
 
+Setup local stack
+
+1. Export keys locally:
+```sh
+# defaults to us-west-2
+export REGION=
+export AWS_ACCESS_KEY_ID=
+export AWS_SECRET_ACCESS_KEY=
+export AWS_SESSION_TOKEN=
+```
+2. Setup OSD plugins
+```sh
+# from OSD root folder
+cd plugins
+# example with security plugin
+git clone --depth 1 --branch $VERSION https://github.com/opensearch-project/security-dashboards-plugin.git
+cd dashboards-assistant
+# NOTE: this will change OSD to be version 2.11.0 for dev purposes
+yarn setup
+cd ../../
+yarn osd bootstrap
+```
+2. Start local environment
+```sh
+# from OSD root folder
+yarn start:assistant
+```
+
 ### Build
 
 To build the plugin's distributable zip simply run `yarn build` in the plugin's directory.

diff --git a/package.json b/package.json
@@ -10,7 +10,8 @@
     "plugin-helpers": "node ../../scripts/plugin_helpers",
     "prepare": "husky install",
     "lint:es": "node ../../scripts/eslint",
-    "lint": "yarn lint:es"
+    "lint": "yarn lint:es",
+    "setup": "scripts/pre_install.sh"
   },
   "lint-staged": {
     "*.{ts,tsx,js,jsx}": [

diff --git a/scripts/assistant/add_model.sh b/scripts/assistant/add_model.sh
@@ -0,0 +1,89 @@
+#!/usr/bin/env bash
+set -e
+
+function add_model() {
+  [ -z "$ACCESS_KEY_ID" ] && ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
+  [ -z "$SECRET_ACCESS_KEY" ] && SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
+  [ -z "$SESSION_TOKEN" ] && SESSION_TOKEN=${AWS_SESSION_TOKEN}
+
+  ENDPOINT=https://admin:${CREDENTIAL}@${BIND_ADDRESS}:${BIND_PORT}
+
+  curl -s -k "${ENDPOINT}/_cluster/settings" -XPUT -H 'Content-Type: application/json' -d '{
+    "persistent" : {
+      "plugins.ml_commons.trusted_connector_endpoints_regex":
+      [ "^https://runtime\\.sagemaker\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
+        "^https://api\\.openai\\.com/.*$",
+        "^https://api\\.cohere\\.ai/.*$",
+        "^https://bedrock-runtime\\.[a-z0-9-]+\\.amazonaws\\.com/.*$"
+      ]
+    }
+  }' | jq
+
+  # shellcheck disable=2016
+  CONNECTOR=$(curl -s -k "${ENDPOINT}/_plugins/_ml/connectors/_create" -XPOST -H 'Content-Type: application/json' -d '{
+    "name": "BedRock test claude Connector",
+    "description": "The connector to BedRock service for claude model",
+    "version": 1,
+    "protocol": "aws_sigv4",
+    "parameters": {
+        "region": "'$REGION'",
+        "service_name": "bedrock",
+        "anthropic_version": "bedrock-2023-05-31",
+        "endpoint": "bedrock.'$REGION'.amazonaws.com",
+        "auth": "Sig_V4",
+        "content_type": "application/json"
+    },
+    "credential": {
+      "access_key": "'$ACCESS_KEY_ID'",
+      "secret_key": "'$SECRET_ACCESS_KEY'",
+      "session_token": "'$SESSION_TOKEN'"
+    },
+    "actions": [
+      {
+        "action_type": "predict",
+        "method": "POST",
+        "url": "https://bedrock-runtime.'"$REGION"'.amazonaws.com/model/anthropic.claude-v1/invoke",
+        "headers": {
+          "content-type": "application/json",
+          "x-amz-content-sha256": "required"
+        },
+        "request_body": "{\"prompt\":\"${parameters.prompt}\", \"max_tokens_to_sample\":${parameters.max_tokens_to_sample}, \"temperature\":${parameters.temperature},  \"anthropic_version\":\"${parameters.anthropic_version}\" }"
+      }
+    ]
+  }' | jq -r '.connector_id')
+  echo "❗connector: ${CONNECTOR}"
+
+  GROUP=$(curl -s -k "${ENDPOINT}/_plugins/_ml/model_groups/_register" -XPOST -H 'Content-Type: application/json' -d '{
+      "name": "test_model_group_public",
+      "description": "This is a public model group"
+  }' | jq | grep -oP '(?<=ID: )(.+)(?=\.)|(?<=model_group_id": ")(.+)(?=",)' | head -n 1)
+  echo "❗group: ${GROUP}"
+
+  EMBEDDINGS_TASK=$(curl -s -k "${ENDPOINT}/_plugins/_ml/models/_register?deploy=true" -XPOST -H 'Content-Type: application/json' -d '{
+      "name": "huggingface/sentence-transformers/all-mpnet-base-v2",
+      "version": "1.0.1",
+      "model_group_id": "'$GROUP'",
+      "model_format": "TORCH_SCRIPT"
+    }' | jq -r '.task_id')
+  echo "❗embeddings_task: ${EMBEDDINGS_TASK}"
+
+  MODEL=$(curl -s -k "${ENDPOINT}/_plugins/_ml/models/_register?deploy=true" -XPOST -H 'Content-Type: application/json' -d '{
+      "name": "Claude model on bedrock",
+      "model_group_id": "'$GROUP'",
+      "function_name": "remote",
+      "version": "1.0.0",
+      "connector_id": "'$CONNECTOR'",
+      "description": "test model"
+    }' | jq -r '.model_id')
+
+  echo "❗model: ${MODEL}"
+  sleep 40
+  EMBEDDINGS_MODEL=$(curl -s -k "${ENDPOINT}/_plugins/_ml/tasks/${EMBEDDINGS_TASK}" | jq -r '.model_id')
+  echo "❗embeddings_model: ${EMBEDDINGS_MODEL}"
+
+  curl -s -k "${ENDPOINT}/.chat-assistant-config/_doc/model-config" -XPOST -H 'Content-Type: application/json' -d '{
+    "model_type":"claude_bedrock",
+    "embeddings_model_id":"'$EMBEDDINGS_MODEL'",
+    "model_id":"'$MODEL'"
+  }' | jq
+}
diff --git a/scripts/assistant/start.sh b/scripts/assistant/start.sh
@@ -0,0 +1,85 @@
+#!/bin/bash
+
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+set -e
+
+. scripts/assistant/utils.sh
+. scripts/assistant/add_model.sh
+
+function usage() {
+    echo ""
+    echo "This script is used to run OpenSearch Assistant"
+    echo "--------------------------------------------------------------------------"
+    echo "Usage: $0 [args]"
+    echo ""
+    echo "Optional arguments:"
+    echo -e "-b BIND_ADDRESS\t, defaults to localhost | 127.0.0.1, can be changed to any IP or domain name for the cluster location."
+    echo -e "-p BIND_PORT\t, defaults to 5601 depends on OpenSearch or Dashboards, can be changed to any port for the cluster location."
+    echo -e "-c CREDENTIAL\t(password), defaults to admin"
+    echo -e "-h\tPrint this message."
+    echo "--------------------------------------------------------------------------"
+}
+
+while getopts ":h:b:p:c:" arg; do
+    case $arg in
+        h)
+            usage
+            exit 1
+            ;;
+        b)
+            BIND_ADDRESS=$OPTARG
+            ;;
+        p)
+            BIND_PORT=$OPTARG
+            ;;
+        c)
+            CREDENTIAL=$OPTARG
+            ;;
+        :)
+            echo "-${OPTARG} requires an argument"
+            usage
+            exit 1
+            ;;
+        ?)
+            echo "Invalid option: -${OPTARG}"
+            exit 1
+            ;;
+    esac
+done
+
+[ -z "$BIND_ADDRESS" ] && BIND_ADDRESS="localhost"
+[ -z "$BIND_PORT" ] && BIND_PORT="9200"
+[ -z "$CREDENTIAL" ] && CREDENTIAL="admin"
+[ -z "$REGION" ] && REGION="us-west-2"
+
+PARENT_PID_LIST=()
+
+PACKAGE_VERSION=$(yarn --silent pkg-version)
+
+# define assistant path
+CWD=$(pwd)
+SNAPSHOT_DIR="$CWD/.opensearch"
+LOGS_DIR="$SNAPSHOT_DIR/$PACKAGE_VERSION/logs"
+
+# Main function
+function execute() {
+  export initialAdminPassword=$CREDENTIAL
+  CLUSTER_SETTINGS="snapshot --assistant --security"
+  CLUSTER_SETTINGS+=" -E plugins.ml_commons.only_run_on_ml_node=true"
+
+  run_opensearch || clean
+  check_opensearch_status
+  echo "Attempting to add models..."
+  echo "(Ensure your environment is exporting your credentials)"
+  (add_model > $LOGS_DIR/add_model.log 2>&1 || clean) & 
+
+  export OPENSEARCH_USERNAME=kibanaserver
+  export OPENSEARCH_PASSWORD=kibanaserver
+  $CWD/scripts/use_node $CWD/scripts/opensearch_dashboards --dev --security || clean
+}
+
+execute
+clean
+exit 0
diff --git a/scripts/assistant/utils.sh b/scripts/assistant/utils.sh
@@ -0,0 +1,86 @@
+#!/bin/bash
+
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+set -e
+
+function open_artifact() {
+  artifact_dir=$1
+  artifact=$2
+  cd $artifact_dir
+
+  # check if artifact provided is URL or attempt if passing by absolute path
+  if curl -I -L $artifact; then
+    curl -L $artifact | tar -xz --strip-components=1
+  else
+    echo "Artifact is not a URL; attempting to unarchive a local file..."
+    tar -xf $artifact --strip-components=1
+  fi
+}
+
+# remove the running opensearch process
+function clean() {
+  echo "Attempt to Terminate Process with PID: ${PARENT_PID_LIST[*]}"
+  for pid_kill in "${PARENT_PID_LIST[@]}"
+  do
+    echo "Closing PID $pid_kill"
+    kill $pid_kill || true
+  done
+  PARENT_PID_LIST=()
+}
+
+function spawn_process_and_save_PID() {
+    echo "Spawn '$@'"
+    eval $@
+    curr_pid=$!
+    echo "PID: $curr_pid"
+    PARENT_PID_LIST+=( $curr_pid )
+}
+
+# Print out a textfile line by line
+function print_txt() {
+  while IFS= read -r line; do
+    echo "text read from $1: $line"
+  done < $1
+}
+
+# this function is used to check the running status of OpenSearch or OpenSearch Dashboards
+# $1 is the path to the tmp file which saves the running status
+# $2 is the error msg to check
+# $3 is the url to curl
+# $4 contains arguments that need to be passed to the curl command
+function check_status() {
+  # Calculate end time as 350s from now
+  check_status_end_time=$(expr 350 + "$(date '+%s')")
+
+  while [ ! -e $1 ] || ! grep -q "$2" $1; do
+    sleep 1
+    # Stop checking after $check_status_end_time
+    if [ $check_status_end_time -lt $(date '+%s') ]; then
+      echo "Error: Status check has timed out"
+      exit 1
+    fi
+  done
+}
+
+# Checks the running status of OpenSearch
+# it calls check_status and passes the OpenSearch tmp file path, error msg, url, and arguments
+# if success, the while loop in the check_status will end and it prints out "OpenSearch is up!"
+function check_opensearch_status() {
+  echo "Checking the status OpenSearch..."
+  # define other paths and tmp files
+  OPENSEARCH_FILE='opensearch.log'
+  OPENSEARCH_LOG_PATH="$LOGS_DIR/$OPENSEARCH_FILE"
+
+  OPENSEARCH_MSG="ML configuration initialized successfully"
+  check_status $OPENSEARCH_LOG_PATH "$OPENSEARCH_MSG" 2>&1
+  echo "OpenSearch is up!" 
+}
+
+# Starts OpenSearch
+function run_opensearch() {
+  echo "[ Attempting to start OpenSearch... ]"
+  rm -rf $LOGS_DIR/opensearch.log 
+  spawn_process_and_save_PID  "$CWD/scripts/use_node $CWD/scripts/opensearch $CLUSTER_SETTINGS > $LOGS_DIR/opensearch.log 2>&1 &"
+}
diff --git a/scripts/pre_install.sh b/scripts/pre_install.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+set -e
+
+[ -z "$SUPPORTED_VERSION" ] && SUPPORTED_VERSION="2.11.0"
+
+cp -r scripts/assistant ../../scripts
+sed -i -E "s|(\"version\": \")[^\"]*|\1${SUPPORTED_VERSION}|" ../../package.json