diff --git a/app.py b/app.py index 9d98f6f..581e7de 100644 --- a/app.py +++ b/app.py @@ -1,17 +1,35 @@ -# -*- coding: utf-8 -*- from fastapi import FastAPI, responses, status, Request from typing import Union import torch +import logging +import os from src.util import request_util from src.data import data_type import controller from submodules.model.business_objects import general -from submodules.model import session +from submodules.model import session, telemetry + + +OTLP_GRPC_ENDPOINT = os.getenv("OTLP_GRPC_ENDPOINT", "tempo:4317") + +app_name = "refinery-embedder" +app = FastAPI(title=app_name) + +if telemetry.ENABLE_TELEMETRY: + print("WARNING: Running telemetry.", flush=True) + telemetry.setting_app_name(app_name) + telemetry.setting_otlp(app, app_name=app_name, endpoint=OTLP_GRPC_ENDPOINT) + app.add_middleware(telemetry.PrometheusMiddleware, app_name=app_name) + app.add_route("/metrics", telemetry.metrics) + + # Filter out /metrics + logging.getLogger("uvicorn.access").addFilter( + lambda record: "GET /metrics" not in record.getMessage() + ) -app = FastAPI() if torch.cuda.is_available(): print( diff --git a/controller.py b/controller.py index 0da0101..22881fa 100644 --- a/controller.py +++ b/controller.py @@ -25,13 +25,13 @@ PrivatemodeAISentenceEmbedder, ) from src.embedders.classification.reduce import PCASentenceReducer # noqa: F401 -from src.util import daemon, request_util +from src.util import request_util from src.util.decorator import param_throttle from src.util.embedders import get_embedder from src.util.notification import send_project_update, embedding_warning_templates from submodules.s3 import controller as s3 -from submodules.model import enums +from submodules.model import enums, daemon from submodules.model.business_objects import ( attribute, embedding, @@ -120,7 +120,7 @@ def get_docbins( def manage_encoding_thread(project_id: str, embedding_id: str) -> int: - daemon.run(prepare_run, project_id, embedding_id) + daemon.run_without_db_token(prepare_run, project_id, embedding_id) return status.HTTP_200_OK diff --git a/src/util/config_handler.py b/src/util/config_handler.py index f9d1a84..024b3b5 100644 --- a/src/util/config_handler.py +++ b/src/util/config_handler.py @@ -1,7 +1,7 @@ from typing import Dict, Any, Optional, Union import requests import time -from src.util import daemon +from submodules.model import daemon __config = None @@ -25,7 +25,7 @@ def refresh_config(): ) global __config __config = response.json() - daemon.run(invalidate_after, 3600) # one hour + daemon.run_without_db_token(invalidate_after, 3600) # one hour def get_config_value( diff --git a/src/util/daemon.py b/src/util/daemon.py deleted file mode 100644 index d87da0e..0000000 --- a/src/util/daemon.py +++ /dev/null @@ -1,11 +0,0 @@ -import threading -from typing import Callable - - -def run(target: Callable, *args, **kwargs) -> None: - threading.Thread( - target=target, - args=args, - kwargs=kwargs, - daemon=True, - ).start() diff --git a/start b/start index e803302..c58dd63 100755 --- a/start +++ b/start @@ -2,11 +2,13 @@ DEBUG_MODE=false DEBUG_PORT=15675 +ENABLE_TELEMETRY=false -while getopts d flag +while getopts dg flag do case "${flag}" in d) DEBUG_MODE=true;; + g) ENABLE_TELEMETRY=true;; esac done @@ -68,6 +70,7 @@ docker run -d --rm \ -e MODEL_PROVIDER=http://refinery-model-provider:80 \ -e WS_NOTIFY_ENDPOINT="http://refinery-websocket:8080" \ -e NEURAL_SEARCH=http://refinery-neural-search:80 \ +-e ENABLE_TELEMETRY=$ENABLE_TELEMETRY \ --mount type=bind,source="$(pwd)"/,target=/app \ -v /var/run/docker.sock:/var/run/docker.sock \ -v "$MODEL_DIR":/models \ diff --git a/submodules/model b/submodules/model index 6691db0..fa52e17 160000 --- a/submodules/model +++ b/submodules/model @@ -1 +1 @@ -Subproject commit 6691db0df45a7ccd30c22a83c3fdc206fdc35e27 +Subproject commit fa52e1725d0691979895644d63c0b61728ea771b