Spaces:

zrxarchit
/

inference

Sleeping

Update start.sh

75cab2d verified about 2 months ago

7.77 kB

#!/usr/bin/env bash
set -euo pipefail

detect_cpu_quota_threads() {
  local quota period threads

  if [ -r /sys/fs/cgroup/cpu.max ]; then
    read -r quota period < /sys/fs/cgroup/cpu.max || true
    if [ "${quota:-max}" != "max" ] && [ "${period:-0}" -gt 0 ] 2>/dev/null; then
      threads=$(( (quota + period - 1) / period ))
      [ "$threads" -gt 0 ] && echo "$threads" && return
    fi
  fi

  if [ -r /sys/fs/cgroup/cpu/cpu.cfs_quota_us ] && [ -r /sys/fs/cgroup/cpu/cpu.cfs_period_us ]; then
    quota=$(cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us)
    period=$(cat /sys/fs/cgroup/cpu/cpu.cfs_period_us)
    if [ "${quota:-0}" -gt 0 ] && [ "${period:-0}" -gt 0 ] 2>/dev/null; then
      threads=$(( (quota + period - 1) / period ))
      [ "$threads" -gt 0 ] && echo "$threads" && return
    fi
  fi

  echo 0
}

detect_memory_kb() {
  local value

  if [ -r /sys/fs/cgroup/memory.max ]; then
    value=$(cat /sys/fs/cgroup/memory.max)
    if [ "$value" != "max" ] && [ "${value:-0}" -gt 0 ] 2>/dev/null; then
      echo $(( value / 1024 ))
      return
    fi
  fi

  if [ -r /sys/fs/cgroup/memory/memory.limit_in_bytes ]; then
    value=$(cat /sys/fs/cgroup/memory/memory.limit_in_bytes)
    if [ "${value:-0}" -gt 0 ] && [ "$value" -lt 9223372036854771712 ] 2>/dev/null; then
      echo $(( value / 1024 ))
      return
    fi
  fi

  if [ -r /proc/meminfo ]; then
    awk '/MemTotal/ {print $2}' /proc/meminfo || echo 0
  else
    echo 0
  fi
}

