Sedna joint inference and federated learning controller optimization-v1

KubeEdge-Sedna

Code repo: https://github.com/kubeedge/sedna

Related PRs:

• initial proposal
• updated proposal
• JointInferenceService controller enhancement
• FederatedLearning controller enhancement

I'd always wanted to wait until the controller part was fully done before writing it up, but I found that later improvements didn't follow the same thinking as the first attempt, so I decided to record the methods by version. v1 is mostly Open Source Promotion Plan (OSPP) work. It will differ from the proposal. As I learn more about Kubernetes, KubeEdge, and Sedna, I hope I can go deeper.

What we needed to fix

• Joint inference and federated learning couldn't cascade-delete properly. On kubectl delete JointInferenceService/FederatedLearningJob **, child resources were not deleted with them.
• On kubectl edit FederatedLearningJob/JointInferenceService **, updating the custom resource didn't update the child pods it managed.
• When a pod was deleted by hand or by mistake, we wanted it recreated.

Cascade delete

Cascade delete in Kubernetes

Owner References give the control plane the association between objects. Through Owner References, Kubernetes gives the control plane and other API clients a way to clean up related resources when something is deleted. In most cases Kubernetes manages Owner References automatically. The garbage collector is what actually does cascade delete.

Kubernetes logic: when a resource is deleted, if other resources' Metadata ownerReference points at it, those resources get cascade-deleted. That behavior is configurable, and defaults to true.

Every resource's metadata has an ownerReferences field, an array of that resource's owners. Whenever an owner is deleted, it's removed from the array. When all owners are gone, GC collects the resource.

Knowing this is enough to fix the problem we had. Kubernetes itself manages cascade delete; if you set the owner reference correctly, cascade delete works. Since it wasn't working, the owner reference wasn't set correctly. Then you just look for the code that creates the job and the pod.

Owner-reference relations for joint inference and federated learning

JointInferenceService

FederatedLearningJob

How the joint-inference controller sets owner reference

Take the joint-inference controller as an example, and walk through how OwnerReference is set.

The whole definition path for owner reference is fine. The problem is actually the variable we pass in.

In pkg/globalmanager/controllers/jointinference/jointinferenceservice.go, we first define the controller name and the CR's kind name. What we pass into OwnerReference should be the kind name, not the controller name.

// Name is this controller name
Name = "JointInference"

// KindName is the kind name of CR this controller controls
KindName = "JointInferenceService"

Kind name of the CR is defined here. Note that the old code passed in Name (JointInference); it should pass Kind Name (JointInferenceService).

// Kind contains the schema.GroupVersionKind for this controller type.
var Kind = sednav1.SchemeGroupVersion.WithKind(KindName)

run starts the controller. First it sets the worker count, then starts worker threads from that count, constantly taking work off the queue and handling it.

for i := 0; i < workers; i++ { go wait.Until(c.worker, time.Second, stopCh) }

In this part of run, wait.Until starts a goroutine that keeps calling c.worker, one second apart, until stopCh closes.

// worker runs a worker thread that just dequeues items, processes them, and marks them done.
// It enforces that the sync is never invoked concurrently with the same key.
func (c *Controller) worker() {
    for c.processNextWorkItem() {
    }
}

worker calls processNextWorkItem(), which takes a task from the queue, handles it, and calls sync for the actual sync, until the queue is closed.

ns, name, err := cache.SplitMetaNamespaceKey(key)
if err != nil {
    return false, err
}
if len(ns) == 0 || len(name) == 0 {
    return false, fmt.Errorf("invalid jointinference service key %q: either namespace or name is missing", key)
}

In sync, cache.SplitMetaNamespaceKey parses the key into namespace and name.

sharedService, err := c.serviceLister.JointInferenceServices(ns).Get(name)
if err != nil {
    if errors.IsNotFound(err) {
        klog.V(4).Infof("JointInferenceService has been deleted: %v", key)
        return true, nil
    }
    return false, err
}
service := *sharedService

Get the JointInferenceService from the lister.

service.SetGroupVersionKind(Kind)

Set GroupVersionKind.

selector, _ := runtime.GenerateSelector(&service)
pods, err := c.podStore.Pods(service.Namespace).List(selector)
if err != nil {
    return false, err
}
klog.V(4).Infof("list jointinference service %v/%v, %v pods: %v", service.Namespace, service.Name, len(pods), pods)

