def test_metrics(app, test_client):
metrics = activate_prometheus_exporter(app)
# test metrics for successful responses
success_labels = {"method": "GET", "status": "200"}
assert (metrics.registry.get_sample_value("mlflow_http_request_total",
labels=success_labels) is None)
resp = test_client.get("/")
assert resp.status_code == 200
assert (metrics.registry.get_sample_value("mlflow_http_request_total",
labels=success_labels) == 1)
# calling the metrics endpoint should not increment the counter
resp = test_client.get("/metrics")
assert resp.status_code == 200
assert (metrics.registry.get_sample_value("mlflow_http_request_total",
labels=success_labels) == 1)
# calling the health endpoint should not increment the counter
resp = test_client.get("/health")
assert resp.status_code == 200
assert (metrics.registry.get_sample_value("mlflow_http_request_total",
labels=success_labels) == 1)
# test metrics for failed responses
failure_labels = {"method": "GET", "status": "404"}
assert (metrics.registry.get_sample_value("mlflow_http_request_total",
labels=failure_labels) is None)
resp = test_client.get("/non-existent-endpoint")
assert resp.status_code == 404
assert (metrics.registry.get_sample_value("mlflow_http_request_total",
labels=failure_labels) == 1)
from flask_cors import CORS
if strtobool(os.getenv(CORS_ENV_VAR)):
print("Enabling CORS")
CORS(app)
for http_path, handler, methods in handlers.get_endpoints():
app.add_url_rule(http_path, handler.__name__, handler, methods=methods)
if os.getenv(PROMETHEUS_EXPORTER_ENV_VAR):
from mlflow.server.prometheus_exporter import activate_prometheus_exporter
prometheus_metrics_path = os.getenv(PROMETHEUS_EXPORTER_ENV_VAR)
if not os.path.exists(prometheus_metrics_path):
os.makedirs(prometheus_metrics_path)
activate_prometheus_exporter(app)
# Provide a health check endpoint to ensure the application is responsive
@app.route("/health")
def health():
return "OK", 200
# Serve the "get-artifact" route.
@app.route(_add_static_prefix("/get-artifact"))
def serve_artifacts():
return get_artifact_handler()
# Serve the "model-versions/get-artifact" route.