1 /* 2 * Copyright (C) 2017 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 package com.android.tradefed.invoker.shard; 17 18 import com.android.tradefed.config.IConfiguration; 19 import com.android.tradefed.invoker.IInvocationContext; 20 import com.android.tradefed.invoker.IRescheduler; 21 import com.android.tradefed.log.LogUtil.CLog; 22 import com.android.tradefed.testtype.IBuildReceiver; 23 import com.android.tradefed.testtype.IDeviceTest; 24 import com.android.tradefed.testtype.IRemoteTest; 25 import com.android.tradefed.testtype.IShardableTest; 26 import com.android.tradefed.testtype.IStrictShardableTest; 27 import com.android.tradefed.testtype.suite.ITestSuite; 28 29 import java.util.ArrayList; 30 import java.util.Collection; 31 import java.util.List; 32 33 /** Sharding strategy to create strict shards that do not report together, */ 34 public class StrictShardHelper extends ShardHelper { 35 36 /** {@inheritDoc} */ 37 @Override 38 public boolean shardConfig( 39 IConfiguration config, IInvocationContext context, IRescheduler rescheduler) { 40 Integer shardCount = config.getCommandOptions().getShardCount(); 41 Integer shardIndex = config.getCommandOptions().getShardIndex(); 42 43 if (shardIndex == null) { 44 return super.shardConfig(config, context, rescheduler); 45 } 46 if (shardCount == null) { 47 throw new RuntimeException("shard-count is null while shard-index is " + shardIndex); 48 } 49 50 // Split tests in place, without actually sharding. 51 if (!config.getCommandOptions().shouldUseTfSharding()) { 52 // TODO: remove when IStrictShardableTest is removed. 53 updateConfigIfSharded(config, shardCount, shardIndex); 54 } else { 55 List<IRemoteTest> listAllTests = getAllTests(config, shardCount, context); 56 config.setTests(splitTests(listAllTests, shardCount, shardIndex)); 57 } 58 return false; 59 } 60 61 // TODO: Retire IStrictShardableTest for IShardableTest and have TF balance the list of tests. 62 private void updateConfigIfSharded(IConfiguration config, int shardCount, int shardIndex) { 63 List<IRemoteTest> testShards = new ArrayList<>(); 64 for (IRemoteTest test : config.getTests()) { 65 if (!(test instanceof IStrictShardableTest)) { 66 CLog.w( 67 "%s is not shardable; the whole test will run in shard 0", 68 test.getClass().getName()); 69 if (shardIndex == 0) { 70 testShards.add(test); 71 } 72 continue; 73 } 74 IRemoteTest testShard = 75 ((IStrictShardableTest) test).getTestShard(shardCount, shardIndex); 76 testShards.add(testShard); 77 } 78 config.setTests(testShards); 79 } 80 81 /** 82 * Helper to return the full list of {@link IRemoteTest} based on {@link IShardableTest} split. 83 * 84 * @param config the {@link IConfiguration} describing the invocation. 85 * @param shardCount the shard count hint to be provided to some tests. 86 * @param context the {@link IInvocationContext} of the parent invocation. 87 * @return the list of all {@link IRemoteTest}. 88 */ 89 private List<IRemoteTest> getAllTests( 90 IConfiguration config, Integer shardCount, IInvocationContext context) { 91 List<IRemoteTest> allTests = new ArrayList<>(); 92 for (IRemoteTest test : config.getTests()) { 93 if (test instanceof IShardableTest) { 94 // Inject current information to help with sharding 95 if (test instanceof IBuildReceiver) { 96 ((IBuildReceiver) test).setBuild(context.getBuildInfos().get(0)); 97 } 98 if (test instanceof IDeviceTest) { 99 ((IDeviceTest) test).setDevice(context.getDevices().get(0)); 100 } 101 // Handling of the ITestSuite is a special case, we do not allow pool of tests 102 // since each shard needs to be independent. 103 if (test instanceof ITestSuite) { 104 ((ITestSuite) test).setShouldMakeDynamicModule(false); 105 } 106 107 Collection<IRemoteTest> subTests = ((IShardableTest) test).split(shardCount); 108 if (subTests == null) { 109 // test did not shard so we add it as is. 110 allTests.add(test); 111 } else { 112 allTests.addAll(subTests); 113 } 114 } else { 115 // if test is not shardable we add it as is. 116 allTests.add(test); 117 } 118 } 119 return allTests; 120 } 121 122 /** 123 * Split the list of tests to run however the implementation see fit. Sharding needs to be 124 * consistent. It is acceptable to return an empty list if no tests can be run in the shard. 125 * 126 * <p>Implement this in order to provide a test suite specific sharding. The default 127 * implementation strictly split by number of tests which is not always optimal. TODO: improve 128 * the splitting criteria. 129 * 130 * @param fullList the initial full list of {@link IRemoteTest} containing all the tests that 131 * need to run. 132 * @param shardCount the total number of shard that need to run. 133 * @param shardIndex the index of the current shard that needs to run. 134 * @return a list of {@link IRemoteTest} that need to run in the current shard. 135 */ 136 private List<IRemoteTest> splitTests( 137 List<IRemoteTest> fullList, int shardCount, int shardIndex) { 138 if (shardCount == 1) { 139 // Not sharded 140 return fullList; 141 } 142 if (shardIndex >= fullList.size()) { 143 // Return empty list when we don't have enough tests for all the shards. 144 return new ArrayList<IRemoteTest>(); 145 } 146 int numPerShard = (int) Math.ceil(fullList.size() / (float) shardCount); 147 if (shardIndex == shardCount - 1) { 148 // last shard take everything remaining. 149 return fullList.subList(shardIndex * numPerShard, fullList.size()); 150 } 151 return fullList.subList(shardIndex * numPerShard, numPerShard + (shardIndex * numPerShard)); 152 } 153 } 154
