aws / aws/sagemaker-mlflow

[BUG] Distributed tracing not working in AWS Sagemaker AI MLflow (3.10.1)

Open
#26 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
26
Forks
19
Avg merge
1d 6h
Merged PRs (30d)
1

Description

### **MLflow version**
dependencies = [
"copilotkit==0.1.87",
"uvicorn[standard]==0.42.0",
"langgraph==1.1.9",
"langchain==1.2.15",
"langchain-chroma==1.1.0",
"langchain-community==0.4.2",
"langchain-openai==1.2.1",
"loguru==0.7.3",
"openpyxl==3.1.5",
"pandas>=2.2.0,<3.0.0",
"pydantic-settings (>=2.10.1,<3.0.0)",
"python-dotenv (>=1.1.1,<2.0.0)",
"pypdf==6.12.2",
"mlflow==3.10.1",
"boto3>=1.43.12",
"sagemaker-mlflow>=0.4.0",
"litellm==1.59.12",
]

### **System information**
OS Platform and Distribution (e.g., Linux Ubuntu 16.04): AWS Sagemaker AI MLflow 3.10.1
Python version: ">=3.11,<3.13"
Describe the problem
There is a distributed tracing issue when changing from local mlflow image (version 3.10.1) to AWS Sagemaker AI MLflow (version 3.10.1). When changing the host to the AWS-provided, all the tracking is registered but the secondary service embedded traces in the main trace.

There are two microservices which generate standalone traces when called independently. However, when the main micro calls the second one and we called the distributed tracing is not working for the AWS version.

### **Code to reproduce issue**
This is the endpoint where distributed tracing is set which locally works (docker MLflow image 3.10.1), though the AWS Sagemaker AI MLflow (3.10.1) is silently failing.

```
@app.post("/chat", response_model=BlueBuddyOutputSchema)
def chat_endpoint(
request: BlueBuddyInputSchema, http_request: Request
) -> BlueBuddyOutputSchema:
headers = dict(http_request.headers)
if "traceparent" in {k.lower() for k in headers}:
with set_tracing_context_from_http_request_headers(headers):
return run_blue_buddy(user_input=request)
return run_blue_buddy(user_input=request)
```
### **Stack trace**
The trace throws no error, but silently fails to log the distributed trace.

### **What component(s) does this bug affect?**

area/tracking: Tracking Service, tracking client APIs, autologging

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the /chat endpoint with the listed Python dependencies against local MLflow 3.10.1 and AWS SageMaker AI MLflow 3.10.1, comparing how the traceparent headers are handled. Done means the secondary service trace is nested in the main trace on AWS, or the incompatibility and affected component are clearly identified.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, python
Domain
cloud, observability-sre
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.