chore: restruct workspace controller code - part 2 (#683)

This PR mainly does the followings: 1) Move tuning package to pkg/workspace 2) Move machine and nodeclaim packages to pkg/utils 3) Move resources/node* and resources/resources* to pkg/resources. I will further clean up the pkg/resource in next change
kaito-project · Nov 9, 2024 · 3c873ec · 3c873ec
1 parent 2ecfdf1
commit 3c873ec
Show file tree

Hide file tree

Showing 19 changed files with 29 additions and 26 deletions.
diff --git a/cmd/workspace/main.go b/cmd/workspace/main.go
@@ -15,8 +15,8 @@ import (
 	awsv1beta1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1beta1"
 	"github.com/kaito-project/kaito/pkg/featuregates"
 	"github.com/kaito-project/kaito/pkg/k8sclient"
-	"github.com/kaito-project/kaito/pkg/nodeclaim"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/nodeclaim"
 	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 	"sigs.k8s.io/karpenter/pkg/apis/v1beta1"
 

diff --git a/pkg/ragengine/controllers/ragengine_controller.go b/pkg/ragengine/controllers/ragengine_controller.go
@@ -18,11 +18,11 @@ import (
 	"github.com/go-logr/logr"
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	"github.com/kaito-project/kaito/pkg/featuregates"
-	"github.com/kaito-project/kaito/pkg/machine"
-	"github.com/kaito-project/kaito/pkg/nodeclaim"
-	"github.com/kaito-project/kaito/pkg/resources"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/machine"
+	"github.com/kaito-project/kaito/pkg/utils/nodeclaim"
+	"github.com/kaito-project/kaito/pkg/utils/resources"
 	"github.com/samber/lo"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"

diff --git a/pkg/ragengine/controllers/ragengine_gc_finalizer.go b/pkg/ragengine/controllers/ragengine_gc_finalizer.go
@@ -8,9 +8,9 @@ import (
 
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	"github.com/kaito-project/kaito/pkg/featuregates"
-	"github.com/kaito-project/kaito/pkg/machine"
-	"github.com/kaito-project/kaito/pkg/nodeclaim"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/machine"
+	"github.com/kaito-project/kaito/pkg/utils/nodeclaim"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"

diff --git a/pkg/machine/machine.go → pkg/utils/machine/machine.go b/pkg/machine/machine.go → pkg/utils/machine/machine.go
@@ -12,8 +12,8 @@ import (
 
 	"github.com/aws/karpenter-core/pkg/apis/v1alpha5"
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
-	"github.com/kaito-project/kaito/pkg/resources"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/resources"
 	"github.com/samber/lo"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"

diff --git a/pkg/machine/machine_test.go → pkg/utils/machine/machine_test.go b/pkg/machine/machine_test.go → pkg/utils/machine/machine_test.go
diff --git a/pkg/nodeclaim/nodeclaim.go → pkg/utils/nodeclaim/nodeclaim.go b/pkg/nodeclaim/nodeclaim.go → pkg/utils/nodeclaim/nodeclaim.go
@@ -15,8 +15,8 @@ import (
 	azurev1alpha2 "github.com/Azure/karpenter-provider-azure/pkg/apis/v1alpha2"
 	awsv1beta1 "github.com/aws/karpenter-provider-aws/pkg/apis/v1beta1"
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
-	"github.com/kaito-project/kaito/pkg/resources"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/resources"
 	"github.com/samber/lo"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"

diff --git a/pkg/nodeclaim/nodeclaim_test.go → pkg/utils/nodeclaim/nodeclaim_test.go b/pkg/nodeclaim/nodeclaim_test.go → pkg/utils/nodeclaim/nodeclaim_test.go
diff --git a/pkg/resources/nodes.go → pkg/utils/resources/nodes.go b/pkg/resources/nodes.go → pkg/utils/resources/nodes.go
diff --git a/pkg/resources/nodes_test.go → pkg/utils/resources/nodes_test.go b/pkg/resources/nodes_test.go → pkg/utils/resources/nodes_test.go
diff --git a/pkg/resources/resources.go → pkg/utils/resources/resources.go b/pkg/resources/resources.go → pkg/utils/resources/resources.go
diff --git a/pkg/resources/resources_test.go → pkg/utils/resources/resources_test.go b/pkg/resources/resources_test.go → pkg/utils/resources/resources_test.go
diff --git a/pkg/workspace/controllers/workspace_controller.go b/pkg/workspace/controllers/workspace_controller.go
@@ -17,19 +17,20 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 
 	"github.com/kaito-project/kaito/pkg/featuregates"
-	"github.com/kaito-project/kaito/pkg/nodeclaim"
-	"github.com/kaito-project/kaito/pkg/tuning"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/nodeclaim"
+	"github.com/kaito-project/kaito/pkg/workspace/tuning"
 	batchv1 "k8s.io/api/batch/v1"
 	"sigs.k8s.io/karpenter/pkg/apis/v1beta1"
 
 	"github.com/aws/karpenter-core/pkg/apis/v1alpha5"
 	"github.com/go-logr/logr"
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
-	"github.com/kaito-project/kaito/pkg/machine"
-	"github.com/kaito-project/kaito/pkg/resources"
+	manifests "github.com/kaito-project/kaito/pkg/resources"
 	"github.com/kaito-project/kaito/pkg/utils"
+	"github.com/kaito-project/kaito/pkg/utils/machine"
 	"github.com/kaito-project/kaito/pkg/utils/plugin"
+	"github.com/kaito-project/kaito/pkg/utils/resources"
 	"github.com/kaito-project/kaito/pkg/workspace/inference"
 	"github.com/samber/lo"
 	appsv1 "k8s.io/api/apps/v1"
@@ -633,12 +634,12 @@ func (c *WorkspaceReconciler) ensureService(ctx context.Context, wObj *kaitov1al
 
 	if presetName := getPresetName(wObj); presetName != "" {
 		model := plugin.KaitoModelRegister.MustGet(presetName)
-		serviceObj := resources.GenerateServiceManifest(ctx, wObj, serviceType, model.SupportDistributedInference())
+		serviceObj := manifests.GenerateServiceManifest(ctx, wObj, serviceType, model.SupportDistributedInference())
 		if err := resources.CreateResource(ctx, serviceObj, c.Client); err != nil {
 			return err
 		}
 		if model.SupportDistributedInference() {
-			headlessService := resources.GenerateHeadlessServiceManifest(ctx, wObj)
+			headlessService := manifests.GenerateHeadlessServiceManifest(ctx, wObj)
 			if err := resources.CreateResource(ctx, headlessService, c.Client); err != nil {
 				return err
 			}
@@ -759,7 +760,7 @@ func (c *WorkspaceReconciler) applyInference(ctx context.Context, wObj *kaitov1a
 							volumes = append(volumes, adapterVolume)
 							volumeMounts = append(volumeMounts, adapterVolumeMount)
 						}
-						initContainers, envs := resources.GenerateInitContainers(wObj, volumeMounts)
+						initContainers, envs := manifests.GenerateInitContainers(wObj, volumeMounts)
 						spec := &deployment.Spec
 
 						spec.Template.Spec.InitContainers = initContainers

diff --git a/pkg/workspace/controllers/workspace_gc_finalizer.go b/pkg/workspace/controllers/workspace_gc_finalizer.go
@@ -8,9 +8,9 @@ import (
 
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	"github.com/kaito-project/kaito/pkg/featuregates"
-	"github.com/kaito-project/kaito/pkg/machine"
-	"github.com/kaito-project/kaito/pkg/nodeclaim"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
+	"github.com/kaito-project/kaito/pkg/utils/machine"
+	"github.com/kaito-project/kaito/pkg/utils/nodeclaim"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"

diff --git a/pkg/workspace/inference/preset-inferences.go b/pkg/workspace/inference/preset-inferences.go
@@ -13,7 +13,8 @@ import (
 
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	"github.com/kaito-project/kaito/pkg/model"
-	"github.com/kaito-project/kaito/pkg/resources"
+	manifests "github.com/kaito-project/kaito/pkg/resources"
+	"github.com/kaito-project/kaito/pkg/utils/resources"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	"k8s.io/apimachinery/pkg/util/intstr"
@@ -156,10 +157,10 @@ func CreatePresetInference(ctx context.Context, workspaceObj *kaitov1alpha1.Work
 
 	var depObj client.Object
 	if supportDistributedInference {
-		depObj = resources.GenerateStatefulSetManifest(ctx, workspaceObj, image, imagePullSecrets, *workspaceObj.Resource.Count, commands,
+		depObj = manifests.GenerateStatefulSetManifest(ctx, workspaceObj, image, imagePullSecrets, *workspaceObj.Resource.Count, commands,
 			containerPorts, livenessProbe, readinessProbe, resourceReq, tolerations, volumes, volumeMounts)
 	} else {
-		depObj = resources.GenerateDeploymentManifest(ctx, workspaceObj, revisionNum, image, imagePullSecrets, *workspaceObj.Resource.Count, commands,
+		depObj = manifests.GenerateDeploymentManifest(ctx, workspaceObj, revisionNum, image, imagePullSecrets, *workspaceObj.Resource.Count, commands,
 			containerPorts, livenessProbe, readinessProbe, resourceReq, tolerations, volumes, volumeMounts)
 	}
 	err = resources.CreateResource(ctx, depObj, kubeClient)

diff --git a/pkg/workspace/inference/template_inference.go b/pkg/workspace/inference/template_inference.go
@@ -6,12 +6,13 @@ import (
 	"context"
 
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
-	"github.com/kaito-project/kaito/pkg/resources"
+	manifests "github.com/kaito-project/kaito/pkg/resources"
+	"github.com/kaito-project/kaito/pkg/utils/resources"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
 func CreateTemplateInference(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, kubeClient client.Client) (client.Object, error) {
-	depObj := resources.GenerateDeploymentManifestWithPodTemplate(ctx, workspaceObj, tolerations)
+	depObj := manifests.GenerateDeploymentManifestWithPodTemplate(ctx, workspaceObj, tolerations)
 	err := resources.CreateResource(ctx, client.Object(depObj), kubeClient)
 	if client.IgnoreAlreadyExists(err) != nil {
 		return nil, err

diff --git a/pkg/tuning/preset-tuning-types.go → pkg/workspace/tuning/preset-tuning-types.go b/pkg/tuning/preset-tuning-types.go → pkg/workspace/tuning/preset-tuning-types.go
diff --git a/pkg/tuning/preset-tuning.go → pkg/workspace/tuning/preset-tuning.go b/pkg/tuning/preset-tuning.go → pkg/workspace/tuning/preset-tuning.go
@@ -16,8 +16,9 @@ import (
 
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	"github.com/kaito-project/kaito/pkg/model"
-	"github.com/kaito-project/kaito/pkg/resources"
+	manifests "github.com/kaito-project/kaito/pkg/resources"
 	"github.com/kaito-project/kaito/pkg/utils"
+	"github.com/kaito-project/kaito/pkg/utils/resources"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/klog/v2"
@@ -356,7 +357,7 @@ func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1alpha1.Workspa
 		Name:  "PYTORCH_CUDA_ALLOC_CONF",
 		Value: "expandable_segments:True",
 	})
-	jobObj := resources.GenerateTuningJobManifest(ctx, workspaceObj, revisionNum, tuningImage, imagePullSecrets, *workspaceObj.Resource.Count, commands,
+	jobObj := manifests.GenerateTuningJobManifest(ctx, workspaceObj, revisionNum, tuningImage, imagePullSecrets, *workspaceObj.Resource.Count, commands,
 		containerPorts, nil, nil, resourceReq, tolerations, initContainers, sidecarContainers, volumes, volumeMounts, envVars)
 
 	err = resources.CreateResource(ctx, jobObj, kubeClient)

diff --git a/pkg/tuning/preset-tuning_test.go → pkg/workspace/tuning/preset-tuning_test.go b/pkg/tuning/preset-tuning_test.go → pkg/workspace/tuning/preset-tuning_test.go
diff --git a/presets/models/falcon/model.go b/presets/models/falcon/model.go
@@ -5,12 +5,11 @@ package falcon
 import (
 	"time"
 
-	"github.com/kaito-project/kaito/pkg/tuning"
-
 	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	"github.com/kaito-project/kaito/pkg/model"
 	"github.com/kaito-project/kaito/pkg/utils/plugin"
 	"github.com/kaito-project/kaito/pkg/workspace/inference"
+	"github.com/kaito-project/kaito/pkg/workspace/tuning"
 )
 
 func init() {