🐛 Bug Report: OpenLLMetry doesn't work with MLFlow

[ThanawaratGithub]

Which component is this bug for?

Traceloop SDK

📜 Description

Right now I'm working with Traceloop integration with mlfow. I'm trying to trace mlflow runs and mlflow evaluate and send the trace to qa tempo. My intention is that I want to add an attribute to the trace which is basically run_id of mlflow, now I tried Traceloop.set_association_properties But it seems that some traces can attach properties that I set some traces not attach (all of them are llmchain.workflow from mlflow evaluation ).

👟 Reproduction steps

Just want to run following code after Traceloop.init()

with mlflow.start_run(
    
    run_name="trace_paralell"+"_v" +run_version,
    tags={
        "mlflow.source.git.commit": run_version,
    }
) as started_run:
    Traceloop.set_association_properties({ "mlflow_run_id": started_run.info.run_id, "mlflow_exp_id": started_run.info.experiment_id})
    avg_tokens = count_tokens([PROMPT])
    for input_var in INPUT_VARIABLES_LIST:
        avg_tokens += count_tokens(dataset[input_var])
        
    model_info = mlflow.langchain.log_model(
        lc_model=create_model(), 
        artifact_path="model",
        signature=SIGNATURE,
    )
    
    if len(dataset.index) <= max_threads:
            max_threads = 1
        
    try:
        split_dataset = numpy.array_split(dataset, max_threads)
        result = []
        print(f"model uri: {model_info.model_uri}")
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
            for chunk in split_dataset:
                executor.submit(
                    evaluate_dataset, 
                    chunk,
                    model_info.model_uri,
                    result
                )
        result_dataset = pd.concat(result)
        
        mlflow.log_metrics({"latency_p50": numpy.percentile(result_dataset.latency_ms, 50)})
        mlflow.log_metrics({"latency_p90": numpy.percentile(result_dataset.latency_ms, 90)})
        
        results = mlflow.evaluate(
            data=result_dataset,
            predictions=RESULT_COLUMN,
            targets=GROUND_TRUTH_COLUMN,
            evaluators="default",
            model_type="classifier",
            validation_thresholds = {
                        "precision_score": MetricThreshold(
                            threshold=PRECISION_THRESHOLD,
                            greater_is_better=True,
                        ),
                    },
            extra_metrics=[ token_count(), latency(), precision_metric]
        )
        
        mlflow.log_metrics(results.metrics)

👍 Expected behavior

I expect that every span after tracing should have mlflow_run_id and mlflow_exp_id as I set in association_properties.

👎 Actual Behavior with Screenshots

There are mlflow_run_id and mlflow_exp_id for some traces not every trace that came from this run. For example , LLMChain.workflow which come from mlflow.evaluate has no assiciation_properties.

This is how all traces look like when I run above code.

This is for POST span with path mlfowserver/api/mlfow/runs/log-inputs

You can see it has association properties
But for the open ai one there is no properties

For OPENAI

🤖 Python Version

3.11

📃 Provide any additional context for the Bug.

From discussing with Nir. It seems like mlflow_runs and mlflow_evaluate may be messing up the context

👀 Have you spent some time to check if this bug has been raised before?

• I checked and didn't find similar issue

Are you willing to submit PR?

None