chore: remove GetAllocationSummary from RM interface (#8846)

determined-ai · Feb 16, 2024 · a5f38cb · a5f38cb
1 parent de28a57
commit a5f38cb
Show file tree

Hide file tree

Showing 7 changed files with 0 additions and 58 deletions.
diff --git a/master/internal/rm/agentrm/agent_resource_manager.go b/master/internal/rm/agentrm/agent_resource_manager.go
@@ -248,20 +248,6 @@ func (a *ResourceManager) GetAllocationSummaries(
 	return summaries, nil
 }
 
-// GetAllocationSummary implements rm.ResourceManager.
-func (a *ResourceManager) GetAllocationSummary(
-	msg sproto.GetAllocationSummary,
-) (*sproto.AllocationSummary, error) {
-	for _, pool := range a.pools {
-		summary, ok := pool.GetAllocationSummary(msg)
-		if !ok {
-			continue
-		}
-		return &summary, nil
-	}
-	return nil, errors.New("allocation not found")
-}
-
 // GetDefaultAuxResourcePool implements rm.ResourceManager.
 func (a *ResourceManager) GetDefaultAuxResourcePool(
 	sproto.GetDefaultAuxResourcePoolRequest,

diff --git a/master/internal/rm/agentrm/agent_resource_manager_intg_test.go b/master/internal/rm/agentrm/agent_resource_manager_intg_test.go
@@ -95,10 +95,6 @@ func TestAgentRMRoutingTaskRelatedMessages(t *testing.T) {
 	require.NoError(t, err)
 
 	// Check the resource pools of the tasks are correct.
-	taskSummary, err := agentRM.GetAllocationSummary(sproto.GetAllocationSummary{ID: cpuTask1.ID})
-	require.NoError(t, err)
-	assert.Equal(t, taskSummary.ResourcePool, cpuTask1.ResourcePool)
-
 	taskSummaries, err = agentRM.GetAllocationSummaries(sproto.GetAllocationSummaries{})
 	require.NoError(t, err)
 	assert.Equal(
@@ -122,10 +118,6 @@ func TestAgentRMRoutingTaskRelatedMessages(t *testing.T) {
 	require.NoError(t, err)
 
 	// Check the resource pools of the tasks are correct.
-	taskSummary, err = agentRM.GetAllocationSummary(sproto.GetAllocationSummary{ID: gpuTask1.ID})
-	require.NoError(t, err)
-	assert.Equal(t, taskSummary.ResourcePool, gpuTask1.ResourcePool)
-
 	taskSummaries, err = agentRM.GetAllocationSummaries(sproto.GetAllocationSummaries{})
 	require.NoError(t, err)
 	assert.Equal(

diff --git a/master/internal/rm/agentrm/resource_pool.go b/master/internal/rm/agentrm/resource_pool.go
@@ -559,17 +559,6 @@ func (rp *resourcePool) NotifyAgentUpdated() {
 	rp.reschedule = true
 }
 
-func (rp *resourcePool) GetAllocationSummary(msg sproto.GetAllocationSummary) (sproto.AllocationSummary, bool) {
-	rp.mu.Lock()
-	defer rp.mu.Unlock()
-
-	resp := rp.taskList.TaskSummary(msg.ID, rp.groups, rp.config.Scheduler.GetType())
-	if resp == nil {
-		return sproto.AllocationSummary{}, false
-	}
-	return *resp, true
-}
-
 func (rp *resourcePool) GetAllocationSummaries(
 	msg sproto.GetAllocationSummaries,
 ) map[model.AllocationID]sproto.AllocationSummary {

diff --git a/master/internal/rm/kubernetesrm/kubernetes_resource_manager.go b/master/internal/rm/kubernetesrm/kubernetes_resource_manager.go
@@ -193,17 +193,6 @@ func (k *ResourceManager) GetAllocationSummaries(
 	return summaries, nil
 }
 
-// GetAllocationSummary implements rm.ResourceManager.
-func (k *ResourceManager) GetAllocationSummary(msg sproto.GetAllocationSummary) (*sproto.AllocationSummary, error) {
-	for _, rp := range k.pools {
-		resp := rp.GetAllocationSummary(msg)
-		if resp != nil {
-			return resp, nil
-		}
-	}
-	return nil, fmt.Errorf("allocation not found: %s", msg.ID)
-}
-
 // GetDefaultAuxResourcePool implements rm.ResourceManager.
 func (k *ResourceManager) GetDefaultAuxResourcePool(
 	sproto.GetDefaultAuxResourcePoolRequest,

diff --git a/master/internal/rm/kubernetesrm/resource_pool.go b/master/internal/rm/kubernetesrm/resource_pool.go
@@ -229,17 +229,6 @@ func (k *kubernetesResourcePool) RecoverJobPosition(msg sproto.RecoverJobPositio
 	k.queuePositions.RecoverJobPosition(msg.JobID, msg.JobPosition)
 }
 
-func (k *kubernetesResourcePool) GetAllocationSummary(msg sproto.GetAllocationSummary) *sproto.AllocationSummary {
-	k.mu.Lock()
-	defer k.mu.Unlock()
-
-	if resp := k.reqList.TaskSummary(
-		msg.ID, k.groups, kubernetesScheduler); resp != nil {
-		return resp
-	}
-	return nil
-}
-
 func (k *kubernetesResourcePool) GetAllocationSummaries(
 	msg sproto.GetAllocationSummaries,
 ) map[model.AllocationID]sproto.AllocationSummary {

diff --git a/master/internal/rm/resource_manager_iface.go b/master/internal/rm/resource_manager_iface.go
@@ -11,7 +11,6 @@ import (
 // ResourceManager is an interface for a resource manager, which can allocate and manage resources.
 type ResourceManager interface {
 	// Basic functionality
-	GetAllocationSummary(sproto.GetAllocationSummary) (*sproto.AllocationSummary, error)
 	GetAllocationSummaries(sproto.GetAllocationSummaries) (map[model.AllocationID]sproto.AllocationSummary, error)
 	Allocate(sproto.AllocateRequest) (*sproto.ResourcesSubscription, error)
 	Release(sproto.ResourcesReleased)

diff --git a/master/internal/sproto/task.go b/master/internal/sproto/task.go
@@ -76,8 +76,6 @@ type (
 		ResourcesID  *ResourcesID
 		ResourcePool string
 	}
-	// GetAllocationSummary returns the summary of the specified task.
-	GetAllocationSummary struct{ ID model.AllocationID }
 	// GetAllocationSummaries returns the summaries of all the tasks in the cluster.
 	GetAllocationSummaries struct{}
 	// AllocationSummary contains information about a task for external display.