normalize_arg_value() {
  local value="$1"

  value="${value//$'\r'/}"
  value="${value//$'\n'/}"
  value="${value#"${value%%[![:space:]]*}"}"
  value="${value%"${value##*[![:space:]]}"}"

  # Hugging Face env values are plain values. A copied shell-style "$q4_0"
  # should be treated as "q4_0" instead of crashing llama-server.
  value="${value#\$}"

  printf '%s' "$value"
}

HOST_THREADS="$(nproc)"
QUOTA_THREADS="$(detect_cpu_quota_threads)"
if [ -z "${THREADS:-}" ]; then
  if [ "$QUOTA_THREADS" -gt 0 ] && [ "$QUOTA_THREADS" -lt "$HOST_THREADS" ]; then
    THREADS="$QUOTA_THREADS"
  else
    THREADS="$HOST_THREADS"
  fi
fi

export OMP_NUM_THREADS="$THREADS"
export OPENBLAS_NUM_THREADS="$THREADS"
export OMP_PROC_BIND="${OMP_PROC_BIND:-FALSE}"
export OMP_PLACES="${OMP_PLACES:-cores}"
export OMP_WAIT_POLICY="${OMP_WAIT_POLICY:-PASSIVE}"
export MKL_NUM_THREADS="$THREADS"
export VECLIB_MAXIMUM_THREADS="$THREADS"
export NUMEXPR_NUM_THREADS="$THREADS"
export HF_HOME="${HF_HOME:-/data/hf-cache}"
export MODEL_DIR="${MODEL_DIR:-/data/models}"
export MODEL_PATH="${MODEL_PATH:-/data/models/model.gguf}"
REQUESTED_CTX_SIZE="${CONTEXT_LENGTH:-${CTX_SIZE:-}}"
export QUANT_PREFERENCE="${QUANT_PREFERENCE:-q4_k_m,q4_k_s,q4_k,q4_0,q4_1}"
export PERF_PROFILE="${PERF_PROFILE:-balanced}"
REQUESTED_CACHE_TYPE_K="${CACHE_TYPE_K:-}"
REQUESTED_CACHE_TYPE_V="${CACHE_TYPE_V:-}"
export REASONING="${REASONING:-auto}"
export TOOLS="${TOOLS:-}"
export ENABLE_TOOLS="${ENABLE_TOOLS:-0}"
export HTTP_THREADS="${HTTP_THREADS:-1}"
export LOG_VERBOSITY="${LOG_VERBOSITY:-3}"
export MMAP="${MMAP:-0}"
export FLASH_ATTN="${FLASH_ATTN:-1}"
export NO_WARMUP="${NO_WARMUP:-1}"
export LANGSEARCH_API_KEY="${LANGSEARCH_API_KEY:-}"

# Ensure data directories exist and are writable. Run as root in container startup.
echo "preparing storage: $HF_HOME and $MODEL_DIR"
install -d -m 0777 "$HF_HOME" "$MODEL_DIR" || mkdir -p "$HF_HOME" "$MODEL_DIR"
chmod 0777 "$HF_HOME" "$MODEL_DIR" || true

# Ensure HF cache logging directories are writable for xet
install -d -m 0777 "$HF_HOME/xet/logs" || true
chmod 0777 "$HF_HOME/xet" "$HF_HOME/xet/logs" 2>/dev/null || true
chown -R 65532:65532 "$HF_HOME/xet" 2>/dev/null || true

# Run downloader as the non-root `appuser` using the venv python when available.
echo "starting model downloader"
if [ -x "/opt/venv/bin/python" ]; then
  su -s /bin/sh appuser -c "/opt/venv/bin/python -u /app/download_model.py"
else
  su -s /bin/sh appuser -c "python3 -u /app/download_model.py"
fi

# If downloader wrote a repo marker with selected filename, update MODEL_PATH accordingly
if [ -f "$MODEL_DIR/model.repo" ]; then
  marker=$(cat "$MODEL_DIR/model.repo" || echo "")
  if [[ "$marker" == *"|"* ]]; then
    IFS='|' read -r _ selected_file _ <<< "$marker"
    if [ -n "$selected_file" ]; then
      MODEL_PATH="$MODEL_DIR/$selected_file"
      echo "using downloaded model path: $MODEL_PATH"
    fi
  fi
fi

LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-/usr/local/bin/llama-server}"

# set server profile defaults
case "$PERF_PROFILE" in
  low_latency)
    DEFAULT_BATCH=64
    DEFAULT_UBATCH=64
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
  balanced)
    DEFAULT_BATCH=512
    DEFAULT_UBATCH=512
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
  throughput)
    DEFAULT_BATCH=512
    DEFAULT_UBATCH=512
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
  *)
    DEFAULT_BATCH=128
    DEFAULT_UBATCH=128
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
esac

# Adjust defaults for low-memory environments (Spaces free tier: ~16GB).
# Use cgroup limits first because /proc/meminfo can report the host machine.
MEM_KB="$(detect_memory_kb)"
echo "detected effective memory (KB): $MEM_KB"
if [ "$MEM_KB" -gt 0 ] && [ "$MEM_KB" -le 18000000 ]; then
  echo "low-memory profile detected: using compact defaults"
  DEFAULT_CTX_SIZE=2048
else
  DEFAULT_CTX_SIZE=4096
fi

CTX_SIZE="${REQUESTED_CTX_SIZE:-$DEFAULT_CTX_SIZE}"
CACHE_TYPE_K="$(normalize_arg_value "${REQUESTED_CACHE_TYPE_K:-f16}")"
CACHE_TYPE_V="$(normalize_arg_value "${REQUESTED_CACHE_TYPE_V:-f16}")"
PARALLEL=1

