| /** |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| package com.datatorrent.stram; |
| |
| import java.util.ArrayList; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| |
| import org.junit.Assert; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| import org.apache.hadoop.yarn.api.records.NodeReport; |
| import org.apache.hadoop.yarn.api.records.NodeState; |
| import org.apache.hadoop.yarn.server.utils.BuilderUtils; |
| |
| import com.google.common.collect.Lists; |
| import com.google.common.collect.Maps; |
| |
| import com.datatorrent.api.AffinityRule; |
| import com.datatorrent.api.AffinityRule.Type; |
| import com.datatorrent.api.AffinityRulesSet; |
| import com.datatorrent.api.Context.DAGContext; |
| import com.datatorrent.api.Context.OperatorContext; |
| import com.datatorrent.api.DAG.Locality; |
| import com.datatorrent.common.partitioner.StatelessPartitioner; |
| import com.datatorrent.stram.StreamingContainerAgent.ContainerStartRequest; |
| import com.datatorrent.stram.engine.GenericTestOperator; |
| import com.datatorrent.stram.engine.TestGeneratorInputOperator; |
| import com.datatorrent.stram.plan.logical.LogicalPlan; |
| import com.datatorrent.stram.plan.physical.PTContainer; |
| import com.datatorrent.stram.plan.physical.PTOperator; |
| import com.datatorrent.stram.support.StramTestSupport.MemoryStorageAgent; |
| import com.datatorrent.stram.support.StramTestSupport.TestMeta; |
| |
| public class AffinityRulesTest |
| { |
| private static final Logger LOG = LoggerFactory.getLogger(AffinityRulesTest.class); |
| |
| @Rule |
| public TestMeta testMeta = new TestMeta(); |
| |
| @Test |
| public void testOperatorPartitionsAntiAffinity() |
| { |
| LogicalPlan dag = new LogicalPlan(); |
| TestGeneratorInputOperator o1 = dag.addOperator("O1", new TestGeneratorInputOperator()); |
| GenericTestOperator o2 = dag.addOperator("O2", new GenericTestOperator()); |
| GenericTestOperator o3 = dag.addOperator("O3", new GenericTestOperator()); |
| dag.addStream("stream1", o1.outport, o2.inport1); |
| dag.addStream("stream2", o2.outport1, o3.inport1); |
| |
| dag.setOperatorAttribute(o2, OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(5)); |
| |
| AffinityRulesSet ruleSet = new AffinityRulesSet(); |
| // Valid case: |
| List<AffinityRule> rules = new ArrayList<>(); |
| ruleSet.setAffinityRules(rules); |
| AffinityRule rule1 = new AffinityRule(Type.ANTI_AFFINITY, Locality.NODE_LOCAL, false, "O2", "O2"); |
| rules.add(rule1); |
| dag.setAttribute(DAGContext.AFFINITY_RULES_SET, ruleSet); |
| dag.validate(); |
| dag.getAttributes().put(com.datatorrent.api.Context.DAGContext.APPLICATION_PATH, testMeta.getAbsolutePath()); |
| dag.setAttribute(OperatorContext.STORAGE_AGENT, new MemoryStorageAgent()); |
| |
| StreamingContainerManager scm = new StreamingContainerManager(dag); |
| |
| // Check Physical Plan assigns anti-affinity preferences correctly |
| |
| for (ContainerStartRequest csr : scm.containerStartRequests) { |
| PTContainer container = csr.container; |
| // Check that for containers for O2 partitions, anti-affinity Preference |
| // are set properly |
| |
| if (container.getOperators().get(0).getName().equals("O2")) { |
| Assert.assertEquals("Anti-affinity containers set should have 4 containers for other partitions ", 4, container.getStrictAntiPrefs().size()); |
| for (PTContainer c : container.getStrictAntiPrefs()) { |
| for (PTOperator operator : c.getOperators()) { |
| Assert.assertEquals("Partion for O2 should be Anti Prefs", "O2", operator.getName()); |
| } |
| } |
| } |
| } |
| |
| // Check resource handler assigns different hosts for each partition |
| |
| ResourceRequestHandler rr = new ResourceRequestHandler(); |
| int containerMem = 1000; |
| Map<String, NodeReport> nodeReports = Maps.newHashMap(); |
| for (int i = 0; i < 10; i++) { |
| String hostName = "host" + i; |
| NodeReport nr = BuilderUtils.newNodeReport(BuilderUtils.newNodeId(hostName, 0), NodeState.RUNNING, "httpAddress", "rackName", BuilderUtils.newResource(0, 0), BuilderUtils.newResource(containerMem * 2, 2), 0, null, 0); |
| nodeReports.put(nr.getNodeId().getHost(), nr); |
| } |
| |
| // set resources |
| rr.updateNodeReports(Lists.newArrayList(nodeReports.values())); |
| |
| Set<String> partitionHostNames = new HashSet<>(); |
| for (ContainerStartRequest csr : scm.containerStartRequests) { |
| String host = rr.getHost(csr, true); |
| csr.container.host = host; |
| if (csr.container.getOperators().get(0).getName().equals("O2")) { |
| Assert.assertNotNull("Host name should not be null", host); |
| LOG.info("Partition {} for operator O2 has host = {} ", csr.container.getId(), host); |
| Assert.assertTrue("Each Partition should have a different host", !partitionHostNames.contains(host)); |
| partitionHostNames.add(host); |
| } |
| } |
| } |
| |
| @Test |
| public void testAntiAffinityInOperators() |
| { |
| LogicalPlan dag = new LogicalPlan(); |
| dag.getAttributes().put(com.datatorrent.api.Context.DAGContext.APPLICATION_PATH, testMeta.getAbsolutePath()); |
| dag.setAttribute(OperatorContext.STORAGE_AGENT, new MemoryStorageAgent()); |
| |
| GenericTestOperator o1 = dag.addOperator("O1", GenericTestOperator.class); |
| dag.setOperatorAttribute(o1, OperatorContext.MEMORY_MB, 256); |
| |
| GenericTestOperator o2 = dag.addOperator("O2", GenericTestOperator.class); |
| dag.setOperatorAttribute(o2, OperatorContext.MEMORY_MB, 256); |
| |
| dag.getMeta(o1).getAttributes().put(OperatorContext.LOCALITY_HOST, "host1"); |
| AffinityRulesSet ruleSet = new AffinityRulesSet(); |
| |
| List<AffinityRule> rules = new ArrayList<>(); |
| ruleSet.setAffinityRules(rules); |
| AffinityRule rule1 = new AffinityRule(Type.ANTI_AFFINITY, Locality.NODE_LOCAL, false, "O1", "O2"); |
| rules.add(rule1); |
| dag.setAttribute(DAGContext.AFFINITY_RULES_SET, ruleSet); |
| |
| dag.addStream("o1_outport1", o1.outport1, o2.inport1);// .setLocality(Locality.NODE_LOCAL); |
| |
| StreamingContainerManager scm = new StreamingContainerManager(dag); |
| |
| ResourceRequestHandler rr = new ResourceRequestHandler(); |
| |
| int containerMem = 1000; |
| Map<String, NodeReport> nodeReports = Maps.newHashMap(); |
| NodeReport nr = BuilderUtils.newNodeReport(BuilderUtils.newNodeId("host1", 0), NodeState.RUNNING, "httpAddress", "rackName", BuilderUtils.newResource(0, 0), BuilderUtils.newResource(containerMem * 2, 2), 0, null, 0); |
| nodeReports.put(nr.getNodeId().getHost(), nr); |
| nr = BuilderUtils.newNodeReport(BuilderUtils.newNodeId("host2", 0), NodeState.RUNNING, "httpAddress", "rackName", BuilderUtils.newResource(0, 0), BuilderUtils.newResource(containerMem * 2, 2), 0, null, 0); |
| nodeReports.put(nr.getNodeId().getHost(), nr); |
| |
| // set resources |
| rr.updateNodeReports(Lists.newArrayList(nodeReports.values())); |
| |
| for (ContainerStartRequest csr : scm.containerStartRequests) { |
| String host = rr.getHost(csr, true); |
| csr.container.host = host; |
| if (csr.container.getOperators().get(0).getName().equals("O1")) { |
| Assert.assertEquals("Hosts set to host1 for Operator O1", "host1", host); |
| } |
| if (csr.container.getOperators().get(0).getName().equals("O2")) { |
| Assert.assertEquals("Hosts set to host2 for Operator O2", "host2", host); |
| } |
| } |
| } |
| } |