@@ -19,22 +19,21 @@ package main
1919import (
2020 "context"
2121 "fmt"
22+ "maps"
2223
23- metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
24+ resourceapi "k8s.io/api/resource/v1beta1"
25+ "k8s.io/apimachinery/pkg/types"
2426 coreclientset "k8s.io/client-go/kubernetes"
2527 "k8s.io/dynamic-resource-allocation/kubeletplugin"
28+ "k8s.io/dynamic-resource-allocation/resourceslice"
2629 "k8s.io/klog/v2"
2730
28- drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
29-
3031 "sigs.k8s.io/dra-example-driver/pkg/consts"
3132)
3233
33- var _ drapbv1.DRAPluginServer = & driver {}
34-
3534type driver struct {
3635 client coreclientset.Interface
37- plugin kubeletplugin.DRAPlugin
36+ helper * kubeletplugin.Helper
3837 state * DeviceState
3938}
4039
@@ -49,87 +48,93 @@ func NewDriver(ctx context.Context, config *Config) (*driver, error) {
4948 }
5049 driver .state = state
5150
52- plugin , err := kubeletplugin .Start (
51+ helper , err := kubeletplugin .Start (
5352 ctx ,
54- [] any { driver } ,
53+ driver ,
5554 kubeletplugin .KubeClient (config .coreclient ),
5655 kubeletplugin .NodeName (config .flags .nodeName ),
5756 kubeletplugin .DriverName (consts .DriverName ),
58- kubeletplugin .RegistrarSocketPath (PluginRegistrationPath ),
59- kubeletplugin .PluginSocketPath (DriverPluginSocketPath ),
60- kubeletplugin .KubeletPluginSocketPath (DriverPluginSocketPath ))
57+ )
6158 if err != nil {
6259 return nil , err
6360 }
64- driver .plugin = plugin
61+ driver .helper = helper
6562
66- var resources kubeletplugin.Resources
67- for _ , device := range state .allocatable {
68- resources .Devices = append (resources .Devices , device )
63+ devices := make ([]resourceapi.Device , 0 , len (state .allocatable ))
64+ for device := range maps .Values (state .allocatable ) {
65+ devices = append (devices , device )
66+ }
67+ resources := resourceslice.DriverResources {
68+ Pools : map [string ]resourceslice.Pool {
69+ config .flags .nodeName : {
70+ Slices : []resourceslice.Slice {
71+ {
72+ Devices : devices ,
73+ },
74+ },
75+ },
76+ },
6977 }
7078
71- if err := plugin .PublishResources (ctx , resources ); err != nil {
79+ if err := helper .PublishResources (ctx , resources ); err != nil {
7280 return nil , err
7381 }
7482
7583 return driver , nil
7684}
7785
7886func (d * driver ) Shutdown (ctx context.Context ) error {
79- d .plugin .Stop ()
87+ d .helper .Stop ()
8088 return nil
8189}
8290
83- func (d * driver ) NodePrepareResources (ctx context.Context , req * drapbv1. NodePrepareResourcesRequest ) (* drapbv1. NodePrepareResourcesResponse , error ) {
84- klog .Infof ("NodePrepareResource is called: number of claims: %d" , len (req . Claims ))
85- preparedResources := & drapbv1. NodePrepareResourcesResponse { Claims : map [string ] * drapbv1. NodePrepareResourceResponse {}}
91+ func (d * driver ) PrepareResourceClaims (ctx context.Context , claims [] * resourceapi. ResourceClaim ) (map [types. UID ]kubeletplugin. PrepareResult , error ) {
92+ klog .Infof ("PrepareResourceClaims is called: number of claims: %d" , len (claims ))
93+ result := make ( map [types. UID ]kubeletplugin. PrepareResult )
8694
87- for _ , claim := range req . Claims {
88- preparedResources . Claims [claim .UID ] = d .nodePrepareResource (ctx , claim )
95+ for _ , claim := range claims {
96+ result [claim .UID ] = d .prepareResourceClaim (ctx , claim )
8997 }
9098
91- return preparedResources , nil
99+ return result , nil
92100}
93101
94- func (d * driver ) nodePrepareResource (ctx context.Context , claim * drapbv1.Claim ) * drapbv1.NodePrepareResourceResponse {
95- resourceClaim , err := d .client .ResourceV1beta1 ().ResourceClaims (claim .Namespace ).Get (
96- ctx ,
97- claim .Name ,
98- metav1.GetOptions {})
102+ func (d * driver ) prepareResourceClaim (_ context.Context , claim * resourceapi.ResourceClaim ) kubeletplugin.PrepareResult {
103+ preparedPBs , err := d .state .Prepare (claim )
99104 if err != nil {
100- return & drapbv1. NodePrepareResourceResponse {
101- Error : fmt .Sprintf ( "failed to fetch ResourceClaim %s in namespace %s "claim .Name , claim . Namespace ),
105+ return kubeletplugin. PrepareResult {
106+ Err : fmt .Errorf ( "error preparing devices for claim %v: %w "claim .UID , err ),
102107 }
103108 }
104-
105- prepared , err := d .state .Prepare (resourceClaim )
106- if err != nil {
107- return & drapbv1.NodePrepareResourceResponse {
108- Error : fmt .Sprintf ("error preparing devices for claim %v: %v" , claim .UID , err ),
109- }
109+ var prepared []kubeletplugin.Device
110+ for _ , preparedPB := range preparedPBs {
111+ prepared = append (prepared , kubeletplugin.Device {
112+ Requests : preparedPB .GetRequestNames (),
113+ PoolName : preparedPB .GetPoolName (),
114+ DeviceName : preparedPB .GetDeviceName (),
115+ CDIDeviceIDs : preparedPB .GetCDIDeviceIDs (),
116+ })
110117 }
111118
112119 klog .Infof ("Returning newly prepared devices for claim '%v': %v" , claim .UID , prepared )
113- return & drapbv1. NodePrepareResourceResponse {Devices : prepared }
120+ return kubeletplugin. PrepareResult {Devices : prepared }
114121}
115122
116- func (d * driver ) NodeUnprepareResources (ctx context.Context , req * drapbv1. NodeUnprepareResourcesRequest ) (* drapbv1. NodeUnprepareResourcesResponse , error ) {
117- klog .Infof ("NodeUnPrepareResource is called: number of claims: %d" , len (req . Claims ))
118- unpreparedResources := & drapbv1. NodeUnprepareResourcesResponse { Claims : map [string ] * drapbv1. NodeUnprepareResourceResponse {}}
123+ func (d * driver ) UnprepareResourceClaims (ctx context.Context , claims []kubeletplugin. NamespacedObject ) (map [types. UID ] error , error ) {
124+ klog .Infof ("UnprepareResourceClaims is called: number of claims: %d" , len (claims ))
125+ result := make ( map [types. UID ] error )
119126
120- for _ , claim := range req . Claims {
121- unpreparedResources . Claims [claim .UID ] = d .nodeUnprepareResource (ctx , claim )
127+ for _ , claim := range claims {
128+ result [claim .UID ] = d .unprepareResourceClaim (ctx , claim )
122129 }
123130
124- return unpreparedResources , nil
131+ return result , nil
125132}
126133
127- func (d * driver ) nodeUnprepareResource (ctx context.Context , claim * drapbv1.Claim ) * drapbv1.NodeUnprepareResourceResponse {
128- if err := d .state .Unprepare (claim .UID ); err != nil {
129- return & drapbv1.NodeUnprepareResourceResponse {
130- Error : fmt .Sprintf ("error unpreparing devices for claim %v: %v" , claim .UID , err ),
131- }
134+ func (d * driver ) unprepareResourceClaim (_ context.Context , claim kubeletplugin.NamespacedObject ) error {
135+ if err := d .state .Unprepare (string (claim .UID )); err != nil {
136+ return fmt .Errorf ("error unpreparing devices for claim %v: %w" , claim .UID , err )
132137 }
133138
134- return & drapbv1. NodeUnprepareResourceResponse {}
139+ return nil
135140}
0 commit comments