echo "effective threads: $THREADS (host=$HOST_THREADS, quota=$QUOTA_THREADS)"
echo "effective llama profile: perf=$PERF_PROFILE single_request=1 ctx=$CTX_SIZE batch=${BATCH_SIZE:-$DEFAULT_BATCH} ubatch=${UBATCH_SIZE:-$DEFAULT_UBATCH} http_threads=$HTTP_THREADS log_verbosity=$LOG_VERBOSITY mmap=$MMAP flash_attn=$FLASH_ATTN"

export CACHE_TYPE_K CACHE_TYPE_V CTX_SIZE PARALLEL

server_args=(
  --model "$MODEL_PATH"
  --host 0.0.0.0
  --port "${PORT:-7860}"
  --api-key "$API_PASSWORD"
  -ngl 0
  -t "$THREADS"
  --threads-batch "${THREADS_BATCH:-$DEFAULT_THREADS_BATCH}"
  --threads-http "$HTTP_THREADS"
  --batch-size "${BATCH_SIZE:-$DEFAULT_BATCH}"
  --ubatch-size "${UBATCH_SIZE:-$DEFAULT_UBATCH}"
  --ctx-size "${CTX_SIZE:-4096}"
  --cache-type-k "${CACHE_TYPE_K}"
  --cache-type-v "${CACHE_TYPE_V}"
  --metrics
  -lv "$LOG_VERBOSITY"
)


if [ "$FLASH_ATTN" = "1" ]; then
  server_args+=(--flash-attn on)
fi

if [ "${NO_WARMUP}" = "1" ]; then
  server_args+=(--no-warmup)
fi

case "${REASONING,,}" in
  true|1|on|yes)
    server_args+=(--reasoning on)
    ;;
  false|0|off|no)
    server_args+=(--reasoning off)
    ;;
  auto|"")
    server_args+=(--reasoning auto)
    ;;
  *)
    echo "invalid REASONING value: $REASONING (expected True/False/auto)"
    exit 1
    ;;
esac

if [ "$ENABLE_TOOLS" = "1" ] && [ -n "$LANGSEARCH_API_KEY" ] && [ -z "$TOOLS" ]; then
  TOOLS="exec_shell_command"
fi

if [[ -n "${TOOLS}" ]]; then
  server_args+=(--tools "$TOOLS")
  echo "enabled tools: $TOOLS"
fi

server_args+=(--parallel "$PARALLEL")

# mmproj (vision projector) support
MMPROJ_PATH=""
if [ -n "${MMPROJ_FILE:-}" ]; then
  MMPROJ_PATH="$MODEL_DIR/$MMPROJ_FILE"
elif [ -f "$MODEL_DIR/model.mmproj" ]; then
  mmproj_name=$(cat "$MODEL_DIR/model.mmproj" || echo "")
  if [ -n "$mmproj_name" ]; then
    MMPROJ_PATH="$MODEL_DIR/$mmproj_name"
  fi
fi

if [ -n "$MMPROJ_PATH" ] && [ -f "$MMPROJ_PATH" ]; then
  server_args+=(--mmproj "$MMPROJ_PATH")
  echo "vision support enabled: $MMPROJ_PATH"
else
  echo "no mmproj found; running text-only"
fi

