From 301911e5b9842bf7aedb133a3a8138b78dca086d Mon Sep 17 00:00:00 2001 From: Fei-Guo Date: Fri, 8 Nov 2024 15:33:12 -0800 Subject: [PATCH 1/2] chore: restruct workspace controller code - part 2 --- cmd/workspace/main.go | 2 +- pkg/ragengine/controllers/ragengine_controller.go | 6 +++--- .../controllers/ragengine_gc_finalizer.go | 4 ++-- pkg/{ => utils}/machine/machine.go | 2 +- pkg/{ => utils}/machine/machine_test.go | 0 pkg/{ => utils}/nodeclaim/nodeclaim.go | 2 +- pkg/{ => utils}/nodeclaim/nodeclaim_test.go | 0 pkg/{ => utils}/resources/nodes.go | 0 pkg/{ => utils}/resources/nodes_test.go | 0 pkg/{ => utils}/resources/resources.go | 0 pkg/{ => utils}/resources/resources_test.go | 0 pkg/workspace/controllers/workspace_controller.go | 15 ++++++++------- .../controllers/workspace_gc_finalizer.go | 4 ++-- pkg/workspace/inference/preset-inferences.go | 7 ++++--- pkg/workspace/inference/template_inference.go | 5 +++-- pkg/{ => workspace}/tuning/preset-tuning-types.go | 0 pkg/{ => workspace}/tuning/preset-tuning.go | 5 +++-- pkg/{ => workspace}/tuning/preset-tuning_test.go | 0 presets/models/falcon/model.go | 3 +-- presets/models/llama2chat/model.go | 2 +- 20 files changed, 30 insertions(+), 27 deletions(-) rename pkg/{ => utils}/machine/machine.go (99%) rename pkg/{ => utils}/machine/machine_test.go (100%) rename pkg/{ => utils}/nodeclaim/nodeclaim.go (99%) rename pkg/{ => utils}/nodeclaim/nodeclaim_test.go (100%) rename pkg/{ => utils}/resources/nodes.go (100%) rename pkg/{ => utils}/resources/nodes_test.go (100%) rename pkg/{ => utils}/resources/resources.go (100%) rename pkg/{ => utils}/resources/resources_test.go (100%) rename pkg/{ => workspace}/tuning/preset-tuning-types.go (100%) rename pkg/{ => workspace}/tuning/preset-tuning.go (99%) rename pkg/{ => workspace}/tuning/preset-tuning_test.go (100%) diff --git a/cmd/workspace/main.go b/cmd/workspace/main.go index 32c846cc6..5dc1b2f40 100644 --- 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 index 54adf7642..9a7036182 100644 --- 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 index c8e729f89..01ae43a37 100644 --- 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 b/pkg/utils/machine/machine.go similarity index 99% rename from pkg/machine/machine.go rename to pkg/utils/machine/machine.go index 7d1cba2cb..187ba2bc2 100644 --- a/pkg/machine/machine.go +++ b/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 b/pkg/utils/machine/machine_test.go similarity index 100% rename from pkg/machine/machine_test.go rename to pkg/utils/machine/machine_test.go diff --git a/pkg/nodeclaim/nodeclaim.go b/pkg/utils/nodeclaim/nodeclaim.go similarity index 99% rename from pkg/nodeclaim/nodeclaim.go rename to pkg/utils/nodeclaim/nodeclaim.go index 242d79e2d..86c797215 100644 --- a/pkg/nodeclaim/nodeclaim.go +++ b/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 b/pkg/utils/nodeclaim/nodeclaim_test.go similarity index 100% rename from pkg/nodeclaim/nodeclaim_test.go rename to pkg/utils/nodeclaim/nodeclaim_test.go diff --git a/pkg/resources/nodes.go b/pkg/utils/resources/nodes.go similarity index 100% rename from pkg/resources/nodes.go rename to pkg/utils/resources/nodes.go diff --git a/pkg/resources/nodes_test.go b/pkg/utils/resources/nodes_test.go similarity index 100% rename from pkg/resources/nodes_test.go rename to pkg/utils/resources/nodes_test.go diff --git a/pkg/resources/resources.go b/pkg/utils/resources/resources.go similarity index 100% rename from pkg/resources/resources.go rename to pkg/utils/resources/resources.go diff --git a/pkg/resources/resources_test.go b/pkg/utils/resources/resources_test.go similarity index 100% rename from pkg/resources/resources_test.go rename to pkg/utils/resources/resources_test.go diff --git a/pkg/workspace/controllers/workspace_controller.go b/pkg/workspace/controllers/workspace_controller.go index bc65a3d63..647b37a31 100644 --- 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 index 61803ec54..351f8fe10 100644 --- 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 index 6bd175fb4..336a4acc2 100644 --- 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 index c40a79f48..b72e251a4 100644 --- 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 b/pkg/workspace/tuning/preset-tuning-types.go similarity index 100% rename from pkg/tuning/preset-tuning-types.go rename to pkg/workspace/tuning/preset-tuning-types.go diff --git a/pkg/tuning/preset-tuning.go b/pkg/workspace/tuning/preset-tuning.go similarity index 99% rename from pkg/tuning/preset-tuning.go rename to pkg/workspace/tuning/preset-tuning.go index d7bf52ce6..112e0016b 100644 --- a/pkg/tuning/preset-tuning.go +++ b/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 b/pkg/workspace/tuning/preset-tuning_test.go similarity index 100% rename from pkg/tuning/preset-tuning_test.go rename to pkg/workspace/tuning/preset-tuning_test.go diff --git a/presets/models/falcon/model.go b/presets/models/falcon/model.go index 83b5a8cac..88487fd63 100644 --- 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() { diff --git a/presets/models/llama2chat/model.go b/presets/models/llama2chat/model.go index 735ecc013..98e974436 100644 --- a/presets/models/llama2chat/model.go +++ b/presets/models/llama2chat/model.go @@ -6,7 +6,7 @@ import ( "time" kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1" - "github.com/kaito-project/kaito/pkg/model" + model "github.com/kaito-project/kaito/pkg/model" "github.com/kaito-project/kaito/pkg/utils/plugin" "github.com/kaito-project/kaito/pkg/workspace/inference" ) From 9f80e7d9ae9ca9bbec511c32a1df8278fb169780 Mon Sep 17 00:00:00 2001 From: Fei-Guo Date: Fri, 8 Nov 2024 15:38:08 -0800 Subject: [PATCH 2/2] revert the unnecessary change --- presets/models/llama2chat/model.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presets/models/llama2chat/model.go b/presets/models/llama2chat/model.go index 98e974436..735ecc013 100644 --- a/presets/models/llama2chat/model.go +++ b/presets/models/llama2chat/model.go @@ -6,7 +6,7 @@ import ( "time" kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1" - model "github.com/kaito-project/kaito/pkg/model" + "github.com/kaito-project/kaito/pkg/model" "github.com/kaito-project/kaito/pkg/utils/plugin" "github.com/kaito-project/kaito/pkg/workspace/inference" )