Generate a selector and list related pods.

When there is no failed worker, if related pods is 0, call createWorkers to create pods.

else {
        if len(pods) == 0 {
            active, manageServiceErr = c.createWorkers(&service)
        }

createWorkers calls createCloudWorker and createEdgeWorker to create cloud and edge worker pods.

In those two functions, runtime.CreatePodWithTemplate creates the pod. Inside, k8scontroller.GetPodFromTemplate sets OwnerReference. What actually does the work is:

if controllerRef != nil {
    pod.OwnerReferences = append(pod.OwnerReferences, *controllerRef)
}

Pod recreation

Automatic pod restart after a fault is decided by RestartPolicy.

In JointInferenceService, RestartPolicy is not set, so it defaults to Always. During a joint-inference job, if something's wrong — e.g. EdgeMesh isn't configured, so the edge can't reach cloud port 5000 for the larger cloud model — the edge pod keeps restarting. In FederatedLearningJob, RestartPolicy is set to OnFailure.

To recreate a pod after it's deleted, you first need the Kubernetes informer mechanism.

k8s informer

Kubernetes uses Informer instead of having the Controller hit the API Server directly. All Controller ops go through Informer, and Informer doesn't hit the API Server every time. Informer uses ListAndWatch: on first start it calls LIST to get the latest resource objects, then WATCH for changes, and keeps events in a read-only cache queue. That speeds up queries and lowers API Server load.

What the pieces in the diagram do:

• Controller: the Informer's carrier. It can create a reflector and drive processLoop. processLoop pops data from DeltaFIFO, first calls Indexer to cache and index, then hands it to the processor.
• Reflector: Informer doesn't talk to k8s-api-server directly; a Reflector does. Reflector ListAndWatches the specified Kubernetes resource. When it changes, e.g. Added, the object is put in the local DeltaFIFO cache.
• DeltaFIFO: a FIFO cache of Watch API events — Added, Updated, Deleted.
• LocalStore: the informer's cache. It caches apiserver objects (some may still be in DeltaFIFO). Callers query the cache instead of apiserver, which takes pressure off apiserver. LocalStore is only accessed by Lister List/Get.
• WorkQueue: after DeltaFIFO gets an event it stores it, then mutates Store, then pops the event onto WorkQueue. The Controller sees WorkQueue events and fires the matching callbacks.

Informer flow

Informer first list/watches apiserver. The Reflector package Informer uses is what connects to apiserver. Reflector ListAndWatch: first list all instances of the resource from apiserver; list gives the latest resourceVersion; then watch all changes after that resourceVersion. If something breaks mid-way, reflector retries watch from the disconnected resourceVersion. Once an instance is created, deleted, or updated, Reflector gets an "event notice". That event plus its API object is a Delta, and it goes into DeltaFIFO.

• Informer keeps reading deltas from DeltaFIFO. For each object, Informer looks at the event type, then creates or updates the local cache (store).
• If the event is Added, Informer uses Indexer to save the API object from the delta into the local cache and index it. On delete, it removes the object from the local cache.
• DeltaFIFO then pops the event to the controller, which calls the ResourceEventHandler that was registered.
• In the ResourceEventHandler, you usually just filter a bit, then put the Object you care about into the workqueue.
• The Controller takes the Object from the workqueue, starts a worker to run its business logic. That logic is usually: how far is the cluster from what the user wants, then tell apiserver to move toward that — e.g. create new pods for a deployment, or scale it.
• In the worker you can use lister to get the resource, without hitting apiserver all the time, because apiserver resource changes show up in the local cache.

Informer has three ResourceEventHandler funcs:

// ResourceEventHandlerFuncs is an adaptor to let you easily specify as many or
// as few of the notification functions as you want while still implementing
// ResourceEventHandler.
type ResourceEventHandlerFuncs struct {
    AddFunc    func(obj interface{})
    UpdateFunc func(oldObj, newObj interface{})
    DeleteFunc func(obj interface{})
}

The handling logic of these three is user-defined. After you register ResourceEventHandler when initializing the controller, create/delete/update on an instance of that object fires the matching handler.

Informer flow in the joint-inference and federated-learning controllers

Take jointinferenceservice.go. New() creates a new JointInferenceService controller so related pods stay in sync with the matching JointInferenceService. In New(), informers are initialized.

podInformer := cc.KubeInformerFactory.Core().V1().Pods()

serviceInformer := cc.SednaInformerFactory.Sedna().V1alpha1().JointInferenceServices()

The service informer uses a custom handler:

serviceInformer.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{
    AddFunc: func(obj interface{}) {
        jc.enqueueController(obj, true)
        jc.syncToEdge(watch.Added, obj)
    },
    UpdateFunc: func(old, cur interface{}) {
        jc.enqueueController(cur, true)
        jc.syncToEdge(watch.Added, cur)
    },

    DeleteFunc: func(obj interface{}) {
        jc.enqueueController(obj, true)
        jc.syncToEdge(watch.Deleted, obj)
    },
})

