Katkat
Back

Debug GPU Resource Allocation

MithrilSoftware Engineer · Onsite
MediumSoftware EngineerOnsiteSoftware Engineering Fundamentals

Last updated

This problem tests your ability to identify and debug concurrency issues and scheduling logic in a custom GPU resource manager.

Problem Statement

You are maintaining a GPU orchestration service that manages leases for training jobs. The service uses a GPUManager class to track availability and allocate GPUs based on a scoring algorithm.

Recently, high-priority jobs have been failing to preempt low-priority jobs despite having higher priority_score values. You are provided with the core allocation logic below:

class GPUManager:
    def __init__(self, total_gpus):
        self.total_gpus = total_gpus
        self.active_leases = {} # map: gpu_id -> {'job_id': id, 'priority': p}

    def request_gpu(self, job_id, priority):
        # 1. Check for idle GPUs
        free_gpus = [i for i in range(self.total_gpus) if i not in self.active_leases]
        if free_gpus:
            self.active_leases[free_gpus[0]] = {'job_id': job_id, 'priority': priority}
            return True
        
        # 2. Attempt preemption
        for gpu_id, lease in self.active_leases.items():
            if priority > lease['priority']:
                self.active_leases[gpu_id] = {'job_id': job_id, 'priority': priority}
                return True
        return False

The Issue

When three jobs with priorities [1, 2, 3] arrive sequentially for a single-GPU system, the GPU is held by job 1. When job 2 arrives, it correctly preempts job 1. However, when job 3 arrives, it fails to preempt job 2.

Your Task

  1. Identify the logical flaw in the current preemption implementation.
  2. Provide a corrected version of the request_gpu method.
  3. Explain how you would write a unit test to catch this regression in the future.

Hints

Hint 1

Analyze the preemption loop: what happens when the first available GPU in the dictionary does not satisfy the priority condition?

Hint 2

The loop exits prematurely if it encounters a lease that it cannot preempt, even if another lease in the dictionary has a lower priority.

Hint 3

Consider tracking the 'weakest' lease currently active rather than preempting the first lease that happens to be lower priority.

Follow-up questions

  • How would you modify the system to handle 'graceful preemption' where the job being preempted has 30 seconds to checkpoint before being terminated?
  • If the cluster grows to thousands of GPUs, how would this implementation impact latency, and how would you optimize the look-up process?
  • How would you incorporate 'anti-affinity' rules, such as ensuring a single user cannot occupy more than 50% of the total GPU pool?