diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 30d3731952a..5c81f67ff83 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -163,6 +163,7 @@ ddtrace/contrib/internal/langgraph @DataDog/ml-observ ddtrace/contrib/internal/crewai @DataDog/ml-observability ddtrace/contrib/internal/openai_agents @DataDog/ml-observability ddtrace/contrib/internal/litellm @DataDog/ml-observability +ddtrace/contrib/internal/mlflow @DataDog/ml-observability ddtrace/contrib/internal/pydantic_ai @DataDog/ml-observability ddtrace/contrib/internal/ray @DataDog/ml-observability ddtrace/contrib/internal/mcp @DataDog/ml-observability @@ -186,6 +187,7 @@ tests/contrib/langgraph @DataDog/ml-observ tests/contrib/crewai @DataDog/ml-observability tests/contrib/openai_agents @DataDog/ml-observability tests/contrib/litellm @DataDog/ml-observability +tests/contrib/mlflow @DataDog/ml-observability tests/contrib/pydantic_ai @DataDog/ml-observability tests/contrib/ray @DataDog/ml-observability tests/contrib/mcp @DataDog/ml-observability diff --git a/.riot/requirements/1b4a60e.txt b/.riot/requirements/1b4a60e.txt new file mode 100644 index 00000000000..4029001701b --- /dev/null +++ b/.riot/requirements/1b4a60e.txt @@ -0,0 +1,100 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --allow-unsafe --no-annotate .riot/requirements/1b4a60e.in +# +alembic==1.18.4 +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +attrs==25.4.0 +blinker==1.9.0 +cachetools==6.2.6 +certifi==2026.2.25 +cffi==2.0.0 +charset-normalizer==3.4.4 +click==8.3.1 +cloudpickle==3.1.2 +contourpy==1.3.3 +coverage[toml]==7.13.4 +cryptography==46.0.5 +cycler==0.12.1 +databricks-sdk==0.96.0 +docker==7.1.0 +fastapi==0.135.1 +flask==3.1.3 +flask-cors==6.0.2 +fonttools==4.61.1 +gitdb==4.0.12 +gitpython==3.1.46 +google-auth==2.48.0 +graphene==3.4.3 +graphql-core==3.2.7 +graphql-relay==3.2.0 +greenlet==3.3.2 +gunicorn==23.0.0 +h11==0.16.0 +huey==2.6.0 +hypothesis==6.45.0 +idna==3.11 +importlib-metadata==8.7.1 +iniconfig==2.3.0 +itsdangerous==2.2.0 +jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.4.9 +mako==1.3.10 +markupsafe==3.0.3 +matplotlib==3.10.8 +mlflow[default]==3.9.0 +mlflow-skinny==3.9.0 +mlflow-tracing==3.9.0 +mock==5.2.0 +numpy==2.4.2 +opentelemetry-api==1.39.1 +opentelemetry-proto==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-semantic-conventions==0.60b1 +opentracing==2.4.0 +packaging==25.0 +pandas==2.3.3 +pillow==12.1.1 +pluggy==1.6.0 +prettytable==3.17.0 +protobuf==6.33.5 +pyarrow==22.0.0 +pyasn1==0.6.2 +pyasn1-modules==0.4.2 +pycparser==3.0 +pydantic==2.12.5 +pydantic-core==2.41.5 +pygments==2.19.2 +pyparsing==3.3.2 +pytest==9.0.2 +pytest-cov==7.0.0 +pytest-mock==3.15.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +pytz==2026.1.post1 +pyyaml==6.0.3 +requests==2.32.5 +rsa==4.9.1 +scikit-learn==1.8.0 +scipy==1.17.1 +six==1.17.0 +skops==0.13.0 +smmap==5.0.2 +sortedcontainers==2.4.0 +sqlalchemy==2.0.48 +sqlparse==0.5.5 +starlette==0.52.1 +threadpoolctl==3.6.0 +typing-extensions==4.15.0 +typing-inspection==0.4.2 +tzdata==2025.3 +urllib3==2.6.3 +uvicorn==0.41.0 +wcwidth==0.6.0 +werkzeug==3.1.6 +zipp==3.23.0 diff --git a/.riot/requirements/1ff5d35.txt b/.riot/requirements/1ff5d35.txt new file mode 100644 index 00000000000..4d7969a590b --- /dev/null +++ b/.riot/requirements/1ff5d35.txt @@ -0,0 +1,100 @@ +# +# This file is autogenerated by pip-compile with Python 3.12 +# by the following command: +# +# pip-compile --allow-unsafe --no-annotate .riot/requirements/1ff5d35.in +# +alembic==1.18.4 +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +attrs==25.4.0 +blinker==1.9.0 +cachetools==6.2.6 +certifi==2026.2.25 +cffi==2.0.0 +charset-normalizer==3.4.4 +click==8.3.1 +cloudpickle==3.1.2 +contourpy==1.3.3 +coverage[toml]==7.13.4 +cryptography==46.0.5 +cycler==0.12.1 +databricks-sdk==0.96.0 +docker==7.1.0 +fastapi==0.135.1 +flask==3.1.3 +flask-cors==6.0.2 +fonttools==4.61.1 +gitdb==4.0.12 +gitpython==3.1.46 +google-auth==2.48.0 +graphene==3.4.3 +graphql-core==3.2.7 +graphql-relay==3.2.0 +greenlet==3.3.2 +gunicorn==23.0.0 +h11==0.16.0 +huey==2.6.0 +hypothesis==6.45.0 +idna==3.11 +importlib-metadata==8.7.1 +iniconfig==2.3.0 +itsdangerous==2.2.0 +jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.4.9 +mako==1.3.10 +markupsafe==3.0.3 +matplotlib==3.10.8 +mlflow[default]==3.9.0 +mlflow-skinny==3.9.0 +mlflow-tracing==3.9.0 +mock==5.2.0 +numpy==2.4.2 +opentelemetry-api==1.39.1 +opentelemetry-proto==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-semantic-conventions==0.60b1 +opentracing==2.4.0 +packaging==25.0 +pandas==2.3.3 +pillow==12.1.1 +pluggy==1.6.0 +prettytable==3.17.0 +protobuf==6.33.5 +pyarrow==22.0.0 +pyasn1==0.6.2 +pyasn1-modules==0.4.2 +pycparser==3.0 +pydantic==2.12.5 +pydantic-core==2.41.5 +pygments==2.19.2 +pyparsing==3.3.2 +pytest==9.0.2 +pytest-cov==7.0.0 +pytest-mock==3.15.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +pytz==2026.1.post1 +pyyaml==6.0.3 +requests==2.32.5 +rsa==4.9.1 +scikit-learn==1.8.0 +scipy==1.17.1 +six==1.17.0 +skops==0.13.0 +smmap==5.0.2 +sortedcontainers==2.4.0 +sqlalchemy==2.0.48 +sqlparse==0.5.5 +starlette==0.52.1 +threadpoolctl==3.6.0 +typing-extensions==4.15.0 +typing-inspection==0.4.2 +tzdata==2025.3 +urllib3==2.6.3 +uvicorn==0.41.0 +wcwidth==0.6.0 +werkzeug==3.1.6 +zipp==3.23.0 diff --git a/.riot/requirements/3b16c7f.txt b/.riot/requirements/3b16c7f.txt new file mode 100644 index 00000000000..d9e2ca0dfb7 --- /dev/null +++ b/.riot/requirements/3b16c7f.txt @@ -0,0 +1,100 @@ +# +# This file is autogenerated by pip-compile with Python 3.13 +# by the following command: +# +# pip-compile --allow-unsafe --no-annotate .riot/requirements/3b16c7f.in +# +alembic==1.18.4 +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +attrs==25.4.0 +blinker==1.9.0 +cachetools==7.0.2 +certifi==2026.2.25 +cffi==2.0.0 +charset-normalizer==3.4.4 +click==8.3.1 +cloudpickle==3.1.2 +contourpy==1.3.3 +coverage[toml]==7.13.4 +cryptography==46.0.5 +cycler==0.12.1 +databricks-sdk==0.96.0 +docker==7.1.0 +fastapi==0.135.1 +flask==3.1.3 +flask-cors==6.0.2 +fonttools==4.61.1 +gitdb==4.0.12 +gitpython==3.1.46 +google-auth==2.48.0 +graphene==3.4.3 +graphql-core==3.2.7 +graphql-relay==3.2.0 +greenlet==3.3.2 +gunicorn==25.1.0 +h11==0.16.0 +huey==2.6.0 +hypothesis==6.45.0 +idna==3.11 +importlib-metadata==8.7.1 +iniconfig==2.3.0 +itsdangerous==2.2.0 +jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.4.9 +mako==1.3.10 +markupsafe==3.0.3 +matplotlib==3.10.8 +mlflow[default]==3.10.0 +mlflow-skinny==3.10.0 +mlflow-tracing==3.10.0 +mock==5.2.0 +numpy==2.4.2 +opentelemetry-api==1.39.1 +opentelemetry-proto==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-semantic-conventions==0.60b1 +opentracing==2.4.0 +packaging==26.0 +pandas==2.3.3 +pillow==12.1.1 +pluggy==1.6.0 +prettytable==3.17.0 +protobuf==6.33.5 +pyarrow==23.0.1 +pyasn1==0.6.2 +pyasn1-modules==0.4.2 +pycparser==3.0 +pydantic==2.12.5 +pydantic-core==2.41.5 +pygments==2.19.2 +pyparsing==3.3.2 +pytest==9.0.2 +pytest-cov==7.0.0 +pytest-mock==3.15.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +pytz==2026.1.post1 +pyyaml==6.0.3 +requests==2.32.5 +rsa==4.9.1 +scikit-learn==1.8.0 +scipy==1.17.1 +six==1.17.0 +skops==0.13.0 +smmap==5.0.2 +sortedcontainers==2.4.0 +sqlalchemy==2.0.48 +sqlparse==0.5.5 +starlette==0.52.1 +threadpoolctl==3.6.0 +typing-extensions==4.15.0 +typing-inspection==0.4.2 +tzdata==2025.3 +urllib3==2.6.3 +uvicorn==0.41.0 +wcwidth==0.6.0 +werkzeug==3.1.6 +zipp==3.23.0 diff --git a/.riot/requirements/48e6e52.txt b/.riot/requirements/48e6e52.txt new file mode 100644 index 00000000000..ae96292fd97 --- /dev/null +++ b/.riot/requirements/48e6e52.txt @@ -0,0 +1,100 @@ +# +# This file is autogenerated by pip-compile with Python 3.13 +# by the following command: +# +# pip-compile --allow-unsafe --no-annotate .riot/requirements/48e6e52.in +# +alembic==1.18.4 +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +attrs==25.4.0 +blinker==1.9.0 +cachetools==6.2.6 +certifi==2026.2.25 +cffi==2.0.0 +charset-normalizer==3.4.4 +click==8.3.1 +cloudpickle==3.1.2 +contourpy==1.3.3 +coverage[toml]==7.13.4 +cryptography==46.0.5 +cycler==0.12.1 +databricks-sdk==0.96.0 +docker==7.1.0 +fastapi==0.135.1 +flask==3.1.3 +flask-cors==6.0.2 +fonttools==4.61.1 +gitdb==4.0.12 +gitpython==3.1.46 +google-auth==2.48.0 +graphene==3.4.3 +graphql-core==3.2.7 +graphql-relay==3.2.0 +greenlet==3.3.2 +gunicorn==23.0.0 +h11==0.16.0 +huey==2.6.0 +hypothesis==6.45.0 +idna==3.11 +importlib-metadata==8.7.1 +iniconfig==2.3.0 +itsdangerous==2.2.0 +jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.4.9 +mako==1.3.10 +markupsafe==3.0.3 +matplotlib==3.10.8 +mlflow[default]==3.9.0 +mlflow-skinny==3.9.0 +mlflow-tracing==3.9.0 +mock==5.2.0 +numpy==2.4.2 +opentelemetry-api==1.39.1 +opentelemetry-proto==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-semantic-conventions==0.60b1 +opentracing==2.4.0 +packaging==25.0 +pandas==2.3.3 +pillow==12.1.1 +pluggy==1.6.0 +prettytable==3.17.0 +protobuf==6.33.5 +pyarrow==22.0.0 +pyasn1==0.6.2 +pyasn1-modules==0.4.2 +pycparser==3.0 +pydantic==2.12.5 +pydantic-core==2.41.5 +pygments==2.19.2 +pyparsing==3.3.2 +pytest==9.0.2 +pytest-cov==7.0.0 +pytest-mock==3.15.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +pytz==2026.1.post1 +pyyaml==6.0.3 +requests==2.32.5 +rsa==4.9.1 +scikit-learn==1.8.0 +scipy==1.17.1 +six==1.17.0 +skops==0.13.0 +smmap==5.0.2 +sortedcontainers==2.4.0 +sqlalchemy==2.0.48 +sqlparse==0.5.5 +starlette==0.52.1 +threadpoolctl==3.6.0 +typing-extensions==4.15.0 +typing-inspection==0.4.2 +tzdata==2025.3 +urllib3==2.6.3 +uvicorn==0.41.0 +wcwidth==0.6.0 +werkzeug==3.1.6 +zipp==3.23.0 diff --git a/.riot/requirements/4cbd638.txt b/.riot/requirements/4cbd638.txt new file mode 100644 index 00000000000..dd074ae6e1a --- /dev/null +++ b/.riot/requirements/4cbd638.txt @@ -0,0 +1,100 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --allow-unsafe --no-annotate .riot/requirements/4cbd638.in +# +alembic==1.18.4 +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +attrs==25.4.0 +blinker==1.9.0 +cachetools==7.0.2 +certifi==2026.2.25 +cffi==2.0.0 +charset-normalizer==3.4.4 +click==8.3.1 +cloudpickle==3.1.2 +contourpy==1.3.3 +coverage[toml]==7.13.4 +cryptography==46.0.5 +cycler==0.12.1 +databricks-sdk==0.96.0 +docker==7.1.0 +fastapi==0.135.1 +flask==3.1.3 +flask-cors==6.0.2 +fonttools==4.61.1 +gitdb==4.0.12 +gitpython==3.1.46 +google-auth==2.48.0 +graphene==3.4.3 +graphql-core==3.2.7 +graphql-relay==3.2.0 +greenlet==3.3.2 +gunicorn==25.1.0 +h11==0.16.0 +huey==2.6.0 +hypothesis==6.45.0 +idna==3.11 +importlib-metadata==8.7.1 +iniconfig==2.3.0 +itsdangerous==2.2.0 +jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.4.9 +mako==1.3.10 +markupsafe==3.0.3 +matplotlib==3.10.8 +mlflow[default]==3.10.0 +mlflow-skinny==3.10.0 +mlflow-tracing==3.10.0 +mock==5.2.0 +numpy==2.4.2 +opentelemetry-api==1.39.1 +opentelemetry-proto==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-semantic-conventions==0.60b1 +opentracing==2.4.0 +packaging==26.0 +pandas==2.3.3 +pillow==12.1.1 +pluggy==1.6.0 +prettytable==3.17.0 +protobuf==6.33.5 +pyarrow==23.0.1 +pyasn1==0.6.2 +pyasn1-modules==0.4.2 +pycparser==3.0 +pydantic==2.12.5 +pydantic-core==2.41.5 +pygments==2.19.2 +pyparsing==3.3.2 +pytest==9.0.2 +pytest-cov==7.0.0 +pytest-mock==3.15.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +pytz==2026.1.post1 +pyyaml==6.0.3 +requests==2.32.5 +rsa==4.9.1 +scikit-learn==1.8.0 +scipy==1.17.1 +six==1.17.0 +skops==0.13.0 +smmap==5.0.2 +sortedcontainers==2.4.0 +sqlalchemy==2.0.48 +sqlparse==0.5.5 +starlette==0.52.1 +threadpoolctl==3.6.0 +typing-extensions==4.15.0 +typing-inspection==0.4.2 +tzdata==2025.3 +urllib3==2.6.3 +uvicorn==0.41.0 +wcwidth==0.6.0 +werkzeug==3.1.6 +zipp==3.23.0 diff --git a/.riot/requirements/53a57db.txt b/.riot/requirements/53a57db.txt new file mode 100644 index 00000000000..8ae3e0cad9c --- /dev/null +++ b/.riot/requirements/53a57db.txt @@ -0,0 +1,100 @@ +# +# This file is autogenerated by pip-compile with Python 3.12 +# by the following command: +# +# pip-compile --allow-unsafe --no-annotate .riot/requirements/53a57db.in +# +alembic==1.18.4 +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +attrs==25.4.0 +blinker==1.9.0 +cachetools==7.0.2 +certifi==2026.2.25 +cffi==2.0.0 +charset-normalizer==3.4.4 +click==8.3.1 +cloudpickle==3.1.2 +contourpy==1.3.3 +coverage[toml]==7.13.4 +cryptography==46.0.5 +cycler==0.12.1 +databricks-sdk==0.96.0 +docker==7.1.0 +fastapi==0.135.1 +flask==3.1.3 +flask-cors==6.0.2 +fonttools==4.61.1 +gitdb==4.0.12 +gitpython==3.1.46 +google-auth==2.48.0 +graphene==3.4.3 +graphql-core==3.2.7 +graphql-relay==3.2.0 +greenlet==3.3.2 +gunicorn==25.1.0 +h11==0.16.0 +huey==2.6.0 +hypothesis==6.45.0 +idna==3.11 +importlib-metadata==8.7.1 +iniconfig==2.3.0 +itsdangerous==2.2.0 +jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.4.9 +mako==1.3.10 +markupsafe==3.0.3 +matplotlib==3.10.8 +mlflow[default]==3.10.0 +mlflow-skinny==3.10.0 +mlflow-tracing==3.10.0 +mock==5.2.0 +numpy==2.4.2 +opentelemetry-api==1.39.1 +opentelemetry-proto==1.39.1 +opentelemetry-sdk==1.39.1 +opentelemetry-semantic-conventions==0.60b1 +opentracing==2.4.0 +packaging==26.0 +pandas==2.3.3 +pillow==12.1.1 +pluggy==1.6.0 +prettytable==3.17.0 +protobuf==6.33.5 +pyarrow==23.0.1 +pyasn1==0.6.2 +pyasn1-modules==0.4.2 +pycparser==3.0 +pydantic==2.12.5 +pydantic-core==2.41.5 +pygments==2.19.2 +pyparsing==3.3.2 +pytest==9.0.2 +pytest-cov==7.0.0 +pytest-mock==3.15.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +pytz==2026.1.post1 +pyyaml==6.0.3 +requests==2.32.5 +rsa==4.9.1 +scikit-learn==1.8.0 +scipy==1.17.1 +six==1.17.0 +skops==0.13.0 +smmap==5.0.2 +sortedcontainers==2.4.0 +sqlalchemy==2.0.48 +sqlparse==0.5.5 +starlette==0.52.1 +threadpoolctl==3.6.0 +typing-extensions==4.15.0 +typing-inspection==0.4.2 +tzdata==2025.3 +urllib3==2.6.3 +uvicorn==0.41.0 +wcwidth==0.6.0 +werkzeug==3.1.6 +zipp==3.23.0 diff --git a/ddtrace/contrib/integration_registry/registry.yaml b/ddtrace/contrib/integration_registry/registry.yaml index 7b48919065d..80d42cbef3c 100644 --- a/ddtrace/contrib/integration_registry/registry.yaml +++ b/ddtrace/contrib/integration_registry/registry.yaml @@ -583,6 +583,16 @@ integrations: min: 1.10.1 max: 1.16.0 +- integration_name: mlflow + is_external_package: true + is_tested: false + dependency_names: + - mlflow + tested_versions_by_dependency: + mlflow: + min: 3.9.0 + max: 3.9.0 + - integration_name: molten is_external_package: true is_tested: true diff --git a/ddtrace/contrib/internal/mlflow/README.md b/ddtrace/contrib/internal/mlflow/README.md new file mode 100644 index 00000000000..b62777c7314 --- /dev/null +++ b/ddtrace/contrib/internal/mlflow/README.md @@ -0,0 +1,19 @@ +## Using Datadog's Authentication Plugin for MLFlow + +1. Set the `DD_API_KEY` and `DD_APP_KEY` environment variables: + + ```bash + export DD_API_KEY="your_api_key_here" + export DD_APP_KEY="your_app_key_here" + export DD_MODEL_LAB_ENABLED="true" + ``` + +2. You will now be able to submit authenticated requests to Datadog's MLFlow tracking server: + + ```python + import mlflow + + mlflow.set_tracking_uri("https://app.dataddoghq.com/api/v1/mlflow-tracking-server"") + ``` + + MLFlow will automatically recognize and use the plugin as long as the ddtrace Python library is installed. diff --git a/ddtrace/contrib/internal/mlflow/__init__.py b/ddtrace/contrib/internal/mlflow/__init__.py new file mode 100644 index 00000000000..867b5102073 --- /dev/null +++ b/ddtrace/contrib/internal/mlflow/__init__.py @@ -0,0 +1,5 @@ +""" +If the environment variables ``DD_API_KEY``, ``DD_APP_KEY`` and ``DD_MODEL_LAB_ENABLED`` +are set, this MLFlow authentication plugin will make MLFlow include the DD-API-KEY +and DD-APPLICATION-KEY HTTP headers in all requests to the MLFlow tracking server. +""" diff --git a/ddtrace/contrib/internal/mlflow/auth_plugin.py b/ddtrace/contrib/internal/mlflow/auth_plugin.py new file mode 100644 index 00000000000..2c9f730d26c --- /dev/null +++ b/ddtrace/contrib/internal/mlflow/auth_plugin.py @@ -0,0 +1,44 @@ +""" +MLFlow authentication plugin for Datadog API key injection. +""" + +from ddtrace.internal.settings._config import config + + +try: + from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider +except ImportError: + # If MLFlow is not installed, create a dummy base class so the module can still be imported + class RequestHeaderProvider: # type: ignore[no-redef] + pass + + +class DatadogHeaderProvider(RequestHeaderProvider): + """ + MLFlow request header provider that injects Datadog API keys. + + This plugin reads the DD_API_KEY and DD_APP_KEY environment variables and injects them + as 'DD-API-KEY' and 'DD-APPLICATION-KEY' headers into all MLFlow HTTP requests, + as long as the DD_MODEL_LAB_ENABLED environment variable is set. + """ + + def in_context(self): + """ + Check if the Datadog API and application keys are configured. + + Returns: + bool: True if DD_API_KEY, DD_APP_KEY and DD_MODEL_LAB_ENABLED env variables are set, else False + """ + return config._dd_api_key and config._dd_app_key and config._model_lab_enabled + + def request_headers(self): + """ + Generate request headers with the Datadog API and application keys. + + Returns: + dict: Dictionary containing the DD-API-KEY and DD-APPLICATION-KEY headers + if DD_API_KEY, DD_APP_KEY and DD_MODEL_LAB_ENABLED are set, empty dictionary otherwise. + """ + if config._dd_api_key and config._dd_app_key and config._model_lab_enabled: + return {"DD-API-KEY": config._dd_api_key, "DD-APPLICATION-KEY": config._dd_app_key} + return {} diff --git a/ddtrace/internal/settings/_config.py b/ddtrace/internal/settings/_config.py index 7f87da906b4..c89e014070f 100644 --- a/ddtrace/internal/settings/_config.py +++ b/ddtrace/internal/settings/_config.py @@ -667,6 +667,8 @@ def __init__(self) -> None: "DD_LLMOBS_INSTRUMENTED_PROXY_URLS", None, lambda x: set(x.strip().split(",")) ) + self._model_lab_enabled = _get_config("DD_MODEL_LAB_ENABLED", False, asbool) + self._inject_force = _get_config("DD_INJECT_FORCE", None, asbool) # Telemetry for whether ssi instrumented an app is tracked by the `instrumentation_source` config self._lib_was_injected = _get_config("_DD_PY_SSI_INJECT", False, asbool, report_telemetry=False) diff --git a/pyproject.toml b/pyproject.toml index ac3a78e4c11..8ca3fe660ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,9 @@ ddcontextvars_context = "ddtrace.internal.opentelemetry.context:DDRuntimeContext [project.entry-points.pytest11] ddtrace = "ddtrace.testing.internal.pytest.entry_point" +[project.entry-points.'mlflow.request_header_provider'] +datadog = "ddtrace.contrib.internal.mlflow.auth_plugin:DatadogHeaderProvider" + [project.entry-points.'ddtrace.products'] "apm-tracing-rc" = "ddtrace.internal.remoteconfig.products.apm_tracing" "code-origin-for-spans" = "ddtrace.debugging._products.code_origin.span" diff --git a/releasenotes/notes/add-auth-plugin-for-mlflow-4aa9aaa6e783078e.yaml b/releasenotes/notes/add-auth-plugin-for-mlflow-4aa9aaa6e783078e.yaml new file mode 100644 index 00000000000..6233ed95cf5 --- /dev/null +++ b/releasenotes/notes/add-auth-plugin-for-mlflow-4aa9aaa6e783078e.yaml @@ -0,0 +1,3 @@ +features: + - | + mlflow: Adds a request header provider (auth plugin) for MLFlow. If the environment variables DD_API_KEY, DD_APP_KEY and DD_MODEL_LAB_ENABLED are set, HTTP requests to the MLFlow tracking server will include the DD-API-KEY and DD-APPLICATION-KEY headers. #16685 diff --git a/riotfile.py b/riotfile.py index 07a9787856a..01c5f06f681 100644 --- a/riotfile.py +++ b/riotfile.py @@ -3136,6 +3136,14 @@ def select_pys(min_version: str = MIN_PYTHON_VERSION, max_version: str = MAX_PYT "ray[default]": ["~=2.46.0", latest], }, ), + Venv( + name="mlflow", + command="pytest {cmdargs} tests/contrib/mlflow", + pys=select_pys(min_version="3.11", max_version="3.13"), + pkgs={ + "mlflow[default]": ["~=3.9.0", latest], + }, + ), Venv( name="logbook", pys=select_pys(), diff --git a/tests/contrib/mlflow/__init__.py b/tests/contrib/mlflow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/contrib/mlflow/test_auth_plugin.py b/tests/contrib/mlflow/test_auth_plugin.py new file mode 100644 index 00000000000..74368e25b8e --- /dev/null +++ b/tests/contrib/mlflow/test_auth_plugin.py @@ -0,0 +1,19 @@ +from ddtrace.contrib.internal.mlflow.auth_plugin import DatadogHeaderProvider +from tests.utils import override_global_config + + +class TestDatadogHeaderProvider: + def test_auth_plugin(self): + test_api_key = "test_api_key_12345" + test_app_key = "test_app_key_12345" + with override_global_config( + dict( + _dd_api_key=test_api_key, + _dd_app_key=test_app_key, + _model_lab_enabled=True, + ) + ): + provider = DatadogHeaderProvider() + assert provider.in_context() + headers = provider.request_headers() + assert headers == {"DD-API-KEY": test_api_key, "DD-APPLICATION-KEY": test_app_key} diff --git a/tests/contrib/suitespec.yml b/tests/contrib/suitespec.yml index 14243505876..027711ba7af 100644 --- a/tests/contrib/suitespec.yml +++ b/tests/contrib/suitespec.yml @@ -116,6 +116,8 @@ components: - ddtrace/contrib/internal/mako/* mariadb: - ddtrace/contrib/internal/mariadb/* + mlflow: + - ddtrace/contrib/internal/mlflow/* molten: - ddtrace/contrib/internal/molten/* mongo: @@ -884,6 +886,17 @@ suites: - mariadb snapshot: true venvs_per_job: 2 + mlflow: + parallelism: 1 + paths: + - '@bootstrap' + - '@core' + - '@contrib' + - '@tracing' + - '@mlflow' + - tests/contrib/mlflow/* + runner: riot + snapshot: true molten: parallelism: 1 paths: diff --git a/tests/telemetry/test_writer.py b/tests/telemetry/test_writer.py index b57083bf342..087cd916707 100644 --- a/tests/telemetry/test_writer.py +++ b/tests/telemetry/test_writer.py @@ -286,6 +286,7 @@ def test_app_started_event_configuration_override(test_agent_session, run_python {"name": "DD_LOGS_INJECTION", "origin": "env_var", "value": True}, {"name": "DD_LOGS_OTEL_ENABLED", "origin": "env_var", "value": True}, {"name": "DD_METRICS_OTEL_ENABLED", "origin": "env_var", "value": True}, + {"name": "DD_MODEL_LAB_ENABLED", "origin": "default", "value": False}, {"name": "DD_PROFILING_AGENTLESS", "origin": "default", "value": False}, {"name": "DD_PROFILING_API_TIMEOUT_MS", "origin": "default", "value": 10000}, {"name": "DD_PROFILING_CAPTURE_PCT", "origin": "env_var", "value": 5.0}, diff --git a/tests/utils.py b/tests/utils.py index 4f7ffe9d41b..d57a1e5b7d7 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -172,6 +172,7 @@ def override_global_config(values): "_data_streams_enabled", "_inferred_proxy_services_enabled", "_lib_was_injected", + "_model_lab_enabled", ] asm_config_keys = asm_config._asm_config_keys