Skip to content

Commit 3ff31f3

Browse files
tkashemtkashem
committed
dra e2e test
1 parent 8ec079c commit 3ff31f3

File tree

9 files changed

+577
-1
lines changed

9 files changed

+577
-1
lines changed

‎pkg/drae2e/dra.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
package drae2e
2+
3+
import (
4+
"context"
5+
)
6+
7+
type Driver interface {
8+
DeviceClassName() string
9+
Setup() error
10+
Cleanup(context.Context) error
11+
Ready() error
12+
}

‎pkg/drae2e/drivertoolkit.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package drae2e
2+
3+
import (
4+
corev1 "k8s.io/api/core/v1"
5+
clientset "k8s.io/client-go/kubernetes"
6+
)
7+
8+
type DriverToolkitProbe struct {
9+
node *corev1.Node
10+
clientset clientset.Interface
11+
}
12+
13+
func (dtk DriverToolkitProbe) Verify() {
14+
15+
}

‎pkg/drae2e/example/example_driver.go

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package example
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"testing"
7+
8+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
9+
"k8s.io/client-go/kubernetes"
10+
11+
"github.com/openshift/origin/pkg/drae2e"
12+
)
13+
14+
func NewDriver(t testing.TB, clientset kubernetes.Interface, p drae2e.HelmParameters) *Driver {
15+
return &Driver{
16+
t: t,
17+
clientset: clientset,
18+
helm: drae2e.NewHelmInstaller(t, p),
19+
namespace: p.Namespace,
20+
}
21+
}
22+
23+
type Driver struct {
24+
t testing.TB
25+
helm *drae2e.Helm
26+
clientset kubernetes.Interface
27+
namespace string
28+
}
29+
30+
func (d *Driver) DeviceClassName() string { return "gpu.example.com" }
31+
func (d *Driver) Setup() error { return d.helm.Install() }
32+
func (d *Driver) Cleanup(ctx context.Context) error { return d.helm.Remove() }
33+
34+
func (d Driver) Ready() error {
35+
client := d.clientset.AppsV1()
36+
ds, err := client.DaemonSets(d.namespace).Get(context.Background(), "dra-example-driver-kubeletplugin", metav1.GetOptions{})
37+
if err != nil {
38+
return fmt.Errorf("failed to get dra-example-driver-kubeletplugin: %w", err)
39+
}
40+
41+
ready, scheduled := ds.Status.NumberReady, ds.Status.DesiredNumberScheduled
42+
if ready != scheduled {
43+
return fmt.Errorf("dra-example-driver-kubeletplugin is not ready, DesiredNumberScheduled: %d, NumberReady: %d", scheduled, ready)
44+
}
45+
return nil
46+
}

‎pkg/drae2e/helm.go

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package drae2e
2+
3+
import (
4+
"fmt"
5+
"os"
6+
"time"
7+
8+
"helm.sh/helm/v3/pkg/action"
9+
"helm.sh/helm/v3/pkg/chart/loader"
10+
"helm.sh/helm/v3/pkg/cli"
11+
"helm.sh/helm/v3/pkg/registry"
12+
)
13+
14+
func NewHelmInstaller(logger logger, p HelmParameters) *Helm {
15+
return &Helm{HelmParameters: p, logger: logger}
16+
}
17+
18+
type HelmParameters struct {
19+
Namespace string
20+
CreateNamespace bool
21+
ChartURL string
22+
ReleaseName string
23+
ChartVersion string
24+
Values map[string]interface{}
25+
}
26+
27+
type logger interface {
28+
Log(args ...any)
29+
Logf(format string, args ...any)
30+
}
31+
32+
type Helm struct {
33+
HelmParameters
34+
logger logger
35+
36+
// save it for uninstal
37+
config *action.Configuration
38+
}
39+
40+
func (h *Helm) Install() error {
41+
os.Setenv("HELM_EXPERIMENTAL_OCI", "1")
42+
settings := cli.New()
43+
44+
registryClient, err := registry.NewClient(
45+
registry.ClientOptDebug(settings.Debug),
46+
registry.ClientOptWriter(os.Stdout),
47+
registry.ClientOptCredentialsFile(settings.RegistryConfig),
48+
)
49+
if err != nil {
50+
return fmt.Errorf("failed to create registry client: %w", err)
51+
}
52+
53+
h.config = &action.Configuration{}
54+
h.config.RegistryClient = registryClient
55+
if err := h.config.Init(settings.RESTClientGetter(), h.Namespace, os.Getenv("HELM_DRIVER"), h.logger.Logf); err != nil {
56+
return fmt.Errorf("failed to initialize Helm action configuration: %w", err)
57+
}
58+
59+
install := action.NewInstall(h.config)
60+
install.ChartPathOptions.Version = h.ChartVersion
61+
localChartPath, err := install.ChartPathOptions.LocateChart(h.ChartURL, settings)
62+
if err != nil {
63+
return fmt.Errorf("failed to locate chart: %w", err)
64+
}
65+
h.logger.Logf("Downloaded chart to: %s", localChartPath)
66+
67+
chart, err := loader.Load(localChartPath)
68+
if err != nil {
69+
return fmt.Errorf("failed to load chart: %w", err)
70+
}
71+
h.logger.Logf("loaded chart: %s", h.ReleaseName)
72+
73+
history := action.NewHistory(h.config)
74+
history.Max = 1
75+
if releases, err := history.Run(h.ReleaseName); err != nil || len(releases) <= 0 {
76+
client := action.NewInstall(h.config)
77+
client.Namespace = h.Namespace
78+
client.ReleaseName = h.ReleaseName
79+
client.CreateNamespace = h.CreateNamespace
80+
client.Wait = true
81+
client.Timeout = 5 * time.Minute
82+
if _, err := client.Run(chart, h.Values); err != nil {
83+
return fmt.Errorf("failed to install chart: %w", err)
84+
}
85+
return nil
86+
}
87+
88+
client := action.NewUpgrade(h.config)
89+
client.Namespace = h.Namespace
90+
client.Wait = true
91+
client.Timeout = 5 * time.Minute
92+
if _, err := client.Run(h.ReleaseName, chart, h.Values); err != nil {
93+
return fmt.Errorf("failed to upgrade chart: %w", err)
94+
}
95+
return nil
96+
}
97+
98+
func (h Helm) Remove() error {
99+
client := action.NewUninstall(h.config)
100+
client.KeepHistory = false
101+
if _, err := client.Run(h.ReleaseName); err != nil {
102+
return fmt.Errorf("failed to uninstall %s: %w", h.ReleaseName, err)
103+
}
104+
return nil
105+
}

