@@ -17,7 +17,7 @@ limitations under the License.
1717package trial
1818
1919import (
20- "sync "
20+ "context "
2121 "testing"
2222 "time"
2323
@@ -48,14 +48,47 @@ import (
4848
4949const (
5050 namespace = "default"
51- trialName = "test-trial"
5251 batchJobName = "test-job"
5352 objectiveMetric = "accuracy"
54- timeout = time .Second * 80
53+ timeout = time .Second * 10
5554)
5655
57- var trialKey = types.NamespacedName {Name : trialName , Namespace : namespace }
58- var batchJobKey = types.NamespacedName {Name : batchJobName , Namespace : namespace }
56+ var (
57+ batchJobKey = types.NamespacedName {Name : batchJobName , Namespace : namespace }
58+ observationLogAvailable = & api_pb.GetObservationLogReply {
59+ ObservationLog : & api_pb.ObservationLog {
60+ MetricLogs : []* api_pb.MetricLog {
61+ {
62+ TimeStamp : "2020-08-10T14:47:38+08:00" ,
63+ Metric : & api_pb.Metric {
64+ Name : objectiveMetric ,
65+ Value : "0.99" ,
66+ },
67+ },
68+ {
69+ TimeStamp : "2020-08-10T14:50:38+08:00" ,
70+ Metric : & api_pb.Metric {
71+ Name : objectiveMetric ,
72+ Value : "0.11" ,
73+ },
74+ },
75+ },
76+ },
77+ }
78+ observationLogUnavailable = & api_pb.GetObservationLogReply {
79+ ObservationLog : & api_pb.ObservationLog {
80+ MetricLogs : []* api_pb.MetricLog {
81+ {
82+ Metric : & api_pb.Metric {
83+ Name : objectiveMetric ,
84+ Value : consts .UnavailableMetricValue ,
85+ },
86+ TimeStamp : time.Time {}.UTC ().Format (time .RFC3339 ),
87+ },
88+ },
89+ },
90+ }
91+ )
5992
6093func init () {
6194 logf .SetLogger (zap .New (zap .UseDevMode (true )))
@@ -112,6 +145,7 @@ func TestReconcileBatchJob(t *testing.T) {
112145 // Try to update status until it be succeeded
113146 for err != nil {
114147 updatedInstance := & trialsv1beta1.Trial {}
148+ trialKey := types.NamespacedName {Name : instance .Name , Namespace : namespace }
115149 if err = c .Get (ctx , trialKey , updatedInstance ); err != nil {
116150 continue
117151 }
@@ -134,59 +168,22 @@ func TestReconcileBatchJob(t *testing.T) {
134168 viper .Set (consts .ConfigTrialResources , trialResources )
135169 g .Expect (add (mgr , recFn )).NotTo (gomega .HaveOccurred ())
136170
137- // Start test manager.
138- wg := & sync. WaitGroup {}
139- wg . Add ( 1 )
171+ // Start test manager
172+ mgrCtx , cancel := context . WithCancel ( context . TODO ())
173+ t . Cleanup ( cancel )
140174 go func () {
141- defer wg .Done ()
142- g .Expect (mgr .Start (ctx )).NotTo (gomega .HaveOccurred ())
175+ g .Expect (mgr .Start (mgrCtx )).NotTo (gomega .HaveOccurred ())
143176 }()
144177
145- // Result for GetTrialObservationLog with some metrics.
146- observationLogAvailable := & api_pb.GetObservationLogReply {
147- ObservationLog : & api_pb.ObservationLog {
148- MetricLogs : []* api_pb.MetricLog {
149- {
150- TimeStamp : "2020-08-10T14:47:38+08:00" ,
151- Metric : & api_pb.Metric {
152- Name : objectiveMetric ,
153- Value : "0.99" ,
154- },
155- },
156- {
157- TimeStamp : "2020-08-10T14:50:38+08:00" ,
158- Metric : & api_pb.Metric {
159- Name : objectiveMetric ,
160- Value : "0.11" ,
161- },
162- },
163- },
164- },
165- }
166- // Empty result for GetTrialObservationLog.
167- // If objective metrics are not parsed, metrics collector reports "unavailable" value to DB.
168- observationLogUnavailable := & api_pb.GetObservationLogReply {
169- ObservationLog : & api_pb.ObservationLog {
170- MetricLogs : []* api_pb.MetricLog {
171- {
172- Metric : & api_pb.Metric {
173- Name : objectiveMetric ,
174- Value : consts .UnavailableMetricValue ,
175- },
176- TimeStamp : time.Time {}.UTC ().Format (time .RFC3339 ),
177- },
178- },
179- },
180- }
181-
182178 t .Run (`Trial run with "Failed" BatchJob.` , func (t * testing.T ) {
183179 g := gomega .NewGomegaWithT (t )
184180 mockManagerClient .EXPECT ().DeleteTrialObservationLog (gomock .Any ()).Return (nil , nil )
185181
186- trial := newFakeTrialBatchJob ()
182+ trial := newFakeTrialBatchJob (commonv1beta1 .StdOutCollector , "test-failed-batch-job" )
183+ trialKey := types.NamespacedName {Name : "test-failed-batch-job" , Namespace : namespace }
187184 batchJob := & batchv1.Job {}
188185
189- // Create the Trial
186+ // Create the Trial with StdOut MC
190187 g .Expect (c .Create (ctx , trial )).NotTo (gomega .HaveOccurred ())
191188
192189 // Expect that BatchJob with appropriate name is created
@@ -239,7 +236,7 @@ func TestReconcileBatchJob(t *testing.T) {
239236 }, timeout ).Should (gomega .BeTrue ())
240237 })
241238
242- t .Run (`Trail with "Complete" BatchJob and Available metrics.` , func (t * testing.T ) {
239+ t .Run (`Trial with "Complete" BatchJob and Available metrics.` , func (t * testing.T ) {
243240 g := gomega .NewGomegaWithT (t )
244241 gomock .InOrder (
245242 mockManagerClient .EXPECT ().GetTrialObservationLog (gomock .Any ()).Return (observationLogAvailable , nil ).MinTimes (1 ),
@@ -262,8 +259,9 @@ func TestReconcileBatchJob(t *testing.T) {
262259 }
263260 g .Expect (c .Status ().Update (ctx , batchJob )).NotTo (gomega .HaveOccurred ())
264261
265- // Create the Trial
266- trial := newFakeTrialBatchJob ()
262+ // Create the Trial with StdOut MC
263+ trial := newFakeTrialBatchJob (commonv1beta1 .StdOutCollector , "test-available-stdout" )
264+ trialKey := types.NamespacedName {Name : "test-available-stdout" , Namespace : namespace }
267265 g .Expect (c .Create (ctx , trial )).NotTo (gomega .HaveOccurred ())
268266
269267 // Expect that Trial status is succeeded and metrics are properly populated
@@ -290,28 +288,71 @@ func TestReconcileBatchJob(t *testing.T) {
290288 }, timeout ).Should (gomega .BeTrue ())
291289 })
292290
293- t .Run (`Trail with "Complete" BatchJob and Unavailable metrics.` , func (t * testing.T ) {
291+ t .Run (`Trial with "Complete" BatchJob and Unavailable metrics(StdOut MC) .` , func (t * testing.T ) {
294292 g := gomega .NewGomegaWithT (t )
295293 gomock .InOrder (
296294 mockManagerClient .EXPECT ().GetTrialObservationLog (gomock .Any ()).Return (observationLogUnavailable , nil ).MinTimes (1 ),
297295 mockManagerClient .EXPECT ().DeleteTrialObservationLog (gomock .Any ()).Return (nil , nil ),
298296 )
299- // Create the Trial
300- trial := newFakeTrialBatchJob ()
297+ // Create the Trial with StdOut MC
298+ trial := newFakeTrialBatchJob (commonv1beta1 .StdOutCollector , "test-unavailable-stdout" )
299+ trialKey := types.NamespacedName {Name : "test-unavailable-stdout" , Namespace : namespace }
301300 g .Expect (c .Create (ctx , trial )).NotTo (gomega .HaveOccurred ())
302301
303302 // Expect that Trial status is succeeded with "false" status and "metrics unavailable" reason.
304303 // Metrics unavailable because GetTrialObservationLog returns "unavailable".
304+ g .Eventually (func (g gomega.Gomega ) {
305+ g .Expect (c .Get (ctx , trialKey , trial )).Should (gomega .Succeed ())
306+ g .Expect (trial .IsMetricsUnavailable ()).Should (gomega .BeTrue ())
307+ g .Expect (trial .Status .Observation .Metrics ).ShouldNot (gomega .HaveLen (0 ))
308+ g .Expect (trial .Status .Observation .Metrics [0 ]).Should (gomega .BeComparableTo (commonv1beta1.Metric {
309+ Name : objectiveMetric ,
310+ Min : consts .UnavailableMetricValue ,
311+ Max : consts .UnavailableMetricValue ,
312+ Latest : consts .UnavailableMetricValue ,
313+ }))
314+ }, timeout ).Should (gomega .Succeed ())
315+
316+ // Delete the Trial
317+ g .Expect (c .Delete (ctx , trial )).NotTo (gomega .HaveOccurred ())
318+
319+ // Expect that Trial is deleted
305320 g .Eventually (func () bool {
306- if err = c .Get (ctx , trialKey , trial ); err != nil {
307- return false
308- }
309- return trial .IsMetricsUnavailable () &&
310- len (trial .Status .Observation .Metrics ) > 0 &&
311- trial .Status .Observation .Metrics [0 ].Min == consts .UnavailableMetricValue &&
312- trial .Status .Observation .Metrics [0 ].Max == consts .UnavailableMetricValue &&
313- trial .Status .Observation .Metrics [0 ].Latest == consts .UnavailableMetricValue
321+ return errors .IsNotFound (c .Get (ctx , trialKey , & trialsv1beta1.Trial {}))
314322 }, timeout ).Should (gomega .BeTrue ())
323+ })
324+
325+ t .Run (`Trial with "Complete" BatchJob and Unavailable metrics(Push MC, failed once).` , func (t * testing.T ) {
326+ mockCtrl .Finish ()
327+ g := gomega .NewGomegaWithT (t )
328+ gomock .InOrder (
329+ mockManagerClient .EXPECT ().GetTrialObservationLog (gomock .Any ()).Return (observationLogUnavailable , nil ),
330+ mockManagerClient .EXPECT ().ReportTrialObservationLog (gomock .Any (), gomock .Any ()).Return (nil , errReportMetricsFailed ),
331+ mockManagerClient .EXPECT ().GetTrialObservationLog (gomock .Any ()).Return (observationLogUnavailable , nil ),
332+ mockManagerClient .EXPECT ().ReportTrialObservationLog (gomock .Any (), gomock .Any ()).Return (nil , nil ),
333+ mockManagerClient .EXPECT ().DeleteTrialObservationLog (gomock .Any ()).Return (nil , nil ),
334+ )
335+ mockManagerClient .EXPECT ().GetTrialObservationLog (gomock .Any ()).Return (observationLogUnavailable , nil ).AnyTimes ()
336+ mockManagerClient .EXPECT ().ReportTrialObservationLog (gomock .Any (), gomock .Any ()).Return (nil , nil ).AnyTimes ()
337+
338+ // Create the Trial with Push MC
339+ trial := newFakeTrialBatchJob (commonv1beta1 .PushCollector , "test-unavailable-push-failed-once" )
340+ trialKey := types.NamespacedName {Name : "test-unavailable-push-failed-once" , Namespace : namespace }
341+ g .Expect (c .Create (ctx , trial )).NotTo (gomega .HaveOccurred ())
342+
343+ // Expect that Trial status is succeeded with "false" status and "metrics unavailable" reason.
344+ // Metrics unavailable because GetTrialObservationLog returns "unavailable".
345+ g .Eventually (func (g gomega.Gomega ) {
346+ g .Expect (c .Get (ctx , trialKey , trial )).Should (gomega .Succeed ())
347+ g .Expect (trial .IsMetricsUnavailable ()).Should (gomega .BeTrue ())
348+ g .Expect (trial .Status .Observation .Metrics ).ShouldNot (gomega .HaveLen (0 ))
349+ g .Expect (trial .Status .Observation .Metrics [0 ]).Should (gomega .BeComparableTo (commonv1beta1.Metric {
350+ Name : objectiveMetric ,
351+ Min : consts .UnavailableMetricValue ,
352+ Max : consts .UnavailableMetricValue ,
353+ Latest : consts .UnavailableMetricValue ,
354+ }))
355+ }, timeout ).Should (gomega .Succeed ())
315356
316357 // Delete the Trial
317358 g .Expect (c .Delete (ctx , trial )).NotTo (gomega .HaveOccurred ())
@@ -386,7 +427,7 @@ func TestGetObjectiveMetricValue(t *testing.T) {
386427 g .Expect (err ).To (gomega .HaveOccurred ())
387428}
388429
389- func newFakeTrialBatchJob () * trialsv1beta1.Trial {
430+ func newFakeTrialBatchJob (mcType commonv1beta1. CollectorKind , trialName string ) * trialsv1beta1.Trial {
390431 primaryContainer := "training-container"
391432
392433 job := & batchv1.Job {
@@ -429,8 +470,13 @@ func newFakeTrialBatchJob() *trialsv1beta1.Trial {
429470 },
430471 Spec : trialsv1beta1.TrialSpec {
431472 PrimaryContainerName : primaryContainer ,
432- SuccessCondition : experimentsv1beta1 .DefaultJobSuccessCondition ,
433- FailureCondition : experimentsv1beta1 .DefaultJobFailureCondition ,
473+ MetricsCollector : commonv1beta1.MetricsCollectorSpec {
474+ Collector : & commonv1beta1.CollectorSpec {
475+ Kind : mcType ,
476+ },
477+ },
478+ SuccessCondition : experimentsv1beta1 .DefaultJobSuccessCondition ,
479+ FailureCondition : experimentsv1beta1 .DefaultJobFailureCondition ,
434480 Objective : & commonv1beta1.ObjectiveSpec {
435481 ObjectiveMetricName : objectiveMetric ,
436482 MetricStrategies : []commonv1beta1.MetricStrategy {
0 commit comments