Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

replace wait.WithTimeout(timeout) with wait.WithContext(ctx)) #463

Merged
merged 1 commit into from
Jul 31, 2024
Merged

replace wait.WithTimeout(timeout) with wait.WithContext(ctx)) #463

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion e2e2/test/cases/neuron/main_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,9 @@ func TestMain(m *testing.M) {
log.Fatalf("failed to initialize test environment: %v", err)
}
testenv = env.NewWithConfig(cfg)
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Minute)
defer cancel()
testenv = testenv.WithContext(ctx)

manifests := [][]byte{
neuronDevicePluginManifest,
Expand All @@ -56,7 +59,7 @@ func TestMain(m *testing.M) {
ObjectMeta: metav1.ObjectMeta{Name: "neuron-device-plugin-daemonset", Namespace: "kube-system"},
}
err := wait.For(fwext.NewConditionExtension(config.Client().Resources()).DaemonSetReady(&ds),
wait.WithTimeout(time.Minute*5))
wait.WithContext(ctx))
if err != nil {
return ctx, err
}
Expand Down
3 changes: 1 addition & 2 deletions e2e2/test/cases/neuron/neuron_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
_ "embed"
"fmt"
"testing"
"time"

fwext "github.com/aws/aws-k8s-tester/e2e2/internal/framework_extensions"
"sigs.k8s.io/e2e-framework/klient/wait"
Expand Down Expand Up @@ -52,7 +51,7 @@ func TestMPIJobPytorchTraining(t *testing.T) {
ObjectMeta: metav1.ObjectMeta{Name: "neuronx-single-node", Namespace: "default"},
}
err := wait.For(fwext.NewConditionExtension(cfg.Client().Resources()).JobSucceeded(job),
wait.WithTimeout(time.Minute*20))
wait.WithContext(ctx))
if err != nil {
t.Fatal(err)
}
Expand Down
9 changes: 6 additions & 3 deletions e2e2/test/cases/nvidia/main_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,9 @@ func TestMain(m *testing.M) {
log.Fatalf("failed to initialize test environment: %v", err)
}
testenv = env.NewWithConfig(cfg)
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Minute)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is a timeout for all suites? should we have a timeout per top-level suite like single, multi, etc so that we don't block all the tests if one hangs?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Having a timeout per top-level suite is very similar to the the previous approach,.

// WithContext provides a way to configure a context that can be used to cancel the wait condition checks. This will enable
// end users to write test in cases where the max timeout is not really predictable or is a factor of a different
// configuration or event.

I think WithContext makes more sense, since some test takes longer than others, e.g. nccl test vs unit test

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should use WithContext regardless, but i meant using a different context per suite

defer cancel()
testenv = testenv.WithContext(ctx)

// all NVIDIA tests require the device plugin and MPI operator
manifests := [][]byte{
Expand All @@ -70,7 +73,7 @@ func TestMain(m *testing.M) {
ObjectMeta: metav1.ObjectMeta{Name: "mpi-operator", Namespace: "mpi-operator"},
}
err := wait.For(conditions.New(config.Client().Resources()).DeploymentConditionMatch(&dep, appsv1.DeploymentAvailable, v1.ConditionTrue),
wait.WithTimeout(time.Minute*5))
wait.WithContext(ctx))
if err != nil {
return ctx, err
}
Expand All @@ -81,7 +84,7 @@ func TestMain(m *testing.M) {
ObjectMeta: metav1.ObjectMeta{Name: "nvidia-device-plugin-daemonset", Namespace: "kube-system"},
}
err := wait.For(fwext.NewConditionExtension(config.Client().Resources()).DaemonSetReady(&ds),
wait.WithTimeout(time.Minute*5))
wait.WithContext(ctx))
if err != nil {
return ctx, err
}
Expand All @@ -105,7 +108,7 @@ func TestMain(m *testing.M) {
ObjectMeta: metav1.ObjectMeta{Name: "efa-device-plugin-daemonset", Namespace: "kube-system"},
}
err = wait.For(fwext.NewConditionExtension(cfg.Client().Resources()).DaemonSetReady(&ds),
wait.WithTimeout(time.Minute*5))
wait.WithContext(ctx))
if err != nil {
return ctx, err
}
Expand Down
7 changes: 2 additions & 5 deletions e2e2/test/cases/nvidia/mpi_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
_ "embed"
"fmt"
"testing"
"time"

fwext "github.com/aws/aws-k8s-tester/e2e2/internal/framework_extensions"
kubeflowv2beta1 "github.com/kubeflow/mpi-operator/pkg/apis/kubeflow/v2beta1"
Expand Down Expand Up @@ -55,9 +54,8 @@ func TestMPIJobPytorchTraining(t *testing.T) {
j := kubeflowv2beta1.MPIJob{
ObjectMeta: metav1.ObjectMeta{Name: "pytorch-training-single-node", Namespace: "default"},
}
timeout := time.Minute * 20
err := wait.For(fwext.NewConditionExtension(rsrc).ResourceMatch(&j, mpiJobSucceeded),
wait.WithTimeout(timeout))
wait.WithContext(ctx))
if err != nil {
t.Fatal(err)
}
Expand Down Expand Up @@ -111,9 +109,8 @@ func TestMPIJobPytorchTraining(t *testing.T) {
j := kubeflowv2beta1.MPIJob{
ObjectMeta: metav1.ObjectMeta{Name: "multi-node-nccl-test", Namespace: "default"},
}
timeout := time.Minute * 10
err := wait.For(conditions.New(rsrc).ResourceMatch(&j, mpiJobSucceeded),
wait.WithTimeout(timeout))
wait.WithContext(ctx))
if err != nil {
t.Fatal(err)
}
Expand Down
3 changes: 1 addition & 2 deletions e2e2/test/cases/nvidia/unit_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
_ "embed"
"fmt"
"testing"
"time"

fwext "github.com/aws/aws-k8s-tester/e2e2/internal/framework_extensions"
"sigs.k8s.io/e2e-framework/klient/wait"
Expand Down Expand Up @@ -52,7 +51,7 @@ func TestSingleNodeUnitTest(t *testing.T) {
ObjectMeta: metav1.ObjectMeta{Name: "unit-test-job", Namespace: "default"},
}
err := wait.For(fwext.NewConditionExtension(cfg.Client().Resources()).JobSucceeded(job),
wait.WithTimeout(time.Minute*20))
wait.WithContext(ctx))
if err != nil {
t.Fatal(err)
}
Expand Down
Loading