chore: Commenting and Renaming

.github/workflows/kind-cluster/docker-job-template.yaml

@@ -3,6 +3,7 @@ kind: Job
 metadata:
   name: docker-build-job-{{JOB_ID}}
 spec:
+  ttlSecondsAfterFinished: 600  # Job and its pods are deleted 10 min after job completion
   backoffLimit: 3  # Number of retries before marking the job as failed
   template:
     spec:

.github/workflows/kind-cluster/main.py

@@ -1,10 +1,11 @@
-import subprocess
 import os
-import shutil
-from pathlib import Path
-import time
 import random
+import shutil
 import string
+import subprocess
+import time
+from pathlib import Path
+
 import yaml
 
 KAITO_REPO_URL = "https://github.com/Azure/kaito.git"
@@ -85,7 +86,8 @@ def write_job_file(job_yaml, job_name):
 def populate_job_template(model, img_tag, job_name, env_vars):
     """Populate the job template with provided values."""
     try:
-        with open("/home/azureuser/docker-job-template.yaml", "r") as file:
+        docker_job_template = Path.cwd() / "repo/.github/workflows/kind-cluster/docker-job-template.yaml"
+        with open(docker_job_template, "r") as file:
             job_template = file.read()
 
         replacements = {
@@ -145,10 +147,10 @@ def check_job_status(job_name):
         return "succeeded"
     elif failed and int(failed) > 0:
         return "failed"
-    else:
+    else: 
         return "running"
 
-def wait_for_jobs_to_complete(job_names, timeout=3600):
+def wait_for_jobs_to_complete(job_names, timeout=10800):
     """Wait for all jobs to complete with a timeout."""
     start_time = time.time()
     while time.time() - start_time < timeout:

pkg/controllers/workspace_controller.go

@@ -420,8 +420,8 @@ func (c *WorkspaceReconciler) ensureService(ctx context.Context, wObj *kaitov1al
 }
 
 func (c *WorkspaceReconciler) updateInferenceParamFromWorkspace(ctx context.Context, wObj *kaitov1alpha1.Workspace, inferenceParam *model.PresetInferenceParam) {
-	inferenceParam.AccessMode = string(wObj.Inference.Preset.PresetMeta.AccessMode)
-	if inferenceParam.AccessMode == "private" && wObj.Inference.Preset.PresetOptions.Image != "" {
+	inferenceParam.ImageAccessMode = string(wObj.Inference.Preset.PresetMeta.AccessMode)
+	if inferenceParam.ImageAccessMode == "private" && wObj.Inference.Preset.PresetOptions.Image != "" {
 		inferenceParam.Image = wObj.Inference.Preset.PresetOptions.Image
 
 		imagePullSecretRefs := []corev1.LocalObjectReference{}

pkg/inference/preset-inference-types.go

@@ -59,15 +59,12 @@ var (
 	presetFalcon40bInstructImage = registryName + fmt.Sprintf("/kaito-%s:0.0.1", PresetFalcon40BInstructModel)
 
 	baseCommandPresetLlama = "cd /workspace/llama/llama-2 && torchrun"
-	// llamaTextInferenceFile       = "inference-api.py" TODO: To support Text Generation Llama Models
-	llamaChatInferenceFile = "inference-api.py"
 	llamaRunParams         = map[string]string{
 		"max_seq_len":    "512",
 		"max_batch_size": "8",
 	}
 
 	baseCommandPresetFalcon = "accelerate launch --use_deepspeed"
-	falconInferenceFile     = "inference-api.py"
 	falconRunParams         = map[string]string{}
 
 	defaultTorchRunParams = map[string]string{
@@ -93,7 +90,7 @@ var (
 		"gpu_ids":       DefaultGPUIds,
 	}
 
-	defaultAccessMode       = "public"
+	defaultImageAccessMode  = "public"
 	defaultImagePullSecrets = []corev1.LocalObjectReference{}
 )
 
@@ -122,7 +119,7 @@ var (
 		"gpu_ids":       DefaultGPUIds,
 	}
 
-	DefaultAccessMode       = "public"
+	DefaultImageAccessMode  = "public"
 	DefaultImagePullSecrets = []corev1.LocalObjectReference{}
 )
 
@@ -132,172 +129,152 @@ var (
 	Llama2PresetInferences = map[string]model.PresetInferenceParam{
 
 		PresetLlama2AChat: {
-			ModelName:                 "LLaMa2",
+			ModelFamilyName:           "LLaMa2",
 			Image:                     "",
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "34Gi",
 			GPUCountRequirement:       "1",
 			TotalGPUMemoryRequirement: "16Gi",
 			TorchRunParams:            defaultTorchRunParams,
 			TorchRunRdzvParams:        defaultTorchRunRdzvParams,
 			ModelRunParams:            llamaRunParams,
-			InferenceFile:             llamaChatInferenceFile,
 			DeploymentTimeout:         time.Duration(10) * time.Minute,
 			BaseCommand:               baseCommandPresetLlama,
 			WorldSize:                 1,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 		PresetLlama2AModel: {
-			ModelName:                 "LLaMa2",
+			ModelFamilyName:           "LLaMa2",
 			Image:                     "",
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "34Gi",
 			GPUCountRequirement:       "1",
 			TotalGPUMemoryRequirement: "16Gi",
 			TorchRunParams:            defaultTorchRunParams,
 			TorchRunRdzvParams:        defaultTorchRunRdzvParams,
 			ModelRunParams:            llamaRunParams,
-			InferenceFile:             llamaChatInferenceFile,
 			DeploymentTimeout:         time.Duration(10) * time.Minute,
 			BaseCommand:               baseCommandPresetLlama,
 			WorldSize:                 1,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 		PresetLlama2BChat: {
-			ModelName:                 "LLaMa2",
+			ModelFamilyName:           "LLaMa2",
 			Image:                     "",
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "46Gi",
 			GPUCountRequirement:       "2",
 			TotalGPUMemoryRequirement: "16Gi",
 			TorchRunParams:            defaultTorchRunParams,
 			TorchRunRdzvParams:        defaultTorchRunRdzvParams,
 			ModelRunParams:            llamaRunParams,
-			InferenceFile:             llamaChatInferenceFile,
 			DeploymentTimeout:         time.Duration(20) * time.Minute,
 			BaseCommand:               baseCommandPresetLlama,
 			WorldSize:                 2,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 		PresetLlama2BModel: {
-			ModelName:                 "LLaMa2",
+			ModelFamilyName:           "LLaMa2",
 			Image:                     "",
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "46Gi",
 			GPUCountRequirement:       "2",
 			TotalGPUMemoryRequirement: "16Gi",
 			TorchRunParams:            defaultTorchRunParams,
 			TorchRunRdzvParams:        defaultTorchRunRdzvParams,
 			ModelRunParams:            llamaRunParams,
-			InferenceFile:             llamaChatInferenceFile,
 			DeploymentTimeout:         time.Duration(20) * time.Minute,
 			BaseCommand:               baseCommandPresetLlama,
 			WorldSize:                 2,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 		PresetLlama2CChat: {
-			ModelName:                 "LLaMa2",
+			ModelFamilyName:           "LLaMa2",
 			Image:                     "",
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "158Gi",
 			GPUCountRequirement:       "8",
 			TotalGPUMemoryRequirement: "19Gi",
 			TorchRunParams:            defaultTorchRunParams,
 			TorchRunRdzvParams:        defaultTorchRunRdzvParams,
 			ModelRunParams:            llamaRunParams,
-			InferenceFile:             llamaChatInferenceFile,
 			DeploymentTimeout:         time.Duration(30) * time.Minute,
 			BaseCommand:               baseCommandPresetLlama,
 			WorldSize:                 8,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 		PresetLlama2CModel: {
-			ModelName:                 "LLaMa2",
+			ModelFamilyName:           "LLaMa2",
 			Image:                     "",
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "158Gi",
 			GPUCountRequirement:       "8",
 			TotalGPUMemoryRequirement: "19Gi",
 			TorchRunParams:            defaultTorchRunParams,
 			TorchRunRdzvParams:        defaultTorchRunRdzvParams,
 			ModelRunParams:            llamaRunParams,
-			InferenceFile:             llamaChatInferenceFile,
 			DeploymentTimeout:         time.Duration(30) * time.Minute,
 			BaseCommand:               baseCommandPresetLlama,
 			WorldSize:                 8,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 	}
 
 	// FalconPresetInferences defines the preset inferences for Falcon.
 	FalconPresetInferences = map[string]model.PresetInferenceParam{
 		PresetFalcon7BModel: {
-			ModelName:                 "Falcon",
+			ModelFamilyName:           "Falcon",
 			Image:                     presetFalcon7bImage,
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "50Gi",
 			GPUCountRequirement:       "1",
 			TotalGPUMemoryRequirement: "14Gi",
 			TorchRunParams:            defaultAccelerateParams,
 			ModelRunParams:            falconRunParams,
-			InferenceFile:             falconInferenceFile,
 			DeploymentTimeout:         time.Duration(30) * time.Minute,
 			BaseCommand:               baseCommandPresetFalcon,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 		PresetFalcon7BInstructModel: {
-			ModelName:                 "Falcon",
+			ModelFamilyName:           "Falcon",
 			Image:                     presetFalcon7bInstructImage,
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "50Gi",
 			GPUCountRequirement:       "1",
 			TotalGPUMemoryRequirement: "14Gi",
 			TorchRunParams:            defaultAccelerateParams,
 			ModelRunParams:            falconRunParams,
-			InferenceFile:             falconInferenceFile,
 			DeploymentTimeout:         time.Duration(30) * time.Minute,
 			BaseCommand:               baseCommandPresetFalcon,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 
 		PresetFalcon40BModel: {
-			ModelName:                 "Falcon",
+			ModelFamilyName:           "Falcon",
 			Image:                     presetFalcon40bImage,
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "400",
 			GPUCountRequirement:       "2",
 			TotalGPUMemoryRequirement: "90Gi",
 			TorchRunParams:            defaultAccelerateParams,
 			ModelRunParams:            falconRunParams,
-			InferenceFile:             falconInferenceFile,
 			DeploymentTimeout:         time.Duration(30) * time.Minute,
 			BaseCommand:               baseCommandPresetFalcon,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 
 		PresetFalcon40BInstructModel: {
-			ModelName:                 "Falcon",
+			ModelFamilyName:           "Falcon",
 			Image:                     presetFalcon40bInstructImage,
 			ImagePullSecrets:          defaultImagePullSecrets,
-			AccessMode:                defaultAccessMode,
+			ImageAccessMode:           defaultImageAccessMode,
 			DiskStorageRequirement:    "400",
 			GPUCountRequirement:       "2",
 			TotalGPUMemoryRequirement: "90Gi",
 			TorchRunParams:            defaultAccelerateParams,
 			ModelRunParams:            falconRunParams,
-			InferenceFile:             falconInferenceFile,
 			DeploymentTimeout:         time.Duration(30) * time.Minute,
 			BaseCommand:               baseCommandPresetFalcon,
-			DefaultVolumeMountPath:    "/dev/shm",
 		},
 	}
 )

pkg/inference/preset-inferences.go

@@ -20,6 +20,8 @@ import (
 const (
 	ProbePath = "/healthz"
 	Port5000  = int32(5000)
+	InferenceFile = "inference-api.py"
+	DefaultVolumeMountPath = "/dev/shm"
 )
 
 var (
@@ -116,10 +118,14 @@ func CreatePresetInference(ctx context.Context, workspaceObj *kaitov1alpha1.Work
 	return depObj, nil
 }
 
+// prepareInferenceParameters builds a PyTorch command:
+// torchrun <TORCH_PARAMS> <OPTIONAL_RDZV_PARAMS> baseCommand <MODEL_PARAMS>
+// and sets the GPU resources required for inference.
+// Returns the command and resource configuration.
 func prepareInferenceParameters(ctx context.Context, inferenceObj *model.PresetInferenceParam) ([]string, corev1.ResourceRequirements) {
 	torchCommand := buildCommandStr(inferenceObj.BaseCommand, inferenceObj.TorchRunParams)
 	torchCommand = buildCommandStr(torchCommand, inferenceObj.TorchRunRdzvParams)
-	modelCommand := buildCommandStr(inferenceObj.InferenceFile, inferenceObj.ModelRunParams)
+	modelCommand := buildCommandStr(InferenceFile, inferenceObj.ModelRunParams)
 	commands := shellCommand(torchCommand + " " + modelCommand)
 
 	resourceRequirements := corev1.ResourceRequirements{
@@ -152,7 +158,7 @@ func configVolume(wObj *kaitov1alpha1.Workspace, inferenceObj *model.PresetInfer
 
 		volumeMount = append(volumeMount, corev1.VolumeMount{
 			Name:      volume[0].Name,
-			MountPath: inferenceObj.DefaultVolumeMountPath,
+			MountPath: DefaultVolumeMountPath,
 		})
 	}
 

pkg/model/interface.go

@@ -13,26 +13,25 @@ type Model interface {
 	SupportDistributedInference() bool //If true, the model workload will be a StatefulSet, using the torch elastic runtime framework.
 }
 
-// PresetInferenceParam defines the preset inference.
+// PresetInferenceParam defines the preset inference parameters for a model.
 type PresetInferenceParam struct {
-	ModelName                 string
-	Image                     string
-	ImagePullSecrets          []corev1.LocalObjectReference
-	AccessMode                string
-	DiskStorageRequirement    string
-	GPUCountRequirement       string
-	TotalGPUMemoryRequirement string
-	PerGPUMemoryRequirement   string
-	TorchRunParams            map[string]string
-	TorchRunRdzvParams        map[string]string
-	ModelRunParams            map[string]string
-	InferenceFile             string
+	ModelFamilyName           string                        // The name of the model family.
+	Image                     string                        // Docker image used for running the inference.
+	ImagePullSecrets          []corev1.LocalObjectReference // Secrets for pulling the image from a private registry.
+	ImageAccessMode           string                        // Defines where the Image is Public or Private.
+	DiskStorageRequirement    string                        // Disk storage requirements for the model.
+	GPUCountRequirement       string                        // Number of GPUs required for the inference.
+	TotalGPUMemoryRequirement string                        // Total GPU memory required for the inference.
+	PerGPUMemoryRequirement   string                        // GPU memory required per GPU.
+	TorchRunParams            map[string]string             // Parameters for configuring the torchrun command.
+	TorchRunRdzvParams        map[string]string             // Optional rendezvous parameters for distributed inference using torchrun (elastic).
+	ModelRunParams            map[string]string             // Parameters for running the model inference.
 	// DeploymentTimeout defines the maximum duration for pulling the Preset image.
-	// This timeout accommodates the size of PresetX, ensuring pull completion
+	// This timeout accommodates the size of the image, ensuring pull completion
 	// even under slower network conditions or unforeseen delays.
 	DeploymentTimeout time.Duration
-	BaseCommand       string
-	// WorldSize defines num of processes required for inference
+	// BaseCommand is the initial command (e.g., 'torchrun', 'accelerate launch') used in the command line.
+	BaseCommand string
+	// WorldSize defines the number of processes required for distributed inference.
 	WorldSize              int
-	DefaultVolumeMountPath string
 }