Add MaxModelToExplore exit strategy to AutoMLExperiment. (#6402)

Author: LittleLittleCloud

src/Microsoft.ML.AutoML/API/BinaryClassificationExperiment.cs

@@ -150,6 +150,7 @@ internal BinaryClassificationExperiment(MLContext context, BinaryExperimentSetti
             {
                 _experiment.SetMaximumMemoryUsageInMegaByte(d);
             }
+            _experiment.SetMaxModelToExplore(settings.MaxModels);
         }
 
         public override ExperimentResult<BinaryClassificationMetrics> Execute(IDataView trainData, ColumnInformation columnInformation, IEstimator<ITransformer> preFeaturizer = null, IProgress<RunDetail<BinaryClassificationMetrics>> progressHandler = null)

src/Microsoft.ML.AutoML/API/MulticlassClassificationExperiment.cs

@@ -142,6 +142,7 @@ internal MulticlassClassificationExperiment(MLContext context, MulticlassExperim
             {
                 _experiment.SetMaximumMemoryUsageInMegaByte(d);
             }
+            _experiment.SetMaxModelToExplore(settings.MaxModels);
         }
 
         public override ExperimentResult<MulticlassClassificationMetrics> Execute(IDataView trainData, ColumnInformation columnInformation, IEstimator<ITransformer> preFeaturizer = null, IProgress<RunDetail<MulticlassClassificationMetrics>> progressHandler = null)

src/Microsoft.ML.AutoML/API/RegressionExperiment.cs

@@ -139,6 +139,7 @@ internal RegressionExperiment(MLContext context, RegressionExperimentSettings se
             }
 
             _experiment.SetTrainingTimeInSeconds(Settings.MaxExperimentTimeInSeconds);
+            _experiment.SetMaxModelToExplore(Settings.MaxModels);
         }
 
         public override ExperimentResult<RegressionMetrics> Execute(IDataView trainData, ColumnInformation columnInformation, IEstimator<ITransformer> preFeaturizer = null, IProgress<RunDetail<RegressionMetrics>> progressHandler = null)

src/Microsoft.ML.AutoML/AutoMLExperiment/AutoMLExperiment.cs

@@ -25,7 +25,6 @@ public class AutoMLExperiment
         private double _bestLoss = double.MaxValue;
         private TrialResult _bestTrialResult = null;
         private readonly IServiceCollection _serviceCollection;
-        private CancellationTokenSource _globalCancellationTokenSource;
 
         public AutoMLExperiment(MLContext context, AutoMLExperimentSettings settings)
         {
@@ -51,14 +50,15 @@ private void InitializeServiceCollection()
             _serviceCollection.TryAddTransient((provider) =>
             {
                 var contextManager = provider.GetRequiredService<IMLContextManager>();
+                var trainingStopManager = provider.GetRequiredService<AggregateTrainingStopManager>();
                 var context = contextManager.CreateMLContext();
-                _globalCancellationTokenSource.Token.Register(() =>
+                trainingStopManager.OnStopTraining += (s, e) =>
                 {
                     // only force-canceling running trials when there's completed trials.
                     // otherwise, wait for the current running trial to be completed.
                     if (_bestTrialResult != null)
                         context.CancelExecution();
-                });
+                };
 
                 return context;
             });
@@ -74,6 +74,29 @@ private void InitializeServiceCollection()
         public AutoMLExperiment SetTrainingTimeInSeconds(uint trainingTimeInSeconds)
         {
             _settings.MaxExperimentTimeInSeconds = trainingTimeInSeconds;
+            _serviceCollection.AddScoped<IStopTrainingManager>((provider) =>
+            {
+                var channel = provider.GetRequiredService<IChannel>();
+                var timeoutManager = new TimeoutTrainingStopManager(TimeSpan.FromSeconds(trainingTimeInSeconds), channel);
+
+                return timeoutManager;
+            });
+
+            return this;
+        }
+
+        public AutoMLExperiment SetMaxModelToExplore(int maxModel)
+        {
+            _context.Assert(maxModel > 0, "maxModel has to be greater than 0");
+            _settings.MaxModels = maxModel;
+            _serviceCollection.AddScoped<IStopTrainingManager>((provider) =>
+            {
+                var channel = provider.GetRequiredService<IChannel>();
+                var maxModelManager = new MaxModelStopManager(maxModel, channel);
+
+                return maxModelManager;
+            });
+
             return this;
         }
 
