Memory leak in text classification pipeline

[IgnasZeb]

System information

• OS version/distro: Microsoft Windows [Version 10.0.18362.418]
• .NET Version (eg., dotnet --info): .NET Core 3.0.100 04339c3a26

Issue

• What did you do? Detected potential memory leak in production application, so I wrote simple application to see if problem persists.
• What happened? Microsoft.ML 1.3.1 and 1.4.0-preview2 both appear to be leaking memory

Source code / logs

Minimal code that consistently reproduces this problem:

    class ModelInput
    {
        [ColumnName("TextColumn"), LoadColumn(0)]
        public string ItemDescription { get; set; }
        [ColumnName("Label"), LoadColumn(1)]
        public int ItemId { get; set; }
    }

    class Program
    {
        static void Main(string[] args)
        {
            Console.WriteLine("Sleeping for 5 seconds. Collect initial memory snapshot...");
            while(true)
            {
                Thread.Sleep(5000);
                BuildAndTrainModel();
                Console.WriteLine("Training done. Collect memory snapshot...");
            }
        }

        static void BuildAndTrainModel()
        {
            MLContext context = new MLContext(seed: 1);

            var dataView = context.Data.LoadFromTextFile<ModelInput>("input.csv", separatorChar: ',');

            var trainingPipeline = context.Transforms.Conversion.MapValueToKey(new[] { new InputOutputColumnPair("Label", "Label") })
                .Append(context.Transforms.Text.FeaturizeText("TextColumn", "TextColumn"))
                .Append(context.MulticlassClassification.Trainers.SdcaMaximumEntropy(labelColumnName: "Label", featureColumnName: "TextColumn", maximumNumberOfIterations: 1));

            var model = trainingPipeline.Fit(dataView);
        }
    }

After each call to BuildAndTrainModel following objects appear to leak:

Amount of objects leaked appears to correlate with number of iterations (increasing maximumNumberOfIterations increases leaked object count).

[ebarsoumMS]

If you put GC.Collect(); before Console.WriteLine, will the issue still repro?

[codemzs]

@antoniovs1029 is looking at this issue.

[codemzs]

@ebarsoumMS Good suggestion that will confirm this is indeed a leak. In the case of the leak calling GC.Collect() will not help.

[antoniovs1029]

Hi, @Ignasz , thanks for noticing this and opening the issue. I ran the sample code you provided, using the wikipedia-detox-250-line-data.tsv file as input. I will post my results down below, but the conclusion is that there's no memory leak, and using GC.Collect() as @ebarsoumMS suggested makes it easier to see it. So I will close this issue, but if you have further comments about this, please feel free to reopen this so we can look into it.

---

If I run your sample code using maximumNumberOfIterations = 1000, I get the following output:

Notice that at first it does seem that in each iteration there's an increase in memory usage, but around the iteration number 6 (at 30s) it seems the GC makes the memory usage to decrease, signaling there's not actually a memory leak.

If I then run your sample, but adding the GC.Collect() suggested by @ebarsoumMS I get the following output, where it is clearer to see that there's in fact no memory leak:

[IgnasZeb]

Hi, thank you for your response. I already tried adding GC.Collect();, but the problem persisted. I just ran the same test with your dataset and while it's difficult to see any leakage in memory chart, I believe it shows clearly in snapshot object diff. This was run with GC.Collect() added before Console.WriteLine:

Comparing snapshot 12 to snapshots 11 and 1, however, show that there is a constant increase in count of certain objects:

Furthermore, I tried running the same code (with GC.Collect() after every BuildAndTrainModel()) for 1000 times and this time memory leak seems to be much more evident in memory chart:

Thanks for looking into it, let me know if you need any more details.

[antoniovs1029]

By running the code 1000 times as you described (including the GC.Collect()) I was able to replicate your issue, and see the memory leak you found in the memory chart. I replicated this by using both ML.NET 1.0, and ML.NET 1.4 (which is the latest stable version), so I will reopen this issue but change its name.

I will try to find the source of this problem.

Thank you for noticing this and taking the time to share your findings.

[eerhardt]

I spent some time investigating this issue today.

The issue is our usage of BufferBlock, as @antoniovs1029 pointed out. There are a few layers to this problem, I’ll try to explain it all.

1. The actual memory leak only happens when the Debugger is attached. I confirmed this is a known issue with the VS debugger. See https://github.com/dotnet/corefx/issues/30582.

See the comments around this Dictionary in the Task class.
https://github.com/dotnet/coreclr/blob/a9f3fc16483eecfc47fb79c362811d870be02249/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs#L182-L188

When the VS Debugger is attached, every Task that gets started gets put into this Dictionary, and then when it gets completed it gets removed from the Dictionary.
When the VS Debugger isn’t attached, the Tasks don’t ever get put in this Dictionary. So even though the Task never completes they still get garbage collected once all references to them go away.

However, we are creating Tasks that aren't getting completed. Thus in VS, they get put in the Dictionary, but never removed.

2. So why aren’t these Tasks getting completed?

The reason these Tasks aren’t getting completed is because we are never calling “Complete()” on the BufferBlocks. You can see in BufferBlock’s constructor that it is doing something on Completion:

https://github.com/dotnet/corefx/blob/5cee7c97d602f294e27c582d4dab81ec388f1d7b/src/System.Threading.Tasks.Dataflow/src/Blocks/BufferBlock.cs#L71

This task (the one created by the call to ContinueWith) is what is not getting completed, because we are not calling Complete() on the BufferBlock in the RowShufflingTransformer.Cursor.

[antoniovs1029]

The issue is our usage of BufferBlock, as @antoniovs1029 pointed out

To your explanation, I just think it's important to add that these BufferBlocks (and the other objects related to this leak) are being created in the constructor of the RowShufflingTransformer.Cursor class (link to constructor). This class gets instantiated multiple times when fitting a SdcaMaximumEntropy trainer; in fact, the number of times it gets instantiated is related to the number of maximumNumberOfIterations parameter, explaining the relation that @Ignasz found between that parameter and the memory leak.

When debugging, the BufferBlocks created by this constructor are not collected by the Garbage Collector because of the reasons explained by @eerhardt . But they do get collected in production releases.