Streaming scheduler resourcemanager (#7070)

2026-07-08 09:51:03 +08:00 · 2020-02-24 19:32:50 +08:00
parent 0ae4fe020d
commit 7e115490d9
20 changed files with 1319 additions and 6 deletions
@@ -1,7 +1,13 @@
 package org.ray.streaming.runtime.graph;

 import com.google.common.collect.Lists;
+
+
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+
+
 import org.ray.streaming.api.context.StreamingContext;
 import org.ray.streaming.api.stream.DataStream;
 import org.ray.streaming.api.stream.DataStreamSource;
@@ -9,6 +15,7 @@ import org.ray.streaming.api.stream.StreamSink;
 import org.ray.streaming.jobgraph.JobGraph;
 import org.ray.streaming.jobgraph.JobGraphBuilder;
 import org.ray.streaming.runtime.BaseUnitTest;
+import org.ray.streaming.runtime.config.StreamingConfig;
 import org.ray.streaming.runtime.core.graph.executiongraph.ExecutionGraph;
 import org.ray.streaming.runtime.core.graph.executiongraph.ExecutionJobVertex;
 import org.ray.streaming.runtime.core.graph.executiongraph.ExecutionVertex;
@@ -26,7 +33,9 @@ public class ExecutionGraphTest extends BaseUnitTest {

  @Test
  public void testBuildExecutionGraph() {
-    GraphManager graphManager = new GraphManagerImpl(new JobRuntimeContext(null));
+    Map<String, String> jobConf = new HashMap<>();
+    StreamingConfig streamingConfig = new StreamingConfig(jobConf);
+    GraphManager graphManager = new GraphManagerImpl(new JobRuntimeContext(streamingConfig));
    JobGraph jobGraph = buildJobGraph();
    ExecutionGraph executionGraph = buildExecutionGraph(graphManager, jobGraph);
    List<ExecutionJobVertex> executionJobVertices = executionGraph.getExecutionJobVertexLices();
@@ -0,0 +1,99 @@
+package org.ray.streaming.runtime.resourcemanager;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.aeonbits.owner.util.Collections;
+import org.ray.api.Ray;
+import org.ray.streaming.jobgraph.JobGraph;
+import org.ray.streaming.runtime.BaseUnitTest;
+import org.ray.streaming.runtime.TestHelper;
+import org.ray.streaming.runtime.config.StreamingConfig;
+import org.ray.streaming.runtime.config.global.CommonConfig;
+import org.ray.streaming.runtime.config.master.ResourceConfig;
+import org.ray.streaming.runtime.core.graph.executiongraph.ExecutionGraph;
+import org.ray.streaming.runtime.core.master.resourcemanager.ResourceManager;
+import org.ray.streaming.runtime.core.master.resourcemanager.ResourceManagerImpl;
+import org.ray.streaming.runtime.core.master.scheduler.strategy.SlotAssignStrategy;
+import org.ray.streaming.runtime.core.master.scheduler.strategy.impl.PipelineFirstStrategy;
+import org.ray.streaming.runtime.core.resource.Container;
+import org.ray.streaming.runtime.core.resource.ContainerID;
+import org.ray.streaming.runtime.core.resource.Slot;
+import org.ray.streaming.runtime.graph.ExecutionGraphTest;
+import org.ray.streaming.runtime.master.JobRuntimeContext;
+import org.ray.streaming.runtime.master.graphmanager.GraphManager;
+import org.ray.streaming.runtime.master.graphmanager.GraphManagerImpl;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class ResourceManagerTest extends BaseUnitTest {
+
+  private static final Logger LOG = LoggerFactory.getLogger(ResourceManagerTest.class);
+
+  @org.testng.annotations.BeforeClass
+  public void setUp() {
+    // ray init
+    Ray.init();
+    TestHelper.setUTFlag();
+  }
+
+  @org.testng.annotations.AfterClass
+  public void tearDown() {
+    TestHelper.clearUTFlag();
+  }
+
+  @Test
+  public void testApi() {
+    Map<String, String> conf = new HashMap<String, String>();
+    conf.put(CommonConfig.JOB_NAME, "testApi");
+    conf.put(ResourceConfig.TASK_RESOURCE_CPU_LIMIT_ENABLE, "true");
+    conf.put(ResourceConfig.TASK_RESOURCE_MEM_LIMIT_ENABLE, "true");
+    conf.put(ResourceConfig.TASK_RESOURCE_MEM, "10");
+    conf.put(ResourceConfig.TASK_RESOURCE_CPU, "2");
+    StreamingConfig config = new StreamingConfig(conf);
+    JobRuntimeContext jobRuntimeContext = new JobRuntimeContext(config);
+    ResourceManager resourceManager = new ResourceManagerImpl(jobRuntimeContext);
+
+    SlotAssignStrategy slotAssignStrategy = resourceManager.getSlotAssignStrategy();
+    Assert.assertTrue(slotAssignStrategy instanceof PipelineFirstStrategy);
+
+    Map<String, Double> containerResource = new HashMap<>();
+    containerResource.put(ResourceConfig.RESOURCE_KEY_CPU, 16.0);
+    containerResource.put(ResourceConfig.RESOURCE_KEY_MEM, 128.0);
+    Container container1 = new Container(null, "testAddress1", "testHostName1");
+    container1.setAvailableResource(containerResource);
+    Container container2 = new Container(null, "testAddress2", "testHostName2");
+    container2.setAvailableResource(new HashMap<>(containerResource));
+    List<Container> containers = Collections.list(container1, container2);
+    resourceManager.getResources().getRegisterContainers().addAll(containers);
+    Assert.assertEquals(resourceManager.getRegisteredContainers().size(), 2);
+
+    //build ExecutionGraph
+    GraphManager graphManager = new GraphManagerImpl(new JobRuntimeContext(config));
+    JobGraph jobGraph = ExecutionGraphTest.buildJobGraph();
+    ExecutionGraph executionGraph = ExecutionGraphTest.buildExecutionGraph(graphManager, jobGraph);
+
+    int slotNumPerContainer = slotAssignStrategy.getSlotNumPerContainer(containers, executionGraph
+    .getMaxParallelism());
+    Assert.assertEquals(slotNumPerContainer, 1);
+
+    slotAssignStrategy.allocateSlot(containers, slotNumPerContainer);
+
+    Map<ContainerID, List<Slot>> allocatingMap = slotAssignStrategy.assignSlot(executionGraph);
+    Assert.assertEquals(allocatingMap.size(), 2);
+
+    executionGraph.getAllAddedExecutionVertices().forEach(vertex -> {
+      Container container = resourceManager.getResources()
+          .getRegisterContainerByContainerId(vertex.getSlot().getContainerID());
+      Map<String, Double> resource = resourceManager.allocateResource(container, vertex.getResources());
+      Assert.assertNotNull(resource);
+    });
+    Assert.assertEquals(container1.getAvailableResource().get(ResourceConfig.RESOURCE_KEY_CPU), 14.0);
+    Assert.assertEquals(container2.getAvailableResource().get(ResourceConfig.RESOURCE_KEY_CPU), 14.0);
+    Assert.assertEquals(container1.getAvailableResource().get(ResourceConfig.RESOURCE_KEY_MEM), 118.0);
+    Assert.assertEquals(container2.getAvailableResource().get(ResourceConfig.RESOURCE_KEY_MEM), 118.0);
+  }
+}
@@ -0,0 +1,108 @@
+package org.ray.streaming.runtime.schedule.strategy;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+
+
+import com.google.common.collect.Lists;
+import org.aeonbits.owner.ConfigFactory;
+import org.ray.api.RayActor;
+import org.ray.api.id.ActorId;
+import org.ray.api.id.ObjectId;
+import org.ray.api.id.UniqueId;
+import org.ray.runtime.actor.LocalModeRayActor;
+import org.ray.streaming.api.context.RuntimeContext;
+import org.ray.streaming.api.context.StreamingContext;
+import org.ray.streaming.api.stream.DataStream;
+import org.ray.streaming.api.stream.DataStreamSink;
+import org.ray.streaming.api.stream.DataStreamSource;
+import org.ray.streaming.jobgraph.JobGraph;
+import org.ray.streaming.jobgraph.JobGraphBuilder;
+import org.ray.streaming.runtime.BaseUnitTest;
+import org.ray.streaming.runtime.config.StreamingConfig;
+import org.ray.streaming.runtime.config.StreamingMasterConfig;
+import org.ray.streaming.runtime.config.master.ResourceConfig;
+import org.ray.streaming.runtime.core.graph.executiongraph.ExecutionGraph;
+import org.ray.streaming.runtime.core.master.scheduler.strategy.SlotAssignStrategy;
+import org.ray.streaming.runtime.core.master.scheduler.strategy.impl.PipelineFirstStrategy;
+import org.ray.streaming.runtime.core.resource.Container;
+import org.ray.streaming.runtime.core.resource.ContainerID;
+import org.ray.streaming.runtime.core.resource.Resources;
+import org.ray.streaming.runtime.core.resource.Slot;
+import org.ray.streaming.runtime.graph.ExecutionGraphTest;
+import org.ray.streaming.runtime.master.JobRuntimeContext;
+import org.ray.streaming.runtime.master.graphmanager.GraphManager;
+import org.ray.streaming.runtime.master.graphmanager.GraphManagerImpl;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+public class PipelineFirstStrategyTest extends BaseUnitTest {
+
+  private Logger LOG = LoggerFactory.getLogger(PipelineFirstStrategyTest.class);
+
+  private SlotAssignStrategy strategy;
+  private List<Container> containers = new ArrayList<>();
+  private JobGraph jobGraph;
+  private ExecutionGraph executionGraph;
+  private int maxParallelism;
+
+  @BeforeClass
+  public void setUp() {
+    strategy = new PipelineFirstStrategy();
+    Map<String, String> conf = new HashMap<>();
+    ResourceConfig resourceConfig = ConfigFactory.create(ResourceConfig.class, conf);
+    Resources resources = new Resources(resourceConfig);
+
+    Map<String, Double> containerResource = new HashMap<>();
+    containerResource.put(ResourceConfig.RESOURCE_KEY_CPU, 16.0);
+    containerResource.put(ResourceConfig.RESOURCE_KEY_MEM, 128.0);
+    for (int i = 0; i < 2; ++i) {
+      UniqueId uniqueId = UniqueId.randomId();
+      Container container = new Container(uniqueId, "1.1.1." + i,  "localhost" + i);
+      container.setAvailableResource(containerResource);
+      containers.add(container);
+      resources.getRegisterContainers().add(container);
+    }
+    strategy.setResources(resources);
+
+    //build ExecutionGraph
+    Map<String, String> jobConf = new HashMap<>();
+    StreamingConfig streamingConfig = new StreamingConfig(jobConf);
+    GraphManager graphManager = new GraphManagerImpl(new JobRuntimeContext(streamingConfig));
+    jobGraph = ExecutionGraphTest.buildJobGraph();
+    executionGraph = ExecutionGraphTest.buildExecutionGraph(graphManager, jobGraph);
+    maxParallelism = executionGraph.getMaxParallelism();
+  }
+
+  @Test
+  public int testSlotNumPerContainer() {
+    int slotNumPerContainer = strategy.getSlotNumPerContainer(containers, maxParallelism);
+    Assert.assertEquals(slotNumPerContainer,
+        (int) Math.ceil(Math.max(maxParallelism, containers.size()) * 1.0 / containers.size()));
+    return slotNumPerContainer;
+  }
+
+  @Test
+  public void testAllocateSlot() {
+    int slotNumPerContainer = testSlotNumPerContainer();
+    strategy.allocateSlot(containers, slotNumPerContainer);
+    for (Container container : containers) {
+      Assert.assertEquals(container.getSlots().size(), slotNumPerContainer);
+    }
+  }
+
+  @Test
+  public void testAssignSlot() {
+    Map<ContainerID, List<Slot>> allocatingMap = strategy.assignSlot(executionGraph);
+    for (Entry<ContainerID, List<Slot>> containerSlotEntry : allocatingMap.entrySet()) {
+      containerSlotEntry.getValue()
+          .forEach(slot -> Assert.assertNotEquals(slot.getExecutionVertexIds().size(), 0));
+    }
+  }
+}