The pod informer uses a custom handler:

podInformer.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{
    AddFunc:    jc.addPod,
    UpdateFunc: jc.updatePod,
    DeleteFunc: jc.deletePod,
})

These EventHandlers (addPod, updatePod, deletePod) really just enqueue the related object; they don't do other work.

podInformer.Lister() creates a Lister for Pods. podInformer.Informer().HasSynced checks whether the Informer cache has synced.

jc.serviceLister = serviceInformer.Lister()
jc.serviceStoreSynced = serviceInformer.Informer().HasSynced
//...
jc.podStore = podInformer.Lister()
jc.podStoreSynced = podInformer.Informer().HasSynced

Syncing from the api server happens in Run() for the jointinferenceservice controller. Run() starts the main goroutine that watches and syncs.

if !cache.WaitForNamedCacheSync(Name, stopCh, c.podStoreSynced, c.serviceStoreSynced) {
    klog.Errorf("failed to wait for %s caches to sync", Name)
    return

}

After starting the Informer, wait for local cache sync, then start workers. When a change event arrives, take the Object from the event, build an object key (namespace/name), put the key in the workerqueue.

worker() calls c.processNextWorkItem().

func (c *Controller) worker() {
    for c.processNextWorkItem() {
    }
}

processNextWorkItem takes a key from the workerqueue, calls sync(). In sync(), the lister gets the real object from the local cache, and the sync work runs.

Federated-learning pod recreation design

Watch delete events. When Informer sees a delete whose OwnerReference is FederatedLearning, start DeletePod, and DeletePod recreates the pod. The new Pod is almost the same as the old one — config and spec kept, identifiers like resourceVersion and UID reset so they can be regenerated.

Code logic

• Checks before recreate
  • Check the pod is owned by a FederatedLearningJob.
  • Check it hasn't already been recreated: c.recreatedPods.Load(pod.Name). If it has, don't create again.

// first check if the pod is owned by a FederatedLearningJob
controllerRef := metav1.GetControllerOf(pod)
if controllerRef == nil || controllerRef.Kind != Kind.Kind {
    return
}

• Recreate the Pod
  • pod.DeepCopy() for a deep copy.
  • Reset unique ids (ResourceVersion, UID) and status fields.
  • c.kubeClient.CoreV1().Pods(pod.Namespace).Create to create the new Pod via the Kubernetes API.
  • On success, log, mark the Pod as recreated.

// Create a deep copy of the old pod
newPod := pod.DeepCopy()
// Reset the resource version and UID as they are unique to each object
newPod.ResourceVersion = ""
newPod.UID = ""
// Clear the status
newPod.Status = v1.PodStatus{}
// Remove the deletion timestamp
newPod.DeletionTimestamp = nil
// Remove the deletion grace period seconds
newPod.DeletionGracePeriodSeconds = nil
_, err := c.kubeClient.CoreV1().Pods(pod.Namespace).Create(context.TODO(), newPod, metav1.CreateOptions{})
if err != nil {
    return
}

• Mark recreated and a timer to clear the record:
  • c.recreatedPods.Store(pod.Name, true) marks this Pod recreated.
  • A timer, 5 seconds later, clears the mark (c.recreatedPods.Delete(pod.Name)), so if the Pod is deleted again later, recreate can fire again.