@@ -204,19 +227,29 @@ public TrialResult Run()
         public async Task<TrialResult> RunAsync(CancellationToken ct = default)
         {
             ValidateSettings();
-            _globalCancellationTokenSource = new CancellationTokenSource();
-            _settings.CancellationToken = ct;
-            // use TimeSpan to avoid overflow.
-            _globalCancellationTokenSource.CancelAfter(TimeSpan.FromSeconds(_settings.MaxExperimentTimeInSeconds));
-            _settings.CancellationToken.Register(() => _globalCancellationTokenSource.Cancel());
+            _serviceCollection.AddScoped((serviceProvider) =>
+            {
+                var logger = serviceProvider.GetRequiredService<IChannel>();
+                var stopServices = serviceProvider.GetServices<IStopTrainingManager>();
+                var cancellationTrainingStopManager = new CancellationTokenStopTrainingManager(ct, logger);
+
+                // always get the most recent added stop service for each type.
+                var mostRecentAddedStopServices = stopServices.GroupBy(s => s.GetType()).Select(g => g.Last()).ToList();
+                mostRecentAddedStopServices.Add(cancellationTrainingStopManager);
+                return new AggregateTrainingStopManager(logger, mostRecentAddedStopServices.ToArray());
+            });
+
             var serviceProvider = _serviceCollection.BuildServiceProvider();
-            var monitor = serviceProvider.GetService<IMonitor>();
+
+            _settings.CancellationToken = ct;
             var logger = serviceProvider.GetRequiredService<IChannel>();
+            var aggregateTrainingStopManager = serviceProvider.GetRequiredService<AggregateTrainingStopManager>();
+            var monitor = serviceProvider.GetService<IMonitor>();
             var trialResultManager = serviceProvider.GetService<ITrialResultManager>();
             var trialNum = trialResultManager?.GetAllTrialResults().Max(t => t.TrialSettings?.TrialId) + 1 ?? 0;
             var tuner = serviceProvider.GetService<ITuner>();
             Contracts.Assert(tuner != null, "tuner can't be null");
-            while (!_globalCancellationTokenSource.Token.IsCancellationRequested)
+            while (!aggregateTrainingStopManager.IsStopTrainingRequested())
             {
                 var setting = new TrialSettings()
                 {
@@ -227,86 +260,95 @@ public async Task<TrialResult> RunAsync(CancellationToken ct = default)
                 setting.Parameter = parameter;
 
                 monitor?.ReportRunningTrial(setting);
-                try
+                using (var trialCancellationTokenSource = new CancellationTokenSource())
                 {
-                    using (var trialCancellationTokenSource = new CancellationTokenSource())
-                    using (var deregisterCallback = _globalCancellationTokenSource.Token.Register(() =>
+                    void handler(object o, EventArgs e)
                     {
                         // only force-canceling running trials when there's completed trials.
                         // otherwise, wait for the current running trial to be completed.
                         if (_bestTrialResult != null)
                             trialCancellationTokenSource.Cancel();
-                    }))
-                    using (var performanceMonitor = serviceProvider.GetService<IPerformanceMonitor>())
-                    using (var runner = serviceProvider.GetRequiredService<ITrialRunner>())
+                    }
+                    try
                     {
-                        performanceMonitor.MemoryUsageInMegaByte += (o, m) =>
+                        using (var performanceMonitor = serviceProvider.GetService<IPerformanceMonitor>())
+                        using (var runner = serviceProvider.GetRequiredService<ITrialRunner>())
                         {
-                            if (_settings.MaximumMemoryUsageInMegaByte is double d && m > d && !trialCancellationTokenSource.IsCancellationRequested)
-                            {
-                                logger.Trace($"cancel current trial {setting.TrialId} because it uses {m} mb memory and the maximum memory usage is {d}");
-                                trialCancellationTokenSource.Cancel();
+                            aggregateTrainingStopManager.OnStopTraining += handler;
 
-                                GC.AddMemoryPressure(Convert.ToInt64(m) * 1024 * 1024);
-                                GC.Collect();
+                            performanceMonitor.MemoryUsageInMegaByte += (o, m) =>
+                            {
+                                if (_settings.MaximumMemoryUsageInMegaByte is double d && m > d && !trialCancellationTokenSource.IsCancellationRequested)
+                                {
+                                    logger.Trace($"cancel current trial {setting.TrialId} because it uses {m} mb memory and the maximum memory usage is {d}");
+                                    trialCancellationTokenSource.Cancel();
+
+                                    GC.AddMemoryPressure(Convert.ToInt64(m) * 1024 * 1024);
+                                    GC.Collect();
+                                }
+                            };
+
+                            performanceMonitor.Start();
+                            logger.Trace($"trial setting - {JsonSerializer.Serialize(setting)}");
+                            var trialResult = await runner.RunAsync(setting, trialCancellationTokenSource.Token);
+
+                            var peakCpu = performanceMonitor?.GetPeakCpuUsage();
+                            var peakMemoryInMB = performanceMonitor?.GetPeakMemoryUsageInMegaByte();
+                            trialResult.PeakCpu = peakCpu;
+                            trialResult.PeakMemoryInMegaByte = peakMemoryInMB;
+
+                            monitor?.ReportCompletedTrial(trialResult);
+                            tuner.Update(trialResult);
+                            trialResultManager?.AddOrUpdateTrialResult(trialResult);
+                            aggregateTrainingStopManager.Update(trialResult);
+
+                            var loss = trialResult.Loss;
+                            if (loss < _bestLoss)
+                            {
+                                _bestTrialResult = trialResult;
+                                _bestLoss = loss;
+                                monitor?.ReportBestTrial(trialResult);
                             }
+                        }
+                    }
+                    catch (OperationCanceledException ex) when (aggregateTrainingStopManager.IsStopTrainingRequested() == false)
+                    {
+                        monitor?.ReportFailTrial(setting, ex);
+                        var result = new TrialResult
+                        {
+                            TrialSettings = setting,
+                            Loss = double.MaxValue,
                         };
 
-                        performanceMonitor.Start();
-                        logger.Trace($"trial setting - {JsonSerializer.Serialize(setting)}");
-                        var trialResult = await runner.RunAsync(setting, trialCancellationTokenSource.Token);
-
-                        var peakCpu = performanceMonitor?.GetPeakCpuUsage();
-                        var peakMemoryInMB = performanceMonitor?.GetPeakMemoryUsageInMegaByte();
-                        trialResult.PeakCpu = peakCpu;
-                        trialResult.PeakMemoryInMegaByte = peakMemoryInMB;
-
-                        monitor?.ReportCompletedTrial(trialResult);
-                        tuner.Update(trialResult);
-                        trialResultManager?.AddOrUpdateTrialResult(trialResult);
+                        tuner.Update(result);
+                        continue;
+                    }
+                    catch (OperationCanceledException) when (aggregateTrainingStopManager.IsStopTrainingRequested())
+                    {
+                        break;
+                    }
+                    catch (Exception ex)
+                    {
+                        monitor?.ReportFailTrial(setting, ex);
 
-                        var loss = trialResult.Loss;
-                        if (loss < _bestLoss)
+                        if (!aggregateTrainingStopManager.IsStopTrainingRequested() && _bestTrialResult == null)
                         {
-                            _bestTrialResult = trialResult;
-                            _bestLoss = loss;
-                            monitor?.ReportBestTrial(trialResult);
+                            // TODO
+                            // it's questionable on whether to abort the entire training process
+                            // for a single fail trial. We should make it an option and only exit
+                            // when error is fatal (like schema mismatch).
+                            throw;
                         }
                     }
-                }
-                catch (OperationCanceledException ex) when (_globalCancellationTokenSource.IsCancellationRequested == false)
-                {
-                    monitor?.ReportFailTrial(setting, ex);
-                    var result = new TrialResult
+                    finally
                     {
-                        TrialSettings = setting,
-                        Loss = double.MaxValue,
-                    };
+                        aggregateTrainingStopManager.OnStopTraining -= handler;
 
-                    tuner.Update(result);
-                    continue;
-                }
-                catch (OperationCanceledException) when (_globalCancellationTokenSource.IsCancellationRequested)
-                {
-                    break;
-                }
-                catch (Exception ex)
-                {
-                    monitor?.ReportFailTrial(setting, ex);
-
-                    if (!_globalCancellationTokenSource.IsCancellationRequested && _bestTrialResult == null)
-                    {
-                        // TODO
-                        // it's questionable on whether to abort the entire training process
-                        // for a single fail trial. We should make it an option and only exit
-                        // when error is fatal (like schema mismatch).
-                        throw;
                     }
                 }
             }
 
             trialResultManager?.Save();
-
             if (_bestTrialResult == null)
             {
                 throw new TimeoutException("Training time finished without completing a trial run");