MLflow logging of drift metrics

[pluflou]

I see a couple of issues with the way MLflow logging is set up. It doesn't seem to happen with the provided examples, but they do with the example set up in the slac-fel branch (maybe due to the higher number of global steps and inner CL loops?)

In that branch I committed a fix that needed to be in there to be able to see drift.detected on MLflow/WandB: logging them as ints not bools.

It's also possible drift.detected should use drift.step (the per-stage step) as its x-axis, not the global step. This way, the chart shows consistent progression independent of how many eval/cl steps happen in between. I did not implement that because I am not sure it resolves the full issue I see below.

Main issue
During a run, the dashboard does not show all the expected drift.detected events (should be 4):

Moreover, after the run ends, MLflow seems to remove some rows from the data, and since we have very few 1s vs 0s, we end up with a flat 0 line:

Exporting the CSV data for drift.detected, for this example, we expected a 1 around step 4500 but it's been removed.

I haven't had the time to create a simple reproducer to put here. Perhaps I or @anagainaru can share the data with whoever is assigned to this (I am also going to upload it to Perlmutter tomorrow).

FWIW, the saved CSV looks fine (4 drifts detected), attached here.
slac-fel.csv

[Zilinghan]

After checking in details, it seems that mlflow does not actually delete data from its backend. The "removal" is a UI-level rendering issue: when a metric has many data points, the chart downsamples for display, and the "Export as CSV" button from the chart tab exports only the rendered (downsampled) subset, not the full history. There's no config flag to disable the UI chart downsampling in mlflow.
Therefore, to make the x-axis with fewer data points, different from wandb, I tried to use per-stage step for different metrics in mlflow logger as the x-axit so that the x-axis becomes less sparse and avoid too much down sampling in the UI. I created a branch on top of the slac-fel branch here to implement the changes in x-axis.

[pluflou]

Thanks for looking into it. I pulled in your changes and ran, and unfortunately it doesn't resolve it (I tried the same implementation earlier as well). But since it sounds like just a rendering issue, maybe it's fine (I was worried it was something on our end).

Here is what I observed with your changes:
right before run ended (peaks would show then disappear as more samples got logged)

after run ended

[Zilinghan]

Thanks @pluflou for your quick testing!@\ It seems that even per-stage step is too sparse to display the detected drift in this example, and it is a limitation for mlflow that we cannot do anything with it. Do you think as long as we have a correct local csv file, we are okay with it?

[pluflou]

The problem with it is that now the chart on MLflow has no value and at worst it might be confusing if the typical user is only checking the dashboard/not aware of the limitations. If it's documented clearly so users know what to expect, that might be fine.

It might also be redundant with drift.score, correct? Is there a case where the number/step loc of scores is not equal to the number/step loc of drifts detected? If they are redundant, then maybe it doesn't have to be logged?

Related: do you know why the drift score persistent across several steps before going back to 0?

[Zilinghan]

@pluflou I created a new fix to the same branch, can you check if it works?
basically, I reduce the log freq for no detection by 90% using log_detected = drift_signal.drift_detected or (np.random.random() <= 0.1)

[Zilinghan]

Also, answer to a continuous high drift.score is it is computed as the mean of past 100 drift histories, so it is no 100% equivalent to drift.detected, if drift was detected multiple times within the 100-length window, the score keeps going high - see an example below

        # Compute drift score based on recent drift frequency
        recent_window = min(100, len(self._drift_history))
        recent_drifts = self._drift_history[-recent_window:]
        drift_score = np.mean(recent_drifts)

[pluflou]

Thank you, reducing the log frequency worked perfectly! I added the fix to my branch.

And thank you for explaining re drift score.

[anagainaru]

Fixed by #89