Use Timer and ctx.CancelExecution() to fix AutoML max-time experiment bug

src/Microsoft.ML.AutoML/Experiment/Experiment.cs

@@ -7,6 +7,8 @@
 using System.Diagnostics;
 using System.IO;
 using System.Linq;
+using System.Threading;
+using Microsoft.ML.Data;
 using Microsoft.ML.Runtime;
 
 namespace Microsoft.ML.AutoML
@@ -25,6 +27,11 @@ internal class Experiment<TRunDetail, TMetrics> where TRunDetail : RunDetail
         private readonly IRunner<TRunDetail> _runner;
         private readonly IList<SuggestedPipelineRunDetail> _history;
         private readonly IChannel _logger;
+        private Timer _maxExperimentTimeTimer;
+        private Timer _mainContextCanceledTimer;
+        private bool _experimentTimerExpired;
+        private MLContext _currentModelMLContext;
+        private Random _newContextSeedGenerator;
 
         public Experiment(MLContext context,
             TaskKind task,
@@ -49,23 +56,75 @@ public Experiment(MLContext context,
             _datasetColumnInfo = datasetColumnInfo;
             _runner = runner;
             _logger = logger;
+            _experimentTimerExpired = false;
+        }
+
+        private void MaxExperimentTimeExpiredEvent(object state)
+        {
+            // If at least one model was run, end experiment immediately.
+            // Else, wait for first model to run before experiment is concluded.
+            _experimentTimerExpired = true;
+            if (_history.Any(r => r.RunSucceeded))
+            {
+                _logger.Warning("Allocated time for Experiment of {0} seconds has elapsed with {1} models run. Ending experiment...",
+                    _experimentSettings.MaxExperimentTimeInSeconds, _history.Count());
+                _currentModelMLContext.CancelExecution();
+            }
+        }
+
+        private void MainContextCanceledEvent(object state)
+        {
+            // If the main MLContext is canceled, cancel the ongoing model training and MLContext.
+            if ((_context.Model.GetEnvironment() as ICancelable).IsCanceled)
+            {
+                _logger.Warning("Main MLContext has been canceled. Ending experiment...");
+                _currentModelMLContext.CancelExecution();
+            }
         }
 
         public IList<TRunDetail> Execute()
         {
-            var stopwatch = Stopwatch.StartNew();
             var iterationResults = new List<TRunDetail>();
+            // Create a timer for the max duration of experiment. When given time has
+            // elapsed, MaxExperimentTimeExpiredEvent is called to interrupt training
+            // of current model. Timer is not used if no experiment time is given, or
+            // is not a positive number.
+            if (_experimentSettings.MaxExperimentTimeInSeconds > 0)
+            {
+                _maxExperimentTimeTimer = new Timer(
+                    new TimerCallback(MaxExperimentTimeExpiredEvent), null,
+                    _experimentSettings.MaxExperimentTimeInSeconds * 1000, Timeout.Infinite
+                );
+            }
+            // If given max duration of experiment is 0, only 1 model will be trained.
+            // _experimentSettings.MaxExperimentTimeInSeconds is of type uint, it is
+            // either 0 or >0.
+            else
+                _experimentTimerExpired = true;
+            // Add second timer to check for the cancelation signal from the main MLContext
+            // to the active child MLContext. This timer will propagate the cancelation
+            // signal from the main to the child MLContexs if the main MLContext is
+            // canceled.
+            _mainContextCanceledTimer = new Timer(new TimerCallback(MainContextCanceledEvent), null, 1000, 1000);
+
+            // Pseudo random number generator to result in deterministic runs with the provided main MLContext's seed and to
+            // maintain variability between training iterations.
+            int? mainContextSeed = ((ISeededEnvironment)_context.Model.GetEnvironment()).Seed;
+            _newContextSeedGenerator = (mainContextSeed.HasValue) ? RandomUtils.Create(mainContextSeed.Value) : RandomUtils.Create();
 
             do
             {
                 var iterationStopwatch = Stopwatch.StartNew();
 
                 // get next pipeline
                 var getPipelineStopwatch = Stopwatch.StartNew();
-                var pipeline = PipelineSuggester.GetNextInferredPipeline(_context, _history, _datasetColumnInfo, _task,
-                    _optimizingMetricInfo.IsMaximizing, _experimentSettings.CacheBeforeTrainer, _trainerAllowList);
 
-                var pipelineInferenceTimeInSeconds = getPipelineStopwatch.Elapsed.TotalSeconds;
+                // A new MLContext is needed per model run. When max experiment time is reached, each used
+                // context is canceled to stop further model training. The cancellation of the main MLContext
+                // a user has instantiated is not desirable, thus additional MLContexts are used.
+                _currentModelMLContext = new MLContext(_newContextSeedGenerator.Next());
+                var pipeline = PipelineSuggester.GetNextInferredPipeline(_currentModelMLContext, _history, _datasetColumnInfo, _task,
+                    _optimizingMetricInfo.IsMaximizing, _experimentSettings.CacheBeforeTrainer, _trainerAllowList);
 
                 // break if no candidates returned, means no valid pipeline available
                 if (pipeline == null)
@@ -101,8 +160,7 @@ public IList<TRunDetail> Execute()
 
             } while (_history.Count < _experimentSettings.MaxModels &&
                     !_experimentSettings.CancellationToken.IsCancellationRequested &&
-                    stopwatch.Elapsed.TotalSeconds < _experimentSettings.MaxExperimentTimeInSeconds);
-
+                    !_experimentTimerExpired);
             return iterationResults;
         }
 

