Resolve trace app name from environment variable (#253)

* Add trace namespace

* Fix app name

* Use arms python trace

* remove setup tracing

* Update arms startup comamnd

* Add instrument

* Update main

* Update aliyun-bootstrap

.pre-commit-config.yaml

@@ -22,13 +22,14 @@ repos:
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
-        exclude: ".*poetry.lock|.*_static|.*__init__.py"
+        exclude: ".*poetry.lock|.*_static|.*__init__.py|.*main.py"
   - repo: https://github.com/psf/black-pre-commit-mirror
     rev: 23.10.1
     hooks:
       - id: black-jupyter
         name: black-src
         alias: black
+        exclude: ".*main.py"
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.0.1
     hooks:

Dockerfile

@@ -11,6 +11,7 @@ WORKDIR /app
 COPY . .
 
 RUN poetry install && rm -rf $POETRY_CACHE_DIR
+RUN aliyun-bootstrap -a install
 
 FROM python:3.11-slim AS prod
 
@@ -20,8 +21,7 @@ ENV VIRTUAL_ENV=/app/.venv \
     PATH="/app/.venv/bin:$PATH"
 
 RUN apt-get update && apt-get install -y libgl1 libglib2.0-0
-
 WORKDIR /app
 COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 COPY . .
-CMD ["pai_rag", "run"]
+CMD ["pai_rag", "serve"]

Dockerfile_gpu

@@ -13,6 +13,7 @@ RUN mv pyproject_gpu.toml pyproject.toml \
     && rm poetry.lock
 
 RUN poetry install && rm -rf $POETRY_CACHE_DIR
+RUN aliyun-bootstrap -a install
 
 FROM python:3.11-slim AS prod
 
@@ -26,4 +27,4 @@ RUN apt-get update && apt-get install -y libgl1 libglib2.0-0
 WORKDIR /app
 COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 COPY . .
-CMD ["pai_rag", "run"]
+CMD ["pai_rag", "serve"]

pyproject.toml

@@ -98,14 +98,14 @@ detectron2 = [
   {markers = "sys_platform != 'win32' and sys_platform != 'linux' ", url = "https://pai-rag.oss-cn-hangzhou.aliyuncs.com/packages/python_wheels/detectron2-0.6%2B864913fpt2.2.2cpu-cp311-cp311-macosx_10_9_universal2.whl"}
 ]
 magic-pdf = {version = "0.7.0b1", extras = ["full"]}
-pai-llm-trace = {url = "https://pai-llm-trace.oss-cn-zhangjiakou.aliyuncs.com/sdk_package/pai_llm_trace-1.0.0-py3-none-any.whl"}
 llama-index-callbacks-arize-phoenix = "0.1.6"
 peft = "^0.12.0"
 duckduckgo-search = "6.2.12"
+aliyun-bootstrap = "^1.0.1"
 docx = "^0.2.4"
 
 [tool.poetry.scripts]
-pai_rag = "pai_rag.main:main"
+pai_rag = "pai_rag.main:run"
 load_data = "pai_rag.tool.load_data_tool:run"
 load_model = "pai_rag.utils.download_models:load_models"
 run_eval_exp = "pai_rag.evaluation.run_evaluation_experiments:run"

pyproject_gpu.toml

@@ -92,12 +92,12 @@ detectron2 = [
   {markers = "sys_platform != 'win32' and sys_platform != 'linux' ", url = "https://pai-rag.oss-cn-hangzhou.aliyuncs.com/packages/python_wheels/detectron2-0.6%2B864913fpt2.2.2cpu-cp311-cp311-macosx_10_9_universal2.whl"}
 ]
 magic-pdf = {version = "0.7.0b1", extras = ["full"]}
-pai-llm-trace = {url = "https://pai-llm-trace.oss-cn-zhangjiakou.aliyuncs.com/sdk_package/pai_llm_trace-1.0.0-py3-none-any.whl"}
 llama-index-callbacks-arize-phoenix = "0.1.6"
 peft = "^0.12.0"
+aliyun-bootstrap = "^1.0.1"
 
 [tool.poetry.scripts]
-pai_rag = "pai_rag.main:main"
+pai_rag = "pai_rag.main:run"
 load_data = "pai_rag.tool.load_data_tool:run"
 load_model = "pai_rag.utils.download_models:load_models"
 run_eval_exp = "pai_rag.evaluation.run_evaluation_experiments:run"

src/pai_rag/core/models/config.py

@@ -70,7 +70,7 @@ class PaiTraceConfig(BaseTraceConfig):
     type: Literal[TraceType.pai_trace] = TraceType.pai_trace
     endpoint: str | None = None
     token: str | None = None
-    app_name: str = "PAI-RAG"
+    app_name: str = "PAIRAG-Service"
 
 
 class ArizeTraceConfig(BaseTraceConfig):

src/pai_rag/core/rag_module.py

@@ -1,7 +1,6 @@
 from typing import Any
 from llama_index.core import Settings
 from llama_index.core.prompts import PromptTemplate
-from pai_rag.core.models.config import ArizeTraceConfig, BaseTraceConfig, PaiTraceConfig
 from pai_rag.core.rag_config import RagConfig
 from pai_rag.core.rag_data_loader import RagDataLoader
 from pai_rag.integrations.agent.pai.pai_agent import PaiAgent
@@ -225,32 +224,3 @@ def resolve_searcher(config: RagConfig) -> BingSearchTool:
     )
 
     return searcher
