1+ /**
2+ * Licensed to the Apache Software Foundation (ASF) under one
3+ * or more contributor license agreements. See the NOTICE file
4+ * distributed with this work for additional information
5+ * regarding copyright ownership. The ASF licenses this file
6+ * to you under the Apache License, Version 2.0 (the
7+ * "License"); you may not use this file except in compliance
8+ * with the License. You may obtain a copy of the License at
9+ *
10+ * http://www.apache.org/licenses/LICENSE-2.0
11+ *
12+ * Unless required by applicable law or agreed to in writing, software
13+ * distributed under the License is distributed on an "AS IS" BASIS,
14+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+ * See the License for the specific language governing permissions and
16+ * limitations under the License.
17+ */
18+ package org .apache .hadoop .yarn .server .resourcemanager .scheduler .fair ;
19+
20+ import org .apache .hadoop .conf .Configuration ;
21+ import org .apache .hadoop .yarn .api .records .ApplicationAttemptId ;
22+ import org .apache .hadoop .yarn .api .records .ApplicationId ;
23+ import org .apache .hadoop .yarn .api .records .ContainerId ;
24+ import org .apache .hadoop .yarn .api .records .Priority ;
25+ import org .apache .hadoop .yarn .api .records .ResourceRequest ;
26+ import org .apache .hadoop .yarn .conf .YarnConfiguration ;
27+ import org .apache .hadoop .yarn .factories .RecordFactory ;
28+ import org .apache .hadoop .yarn .factory .providers .RecordFactoryProvider ;
29+ import org .apache .hadoop .yarn .server .resourcemanager .ResourceManager ;
30+ import org .apache .hadoop .yarn .server .resourcemanager .scheduler .ResourceScheduler ;
31+ import org .apache .hadoop .yarn .server .utils .BuilderUtils ;
32+ import org .apache .hadoop .yarn .util .Clock ;
33+
34+ import java .io .File ;
35+ import java .util .ArrayList ;
36+ import java .util .List ;
37+
38+ public class FairSchedulerTestBase {
39+ protected static class MockClock implements Clock {
40+ private long time = 0 ;
41+ @ Override
42+ public long getTime () {
43+ return time ;
44+ }
45+
46+ public void tick (int seconds ) {
47+ time = time + seconds * 1000 ;
48+ }
49+ }
50+
51+ protected final static String TEST_DIR =
52+ new File (System .getProperty ("test.build.data" , "/tmp" )).getAbsolutePath ();
53+
54+ private static RecordFactory
55+ recordFactory = RecordFactoryProvider .getRecordFactory (null );
56+
57+ protected int APP_ID = 1 ; // Incrementing counter for scheduling apps
58+ protected int ATTEMPT_ID = 1 ; // Incrementing counter for scheduling attempts
59+
60+ protected Configuration conf ;
61+ protected FairScheduler scheduler ;
62+ protected ResourceManager resourceManager ;
63+
64+ // Helper methods
65+ protected Configuration createConfiguration () {
66+ Configuration conf = new YarnConfiguration ();
67+ conf .setClass (YarnConfiguration .RM_SCHEDULER , FairScheduler .class ,
68+ ResourceScheduler .class );
69+ conf .setInt (YarnConfiguration .RM_SCHEDULER_MINIMUM_ALLOCATION_MB , 0 );
70+ conf .setInt (FairSchedulerConfiguration .RM_SCHEDULER_INCREMENT_ALLOCATION_MB ,
71+ 1024 );
72+ conf .setInt (YarnConfiguration .RM_SCHEDULER_MAXIMUM_ALLOCATION_MB , 10240 );
73+ conf .setBoolean (FairSchedulerConfiguration .ASSIGN_MULTIPLE , false );
74+ conf .setFloat (FairSchedulerConfiguration .PREEMPTION_THRESHOLD , 0f );
75+ return conf ;
76+ }
77+
78+ protected ApplicationAttemptId createAppAttemptId (int appId , int attemptId ) {
79+ ApplicationId appIdImpl = ApplicationId .newInstance (0 , appId );
80+ return ApplicationAttemptId .newInstance (appIdImpl , attemptId );
81+ }
82+
83+ protected ResourceRequest createResourceRequest (
84+ int memory , String host , int priority , int numContainers ,
85+ boolean relaxLocality ) {
86+ return createResourceRequest (memory , 1 , host , priority , numContainers ,
87+ relaxLocality );
88+ }
89+
90+ protected ResourceRequest createResourceRequest (
91+ int memory , int vcores , String host , int priority , int numContainers ,
92+ boolean relaxLocality ) {
93+ ResourceRequest request = recordFactory .newRecordInstance (ResourceRequest .class );
94+ request .setCapability (BuilderUtils .newResource (memory , vcores ));
95+ request .setResourceName (host );
96+ request .setNumContainers (numContainers );
97+ Priority prio = recordFactory .newRecordInstance (Priority .class );
98+ prio .setPriority (priority );
99+ request .setPriority (prio );
100+ request .setRelaxLocality (relaxLocality );
101+ return request ;
102+ }
103+
104+ /**
105+ * Creates a single container priority-1 request and submits to
106+ * scheduler.
107+ */
108+ protected ApplicationAttemptId createSchedulingRequest (
109+ int memory , String queueId , String userId ) {
110+ return createSchedulingRequest (memory , queueId , userId , 1 );
111+ }
112+
113+ protected ApplicationAttemptId createSchedulingRequest (
114+ int memory , int vcores , String queueId , String userId ) {
115+ return createSchedulingRequest (memory , vcores , queueId , userId , 1 );
116+ }
117+
118+ protected ApplicationAttemptId createSchedulingRequest (
119+ int memory , String queueId , String userId , int numContainers ) {
120+ return createSchedulingRequest (memory , queueId , userId , numContainers , 1 );
121+ }
122+
123+ protected ApplicationAttemptId createSchedulingRequest (
124+ int memory , int vcores , String queueId , String userId , int numContainers ) {
125+ return createSchedulingRequest (memory , vcores , queueId , userId , numContainers , 1 );
126+ }
127+
128+ protected ApplicationAttemptId createSchedulingRequest (
129+ int memory , String queueId , String userId , int numContainers , int priority ) {
130+ return createSchedulingRequest (memory , 1 , queueId , userId , numContainers ,
131+ priority );
132+ }
133+
134+ protected ApplicationAttemptId createSchedulingRequest (
135+ int memory , int vcores , String queueId , String userId , int numContainers ,
136+ int priority ) {
137+ ApplicationAttemptId id = createAppAttemptId (this .APP_ID ++, this .ATTEMPT_ID ++);
138+ scheduler .addApplication (id .getApplicationId (), queueId , userId );
139+ // This conditional is for testAclSubmitApplication where app is rejected
140+ // and no app is added.
141+ if (scheduler .getSchedulerApplications ().containsKey (id .getApplicationId ())) {
142+ scheduler .addApplicationAttempt (id , false );
143+ }
144+ List <ResourceRequest > ask = new ArrayList <ResourceRequest >();
145+ ResourceRequest request = createResourceRequest (memory , vcores , ResourceRequest .ANY ,
146+ priority , numContainers , true );
147+ ask .add (request );
148+ scheduler .allocate (id , ask , new ArrayList <ContainerId >(), null , null );
149+ return id ;
150+ }
151+
152+ protected void createSchedulingRequestExistingApplication (
153+ int memory , int priority , ApplicationAttemptId attId ) {
154+ ResourceRequest request = createResourceRequest (memory , ResourceRequest .ANY ,
155+ priority , 1 , true );
156+ createSchedulingRequestExistingApplication (request , attId );
157+ }
158+
159+ protected void createSchedulingRequestExistingApplication (
160+ int memory , int vcores , int priority , ApplicationAttemptId attId ) {
161+ ResourceRequest request = createResourceRequest (memory , vcores , ResourceRequest .ANY ,
162+ priority , 1 , true );
163+ createSchedulingRequestExistingApplication (request , attId );
164+ }
165+
166+ protected void createSchedulingRequestExistingApplication (
167+ ResourceRequest request , ApplicationAttemptId attId ) {
168+ List <ResourceRequest > ask = new ArrayList <ResourceRequest >();
169+ ask .add (request );
170+ scheduler .allocate (attId , ask , new ArrayList <ContainerId >(), null , null );
171+ }
172+ }
0 commit comments