server_cmd=("$LLAMA_SERVER_BIN" "${server_args[@]}")
echo "server command: $LLAMA_SERVER_BIN ${server_args[*]}"
echo "starting llama-server"
exec su -s /bin/sh appuser -c "$(printf '%q ' "${server_cmd[@]}")"

	#!/usr/bin/env bash
	set -euo pipefail

	detect_cpu_quota_threads() {
	local quota period threads

	if [ -r /sys/fs/cgroup/cpu.max ]; then
	read -r quota period < /sys/fs/cgroup/cpu.max \|\| true
	if [ "${quota:-max}" != "max" ] && [ "${period:-0}" -gt 0 ] 2>/dev/null; then
	threads=$(( (quota + period - 1) / period ))
	[ "$threads" -gt 0 ] && echo "$threads" && return
	fi
	fi

	if [ -r /sys/fs/cgroup/cpu/cpu.cfs_quota_us ] && [ -r /sys/fs/cgroup/cpu/cpu.cfs_period_us ]; then
	quota=$(cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us)
	period=$(cat /sys/fs/cgroup/cpu/cpu.cfs_period_us)
	if [ "${quota:-0}" -gt 0 ] && [ "${period:-0}" -gt 0 ] 2>/dev/null; then
	threads=$(( (quota + period - 1) / period ))
	[ "$threads" -gt 0 ] && echo "$threads" && return
	fi
	fi

	echo 0
	}

	detect_memory_kb() {
	local value

	if [ -r /sys/fs/cgroup/memory.max ]; then
	value=$(cat /sys/fs/cgroup/memory.max)
	if [ "$value" != "max" ] && [ "${value:-0}" -gt 0 ] 2>/dev/null; then
	echo $(( value / 1024 ))
	return
	fi
	fi

	if [ -r /sys/fs/cgroup/memory/memory.limit_in_bytes ]; then
	value=$(cat /sys/fs/cgroup/memory/memory.limit_in_bytes)
	if [ "${value:-0}" -gt 0 ] && [ "$value" -lt 9223372036854771712 ] 2>/dev/null; then
	echo $(( value / 1024 ))
	return
	fi
	fi

	if [ -r /proc/meminfo ]; then
	awk '/MemTotal/ {print $2}' /proc/meminfo \|\| echo 0
	else
	echo 0
	fi
	}

	normalize_arg_value() {
	local value="$1"

	value="${value//$'\r'/}"
	value="${value//$'\n'/}"
	value="${value#"${value%%[![:space:]]*}"}"
	value="${value%"${value##*[![:space:]]}"}"

	# Hugging Face env values are plain values. A copied shell-style "$q4_0"
	# should be treated as "q4_0" instead of crashing llama-server.
	value="${value#\$}"

	printf '%s' "$value"
	}

	HOST_THREADS="$(nproc)"
	QUOTA_THREADS="$(detect_cpu_quota_threads)"
	if [ -z "${THREADS:-}" ]; then
	if [ "$QUOTA_THREADS" -gt 0 ] && [ "$QUOTA_THREADS" -lt "$HOST_THREADS" ]; then
	THREADS="$QUOTA_THREADS"
	else
	THREADS="$HOST_THREADS"
	fi
	fi

	export OMP_NUM_THREADS="$THREADS"
	export OPENBLAS_NUM_THREADS="$THREADS"
	export OMP_PROC_BIND="${OMP_PROC_BIND:-FALSE}"
	export OMP_PLACES="${OMP_PLACES:-cores}"
	export OMP_WAIT_POLICY="${OMP_WAIT_POLICY:-PASSIVE}"
	export MKL_NUM_THREADS="$THREADS"
	export VECLIB_MAXIMUM_THREADS="$THREADS"
	export NUMEXPR_NUM_THREADS="$THREADS"
	export HF_HOME="${HF_HOME:-/data/hf-cache}"
	export MODEL_DIR="${MODEL_DIR:-/data/models}"
	export MODEL_PATH="${MODEL_PATH:-/data/models/model.gguf}"
	REQUESTED_CTX_SIZE="${CONTEXT_LENGTH:-${CTX_SIZE:-}}"
	export QUANT_PREFERENCE="${QUANT_PREFERENCE:-q4_k_m,q4_k_s,q4_k,q4_0,q4_1}"
	export PERF_PROFILE="${PERF_PROFILE:-balanced}"
	REQUESTED_CACHE_TYPE_K="${CACHE_TYPE_K:-}"
	REQUESTED_CACHE_TYPE_V="${CACHE_TYPE_V:-}"
	export REASONING="${REASONING:-auto}"
	export TOOLS="${TOOLS:-}"
	export ENABLE_TOOLS="${ENABLE_TOOLS:-0}"
	export HTTP_THREADS="${HTTP_THREADS:-1}"
	export LOG_VERBOSITY="${LOG_VERBOSITY:-3}"
	export MMAP="${MMAP:-0}"
	export FLASH_ATTN="${FLASH_ATTN:-1}"
	export NO_WARMUP="${NO_WARMUP:-1}"
	export LANGSEARCH_API_KEY="${LANGSEARCH_API_KEY:-}"

	# Ensure data directories exist and are writable. Run as root in container startup.
	echo "preparing storage: $HF_HOME and $MODEL_DIR"
	install -d -m 0777 "$HF_HOME" "$MODEL_DIR" \|\| mkdir -p "$HF_HOME" "$MODEL_DIR"
	chmod 0777 "$HF_HOME" "$MODEL_DIR" \|\| true

	# Ensure HF cache logging directories are writable for xet
	install -d -m 0777 "$HF_HOME/xet/logs" \|\| true
	chmod 0777 "$HF_HOME/xet" "$HF_HOME/xet/logs" 2>/dev/null \|\| true
	chown -R 65532:65532 "$HF_HOME/xet" 2>/dev/null \|\| true

	# Run downloader as the non-root `appuser` using the venv python when available.
	echo "starting model downloader"
	if [ -x "/opt/venv/bin/python" ]; then
	su -s /bin/sh appuser -c "/opt/venv/bin/python -u /app/download_model.py"
	else
	su -s /bin/sh appuser -c "python3 -u /app/download_model.py"
	fi

	# If downloader wrote a repo marker with selected filename, update MODEL_PATH accordingly
	if [ -f "$MODEL_DIR/model.repo" ]; then
	marker=$(cat "$MODEL_DIR/model.repo" \|\| echo "")
	if [[ "$marker" == "\|" ]]; then
	IFS='\|' read -r _ selected_file _ <<< "$marker"
	if [ -n "$selected_file" ]; then
	MODEL_PATH="$MODEL_DIR/$selected_file"
	echo "using downloaded model path: $MODEL_PATH"
	fi
	fi
	fi

	LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-/usr/local/bin/llama-server}"

	# set server profile defaults
	case "$PERF_PROFILE" in
	low_latency)
	DEFAULT_BATCH=64
	DEFAULT_UBATCH=64
	DEFAULT_THREADS_BATCH=${THREADS}
	;;
	balanced)
	DEFAULT_BATCH=512
	DEFAULT_UBATCH=512
	DEFAULT_THREADS_BATCH=${THREADS}
	;;
	throughput)
	DEFAULT_BATCH=512
	DEFAULT_UBATCH=512
	DEFAULT_THREADS_BATCH=${THREADS}
	;;
	*)
	DEFAULT_BATCH=128
	DEFAULT_UBATCH=128
	DEFAULT_THREADS_BATCH=${THREADS}
	;;
	esac

	# Adjust defaults for low-memory environments (Spaces free tier: ~16GB).
	# Use cgroup limits first because /proc/meminfo can report the host machine.
	MEM_KB="$(detect_memory_kb)"
	echo "detected effective memory (KB): $MEM_KB"
	if [ "$MEM_KB" -gt 0 ] && [ "$MEM_KB" -le 18000000 ]; then
	echo "low-memory profile detected: using compact defaults"
	DEFAULT_CTX_SIZE=2048
	else
	DEFAULT_CTX_SIZE=4096
	fi

	CTX_SIZE="${REQUESTED_CTX_SIZE:-$DEFAULT_CTX_SIZE}"
	CACHE_TYPE_K="$(normalize_arg_value "${REQUESTED_CACHE_TYPE_K:-f16}")"
	CACHE_TYPE_V="$(normalize_arg_value "${REQUESTED_CACHE_TYPE_V:-f16}")"
	PARALLEL=1

	echo "effective threads: $THREADS (host=$HOST_THREADS, quota=$QUOTA_THREADS)"
	echo "effective llama profile: perf=$PERF_PROFILE single_request=1 ctx=$CTX_SIZE batch=${BATCH_SIZE:-$DEFAULT_BATCH} ubatch=${UBATCH_SIZE:-$DEFAULT_UBATCH} http_threads=$HTTP_THREADS log_verbosity=$LOG_VERBOSITY mmap=$MMAP flash_attn=$FLASH_ATTN"

	export CACHE_TYPE_K CACHE_TYPE_V CTX_SIZE PARALLEL

	server_args=(
	--model "$MODEL_PATH"
	--host 0.0.0.0
	--port "${PORT:-7860}"
	--api-key "$API_PASSWORD"
	-ngl 0
	-t "$THREADS"
	--threads-batch "${THREADS_BATCH:-$DEFAULT_THREADS_BATCH}"
	--threads-http "$HTTP_THREADS"
	--batch-size "${BATCH_SIZE:-$DEFAULT_BATCH}"
	--ubatch-size "${UBATCH_SIZE:-$DEFAULT_UBATCH}"
	--ctx-size "${CTX_SIZE:-4096}"
	--cache-type-k "${CACHE_TYPE_K}"
	--cache-type-v "${CACHE_TYPE_V}"
	--metrics
	-lv "$LOG_VERBOSITY"
	)


	if [ "$FLASH_ATTN" = "1" ]; then
	server_args+=(--flash-attn on)
	fi

	if [ "${NO_WARMUP}" = "1" ]; then
	server_args+=(--no-warmup)
	fi

	case "${REASONING,,}" in
	true\|1\|on\|yes)
	server_args+=(--reasoning on)
	;;
	false\|0\|off\|no)
	server_args+=(--reasoning off)
	;;
	auto\|"")
	server_args+=(--reasoning auto)
	;;
	*)
	echo "invalid REASONING value: $REASONING (expected True/False/auto)"
	exit 1
	;;
	esac

	if [ "$ENABLE_TOOLS" = "1" ] && [ -n "$LANGSEARCH_API_KEY" ] && [ -z "$TOOLS" ]; then
	TOOLS="exec_shell_command"
	fi

	if [[ -n "${TOOLS}" ]]; then
	server_args+=(--tools "$TOOLS")
	echo "enabled tools: $TOOLS"
	fi

	server_args+=(--parallel "$PARALLEL")

	# mmproj (vision projector) support
	MMPROJ_PATH=""
	if [ -n "${MMPROJ_FILE:-}" ]; then
	MMPROJ_PATH="$MODEL_DIR/$MMPROJ_FILE"
	elif [ -f "$MODEL_DIR/model.mmproj" ]; then
	mmproj_name=$(cat "$MODEL_DIR/model.mmproj" \|\| echo "")
	if [ -n "$mmproj_name" ]; then
	MMPROJ_PATH="$MODEL_DIR/$mmproj_name"
	fi
	fi

	if [ -n "$MMPROJ_PATH" ] && [ -f "$MMPROJ_PATH" ]; then
	server_args+=(--mmproj "$MMPROJ_PATH")
	echo "vision support enabled: $MMPROJ_PATH"
	else
	echo "no mmproj found; running text-only"
	fi

	server_cmd=("$LLAMA_SERVER_BIN" "${server_args[@]}")
	echo "server command: $LLAMA_SERVER_BIN ${server_args[*]}"
	echo "starting llama-server"
	exec su -s /bin/sh appuser -c "$(printf '%q ' "${server_cmd[@]}")"