// mark the pod as recreated
c.recreatedPods.Store(newPod.Name, true)
// set a timer to delete the record from the map after a while
go func() {
    time.Sleep(5 * time.Second)
    c.recreatedPods.Delete(pod.Name)
}()

We add a sync.Map recreatedPods on the Controller struct, to avoid creating the same pod twice in one delete event. After a manual delete and a successful create, the pod name goes into recreatedPods. When deletePod is called again for the same delete event, the name is already there, so the manually deleted pod isn't deleted and created again. A timer clears the mark after 5s, so a later manual delete can trigger recreate again.

type Controller struct{
//...
preventRecreation bool
//...
}

Joint-inference pod recreation design

Inference itself is a stateless workload, so we can use the native k8s deployment for pod self-heal. We need a full watch-and-handle path for resource changes.

• Get a deployment informer from the informer factory.
• Register event handlers on the informer (addDeployment, updateDeployment, deleteDeployment).
• Start the informer and sync with the k8s api.
• Before handling events, wait for the local cache to sync with API Server.
• When a Deployment in the cluster changes, Informer fires the matching handler.

Updating pods when the federated-learning CRD changes

Flow for the federated-learning controller updating pods on CRD change:

Watch update events. When Informer sees an update on FederatedLearningJob, if the CRD changed, delete the old pods and create pods from the new CRD params.

Code logic

Start updateJob. updateJob first decides whether an update is needed.

• If old and cur can't convert to sednav1.FederatedLearningJob, return.
• If oldJob and curJob have the same ResourceVersion, nothing changed; return, no update to handle.
• Set preventRecreation to true, so we don't recreate Pods in the middle of the update.

Then compare oldJob and curJob params.

• Compare old.Generation and cur.Generation. CRDs have a Generation field. It's auto-generated: every create/modify of the crd object changes it. Create starts at 1, each later modify +1. Only spec changes bump Generation; status changes don't. So you can use it to tell spec vs status. If they differ, FederatedLearningJob params changed.
• Walk the Pod list and delete each Pod.
• Recreate AggWorker and TrainWorker from the updated curJob.Spec.
• Reset preventRecreation to false so later pod self-heal still works.

Updating pods when the joint-inference CRD changes

Flow for the joint-inference controller updating pods on CRD change:

Watch update events. When Informer sees an update on JointInferenceService, if the CRD changed, delete the old pods and create pods from the new CRD params.

Code logic

On CRD-change pod update, joint inference is basically the same as federated learning.

Start updateService, compare old.Generation and cur.Generation. If they differ, JointInferenceService params changed.

• Walk the Pod list, delete each Pod.
• Recreate cloudWorker and edgeWorker from the updated curService.Spec.

Tests

Federated-learning unit tests

Unit tests, not e2e. Focus on the two changed functions, deletePod() and updateJob().

Test_deletePod()

• fake.NewSimpleClientset() for a fake k8s client.
• Create a test pod via fakeclient.
• Create a controller, register fakeclient.
• Pass the test pod into controller.deletePod(), then fakeClient.CoreV1().Pods("default").Get(context.TODO(), "test-pod", metav1.GetOptions{}) to check the pod was recreated. If not, fail.

• Create a fake client.
• Create a controller.
• Call controller.deletePod() on a missing pod.
• Confirm whether an error happens.

Test_updateJob()

• Mock the pod list.
• Create a fake client.
• Create dataset, model, and related job and pod resources.
• Init the controller with the fake client, the test job, the mock pod list, and an event broadcaster and other deps.
• Define a new job, update some params (change TrainingWorker batch_size from 32 to 16).
• Call updateJob to update the old job to the new one, simulating a real job update.
• Check the result: if updated params match expectations, pass.

Joint-inference unit tests

Test_UpdateService()

• Create fake Sedna and Kubernetes clients.
• Create an old service with cloud-worker and edge-worker config. Create two model resources.
• From the old service, create deployment and pod resources.
• Init controller, set fake clients, pod list, deployment list.
• Controller sendToEdgeFunc is set to an empty func (no real edge communication).
• Copy the old joint-inference service, change the hard-example-mining param on the edge worker from value1 to value2, and bump Generation.
• Call updateService to trigger the update.
• The test checks that the updated joint-inference service can be fetched via the fake client.
• Check the updated HEM param went from value1 to value2, so the service-update logic ran correctly.