91 lines
4.0 KiB
Bash
91 lines
4.0 KiB
Bash
# shellcheck shell=bash
|
|
# Central environment for the Ollama x Mooncake KVCache Bus integration.
|
|
#
|
|
# Layout philosophy:
|
|
# * WS -> the *source* tree (this integration dir inside the Mooncake
|
|
# repo). Resolved from this file's location, so the checkout
|
|
# can live at any path without edits. Never written to by the
|
|
# build/run (keeps the git tree clean for a PR).
|
|
# * OMB_STATE -> all *generated* state: toolchains, venv, caches, llama.cpp /
|
|
# ollama checkouts, model weights, logs, run dirs. Defaults to
|
|
# a sibling ".omb-state" OUTSIDE the repo so heavy artifacts
|
|
# never pollute the tree and never land on '/'. Override with
|
|
# OMB_STATE_DIR to point at a large data volume.
|
|
#
|
|
# Source at the top of every script: source "<dir>/env.sh"
|
|
|
|
# --- Source tree root (this integration dir) ---
|
|
export WS="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)"
|
|
|
|
# --- Generated-state root (kept OUTSIDE the source tree) ---
|
|
# Default: a ".omb-state" directory next to the top of the git repo if we can
|
|
# find it, else next to WS. Override via OMB_STATE_DIR for a big data volume.
|
|
if [ -z "${OMB_STATE_DIR:-}" ]; then
|
|
_repo_root="$(cd "$WS" && git rev-parse --show-toplevel 2>/dev/null || true)"
|
|
if [ -n "${_repo_root:-}" ]; then
|
|
OMB_STATE_DIR="$(dirname "$_repo_root")/.omb-state"
|
|
else
|
|
OMB_STATE_DIR="$(dirname "$WS")/.omb-state"
|
|
fi
|
|
fi
|
|
export OMB_STATE="$OMB_STATE_DIR"
|
|
|
|
# --- Keep ALL caches off '/' and out of the repo ---
|
|
export TMPDIR="$OMB_STATE/tmp"
|
|
export PIP_CACHE_DIR="$OMB_STATE/.cache/pip"
|
|
export XDG_CACHE_HOME="$OMB_STATE/.cache/xdg"
|
|
export HF_HOME="$OMB_STATE/.cache/hf"
|
|
export HUGGINGFACE_HUB_CACHE="$OMB_STATE/.cache/hf/hub"
|
|
|
|
# --- Go toolchain (installed locally, never system-wide) ---
|
|
export GOROOT="$OMB_STATE/go"
|
|
export GOPATH="$OMB_STATE/.cache/gopath"
|
|
export GOMODCACHE="$OMB_STATE/.cache/gomod"
|
|
export GOCACHE="$OMB_STATE/.cache/go-build"
|
|
export GOBIN="$OMB_STATE/.cache/gopath/bin"
|
|
export GOFLAGS="-mod=mod"
|
|
export GOTOOLCHAIN="local" # never auto-download a different toolchain
|
|
|
|
# --- Python venv ---
|
|
export VENV="$OMB_STATE/.venv"
|
|
|
|
# --- CUDA / GPU. Restrict our processes to a subset of devices. ---
|
|
export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda-12.4}"
|
|
export OMB_GPUS="${OMB_GPUS:-2,3,4,5,6,7}"
|
|
|
|
# --- llama.cpp build/run ---
|
|
export LLAMA_DIR="$OMB_STATE/llama.cpp"
|
|
export LLAMA_BUILD="$LLAMA_DIR/build"
|
|
export MODELS_DIR="$OMB_STATE/models"
|
|
export OLLAMA_DIR="$OMB_STATE/ollama"
|
|
|
|
# --- Runtime dirs ---
|
|
export OMB_RUN="$OMB_STATE/run"
|
|
export OMB_LOGS="$OMB_STATE/run/logs"
|
|
# Slot save files live on tmpfs (RAM) so the GPU<->host<->store path never
|
|
# touches a spinning/NVMe disk. Scoped + cleaned after use.
|
|
export OMB_SLOTS="${OMB_SLOTS_OVERRIDE:-/dev/shm/omb_slots}"
|
|
export OMB_SOCK="$OMB_STATE/run/sockets"
|
|
export OMB_STORE_DATA="$OMB_STATE/run/store" # mooncake local store spill
|
|
|
|
# --- PATH ---
|
|
export PATH="$GOROOT/bin:$GOBIN:$VENV/bin:$CUDA_HOME/bin:$PATH"
|
|
export LD_LIBRARY_PATH="$CUDA_HOME/lib64:$LLAMA_BUILD/bin:${LD_LIBRARY_PATH:-}"
|
|
|
|
# --- Default service ports (loopback only, high range to avoid clashes) ---
|
|
export OMB_BRIDGE_GRPC_PORT="${OMB_BRIDGE_GRPC_PORT:-52051}" # Go sidecar gRPC
|
|
export OMB_BRIDGE_HTTP_PORT="${OMB_BRIDGE_HTTP_PORT:-52052}" # Go sidecar metrics/HTTP
|
|
export OMB_STORE_PROXY_PORT="${OMB_STORE_PROXY_PORT:-52060}" # Python store proxy gRPC
|
|
export OMB_MASTER_PORT="${OMB_MASTER_PORT:-52061}" # mooncake master
|
|
export OMB_META_PORT="${OMB_META_PORT:-52062}" # mooncake metadata (http)
|
|
export OMB_LLAMA_PORT="${OMB_LLAMA_PORT:-52070}" # llama-server base port
|
|
|
|
mkdir -p "$TMPDIR" "$OMB_LOGS" "$OMB_SLOTS" "$OMB_SOCK" "$OMB_STORE_DATA" \
|
|
"$PIP_CACHE_DIR" "$GOPATH" "$GOMODCACHE" "$GOCACHE" "$HF_HOME" "$XDG_CACHE_HOME" 2>/dev/null
|
|
|
|
# Helper: activate python venv if present
|
|
if [ -f "$VENV/bin/activate" ]; then
|
|
# shellcheck disable=SC1091
|
|
source "$VENV/bin/activate"
|
|
fi
|