Working version final

Author: adityapb

examples/v2beta1/charm/generate_jobs.py

@@ -1,11 +1,20 @@
-from random import randint, seed
+from random import randint, seed, choice, shuffle
 import sys
 import os
 import time
 
 seed(42)
 
-def create_job(job_index, priority, problem_size, min_replicas, max_replicas, timesteps):
+size_index = [0, 1, 2]
+sizes_per_pe = [1024, 1024, 2048]
+min_pes = [1, 4, 4]
+timesteps_per_job = [1000, 2000, 500]
+job_prefixes = ["small", "medium", "large"]
+
+njobs = 20
+
+
+def create_job(prefix, job_index, priority, problem_size, min_replicas, max_replicas, timesteps):
     with open("charm-template.yaml", "r") as file:
         template = file.read()
 
@@ -15,6 +24,7 @@ def create_job(job_index, priority, problem_size, min_replicas, max_replicas, ti
     chare_size = problem_size // num_chares
 
     job_yaml = template.format(
+        prefix=prefix,
         job_index=job_index,
         priority=priority,
         problem_size=problem_size,
@@ -29,17 +39,24 @@ def create_job(job_index, priority, problem_size, min_replicas, max_replicas, ti
 
 
 def generate_jobs():
-    for job_index in range(10):
-        priority = randint(1, 5)
-        problem_size = 2 ** (randint(9, 12))
-        min_replicas = randint(1, 3)
-        max_replicas = randint(min_replicas+1, 8)
-        timesteps = 100 * randint(20, 30)
-        create_job(job_index, priority, problem_size, min_replicas, max_replicas, timesteps)
+    jobs = [0] * int(njobs * 0.25)
+    jobs += [1] * int(njobs * 0.5)
+    jobs += [2] * int(njobs * 0.25)
+    shuffle(jobs)
+    print(jobs)
+    for i, job_index in enumerate(jobs):
+        idx = job_index
+        priority = 3 - idx + randint(0, 3)
+        min_replicas = min_pes[idx]
+        max_replicas = 4 * min_replicas
+        problem_size = min_replicas * sizes_per_pe[idx]
+        timesteps = timesteps_per_job[idx] + 100 * randint(0, 10)
+        prefix = job_prefixes[idx]
+        create_job(prefix, i, priority, problem_size, min_replicas, max_replicas, timesteps)
 
 
 def submit_jobs():
-    for job_index in range(10):
+    for job_index in range(njobs):
         job_file = f"jobs/charm-job-{job_index}.yaml"
         print(f"Submitting {job_file}")
         # Here you would submit the job using your cluster's job submission command

examples/v2beta1/charm/track_utilization.py

@@ -1,7 +1,9 @@
 import subprocess
 import time
 
-output_file = "/home/aditya/mpi-operator/examples/v2beta1/charm/pod_utilization.log"
+output_file = "/home/aditya/mpi-operator/examples/v2beta1/charm/pod_utilization_long2.log"
+
+open(output_file, "w").close()  # Clear the file before starting
 
 while True:
     result = subprocess.run(["kubectl", "get", "pods"], capture_output=True, text=True)

pkg/apis/kubeflow/v2beta1/default.go

@@ -29,7 +29,7 @@ func setDefaultsTypeLauncher(spec *ReplicaSpec) {
 		return
 	}
 	if spec.RestartPolicy == "" {
-		spec.RestartPolicy = DefaultLauncherRestartPolicy
+		spec.RestartPolicy = DefaultRestartPolicy
 	}
 	if spec.MaxReplicas == nil {
 		spec.MaxReplicas = ptr.To[int32](1)

pkg/controller/mpi_job_controller.go

@@ -276,16 +276,17 @@ func (pq PriorityQueue) DeepCopy() PriorityQueue {
 
 func (pq PriorityQueue) Len() int { return len(pq) }
 
-func (pq PriorityQueue) Less(i, j int) bool {
-	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
-	if pq[i].priority == pq[j].priority {
-		return pq[j].timestamp.Before(pq[i].timestamp) // earlier timestamp has higher priority
-	}
-	return pq[i].priority < pq[j].priority
-}
-
 func compare(a *Item, b *Item) int {
-	return cmp.Compare(a.priority, b.priority)
+	if a.priority == b.priority {
+		if a.timestamp.Before(b.timestamp) {
+			return -1 // earlier timestamp has higher priority
+		} else if a.timestamp.After(b.timestamp) {
+			return 1 // earlier timestamp has higher priority
+		} else {
+			return 0 // equal timestamps
+		} // earlier timestamp has higher priority
+	}
+	return cmp.Compare(b.priority, a.priority)
 }
 
 func (pq *PriorityQueue) Push(x any) {
@@ -780,8 +781,8 @@ func (c *MPIJobController) assignFreeSlots() error {
 		}
 
 		if runPriority < queuePriority {
-			idxQueued += 1
 			it = itQueued
+			c.queuedJobs = append(c.queuedJobs[:idxQueued], c.queuedJobs[idxQueued+1:]...)
 			launcherCount = 1
 			//action = create
 		} else {
@@ -799,10 +800,15 @@ func (c *MPIJobController) assignFreeSlots() error {
 		jobMinReplicas := *it.mpiJob.Spec.MPIReplicaSpecs[kubeflow.MPIReplicaTypeWorker].MinReplicas
 		if int32(len(workerPodList)) < jobMaxReplicas {
 			newReplicas = int32(math.Min(float64(jobMaxReplicas),
-				float64(int(c.latestReplicas[getJobKey(&it.mpiJob)])+c.freeSlots)))
+				float64(int(c.latestReplicas[getJobKey(&it.mpiJob)])+c.freeSlots-int(launcherCount))))
 			klog.Infof("Expanding %s to %d, freecount = %d", getJobKey(&it.mpiJob),
 				newReplicas, c.freeSlots)
+
 			if newReplicas < jobMinReplicas {
+				if launcherCount == 1 {
+					c.queuedJobs.Push(it)
+					idxQueued += 1
+				}
 				continue
 			}
 
@@ -811,6 +817,10 @@ func (c *MPIJobController) assignFreeSlots() error {
 				continue
 			}
 
+			if c.jobStatus[getJobKey(&it.mpiJob)] == expanding {
+				continue
+			}
+
 			c.latestReplicas[getJobKey(&it.mpiJob)] = newReplicas
 			c.freeSlots -= (int(newReplicas) - len(workerPodList) + int(launcherCount))
 
@@ -823,7 +833,7 @@ func (c *MPIJobController) assignFreeSlots() error {
 		}
 	}
 	//c.runningJobs = c.runningJobs[idxRunning:]
-	c.queuedJobs = c.queuedJobs[idxQueued:]
+	//c.queuedJobs = c.queuedJobs[idxQueued:]
 
 	if minTime < math.MaxInt64 {
 		c.queue.AddAfter(assignFreeSlotsFlag, minTime)
@@ -980,7 +990,7 @@ func (c *MPIJobController) syncHandler(key string) error {
 			//action, _, newPods, err = c.getAction(mpiJob)
 
 			isExpand := false
-			if status, ok := c.jobStatus[getJobKey(mpiJob)]; ok && status == running {
+			if status, ok := c.jobStatus[getJobKey(mpiJob)]; ok && status == expanding {
 				selector, err := workerSelector(mpiJob.Name)
 				if err != nil {
 					return err
@@ -990,7 +1000,7 @@ func (c *MPIJobController) syncHandler(key string) error {
 					return err
 				}
 				if len(podFullList) < int(lastReplicas) {
-					isExpand = true
+					isExpand = true // this will be true only during the first pass through synchandler
 				}
 			}
 
@@ -1463,6 +1473,8 @@ func (c *MPIJobController) checkJobQueue() error {
 			index += 1
 			continue
 		}
+		c.freeSlots -= int(c.latestReplicas[getJobKey(&mpiJob)]) // This is for the workers
+		c.freeSlots -= 1                                         // This one is for the launcher
 		c.queue.AddRateLimited(getJobKey(&mpiJob))
 		if index < len(c.queuedJobs)-1 {
 			c.queuedJobs = append(c.queuedJobs[:index], c.queuedJobs[index+1:]...)
@@ -1529,6 +1541,7 @@ func (c *MPIJobController) calculateWorkerReplicas(mpiJob *kubeflow.MPIJob) (int
 				}
 
 				it := c.runningJobs[index]
+				index -= 1
 
 				// if the running job priority is higher than the new job
 				// don't shrink it
@@ -1539,7 +1552,7 @@ func (c *MPIJobController) calculateWorkerReplicas(mpiJob *kubeflow.MPIJob) (int
 				if c.jobStatus[getJobKey(&it.mpiJob)] != running {
 					continue
 				}
-				index -= 1
+
 				workerPodList, err := c.getRunningWorkerPods(&it.mpiJob)
 				if err != nil {
 					return -1, err

pkg/controller/rescale_client.C

@@ -21,8 +21,9 @@ int main (int argc, char **argv)
 
     // Create a CcsServer and connect to the given hostname and port
     CcsServer server;
-    char host[BUF], *bitmap;
-    int i, port, cmdLen, mode;
+    char host[BUF], *msg;
+    int i, port, cmdLen;
+    bool isExpand;
 
     sprintf(host, "%s", argv[1]);
     sscanf(argv[2], "%d", &port);
@@ -32,11 +33,11 @@ int main (int argc, char **argv)
     //printf("Rescaling from %i to %i\n", OLDNPROCS, NEWNPROCS);
 
     if( NEWNPROCS > OLDNPROCS)
-        mode = EXPAND;
+        isExpand = true;
     else if(OLDNPROCS > NEWNPROCS)
-        mode = SHRINK;
+        isExpand = false;
     else{
-        printf("0");
+        printf("1");
         return 0;
     }
     //printf("Connecting to server %s %d\n", host, port);
@@ -46,33 +47,18 @@ int main (int argc, char **argv)
     }
     //printf("Connected to server\n");
 
-    cmdLen = OLDNPROCS * sizeof(char) + sizeof(int) + sizeof(char);
-    bitmap = (char *) malloc(cmdLen);
+    cmdLen = sizeof(int) + sizeof(bool);
+    msg = (char *) malloc(cmdLen);
+    memcpy(msg, &isExpand, sizeof(bool));
+    memcpy(&msg[sizeof(bool)], &NEWNPROCS, sizeof(int));
 
-    if (mode == EXPAND) {
-        //printf("Sending expand command.\n");
-        for (i = 0; i < OLDNPROCS; i++) {
-            bitmap[i] = 1;
-        }
-    }
-    else {
-        //printf("Sending shrink command.\n");
-        for (i = 0; i < OLDNPROCS; i++) {
-            if (i < NEWNPROCS)
-                bitmap[i] = 1;
-            else
-                bitmap[i] = 0;
-        }
-    }
-    memcpy(&bitmap[OLDNPROCS], &NEWNPROCS, sizeof(int));
-    bitmap[OLDNPROCS+sizeof(int)] = '\0';
-    if (CcsSendRequest(&server, "set_bitmap", 0, cmdLen, bitmap) == -1) {
+    if (CcsSendRequest(&server, "realloc", 0, cmdLen, msg) == -1) {
         printf("0");
         return 0;
     }
 
     //printf("Waiting for reply...\n" );
-    if (CcsRecvResponse(&server, cmdLen, bitmap , 30) == -1) {
+    if (CcsRecvResponse(&server, cmdLen, msg , 15) == -1) {
         printf("0");
         return 0;
     }