src/Microsoft.ML.AutoML/Experiment/Runners/CrossValSummaryRunner.cs

@@ -58,7 +58,7 @@ public CrossValSummaryRunner(MLContext context,
             for (var i = 0; i < _trainDatasets.Length; i++)
             {
                 var modelFileInfo = RunnerUtil.GetModelFileInfo(modelDirectory, iterationNum, i + 1);
-                var trainResult = RunnerUtil.TrainAndScorePipeline(_context, pipeline, _trainDatasets[i], _validDatasets[i],
+                var trainResult = RunnerUtil.TrainAndScorePipeline(pipeline.GetContext(), pipeline, _trainDatasets[i], _validDatasets[i],
                     _groupIdColumn, _labelColumn, _metricsAgent, _preprocessorTransforms?.ElementAt(i), modelFileInfo, _modelInputSchema,
                     _logger);
                 trainResults.Add(trainResult);

src/Microsoft.ML.AutoML/Experiment/SuggestedPipeline.cs

@@ -52,6 +52,11 @@ public override int GetHashCode()
             return ToString().GetHashCode();
         }
 
+        public MLContext GetContext()
+		{
+            return _context;
+		}
+
         public Pipeline ToPipeline()
         {
             var pipelineElements = new List<PipelineNode>();

src/Microsoft.ML.Core/Data/IHostEnvironment.cs

@@ -72,7 +72,7 @@ public interface IHostEnvironment : IChannelProvider, IProgressChannelProvider
     internal interface ICancelable
     {
         /// <summary>
-        /// Signal to stop exection in all the hosts.
+        /// Signal to stop execution in all the hosts.
         /// </summary>
         void CancelExecution();
 

test/Microsoft.ML.AutoML.Tests/AutoFitTests.cs

@@ -2,12 +2,14 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System;
 using System.Linq;
+using System.Reflection;
 using Microsoft.ML.Data;
+using Microsoft.ML.Runtime;
 using Microsoft.ML.TestFramework;
 using Microsoft.ML.TestFramework.Attributes;
 using Microsoft.ML.TestFrameworkCommon;
-using Microsoft.ML.Trainers.LightGbm;
 using Xunit;
 using Xunit.Abstractions;
 using static Microsoft.ML.DataOperationsCatalog;
@@ -117,7 +119,7 @@ public void AutoFitRegressionTest()
                 .Execute(trainData, validationData,
                     new ColumnInformation() { LabelColumnName = DatasetUtil.MlNetGeneratedRegressionLabel });
 
-            Assert.True(result.RunDetails.Max(i => i.ValidationMetrics.RSquared > 0.9));
+            Assert.True(result.RunDetails.Max(i => i?.ValidationMetrics?.RSquared) > 0.9);
         }
 
         [LightGBMFact]
@@ -246,7 +248,7 @@ public void AutoFitRecommendationTest()
             RunDetail<RegressionMetrics> bestRun = experimentResult.BestRun;
             Assert.True(experimentResult.RunDetails.Count() > 1);
             Assert.NotNull(bestRun.ValidationMetrics);
-            Assert.True(experimentResult.RunDetails.Max(i => i.ValidationMetrics.RSquared != 0));
+            Assert.True(experimentResult.RunDetails.Max(i => i?.ValidationMetrics?.RSquared* i?.ValidationMetrics?.RSquared) > 0.5);
 
             var outputSchema = bestRun.Model.GetOutputSchema(trainDataView.Schema);
             var expectedOutputNames = new string[] { labelColumnName, userColumnName, userColumnName, itemColumnName, itemColumnName, scoreColumnName };
@@ -320,6 +322,33 @@ public void AutoFitWithPresplittedData()
 
         }
 
+        [Fact]
+        public void AutoFitMaxExperimentTimeTest()
+		{
+            // A single binary classification experiment takes less than 5 seconds.
+            // System.OperationCanceledException is thrown when ongoing experiment
+            // is canceled and at least one model has been generated.
+            var context = new MLContext(1);
+            var dataPath = DatasetUtil.GetUciAdultDataset();
+            var columnInference = context.Auto().InferColumns(dataPath, DatasetUtil.UciAdultLabel);
+            var textLoader = context.Data.CreateTextLoader(columnInference.TextLoaderOptions);
+            var trainData = textLoader.Load(dataPath);
+            var experiment = context.Auto()
+                .CreateBinaryClassificationExperiment(10)
+                .Execute(trainData, new ColumnInformation() { LabelColumnName = DatasetUtil.UciAdultLabel });
+
+            // Ensure the (last) model that was training when maximum experiment time was reached has been stopped,
+            // and that its MLContext has been canceled.
+            Assert.True(experiment.RunDetails.Last().Exception.Message.Contains("Operation was canceled"),
+                        "Training process was not successfully canceled after maximum experiment time was reached.");
+
+            // Ensure that the best found model can still run after maximum experiment time was reached.
+            var refitModel = experiment.BestRun.Estimator.Fit(trainData);
+            IDataView predictions = refitModel.Transform(trainData);
+            var metrics = context.BinaryClassification.Evaluate(predictions, labelColumnName: DatasetUtil.UciAdultLabel);
+            Assert.True(metrics?.Accuracy > 0.5);
+        }
+
         private TextLoader.Options GetLoaderArgs(string labelColumnName, string userIdColumnName, string itemIdColumnName)
         {
             return new TextLoader.Options()