Fix test cases

intelligent-machine-learning · Jul 31, 2023 · 8678770 · 8678770
1 parent 73ad550
commit 8678770
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/dlrover/python/master/resource/job.py b/dlrover/python/master/resource/job.py
@@ -351,7 +351,7 @@ def _adjust_oom_ps_resource(self, node: Node):
         plan = self._resource_optimizer.generate_oom_recovery_plan(
             [node], JobOptStage.PS_INITIAL
         )
-        if plan and not plan.empty():
+        if plan and not plan.empty() and node.name in plan.node_resources:
             resource = plan.node_resources[node.name]
             self._ps_resource.node_resource.memory = max(
                 self._ps_resource.node_resource.memory,

diff --git a/dlrover/python/tests/test_resource_optimizer.py b/dlrover/python/tests/test_resource_optimizer.py
@@ -43,9 +43,8 @@ class MockStub(object):
     def optimize(self, request):
         res = brain_pb2.OptimizeResponse()
         res.job_optimize_plans.add()
-        group_resources = res.job_optimize_plans[
-            0
-        ].resource.task_group_resources
+        plan = res.job_optimize_plans[0]
+        group_resources = plan.resource.task_group_resources
         group_resources[NodeType.WORKER].count = 5
         group_resources[NodeType.WORKER].resource.memory = (
             _MEMORY * MemoryUnit.MB
@@ -55,6 +54,7 @@ def optimize(self, request):
         group_resources[NodeType.PS].count = 2
         group_resources[NodeType.PS].resource.memory = _MEMORY * MemoryUnit.MB
         group_resources[NodeType.PS].resource.cpu = 16
+        plan.resource.pod_resources["ps-0"].memory = _MEMORY * MemoryUnit.MB
         return res