-
-
-def setup_tracing(trace_config: BaseTraceConfig):
-    from pai.llm_trace.instrumentation import init_opentelemetry
-    from pai.llm_trace.instrumentation.llama_index import LlamaIndexInstrumentor
-    from llama_index.core import set_global_handler
-
-    if isinstance(trace_config, PaiTraceConfig):
-        if not trace_config.token:
-            logger.info("Trace is not enabled since PaiTraceConfig.token is not set.")
-            return
-
-        init_opentelemetry(
-            LlamaIndexInstrumentor,
-            grpc_endpoint=trace_config.endpoint,
-            token=trace_config.token,
-            service_name=trace_config.app_name,
-            service_version="0.1.0",
-            service_id="",
-            deployment_environment="",
-            service_owner_id="",
-            service_owner_sub_id="",
-        )
-        logger.info(f"Pai-LLM-Trace enabled with endpoint: '{trace_config.endpoint}'.")
-    elif isinstance(trace_config, ArizeTraceConfig):
-        set_global_handler("arize_phoenix")
-        logger.info("Arize trace enabled.")
-    else:
-        logger.warning("Trace is not configured.")

src/pai_rag/main.py

@@ -1,3 +1,5 @@
+from aliyun.opentelemetry.instrumentation.auto_instrumentation import sitecustomize # ByPass unused import: use for arms trace initialization
+
 import click
 import uvicorn
 from fastapi import FastAPI
@@ -68,14 +70,14 @@ def init_log():
 @click.group(invoke_without_command=True)
 @click.pass_context
 @click.option("-V", "--version", is_flag=True, help="Show version and exit.")
-def main(ctx, version):
+def run(ctx, version):
     if version:
         click.echo(version)
     elif ctx.invoked_subcommand is None:
         click.echo(ctx.get_help())
 
 
-@main.command()
+@run.command()
 @click.option(
     "-h",
     "--host",
@@ -105,7 +107,7 @@ def ui(host, port, rag_url):
     uvicorn.run(app, host=host, port=port, loop="asyncio")
 
 
-@main.command()
+@run.command()
 @click.option(
     "-h",
     "--host",
@@ -157,12 +159,10 @@ def ui(host, port, rag_url):
 )
 def serve(host, port, config_file, workers, enable_example, skip_download_models):
     from pai_rag.app.api.service import configure_app
-    from pai_rag.core.rag_module import setup_tracing
     from pai_rag.utils.download_models import ModelScopeDownloader
 
     rag_configuration = RagConfigManager.from_file(config_file)
     rag_configuration.persist()
-    setup_tracing(rag_configuration.get_value().trace)
 
     if not skip_download_models and DEFAULT_MODEL_DIR != EAS_DEFAULT_MODEL_DIR:
         logger.info("Start to download models.")
@@ -178,3 +178,7 @@ def serve(host, port, config_file, workers, enable_example, skip_download_models
     app = FastAPI()
     configure_app(app, rag_configuration)
     uvicorn.run(app=app, host=host, port=port, loop="asyncio", workers=workers)
+
+
+if __name__ == "__main__":
+    serve()

src/pai_rag/tools/query_tool.py

@@ -2,7 +2,7 @@
 import os
 from pathlib import Path
 from pai_rag.core.rag_config_manager import RagConfigManager
-from pai_rag.core.rag_module import resolve_query_engine, setup_tracing
+from pai_rag.core.rag_module import resolve_query_engine
 from pai_rag.integrations.synthesizer.pai_synthesizer import PaiQueryBundle
 import logging
 
@@ -44,7 +44,6 @@ def run(
     logging.basicConfig(level=logging.INFO)
 
     config = RagConfigManager.from_file(config_file).get_value()
-    setup_tracing(config.trace)
 
     query_engine = resolve_query_engine(config)