diff --git a/lang/cs/Org.Apache.REEF.Examples/MachineLearning/KMeans/KMeansDriverHandlers.cs b/lang/cs/Org.Apache.REEF.Examples/MachineLearning/KMeans/KMeansDriverHandlers.cs
index 6ae60bd..b01065d 100644
--- a/lang/cs/Org.Apache.REEF.Examples/MachineLearning/KMeans/KMeansDriverHandlers.cs
+++ b/lang/cs/Org.Apache.REEF.Examples/MachineLearning/KMeans/KMeansDriverHandlers.cs
@@ -25,6 +25,7 @@
 using Org.Apache.REEF.Driver.Bridge;
 using Org.Apache.REEF.Driver.Context;
 using Org.Apache.REEF.Driver.Evaluator;
+using Org.Apache.REEF.Driver.Task;
 using Org.Apache.REEF.Examples.MachineLearning.KMeans.codecs;
 using Org.Apache.REEF.Network.Group.Config;
 using Org.Apache.REEF.Network.Group.Driver;
@@ -45,7 +46,8 @@
     public class KMeansDriverHandlers : 
         IObserver<IAllocatedEvaluator>,
         IObserver<IActiveContext>, 
-        IObserver<IDriverStarted>
+        IObserver<IDriverStarted>,
+        IObserver<ICompletedTask>
     {
         private static readonly Logger _Logger = Logger.GetLogger(typeof(KMeansDriverHandlers));
         private readonly object _lockObj = new object();
@@ -192,12 +194,15 @@
 
         public void OnError(Exception error)
         {
-            throw new NotImplementedException();
         }
 
         public void OnCompleted()
         {
-            throw new NotImplementedException();
+        }
+
+        public void OnNext(ICompletedTask value)
+        {
+            value.ActiveContext.Dispose();
         }
     }
 
diff --git a/lang/cs/Org.Apache.REEF.Tests/Functional/ML/KMeans/TestKMeans.cs b/lang/cs/Org.Apache.REEF.Tests/Functional/ML/KMeans/TestKMeans.cs
index a1cc91c..4cc91ff 100644
--- a/lang/cs/Org.Apache.REEF.Tests/Functional/ML/KMeans/TestKMeans.cs
+++ b/lang/cs/Org.Apache.REEF.Tests/Functional/ML/KMeans/TestKMeans.cs
@@ -123,7 +123,7 @@
             string dataFilePath = GenerateDataFileAndGetPath();
 
             TestRun(DriverConfiguration(dataFilePath), typeof(KMeansDriverHandlers), Partitions + 1, "KMeansDriverHandlers", "local", testFolder);
-            ValidateSuccessForLocalRuntime(Partitions + 1, testFolder: testFolder);
+            ValidateSuccessForLocalRuntime(Partitions + 1, testFolder: testFolder, retryCount: 120);
             CleanUp(testFolder);
             try
             {
@@ -156,11 +156,12 @@
 
             IConfiguration driverConfig = TangFactory.GetTang().NewConfigurationBuilder(
                 Org.Apache.REEF.Driver.DriverConfiguration.ConfigurationModule
-                    .Set(Org.Apache.REEF.Driver.DriverConfiguration.OnDriverStarted, GenericType<KMeansDriverHandlers>.Class)
-                    .Set(Org.Apache.REEF.Driver.DriverConfiguration.OnEvaluatorAllocated, GenericType<KMeansDriverHandlers>.Class)
-                    .Set(Org.Apache.REEF.Driver.DriverConfiguration.OnContextActive, GenericType<KMeansDriverHandlers>.Class)
-                    .Set(Org.Apache.REEF.Driver.DriverConfiguration.CommandLineArguments, dataFilePath)
-                    .Set(Org.Apache.REEF.Driver.DriverConfiguration.CustomTraceLevel, Level.Info.ToString())
+                    .Set(REEF.Driver.DriverConfiguration.OnDriverStarted, GenericType<KMeansDriverHandlers>.Class)
+                    .Set(REEF.Driver.DriverConfiguration.OnEvaluatorAllocated, GenericType<KMeansDriverHandlers>.Class)
+                    .Set(REEF.Driver.DriverConfiguration.OnContextActive, GenericType<KMeansDriverHandlers>.Class)
+                    .Set(REEF.Driver.DriverConfiguration.OnTaskCompleted, GenericType<KMeansDriverHandlers>.Class)
+                    .Set(REEF.Driver.DriverConfiguration.CommandLineArguments, dataFilePath)
+                    .Set(REEF.Driver.DriverConfiguration.CustomTraceLevel, Level.Info.ToString())
                     .Build())
                 .BindIntNamedParam<NumPartitions>(Partitions.ToString())
                 .Build();