‎pkg/drae2e/helper.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package drae2e
2+
3+
import (
4+
corev1 "k8s.io/api/core/v1"
5+
resourceapi "k8s.io/api/resource/v1beta1"
6+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
7+
e2epodutil "k8s.io/kubernetes/test/e2e/framework/pod"
8+
"k8s.io/utils/ptr"
9+
)
10+
11+
func NewHelper(namespace, deviceclass string) Helper {
12+
return Helper{namespace: namespace, deviceclass: deviceclass}
13+
}
14+
15+
type Helper struct {
16+
namespace string
17+
deviceclass string
18+
}
19+
20+
func (h Helper) NewResourceClaimTemplate(name string) *resourceapi.ResourceClaimTemplate {
21+
return &resourceapi.ResourceClaimTemplate{
22+
ObjectMeta: metav1.ObjectMeta{
23+
Namespace: h.namespace,
24+
Name: name,
25+
},
26+
Spec: resourceapi.ResourceClaimTemplateSpec{
27+
Spec: resourceapi.ResourceClaimSpec{
28+
Devices: resourceapi.DeviceClaim{
29+
Requests: []resourceapi.DeviceRequest{
30+
{
31+
Name: "gpu",
32+
DeviceClassName: h.deviceclass,
33+
},
34+
},
35+
},
36+
},
37+
},
38+
}
39+
}
40+
41+
func (h Helper) NewPod(name string) *corev1.Pod {
42+
pod := &corev1.Pod{
43+
TypeMeta: metav1.TypeMeta{
44+
Kind: "Pod",
45+
APIVersion: "v1",
46+
},
47+
ObjectMeta: metav1.ObjectMeta{
48+
Name: name,
49+
Namespace: h.namespace,
50+
Labels: make(map[string]string),
51+
Annotations: make(map[string]string),
52+
},
53+
Spec: corev1.PodSpec{
54+
RestartPolicy: corev1.RestartPolicyNever,
55+
TerminationGracePeriodSeconds: ptr.To[int64](1),
56+
SecurityContext: e2epodutil.GetRestrictedPodSecurityContext(),
57+
},
58+
}
59+
return pod
60+
}
61+
62+
func NewContainer(name string) corev1.Container {
63+
return corev1.Container{
64+
Name: name,
65+
Image: e2epodutil.GetDefaultTestImage(),
66+
Command: e2epodutil.GenerateScriptCmd("env && sleep 100000"),
67+
SecurityContext: e2epodutil.GetRestrictedContainerSecurityContext(),
68+
}
69+
}
70+
71+
func NewNvidiaSMIContainer(name string) corev1.Container {
72+
return corev1.Container{
73+
Name: name,
74+
Image: "ubuntu:22.04",
75+
Command: []string{"bash", "-c"},
76+
Args: []string{"nvidia-smi -L; trap 'exit 0' TERM; sleep 9999 & wait"},
77+
}
78+
}

‎pkg/drae2e/nvidia/nvidia.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
package nvidia
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"k8s.io/client-go/kubernetes"
8+
9+
"github.com/openshift/origin/pkg/drae2e"
10+
)
11+
12+
func NewDriver(t testing.TB, clientset kubernetes.Interface, p drae2e.HelmParameters) *Driver {
13+
return &Driver{
14+
t: t,
15+
namespace: p.Namespace,
16+
clientset: clientset,
17+
helm: drae2e.NewHelmInstaller(t, p),
18+
}
19+
}
20+
21+
type Driver struct {
22+
t testing.TB
23+
helm *drae2e.Helm
24+
clientset kubernetes.Interface
25+
namespace string
26+
}
27+
28+
func (d *Driver) DeviceClassName() string { return "gpu.nvidia.com" }
29+
func (d *Driver) Setup() error { return d.helm.Install() }
30+
func (d *Driver) Cleanup(ctx context.Context) error { return d.helm.Remove() }
31+
func (d Driver) Ready() error { return nil }

0 commit comments

Comments
 (0)