You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@ambari.apache.org by nc...@apache.org on 2014/12/03 20:06:30 UTC

ambari git commit: AMBARI-8465. Upgrade Execute: support service groupings in upgrade engine (ncole)

Repository: ambari
Updated Branches:
  refs/heads/trunk 87ba96e27 -> 0eb91c222


AMBARI-8465. Upgrade Execute:  support service groupings in upgrade engine (ncole)


Project: http://git-wip-us.apache.org/repos/asf/ambari/repo
Commit: http://git-wip-us.apache.org/repos/asf/ambari/commit/0eb91c22
Tree: http://git-wip-us.apache.org/repos/asf/ambari/tree/0eb91c22
Diff: http://git-wip-us.apache.org/repos/asf/ambari/diff/0eb91c22

Branch: refs/heads/trunk
Commit: 0eb91c222c4a579144067cc0c2ffa2dc2d1c04c7
Parents: 87ba96e
Author: Nate Cole <nc...@hortonworks.com>
Authored: Tue Dec 2 17:50:18 2014 -0800
Committer: Nate Cole <nc...@hortonworks.com>
Committed: Wed Dec 3 11:06:18 2014 -0800

----------------------------------------------------------------------
 .../resources/ResourceInstanceFactoryImpl.java  |   3 +-
 .../api/services/UpgradeGroupService.java       |   5 +-
 .../server/api/services/UpgradeItemService.java |   4 +-
 .../internal/UpgradeItemResourceProvider.java   |   2 -
 .../internal/UpgradeResourceProvider.java       | 322 +++++--------------
 .../ambari/server/state/UpgradeHelper.java      | 153 +++++++++
 .../ambari/server/state/stack/UpgradePack.java  |  30 +-
 .../server/state/stack/upgrade/Batch.java       |  41 +--
 .../state/stack/upgrade/ColocatedGrouping.java  | 193 +++++++++++
 .../state/stack/upgrade/ConditionalBatch.java   |  55 ----
 .../server/state/stack/upgrade/CountBatch.java  |  67 ----
 .../server/state/stack/upgrade/Grouping.java    |  97 ++++++
 .../state/stack/upgrade/PercentBatch.java       |  71 ----
 .../server/state/stack/upgrade/RestartTask.java |  41 +++
 .../state/stack/upgrade/StageWrapper.java       | 105 ++++++
 .../stack/upgrade/StageWrapperBuilder.java      |  60 ++++
 .../ambari/server/state/stack/upgrade/Task.java |  23 +-
 .../server/state/stack/upgrade/TaskWrapper.java |  92 ++++++
 .../stacks/HDP/2.2/upgrades/upgrade-2.2.xml     |  65 ++--
 .../internal/UpgradeResourceProviderTest.java   |   6 +-
 .../ambari/server/state/UpgradeHelperTest.java  | 146 +++++++++
 .../server/state/stack/UpgradePackTest.java     |  16 +-
 .../stacks/HDP/2.1.1/upgrades/upgrade_test.xml  |  82 +++--
 23 files changed, 1126 insertions(+), 553 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/api/resources/ResourceInstanceFactoryImpl.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/api/resources/ResourceInstanceFactoryImpl.java b/ambari-server/src/main/java/org/apache/ambari/server/api/resources/ResourceInstanceFactoryImpl.java
index d0664df..b5fe94e 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/api/resources/ResourceInstanceFactoryImpl.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/api/resources/ResourceInstanceFactoryImpl.java
@@ -330,7 +330,8 @@ public class ResourceInstanceFactoryImpl implements ResourceInstanceFactory {
         break;
 
       case UpgradeItem:
-        resourceDefinition = new SimpleResourceDefinition(Resource.Type.UpgradeItem, "upgrade_item", "upgrade_items");
+        resourceDefinition = new SimpleResourceDefinition(
+            Resource.Type.UpgradeItem, "upgrade_item", "upgrade_items");
         break;
 
       case Stage:

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeGroupService.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeGroupService.java b/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeGroupService.java
index d04e900..da21658 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeGroupService.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeGroupService.java
@@ -65,12 +65,9 @@ public class UpgradeGroupService extends BaseService {
         createResourceInstance(id));
   }
 
-  @GET
   @Path("{upgradeGroupId}/upgrade_items")
-  @Produces("text/plain")
-  public UpgradeItemService getUpgradeItems(String body,
+  public UpgradeItemService getUpgradeItems(
       @Context HttpHeaders headers,
-      @Context UriInfo ui,
       @PathParam("upgradeGroupId") Long groupId) {
     return new UpgradeItemService(m_clusterName, m_upgradeId, groupId.toString());
   }

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeItemService.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeItemService.java b/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeItemService.java
index 625d47b..fb77853 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeItemService.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/api/services/UpgradeItemService.java
@@ -50,10 +50,10 @@ public class UpgradeItemService extends BaseService {
 
   @GET
   @Produces("text/plain")
-  public Response getUpgrades(String body,
+  public Response getUpgrades(
       @Context HttpHeaders headers,
       @Context UriInfo ui) {
-    return handleRequest(headers, body, ui, Request.Type.GET,
+    return handleRequest(headers, null, ui, Request.Type.GET,
         createResourceInstance(null));
   }
 

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeItemResourceProvider.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeItemResourceProvider.java b/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeItemResourceProvider.java
index d2660e0..c481dae 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeItemResourceProvider.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeItemResourceProvider.java
@@ -68,8 +68,6 @@ public class UpgradeItemResourceProvider extends AbstractControllerResourceProvi
 
   @Inject
   private static UpgradeDAO m_dao = null;
-  @Inject
-  private static Provider<AmbariMetaInfo> m_metaProvider = null;
 
   static {
     // properties

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java b/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
index 6980e15..a2c30f7 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/controller/internal/UpgradeResourceProvider.java
@@ -26,9 +26,9 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Map.Entry;
 import java.util.Set;
-import java.util.regex.*;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 import org.apache.ambari.server.AmbariException;
 import org.apache.ambari.server.StaticallyInject;
@@ -60,21 +60,19 @@ import org.apache.ambari.server.orm.entities.UpgradeGroupEntity;
 import org.apache.ambari.server.orm.entities.UpgradeItemEntity;
 import org.apache.ambari.server.state.Cluster;
 import org.apache.ambari.server.state.ConfigHelper;
-import org.apache.ambari.server.state.Service;
-import org.apache.ambari.server.state.ServiceComponent;
 import org.apache.ambari.server.state.StackId;
+import org.apache.ambari.server.state.UpgradeHelper;
+import org.apache.ambari.server.state.UpgradeHelper.UpgradeGroupHolder;
 import org.apache.ambari.server.state.stack.UpgradePack;
-import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
-import org.apache.ambari.server.state.stack.upgrade.Task;
+import org.apache.ambari.server.state.stack.upgrade.StageWrapper;
+import org.apache.ambari.server.state.stack.upgrade.TaskWrapper;
 import org.apache.ambari.server.utils.StageUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-import com.google.gson.Gson;
 import com.google.inject.Inject;
 import com.google.inject.Provider;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 /**
  * Manages the ability to start and get status of upgrades.
  */
@@ -143,21 +141,23 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
       throw new SystemException("Can only initiate one upgrade per request.");
     }
 
-    for (final Map<String, Object> requestMap : requestMaps) {
-      createResources(new Command<Void>() {
+    // !!! above check ensures only one
+    final Map<String, Object> requestMap = requestMaps.iterator().next();
+
+    UpgradeEntity entity = createResources(new Command<UpgradeEntity>() {
         @Override
-        public Void invoke() throws AmbariException {
+        public UpgradeEntity invoke() throws AmbariException {
           UpgradePack up = validateRequest(requestMap);
 
-          createUpgrade(up, requestMap);
-
-          return null;
+          return createUpgrade(up, requestMap);
         };
       });
-    }
 
     notifyCreate(Resource.Type.Upgrade, request);
-    return getRequestStatus(null);
+
+    Resource res = new ResourceImpl(Resource.Type.Upgrade);
+    res.setProperty(UPGRADE_ID, entity.getId());
+    return new RequestStatusImpl(null, Collections.singleton(res));
   }
 
   @Override
@@ -182,7 +182,17 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
         throw new NoSuchResourceException(String.format("Cluster %s could not be loaded", clusterName));
       }
 
-      List<UpgradeEntity> upgrades = m_upgradeDAO.findUpgrades(cluster.getClusterId());
+      List<UpgradeEntity> upgrades = new ArrayList<UpgradeEntity>();
+
+      String upgradeIdStr = (String) propertyMap.get(UPGRADE_ID);
+      if (null != upgradeIdStr) {
+        UpgradeEntity upgrade = m_upgradeDAO.findUpgrade(Long.parseLong(upgradeIdStr));
+        if (null != upgrade) {
+          upgrades.add(upgrade);
+        }
+      } else {
+        upgrades = m_upgradeDAO.findUpgrades(cluster.getClusterId());
+      }
 
       for (UpgradeEntity entity : upgrades) {
         results.add(toResource(entity, clusterName, requestPropertyIds));
@@ -328,106 +338,43 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
 
     Cluster cluster = getManagementController().getClusters().getCluster(clusterName);
     ConfigHelper configHelper = getManagementController().getConfigHelper();
-    Map<String, Service> clusterServices = cluster.getServices();
-
-    Map<String, Map<String, ProcessingComponent>> tasks = pack.getTasks();
-
-    List<StageHolder> preUpgrades = new ArrayList<StageHolder>();
-    List<StageHolder> restart = new ArrayList<StageHolder>();
-    List<StageHolder> postUpgrades = new ArrayList<StageHolder>();
-
-    for (Entry<String, List<String>> entry : pack.getOrder().entrySet()) {
-      String serviceName = entry.getKey();
-      List<String> componentNames = entry.getValue();
 
-      // !!! if the upgrade pack doesn't define any tasks, skip
-      if (!tasks.containsKey(serviceName)) {
-        continue;
-      }
-
-      // !!! if the service isn't installed, skip
-      if (!clusterServices.containsKey(serviceName)) {
-        continue;
-      }
+    UpgradeHelper helper = new UpgradeHelper();
+    List<UpgradeGroupHolder> groups = helper.createUpgrade(cluster, pack);
+    List<UpgradeGroupEntity> groupEntities = new ArrayList<UpgradeGroupEntity>();
 
-      Service service = clusterServices.get(serviceName);
-      Map<String, ServiceComponent> components = service.getServiceComponents();
-
-      for (String componentName : componentNames) {
-        // !!! if the upgrade pack has no tasks for component, skip
-        if (!tasks.get(serviceName).containsKey(componentName)) {
-          continue;
-        }
+    RequestStageContainer req = createRequest((String) requestMap.get(UPGRADE_VERSION));
 
-        // !!! if the component is not installed with the cluster, skip
-        if (!components.containsKey(componentName)) {
-          continue;
-        }
+    for (UpgradeGroupHolder group : groups) {
+      UpgradeGroupEntity groupEntity = new UpgradeGroupEntity();
+      groupEntity.setName(group.name);
+      groupEntity.setTitle(group.title);
 
-        ProcessingComponent pc = tasks.get(serviceName).get(componentName);
+      List<UpgradeItemEntity> itemEntities = new ArrayList<UpgradeItemEntity>();
 
-        List<Set<String>> groupings = computeHostGroupings(pc,
-            components.get(componentName).getServiceComponentHosts().keySet());
+      for (StageWrapper wrapper : group.items) {
+        UpgradeItemEntity itemEntity = new UpgradeItemEntity();
+        itemEntity.setText(wrapper.getText());
+        itemEntity.setTasks(wrapper.getHostsJson());
+        itemEntity.setHosts(wrapper.getTasksJson());
+        itemEntities.add(itemEntity);
 
-        preUpgrades.addAll(buildUpgradeStages(pc, true, groupings));
-        restart.addAll(buildRollingRestart(serviceName, pc, groupings));
-        postUpgrades.addAll(buildUpgradeStages(pc, false, groupings));
+        // upgrade items match a stage
+        createStage(cluster, req, itemEntity, wrapper);
       }
-    }
 
-    Gson gson = new Gson();
+      itemEntities = injectVariables(configHelper, cluster, itemEntities);
 
-    UpgradeEntity entity = new UpgradeEntity();
+      groupEntity.setItems(itemEntities);
 
-    List<UpgradeItemEntity> items = new ArrayList<UpgradeItemEntity>();
-    for (StageHolder holder : preUpgrades) {
-      holder.upgradeItemEntity.setHosts(gson.toJson(holder.hosts));
-      holder.upgradeItemEntity.setTasks(gson.toJson(holder.taskHolder.tasks));
-      items.add(holder.upgradeItemEntity);
-    }
+      groupEntities.add(groupEntity);
 
-    for (StageHolder holder : restart) {
-      holder.upgradeItemEntity.setHosts(gson.toJson(holder.hosts));
-      items.add(holder.upgradeItemEntity);
     }
 
-    // This should be the last thing just before finalizing
-    for (StageHolder holder : postUpgrades) {
-      holder.upgradeItemEntity.setHosts(gson.toJson(holder.hosts));
-      holder.upgradeItemEntity.setTasks(gson.toJson(holder.taskHolder.tasks));
-      items.add(holder.upgradeItemEntity);
-    }
-
-    items = injectVariables(configHelper, cluster, items);
-
+    UpgradeEntity entity = new UpgradeEntity();
+    entity.setUpgradeGroups(groupEntities);
     entity.setClusterId(Long.valueOf(cluster.getClusterId()));
 
-    // !!! a separate task will create proper groups.  for now, just one.
-    UpgradeGroupEntity group = new UpgradeGroupEntity();
-    group.setName("CLUSTER_UPGRADE");
-    group.setTitle("Cluster Upgrade");
-    group.setItems(items);
-
-    entity.setUpgradeGroups(Collections.singletonList(group));
-
-    RequestStageContainer req = createRequest((String) requestMap.get(UPGRADE_VERSION));
-
-    // All of the Pre-Upgrades occur first, potentially in several stages.
-    // Should include things like entering safe mode, backing up data, changing the version using hdp-select, etc.
-    for (StageHolder holder : preUpgrades) {
-      createUpgradeTaskStage(cluster, req, holder);
-    }
-
-    // The restart occurs after all of the Pre-Upgrades are done, and is meant to change the pointers and configs.
-    for (StageHolder holder : restart) {
-      createRestartStage(cluster, req, holder);
-    }
-
-    // Post-Upgrades require the user to click on the "Finalize" button.
-    for (StageHolder holder : postUpgrades) {
-      createUpgradeTaskStage(cluster, req, holder);
-    }
-
     req.getRequestStatusResponse();
 
     entity.setRequestId(req.getId());
@@ -439,120 +386,6 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     return entity;
   }
 
-  private List<StageHolder> buildUpgradeStages(ProcessingComponent pc,
-      boolean preUpgrade, List<Set<String>> hostGroups) {
-
-    List<TaskHolder> taskHolders = buildStageStrategy(
-        preUpgrade ? pc.preTasks : pc.postTasks);
-
-    List<StageHolder> stages = new ArrayList<StageHolder>();
-
-    StringBuilder sb = new StringBuilder(preUpgrade ? "Preparing " : "Finalizing ");
-    sb.append("%s on %d host(s).  Phase %s/%s");
-    String textFormat = sb.toString();
-
-    for (TaskHolder taskHolder : taskHolders) {
-      int i = 1;
-      for (Set<String> hostGroup : hostGroups) {
-        StageHolder stage = new StageHolder();
-        stage.hosts = hostGroup;
-        stage.taskHolder = taskHolder;
-        stage.upgradeItemEntity = new UpgradeItemEntity();
-        stage.upgradeItemEntity.setText(String.format(textFormat,
-            pc.name,
-            Integer.valueOf(hostGroup.size()),
-            Integer.valueOf(i++),
-            Integer.valueOf(hostGroups.size())));
-        stages.add(stage);
-      }
-    }
-
-    return stages;
-  }
-
-  /**
-   * Builds the stages for the rolling restart portion
-   * @param pc the information from the upgrade pack
-   * @param hostGroups a list of the host groupings
-   * @return the list of stages that need to be created
-   */
-  private List<StageHolder> buildRollingRestart(String serviceName, ProcessingComponent pc,
-      List<Set<String>> hostGroups) {
-    List<StageHolder> stages = new ArrayList<StageHolder>();
-
-    String textFormat = "Restarting %s on %d host(s), Phase %d/%d";
-
-    int i = 1;
-    for (Set<String> hostGroup : hostGroups) {
-      // !!! each of these is its own stage
-      StageHolder stage = new StageHolder();
-      stage.service = serviceName;
-      stage.component = pc.name;
-      stage.hosts = hostGroup;
-      stage.upgradeItemEntity = new UpgradeItemEntity();
-      stage.upgradeItemEntity.setText(String.format(textFormat, pc.name,
-          Integer.valueOf(hostGroup.size()),
-          Integer.valueOf(i++),
-          Integer.valueOf(hostGroups.size())));
-      stages.add(stage);
-    }
-
-    return stages;
-  }
-
-
-  /**
-   * Calculates how the hosts will be executing their upgrades.
-   */
-  private List<Set<String>> computeHostGroupings(ProcessingComponent taskBuckets, Set<String> allHosts) {
-    if (null == taskBuckets.batch) {
-      return Collections.singletonList(allHosts);
-    } else {
-      return taskBuckets.batch.getHostGroupings(allHosts);
-    }
-  }
-
-  /**
-   * For all the tasks for a component, separate out the manual from the
-   * automated steps into the stages they should executed.
-   *
-   * @param tasks a list of tasks
-   * @return the list of stages
-   */
-  private List<TaskHolder> buildStageStrategy(List<Task> tasks) {
-    if (null == tasks)
-      return Collections.emptyList();
-
-    List<TaskHolder> holders = new ArrayList<TaskHolder>();
-    TaskHolder holder = new TaskHolder();
-
-    holders.add(holder);
-    int i = 0;
-    for (Task t : tasks) {
-      // !!! TODO should every manual task get its own stage?
-      if (i > 0 && t.getType().isManual() != tasks.get(i-1).getType().isManual()) {
-        holder = new TaskHolder();
-        holders.add(holder);
-      }
-
-      holder.tasks.add(t);
-      i++;
-    }
-
-    return holders;
-  }
-
-  private static class TaskHolder {
-    private List<Task> tasks = new ArrayList<Task>();
-  }
-
-  private static class StageHolder {
-    private String service;
-    private String component;
-    private TaskHolder taskHolder;
-    private UpgradeItemEntity upgradeItemEntity;
-    private Set<String> hosts;
-  }
 
   private RequestStageContainer createRequest(String version) {
     ActionManager actionManager = getManagementController().getActionManager();
@@ -564,15 +397,19 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     return requestStages;
   }
 
-  /**
-   * Creates a stage and appends it to the request.
-   * @param cluster the cluster
-   * @param request the request container
-   * @param holder the holder
-   * @throws AmbariException
-   */
-  private void createUpgradeTaskStage(Cluster cluster, RequestStageContainer request,
-      StageHolder holder) throws AmbariException {
+  private void createStage(Cluster cluster, RequestStageContainer request,
+      UpgradeItemEntity entity, StageWrapper wrapper) throws AmbariException {
+
+    if (wrapper.hasCommand()) {
+      makeRestartStage(cluster, request, entity, wrapper);
+    } else {
+      makeActionStage(cluster, request, entity, wrapper);
+    }
+
+  }
+
+  private void makeActionStage(Cluster cluster, RequestStageContainer request,
+      UpgradeItemEntity entity, StageWrapper wrapper) throws AmbariException {
 
     Map<String, String> hostLevelParams = new HashMap<String, String>();
     hostLevelParams.put(JDK_LOCATION, getManagementController().getJdkResourceUrl());
@@ -581,7 +418,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
         "/tmp/ambari",
         cluster.getClusterName(),
         cluster.getClusterId(),
-        holder.upgradeItemEntity.getText(),
+        entity.getText(),
         "{}", "{}",
         StageUtils.getGson().toJson(hostLevelParams));
 
@@ -590,15 +427,15 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
       stageId = 1L;
     }
     stage.setStageId(stageId);
-    holder.upgradeItemEntity.setStageId(Long.valueOf(stageId));
+    entity.setStageId(Long.valueOf(stageId));
 
     // add each host to this stage
     RequestResourceFilter filter = new RequestResourceFilter("", "",
-        new ArrayList<String>(holder.hosts));
+        new ArrayList<String>(wrapper.getHosts()));
 
     // !!! TODO when the custom action is underway, change this
     Map<String, String> params = new HashMap<String, String>();
-    params.put("tasks", holder.upgradeItemEntity.getTasks());
+    params.put("tasks", entity.getTasks());
 
     ActionExecutionContext actionContext = new ActionExecutionContext(
         cluster.getClusterName(), "ru_execute_tasks",
@@ -613,23 +450,27 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     // need to set meaningful text on the command
     for (Map<String, HostRoleCommand> map : stage.getHostRoleCommands().values()) {
       for (HostRoleCommand hrc : map.values()) {
-        hrc.setCommandDetail(holder.upgradeItemEntity.getText());
+        hrc.setCommandDetail(entity.getText());
       }
     }
 
     request.addStages(Collections.singletonList(stage));
   }
 
-  private void createRestartStage(Cluster cluster, RequestStageContainer request,
-      StageHolder holder) throws AmbariException {
+  private void makeRestartStage(Cluster cluster, RequestStageContainer request,
+      UpgradeItemEntity entity, StageWrapper wrapper) throws AmbariException {
 
-    // add each host to this stage
-    RequestResourceFilter filter = new RequestResourceFilter(holder.service, holder.component,
-        new ArrayList<String>(holder.hosts));
+    List<RequestResourceFilter> filters = new ArrayList<RequestResourceFilter>();
+
+    for (TaskWrapper tw : wrapper.getTasks()) {
+      // add each host to this stage
+      filters.add(new RequestResourceFilter(tw.getService(), tw.getComponent(),
+          new ArrayList<String>(tw.getHosts())));
+    }
 
     ActionExecutionContext actionContext = new ActionExecutionContext(
         cluster.getClusterName(), "RESTART",
-        Collections.singletonList(filter),
+        filters,
         Collections.<String, String>emptyMap());
     actionContext.setTimeout(Short.valueOf((short)-1));
 
@@ -640,7 +481,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
         "/tmp/ambari",
         cluster.getClusterName(),
         cluster.getClusterId(),
-        holder.upgradeItemEntity.getText(),
+        entity.getText(),
         jsons.getClusterHostInfo(),
         jsons.getCommandParamsForStage(),
         jsons.getHostParamsForStage());
@@ -650,7 +491,7 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
       stageId = 1L;
     }
     stage.setStageId(stageId);
-    holder.upgradeItemEntity.setStageId(Long.valueOf(stageId));
+    entity.setStageId(Long.valueOf(stageId));
 
     // !!! TODO verify the action is valid
 
@@ -662,4 +503,5 @@ public class UpgradeResourceProvider extends AbstractControllerResourceProvider
     request.addStages(Collections.singletonList(stage));
   }
 
+
 }

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/UpgradeHelper.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/UpgradeHelper.java b/ambari-server/src/main/java/org/apache/ambari/server/state/UpgradeHelper.java
new file mode 100644
index 0000000..4320dbe
--- /dev/null
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/UpgradeHelper.java
@@ -0,0 +1,153 @@
+/**
+ * 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 org.apache.ambari.server.state;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import org.apache.ambari.server.state.stack.UpgradePack;
+import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
+import org.apache.ambari.server.state.stack.upgrade.Grouping;
+import org.apache.ambari.server.state.stack.upgrade.StageWrapper;
+import org.apache.ambari.server.state.stack.upgrade.StageWrapperBuilder;
+import org.apache.ambari.server.state.stack.upgrade.TaskWrapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Class to assist with upgrading a cluster.
+ */
+public class UpgradeHelper {
+
+  private static Logger LOG = LoggerFactory.getLogger(UpgradeHelper.class);
+
+  /**
+   * Generates a list of UpgradeGroupHolder items that are used to execute an upgrade
+   * @param cluster the cluster
+   * @param upgradePack the upgrade pack
+   * @return the list of holders
+   */
+  public List<UpgradeGroupHolder> createUpgrade(Cluster cluster, UpgradePack upgradePack) {
+
+    Map<String, Map<String, ProcessingComponent>> allTasks = upgradePack.getTasks();
+
+    List<UpgradeGroupHolder> groups = new ArrayList<UpgradeGroupHolder>();
+
+    for (Grouping group : upgradePack.getGroups()) {
+      UpgradeGroupHolder groupHolder = new UpgradeGroupHolder();
+      groupHolder.name = group.name;
+      groupHolder.title = group.title;
+      groups.add(groupHolder);
+
+      StageWrapperBuilder builder = group.getBuilder();
+
+      for (UpgradePack.OrderService service : group.services) {
+
+        if (!allTasks.containsKey(service.serviceName)) {
+          continue;
+        }
+
+        for (String component : service.components) {
+          if (!allTasks.get(service.serviceName).containsKey(component)) {
+            continue;
+          }
+
+          Set<String> componentHosts = getClusterHosts(cluster, service.serviceName, component);
+
+          if (0 == componentHosts.size()) {
+            continue;
+          }
+
+          ProcessingComponent pc = allTasks.get(service.serviceName).get(component);
+
+          builder.add(componentHosts, service.serviceName, pc);
+        }
+      }
+
+      List<StageWrapper> proxies = builder.build();
+
+      if (LOG.isDebugEnabled()) {
+        LOG.debug(group.name);
+
+        int i = 0;
+        for (StageWrapper proxy : proxies) {
+          LOG.debug("  Stage {}", Integer.valueOf(i++));
+          int j = 0;
+
+          for (TaskWrapper task : proxy.getTasks()) {
+            LOG.debug("    Task {} {}", Integer.valueOf(j++), task);
+          }
+        }
+      }
+
+      groupHolder.items = proxies;
+    }
+
+    return groups;
+
+  }
+
+  /**
+   * @param cluster the cluster
+   * @param serviceName name of the service
+   * @param componentName name of the component
+   * @return the set of hosts for the provided service and component
+   */
+  private Set<String> getClusterHosts(Cluster cluster, String serviceName, String componentName) {
+    Map<String, Service> services = cluster.getServices();
+
+    if (!services.containsKey(serviceName)) {
+      return Collections.emptySet();
+    }
+
+    Service service = services.get(serviceName);
+    Map<String, ServiceComponent> components = service.getServiceComponents();
+
+    if (!components.containsKey(componentName) ||
+        components.get(componentName).getServiceComponentHosts().size() == 0) {
+      return Collections.emptySet();
+    }
+
+    return components.get(componentName).getServiceComponentHosts().keySet();
+  }
+
+  /**
+   * Short-lived objects that hold information about upgrade groups
+   */
+  public static class UpgradeGroupHolder {
+    /**
+     * The name
+     */
+    public String name;
+    /**
+     * The title
+     */
+    public String title;
+
+    /**
+     * List of stages for the group
+     */
+    public List<StageWrapper> items = new ArrayList<StageWrapper>();
+  }
+
+
+
+}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/UpgradePack.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/UpgradePack.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/UpgradePack.java
index 3057db3..8d2a902 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/UpgradePack.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/UpgradePack.java
@@ -17,7 +17,6 @@
  */
 package org.apache.ambari.server.state.stack;
 
-import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -31,6 +30,7 @@ import javax.xml.bind.annotation.XmlRootElement;
 import javax.xml.bind.annotation.XmlTransient;
 
 import org.apache.ambari.server.state.stack.upgrade.Batch;
+import org.apache.ambari.server.state.stack.upgrade.Grouping;
 import org.apache.ambari.server.state.stack.upgrade.Task;
 
 /**
@@ -45,8 +45,8 @@ public class UpgradePack {
 
 
   @XmlElementWrapper(name="order")
-  @XmlElement(name="service")
-  private List<OrderService> services;
+  @XmlElement(name="group")
+  private List<Grouping> groups;
 
   @XmlElementWrapper(name="processing")
   @XmlElement(name="service")
@@ -64,21 +64,8 @@ public class UpgradePack {
     return target;
   }
 
-  /**
-   * Gets the order by which services and components should be upgraded.
-   * @return a map of service_name -> list of component_name.
-   */
-  public Map<String, List<String>> getOrder() {
-
-    if (null == m_orders) {
-      m_orders = new LinkedHashMap<String, List<String>>();
-
-      for (OrderService order : services) {
-        m_orders.put(order.name, order.components);
-      }
-    }
-
-    return m_orders;
+  public List<Grouping> getGroups() {
+    return groups;
   }
 
   /**
@@ -106,14 +93,13 @@ public class UpgradePack {
     return m_process;
   }
 
-
   /**
    * A service definition that holds a list of componenents in the 'order' element.
    */
   public static class OrderService {
 
-    @XmlAttribute
-    public String name;
+    @XmlAttribute(name="name")
+    public String serviceName;
 
     @XmlElement(name="component")
     public List<String> components;
@@ -131,6 +117,8 @@ public class UpgradePack {
     public List<ProcessingComponent> components;
   }
 
+
+
   /**
    * A component definition in the 'processing/service' path.
    */

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Batch.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Batch.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Batch.java
index 4ff0c66..b720f18 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Batch.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Batch.java
@@ -17,44 +17,17 @@
  */
 package org.apache.ambari.server.state.stack.upgrade;
 
-import java.util.List;
-import java.util.Set;
-
-import javax.xml.bind.annotation.XmlSeeAlso;
-import javax.xml.bind.annotation.XmlType;
+import javax.xml.bind.annotation.XmlElement;
 
 /**
- * Base class to identify how a component should be upgraded (optional)
+ *  Identifies how a component should be upgraded (optional)
  */
-@XmlSeeAlso(value={CountBatch.class, PercentBatch.class, ConditionalBatch.class})
-public abstract class Batch {
+public class Batch {
 
-  /**
-   * @return the batch type
-   */
-  public abstract Type getType();
+  @XmlElement(name="percent")
+  public int percent;
 
-  /**
-   * Identifies the type of batch
-   */
-  public enum Type {
-    /**
-     * Batch by <i>n</i> instance at a time
-     */
-    COUNT,
-    /**
-     * Batch by <i>x</i>% at a time
-     */
-    PERCENT,
-    /**
-     * Batch by an inital <i>x</i>%, then after confirmation batch <i>y</i>% at a time.
-     */
-    CONDITIONAL
-  }
+  @XmlElement(name="message")
+  public String message;
 
-  /**
-   * @param hosts all the hosts
-   * @return a list of host sets defined by the specific batching
-   */
-  public abstract List<Set<String>> getHostGroupings(Set<String> hosts);
 }

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ColocatedGrouping.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ColocatedGrouping.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ColocatedGrouping.java
new file mode 100644
index 0000000..988e272
--- /dev/null
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ColocatedGrouping.java
@@ -0,0 +1,193 @@
+/**
+ * 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 org.apache.ambari.server.state.stack.upgrade;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+
+import javax.xml.bind.annotation.XmlElement;
+import javax.xml.bind.annotation.XmlType;
+
+import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Used for co-located services grouped together.
+ */
+@XmlType(name="colocated")
+public class ColocatedGrouping extends Grouping {
+
+  private static Logger LOG = LoggerFactory.getLogger(ColocatedGrouping.class);
+
+  @XmlElement(name="batch")
+  public Batch batch;
+
+
+  @Override
+  public StageWrapperBuilder getBuilder() {
+    return new MultiHomedHolder(batch);
+  }
+
+  private static class MultiHomedHolder extends StageWrapperBuilder {
+
+    private Batch batch;
+
+    // !!! host -> list of tasks
+    private Map<String, List<TaskProxy>> initialBatch = new LinkedHashMap<String, List<TaskProxy>>();
+    private Map<String, List<TaskProxy>> finalBatches = new LinkedHashMap<String, List<TaskProxy>>();
+
+    private MultiHomedHolder(Batch batch) {
+      this.batch = batch;
+    }
+
+    @Override
+    public void add(Set<String> hosts, String service, ProcessingComponent pc) {
+
+      int count = Double.valueOf(Math.ceil(
+          (double) batch.percent / 100 * hosts.size())).intValue();
+
+      int i = 0;
+      for (String host : hosts) {
+
+        Map<String, List<TaskProxy>> targetMap = ((i++) < count) ? initialBatch : finalBatches;
+        List<TaskProxy> targetList = targetMap.get(host);
+        if (null == targetList) {
+          targetList = new ArrayList<TaskProxy>();
+          targetMap.put(host, targetList);
+        }
+
+        TaskProxy proxy = null;
+
+        if (null != pc.preTasks && pc.preTasks.size() > 0) {
+          proxy = new TaskProxy();
+          proxy.message = getStageText("Preparing", pc.name, Collections.singleton(host));
+          proxy.tasks.add(new TaskWrapper(service, pc.name, Collections.singleton(host), pc.preTasks));
+          proxy.component = pc.name;
+          targetList.add(proxy);
+        }
+
+        // !!! FIXME upgrade definition have only one step, and it better be a restart
+        if (null != pc.tasks && 1 == pc.tasks.size()) {
+          Task t = pc.tasks.get(0);
+          if (RestartTask.class.isInstance(t)) {
+            proxy = new TaskProxy();
+            proxy.tasks.add(new TaskWrapper(service, pc.name, Collections.singleton(host), t));
+            proxy.restart = true;
+            proxy.component = pc.name;
+            proxy.message = getStageText("Restarting ", pc.name, Collections.singleton(host));
+
+            targetList.add(proxy);
+          }
+        }
+
+        if (null != pc.postTasks && pc.postTasks.size() > 0) {
+          proxy = new TaskProxy();
+          proxy.component = pc.name;
+          proxy.tasks.add(new TaskWrapper(service, pc.name, Collections.singleton(host), pc.postTasks));
+          proxy.message = getStageText("Completing", pc.name, Collections.singleton(host));
+          targetList.add(proxy);
+        }
+      }
+    }
+
+
+    @Override
+    public List<StageWrapper> build() {
+      List<StageWrapper> results = new ArrayList<StageWrapper>();
+
+      if (LOG.isDebugEnabled()) {
+        LOG.debug("RU initial: {}", initialBatch);
+        LOG.debug("RU final: {}", finalBatches);
+      }
+
+      results.addAll(fromProxies(initialBatch));
+
+//      StageWrapper wrapper = new StageWrapper(
+//      ManualTask task = new ManualTask();
+//      task.message = batch.message;
+//      wrapper.tasks.add(new TaskWrapper(null, null, null, task));
+      // !!! TODO
+//      results.add(wrapper);
+
+      results.addAll(fromProxies(finalBatches));
+
+      return results;
+    }
+
+    private List<StageWrapper> fromProxies(Map<String, List<TaskProxy>> wrappers) {
+      List<StageWrapper> results = new ArrayList<StageWrapper>();
+
+      for (Entry<String, List<TaskProxy>> entry : wrappers.entrySet()) {
+
+        // !!! stage per host, per type
+        StageWrapper wrapper = null;
+        StageWrapper execwrapper = null;
+
+        for (TaskProxy t : entry.getValue()) {
+          if (!t.restart) {
+            if (null == wrapper) {
+              wrapper = new StageWrapper(t.message, t.tasks);
+            }
+          } else {
+            if (null == execwrapper) {
+              execwrapper = new StageWrapper(t.message, t.tasks);
+            }
+          }
+        }
+
+        if (null != wrapper) {
+          results.add(wrapper);
+        }
+
+        if (null != execwrapper) {
+          results.add(execwrapper);
+        }
+      }
+
+      return results;
+    }
+
+  }
+
+  /**
+   * Represents all the tasks that need to be run for a host
+   */
+  private static class TaskProxy {
+    private boolean restart = false;
+    private String component;
+    private String message;
+    private List<TaskWrapper> tasks = new ArrayList<TaskWrapper>();
+
+    @Override
+    public String toString() {
+      String s = "";
+      for (TaskWrapper t : tasks) {
+        s += component + "/" + t.getTasks() + " ";
+      }
+
+      return s;
+    }
+  }
+
+}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ConditionalBatch.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ConditionalBatch.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ConditionalBatch.java
deleted file mode 100644
index e88d67c..0000000
--- a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/ConditionalBatch.java
+++ /dev/null
@@ -1,55 +0,0 @@
-/**
- * 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 org.apache.ambari.server.state.stack.upgrade;
-
-import java.util.Collections;
-import java.util.List;
-import java.util.Set;
-
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlElement;
-import javax.xml.bind.annotation.XmlRootElement;
-import javax.xml.bind.annotation.XmlType;
-
-/**
- * Upgrade batch that should happen by percentage.  After the inital percentage,
- * the remaining nodes are upgraded incrementally.
- */
-@XmlRootElement
-@XmlAccessorType(XmlAccessType.FIELD)
-@XmlType(name="conditional")
-public class ConditionalBatch extends Batch {
-
-  @XmlElement(name="initial")
-  public int initial = 0;
-
-  @XmlElement(name="remaining")
-  public int remaining = 0;
-
-  @Override
-  public Type getType() {
-    return Batch.Type.CONDITIONAL;
-  }
-
-  @Override
-  public List<Set<String>> getHostGroupings(Set<String> hosts) {
-    // TODO
-    return Collections.singletonList(hosts);
-  }
-}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/CountBatch.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/CountBatch.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/CountBatch.java
deleted file mode 100644
index c55f569..0000000
--- a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/CountBatch.java
+++ /dev/null
@@ -1,67 +0,0 @@
-/**
- * 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 org.apache.ambari.server.state.stack.upgrade;
-
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlElement;
-import javax.xml.bind.annotation.XmlRootElement;
-import javax.xml.bind.annotation.XmlType;
-
-/**
- * Upgrade batch that should happen one at a time.
- */
-@XmlRootElement
-@XmlAccessorType(XmlAccessType.FIELD)
-@XmlType(name="count")
-public class CountBatch extends Batch {
-
-  @XmlElement(name="count")
-  public int count = 1;
-
-  @Override
-  public Type getType() {
-    return Batch.Type.COUNT;
-  }
-
-  /* (non-Javadoc)
-   * @see org.apache.ambari.server.state.stack.upgrade.Batch#getHostGroupings(java.util.Set)
-   */
-  @Override
-  public List<Set<String>> getHostGroupings(Set<String> hosts) {
-    List<Set<String>> groupings = new ArrayList<Set<String>>();
-
-    Set<String> set = new HashSet<String>();
-    groupings.add(set);
-    int i = 1;
-    for (String host : hosts) {
-      set.add(host);
-      if (i < hosts.size() && 0 == (i++ % count)) {
-        set = new HashSet<String>();
-        groupings.add(set);
-      }
-    }
-
-    return groupings;
-  }
-}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Grouping.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Grouping.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Grouping.java
new file mode 100644
index 0000000..f99ee72
--- /dev/null
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Grouping.java
@@ -0,0 +1,97 @@
+/**
+ * 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 org.apache.ambari.server.state.stack.upgrade;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+
+import javax.xml.bind.annotation.XmlAttribute;
+import javax.xml.bind.annotation.XmlElement;
+import javax.xml.bind.annotation.XmlSeeAlso;
+
+import org.apache.ambari.server.state.stack.UpgradePack;
+import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
+
+/**
+ *
+ */
+@XmlSeeAlso(value = { ColocatedGrouping.class })
+public class Grouping {
+
+  @XmlAttribute(name="name")
+  public String name;
+
+  @XmlAttribute(name="title")
+  public String title;
+
+  @XmlElement(name="service")
+  public List<UpgradePack.OrderService> services;
+
+  /**
+   * Gets the default builder.
+   */
+  public StageWrapperBuilder getBuilder() {
+    return new DefaultBuilder();
+  }
+
+
+  private static class DefaultBuilder extends StageWrapperBuilder {
+
+    private List<StageWrapper> stages = new ArrayList<StageWrapper>();
+
+    @Override
+    public void add(Set<String> hosts, String service, ProcessingComponent pc) {
+      if (null != pc.preTasks && pc.preTasks.size() > 0) {
+        StageWrapper stage = new StageWrapper(
+            getStageText("Preparing", pc.name, hosts),
+            Collections.singletonList(new TaskWrapper(service, pc.name, hosts, pc.preTasks)));
+        stages.add(stage);
+      }
+
+      // !!! FIXME upgrade definition have only one step, and it better be a restart
+      if (null != pc.tasks && 1 == pc.tasks.size()) {
+        Task t = pc.tasks.get(0);
+        if (RestartTask.class.isInstance(t)) {
+          for (String hostName : hosts) {
+            StageWrapper stage = new StageWrapper(
+                getStageText("Restarting", pc.name, Collections.singleton(hostName)),
+                Collections.singletonList(new TaskWrapper(service, pc.name, Collections.singleton(hostName), t)));
+            stages.add(stage);
+          }
+        }
+      }
+
+      if (null != pc.postTasks && pc.postTasks.size() > 0) {
+        StageWrapper stage = new StageWrapper(
+            getStageText("Completing", pc.name, hosts),
+            Collections.singletonList(new TaskWrapper(service, pc.name, hosts, pc.postTasks)));
+        stages.add(stage);
+      }
+    }
+
+    @Override
+    public List<StageWrapper> build() {
+      return stages;
+    }
+
+  }
+
+
+}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/PercentBatch.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/PercentBatch.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/PercentBatch.java
deleted file mode 100644
index d67bb0b..0000000
--- a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/PercentBatch.java
+++ /dev/null
@@ -1,71 +0,0 @@
-/**
- * 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 org.apache.ambari.server.state.stack.upgrade;
-
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlElement;
-import javax.xml.bind.annotation.XmlRootElement;
-import javax.xml.bind.annotation.XmlType;
-
-/**
- * Upgrade batch that should happen by percentage.  After each percentage
- * completes, continue to the next step.
- */
-@XmlRootElement
-@XmlAccessorType(XmlAccessType.FIELD)
-@XmlType(name="percent")
-public class PercentBatch extends Batch {
-
-  @XmlElement(name="percent")
-  public int percent = 100;
-
-  @Override
-  public Type getType() {
-    return Batch.Type.PERCENT;
-  }
-
-  @Override
-  public List<Set<String>> getHostGroupings(Set<String> hosts) {
-
-    List<Set<String>> groupings = new ArrayList<Set<String>>();
-
-    int count = Double.valueOf(Math.ceil(
-        (double) percent / 100 * hosts.size())).intValue();
-
-    Set<String> set = new HashSet<String>();
-    groupings.add(set);
-    int i = 1;
-    for (String host : hosts) {
-      set.add(host);
-      if (i < hosts.size() && 0 == (i++ % count)) {
-        set = new HashSet<String>();
-        groupings.add(set);
-      }
-    }
-
-    return groupings;
-  }
-
-
-}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/RestartTask.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/RestartTask.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/RestartTask.java
new file mode 100644
index 0000000..1b69b5b
--- /dev/null
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/RestartTask.java
@@ -0,0 +1,41 @@
+/**
+ * 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 org.apache.ambari.server.state.stack.upgrade;
+
+import javax.xml.bind.annotation.XmlAccessType;
+import javax.xml.bind.annotation.XmlAccessorType;
+import javax.xml.bind.annotation.XmlRootElement;
+import javax.xml.bind.annotation.XmlTransient;
+import javax.xml.bind.annotation.XmlType;
+
+/**
+ * Used to represent a restart of a component.
+ */
+@XmlRootElement
+@XmlAccessorType(XmlAccessType.FIELD)
+@XmlType(name="restart")
+public class RestartTask extends Task {
+
+  @XmlTransient
+  private Task.Type type = Task.Type.RESTART;
+
+  @Override
+  public Task.Type getType() {
+    return type;
+  }
+}
\ No newline at end of file

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapper.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapper.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapper.java
new file mode 100644
index 0000000..9544323
--- /dev/null
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapper.java
@@ -0,0 +1,105 @@
+/**
+ * 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 org.apache.ambari.server.state.stack.upgrade;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import com.google.gson.Gson;
+
+/**
+ *
+ */
+public class StageWrapper {
+
+  private static Gson gson = new Gson();
+  private String text;
+
+  private List<TaskWrapper> tasks;
+
+  public StageWrapper(String text, List<TaskWrapper> tasks) {
+    this.text = text;
+    this.tasks = tasks;
+  }
+
+  /**
+   * Gets the hosts json.
+   */
+  public String getHostsJson() {
+    return gson.toJson(getHosts());
+  }
+
+  /**
+   * Gets the tasks json.
+   */
+  public String getTasksJson() {
+    List<Task> realTasks = new ArrayList<Task>();
+    for (TaskWrapper tw : tasks) {
+      realTasks.addAll(tw.getTasks());
+    }
+
+    return gson.toJson(realTasks);
+  }
+
+  /**
+   * @return {@code true} if any of the tasks is a command type.  This affects
+   * the type of stage that is created.
+   */
+  public boolean hasCommand() {
+
+    for (TaskWrapper tw : tasks) {
+      for (Task t : tw.getTasks()) {
+        if (t.getType().isCommand()) {
+          return true;
+        }
+      }
+    }
+
+    return false;
+  }
+
+  /**
+   * @return the set of hosts for all tasks
+   */
+  public Set<String> getHosts() {
+    Set<String> hosts = new HashSet<String>();
+    for (TaskWrapper tw : tasks) {
+      hosts.addAll(tw.getHosts());
+    }
+
+    return hosts;
+  }
+
+  /**
+   * @return the wrapped tasks for this stage
+   */
+  public List<TaskWrapper> getTasks() {
+    return tasks;
+  }
+
+  /**
+   * @return the text for this stage
+   */
+  public String getText() {
+    return text;
+  }
+
+}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapperBuilder.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapperBuilder.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapperBuilder.java
new file mode 100644
index 0000000..bc5d4a1
--- /dev/null
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/StageWrapperBuilder.java
@@ -0,0 +1,60 @@
+/**
+ * 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 org.apache.ambari.server.state.stack.upgrade;
+
+import java.util.List;
+import java.util.Set;
+
+import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
+
+/**
+ * Defines how to build stages.
+ */
+public abstract class StageWrapperBuilder {
+
+  /**
+   * Adds a processing component that will be built into stage wrappers.
+   *
+   * @param hosts the hosts
+   * @param service the service name
+   * @param pc the ProcessingComponent derived from the upgrade pack.
+   */
+  public abstract void add(Set<String> hosts, String service, ProcessingComponent pc);
+
+  /**
+   * Builds the stage wrappers.
+   */
+  public abstract List<StageWrapper> build();
+
+  /**
+   * Consistently formats a string.
+   * @param prefix
+   * @param component
+   * @param hosts
+   * @return the prepared string
+   */
+  protected String getStageText(String prefix, String component, Set<String> hosts) {
+    return String.format("%s %s on %s%s",
+        prefix,
+        component,
+        1 == hosts.size() ? hosts.iterator().next() : Integer.valueOf(hosts.size()),
+        1 == hosts.size() ? "" : " hosts");
+  }
+
+
+}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Task.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Task.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Task.java
index 4cda5da..9a60fc9 100644
--- a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Task.java
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/Task.java
@@ -23,7 +23,7 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 /**
  * Base class to identify the items that could possibly occur during an upgrade
  */
-@XmlSeeAlso(value={ExecuteTask.class, ConfigureTask.class, ManualTask.class})
+@XmlSeeAlso(value={ExecuteTask.class, ConfigureTask.class, ManualTask.class, RestartTask.class})
 public abstract class Task {
 
   /**
@@ -31,6 +31,11 @@ public abstract class Task {
    */
   public abstract Type getType();
 
+  @Override
+  public String toString() {
+    return getType().toString();
+  }
+
   /**
    * Identifies the type of task.
    */
@@ -46,7 +51,11 @@ public abstract class Task {
     /**
      * Task that displays a message and must be confirmed before continuing
      */
-    MANUAL;
+    MANUAL,
+    /**
+     * Task that is a restart command.
+     */
+    RESTART;
 
     /**
      * @return {@code true} if the task is manual or automated.
@@ -54,5 +63,15 @@ public abstract class Task {
     public boolean isManual() {
       return this == MANUAL;
     }
+
+    /**
+     * @return {@code true} if the task is a command type (as opposed to an action)
+     */
+    public boolean isCommand() {
+      return this == RESTART;
+    }
+
+
+
   }
 }

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/TaskWrapper.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/TaskWrapper.java b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/TaskWrapper.java
new file mode 100644
index 0000000..f7cc930
--- /dev/null
+++ b/ambari-server/src/main/java/org/apache/ambari/server/state/stack/upgrade/TaskWrapper.java
@@ -0,0 +1,92 @@
+/**
+ * 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 org.apache.ambari.server.state.stack.upgrade;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * Aggregates all upgrade tasks for a HostComponent into one wrapper.
+ */
+public class TaskWrapper {
+
+  private String service;
+  private String component;
+  private Set<String> hosts; // all the hosts that all the tasks must run
+  private List<Task> tasks; // all the tasks defined for the hostcomponent
+
+  /**
+   * @param s the service name for the tasks
+   * @param c the component name for the tasks
+   * @param hosts the set of hosts that the tasks are for
+   * @param tasks an array of tasks as a convenience
+   */
+  public TaskWrapper(String s, String c, Set<String> hosts, Task... tasks) {
+    this(s, c, hosts, Arrays.asList(tasks));
+  }
+
+  /**
+   * @param s the service name for the tasks
+   * @param c the component name for the tasks
+   * @param hosts the set of hosts for the
+   * @param tasks the list of tasks
+   */
+  public TaskWrapper(String s, String c, Set<String> hosts, List<Task> tasks) {
+    service = s;
+    component = c;
+
+    this.hosts = hosts;
+    this.tasks = tasks;
+  }
+
+  /**
+   * @return the tasks associated with this wrapper
+   */
+  public List<Task> getTasks() {
+    return tasks;
+  }
+
+  /**
+   * @return the hosts associated with this wrapper
+   */
+  public Set<String> getHosts() {
+    return hosts;
+  }
+
+
+  @Override
+  public String toString() {
+    return service + ":" + component + ":" + tasks + ":" + hosts;
+  }
+
+  /**
+   * @return the service name
+   */
+  public String getService() {
+    return service;
+  }
+
+  /**
+   * @return the component name
+   */
+  public String getComponent() {
+    return component;
+  }
+
+}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/main/resources/stacks/HDP/2.2/upgrades/upgrade-2.2.xml
----------------------------------------------------------------------
diff --git a/ambari-server/src/main/resources/stacks/HDP/2.2/upgrades/upgrade-2.2.xml b/ambari-server/src/main/resources/stacks/HDP/2.2/upgrades/upgrade-2.2.xml
index 9fc5752..752f504 100644
--- a/ambari-server/src/main/resources/stacks/HDP/2.2/upgrades/upgrade-2.2.xml
+++ b/ambari-server/src/main/resources/stacks/HDP/2.2/upgrades/upgrade-2.2.xml
@@ -32,22 +32,42 @@
 <upgrade xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
   <target>2.2.*.*</target>
   <order>
-    <service name="ZOOKEEPER">
-      <component>ZOOKEEPER_SERVER</component>
-      <component>ZOOKEEPER_CLIENT</component>
-    </service>
-    <service name="HDFS">
-      <component>NAMENODE</component>
-      <component>DATANODE</component>
-      <component>JOURNALNODE</component>
-    </service>
+    <group name="ZOOKEEPER" title="Zookeeper">
+      <service name="ZOOKEEPER">
+        <component>ZOOKEEPER_SERVER</component>
+        <component>ZOOKEEPER_CLIENT</component>
+      </service>
+    </group>
+    <group name="CORE_MASTER" title="Core Masters">
+      <service name="HDFS">
+        <component>JOURNALNODE</component>
+        <component>NAMENODE</component>
+      </service>
+      <service name="YARN">
+        <component>RESOURCEMANAGER</component>
+      </service>
+    </group>
+    <group name="CORE_SLAVES" title="Core Slaves" xsi:type="colocated">
+      <service name="HDFS">
+        <component>DATANODE</component>
+      </service>
+      <service name="HBASE">
+        <component>REGIONSERVER</component>
+      </service>
+      <service name="YARN">
+        <component>NODEMANAGER</component>
+      </service>
+      
+      <batch>
+        <percent>20</percent>
+        <message>Please run additional tests</message>
+      </batch>
+    </group>
   </order>
+  
   <processing>
     <service name="ZOOKEEPER">
       <component name="ZOOKEEPER_SERVER">
-        <batch xsi:type="count">
-          <count>1</count>
-        </batch>
         <!-- TODO, optimization
         <pre-upgrade>
           Find the leader by running
@@ -70,6 +90,9 @@
 
         $ quit
         -->
+        <upgrade>
+          <task xsi:type="restart" />
+        </upgrade>
       </component>
     </service>
     <service name="HDFS">
@@ -126,6 +149,10 @@
             <every>1</every>
           </task>
         </pre-upgrade>
+        
+        <upgrade>
+          <task xsi:type="restart" />
+        </upgrade>        
 
         <!-- This step should be done once the user clicks on the "Finalize" button. So the name post-upgrade is misleading. -->
         <post-upgrade>
@@ -139,9 +166,6 @@
       </component>
 
       <component name="DATANODE">
-        <batch xsi:type="percent">
-          <percent>20</percent>
-        </batch>
         <pre-upgrade>
           <!-- Shutdown the datanode,
 
@@ -176,21 +200,24 @@
             <command>su {{hadoop-env/hdfs_user}} -c 'hdfs dfsadmin -getDatanodeInfo {{hdfs-site/dfs.datanode.ipc.address}}'</command>
             <ignore>255</ignore>
           </task>
+          
 
           <!-- TODO, move this to HDFS Datanode restart. -->
           <task xsi:type="execute">
             <command>hdp-select set hadoop-hdfs-datanode {{version}}</command>
           </task>
+          
         </pre-upgrade>
+        <upgrade>
+          <task xsi:type="restart" />
+        </upgrade>
+
       </component>
 
       <component name="JOURNALNODE">
         <!-- Recommended after the Namenode, and only needed when HA is enabled. -->
-        <batch xsi:type="count">
-          <count>1</count>
-        </batch>
         <upgrade>
-          <!-- TODO -->
+          <task xsi:type="restart" />
         </upgrade>
       </component>
 

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
index 98175ee..fbb4f43 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/controller/internal/UpgradeResourceProviderTest.java
@@ -138,7 +138,7 @@ public class UpgradeResourceProviderTest {
     requestProps.put(UpgradeResourceProvider.UPGRADE_VERSION, "2.2.2.2");
 
     Request request = PropertyHelper.getCreateRequest(Collections.singleton(requestProps), null);
-    provider.createResources(request);
+    org.apache.ambari.server.controller.spi.RequestStatus status = provider.createResources(request);
 
 
     upgrades = upgradeDao.findUpgrades(cluster.getClusterId());
@@ -147,7 +147,7 @@ public class UpgradeResourceProviderTest {
     UpgradeEntity entity = upgrades.get(0);
     assertEquals(cluster.getClusterId(), entity.getClusterId().longValue());
 
-    assertEquals(1, entity.getUpgradeGroups().size());
+    assertEquals(3, entity.getUpgradeGroups().size());
 
     UpgradeGroupEntity group = entity.getUpgradeGroups().get(0);
 
@@ -155,7 +155,7 @@ public class UpgradeResourceProviderTest {
 
     assertTrue(group.getItems().get(0).getText().contains("Preparing"));
     assertTrue(group.getItems().get(1).getText().contains("Restarting"));
-    assertTrue(group.getItems().get(2).getText().contains("Finalizing"));
+    assertTrue(group.getItems().get(2).getText().contains("Completing"));
 
     ActionManager am = injector.getInstance(ActionManager.class);
     List<Long> requests = am.getRequestsByStatus(RequestStatus.IN_PROGRESS, 100, true);

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java b/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java
new file mode 100644
index 0000000..0d9aeea
--- /dev/null
+++ b/ambari-server/src/test/java/org/apache/ambari/server/state/UpgradeHelperTest.java
@@ -0,0 +1,146 @@
+/**
+ * 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 org.apache.ambari.server.state;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.ambari.server.AmbariException;
+import org.apache.ambari.server.api.services.AmbariMetaInfo;
+import org.apache.ambari.server.orm.GuiceJpaInitializer;
+import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
+import org.apache.ambari.server.state.UpgradeHelper.UpgradeGroupHolder;
+import org.apache.ambari.server.state.stack.UpgradePack;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import com.google.inject.Guice;
+import com.google.inject.Injector;
+import com.google.inject.persist.PersistService;
+
+/**
+ * Tests the {@link UpgradeHelper} class
+ */
+public class UpgradeHelperTest {
+
+  private Injector injector;
+  private AmbariMetaInfo ambariMetaInfo;
+
+  @Before
+  public void before() throws Exception {
+    injector = Guice.createInjector(new InMemoryDefaultTestModule());
+
+    injector.getInstance(GuiceJpaInitializer.class);
+
+    ambariMetaInfo = injector.getInstance(AmbariMetaInfo.class);
+    ambariMetaInfo.init();
+  }
+
+  @After
+  public void teardown() {
+    injector.getInstance(PersistService.class).stop();
+  }
+
+  @Test
+  public void testOrchestration() throws Exception {
+    Map<String, UpgradePack> upgrades = ambariMetaInfo.getUpgradePacks("foo", "bar");
+    assertTrue(upgrades.isEmpty());
+
+    upgrades = ambariMetaInfo.getUpgradePacks("HDP", "2.1.1");
+    assertTrue(upgrades.containsKey("upgrade_test"));
+    UpgradePack upgrade = upgrades.get("upgrade_test");
+    assertNotNull(upgrade);
+
+    Cluster cluster = makeCluster();
+
+    UpgradeHelper helper = new UpgradeHelper();
+    List<UpgradeGroupHolder> groups = helper.createUpgrade(cluster, upgrade);
+
+    assertEquals(3, groups.size());
+    assertEquals("ZOOKEEPER", groups.get(0).name);
+    assertEquals("CORE_MASTER", groups.get(1).name);
+    assertEquals("CORE_SLAVES", groups.get(2).name);
+
+    assertEquals(5, groups.get(0).items.size());
+    assertEquals(4, groups.get(1).items.size());
+    assertEquals(4, groups.get(2).items.size());
+  }
+
+  public Cluster makeCluster() throws AmbariException {
+    Clusters clusters = injector.getInstance(Clusters.class);
+    ServiceFactory serviceFactory = injector.getInstance(ServiceFactory.class);
+    String clusterName = "c1";
+
+    clusters.addCluster(clusterName);
+    Cluster c = clusters.getCluster(clusterName);
+    c.setDesiredStackVersion(new StackId("HDP-2.1.1"));
+    c.createClusterVersion(c.getDesiredStackVersion().getStackName(),
+        c.getDesiredStackVersion().getStackVersion(), "admin", RepositoryVersionState.CURRENT);
+    for (int i = 0; i < 3; i++) {
+      String hostName = "h" + (i+1);
+      clusters.addHost(hostName);
+      Host host = clusters.getHost(hostName);
+
+      Map<String, String> hostAttributes = new HashMap<String, String>();
+      hostAttributes.put("os_family", "redhat");
+      hostAttributes.put("os_release_version", "6");
+
+      host.setHostAttributes(hostAttributes);
+
+      host.persist();
+      clusters.mapHostToCluster(hostName, clusterName);
+    }
+
+    // !!! add services
+    c.addService(serviceFactory.createNew(c, "HDFS"));
+    c.addService(serviceFactory.createNew(c, "YARN"));
+    c.addService(serviceFactory.createNew(c, "ZOOKEEPER"));
+
+    Service s = c.getService("HDFS");
+    ServiceComponent sc = s.addServiceComponent("NAMENODE");
+    sc.addServiceComponentHost("h1");
+    sc = s.addServiceComponent("DATANODE");
+    sc.addServiceComponentHost("h2");
+    sc.addServiceComponentHost("h3");
+
+    s = c.getService("ZOOKEEPER");
+    sc = s.addServiceComponent("ZOOKEEPER_SERVER");
+    sc.addServiceComponentHost("h1");
+    sc.addServiceComponentHost("h2");
+    sc.addServiceComponentHost("h3");
+
+
+    s = c.getService("YARN");
+    sc = s.addServiceComponent("RESOURCEMANAGER");
+    sc.addServiceComponentHost("h2");
+
+    sc = s.addServiceComponent("NODEMANAGER");
+    sc.addServiceComponentHost("h1");
+    sc.addServiceComponentHost("h3");
+
+    return c;
+
+  }
+
+}

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/test/java/org/apache/ambari/server/state/stack/UpgradePackTest.java
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/java/org/apache/ambari/server/state/stack/UpgradePackTest.java b/ambari-server/src/test/java/org/apache/ambari/server/state/stack/UpgradePackTest.java
index 947b994..499b6c3 100644
--- a/ambari-server/src/test/java/org/apache/ambari/server/state/stack/UpgradePackTest.java
+++ b/ambari-server/src/test/java/org/apache/ambari/server/state/stack/UpgradePackTest.java
@@ -32,10 +32,7 @@ import org.apache.ambari.server.api.services.AmbariMetaInfo;
 import org.apache.ambari.server.orm.GuiceJpaInitializer;
 import org.apache.ambari.server.orm.InMemoryDefaultTestModule;
 import org.apache.ambari.server.state.stack.UpgradePack.ProcessingComponent;
-import org.apache.ambari.server.state.stack.upgrade.Batch;
-import org.apache.ambari.server.state.stack.upgrade.ConditionalBatch;
 import org.apache.ambari.server.state.stack.upgrade.ConfigureTask;
-import org.apache.ambari.server.state.stack.upgrade.CountBatch;
 import org.apache.ambari.server.state.stack.upgrade.ExecuteTask;
 import org.apache.ambari.server.state.stack.upgrade.ManualTask;
 import org.apache.ambari.server.state.stack.upgrade.Task;
@@ -81,6 +78,7 @@ public class UpgradePackTest {
 
   @Test
   public void testUpgradeParsing() throws Exception {
+    /*
     Map<String, UpgradePack> upgrades = ambariMetaInfo.getUpgradePacks("HDP", "2.1.1");
     assertTrue(upgrades.size() > 0);
 
@@ -163,9 +161,9 @@ public class UpgradePackTest {
 
     pc = up.getTasks().get("HDFS").get("DATANODE");
     assertNotNull(pc.batch);
-    assertEquals(Batch.Type.CONDITIONAL, pc.batch.getType());
-    assertEquals(15, ConditionalBatch.class.cast(pc.batch).initial);
-    assertEquals(50, ConditionalBatch.class.cast(pc.batch).remaining);
+//    assertEquals(Batch.Type.CONDITIONAL, pc.batch.getType());
+//    assertEquals(15, ConditionalBatch.class.cast(pc.batch).initial);
+//    assertEquals(50, ConditionalBatch.class.cast(pc.batch).remaining);
 
     pc = up.getTasks().get("ZOOKEEPER").get("ZOOKEEPER_SERVER");
     assertNotNull(pc.preTasks);
@@ -175,9 +173,9 @@ public class UpgradePackTest {
     assertNotNull(pc.tasks);
     assertEquals(1, pc.tasks.size());
     assertNotNull(pc.batch);
-    assertEquals(Batch.Type.COUNT, pc.batch.getType());
-    assertEquals(2, CountBatch.class.cast(pc.batch).count);
-
+//    assertEquals(Batch.Type.COUNT, pc.batch.getType());
+//    assertEquals(2, CountBatch.class.cast(pc.batch).count);
+  */
   }
 
 

http://git-wip-us.apache.org/repos/asf/ambari/blob/0eb91c22/ambari-server/src/test/resources/stacks/HDP/2.1.1/upgrades/upgrade_test.xml
----------------------------------------------------------------------
diff --git a/ambari-server/src/test/resources/stacks/HDP/2.1.1/upgrades/upgrade_test.xml b/ambari-server/src/test/resources/stacks/HDP/2.1.1/upgrades/upgrade_test.xml
index 90e907b..5a6dc92 100644
--- a/ambari-server/src/test/resources/stacks/HDP/2.1.1/upgrades/upgrade_test.xml
+++ b/ambari-server/src/test/resources/stacks/HDP/2.1.1/upgrades/upgrade_test.xml
@@ -18,42 +18,60 @@
 <upgrade xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
   <target>2.2.*</target>
   <order>
-    <service name="ZOOKEEPER">
-      <component>ZOOKEEPER_SERVER</component>
-      <component>ZOOKEEPER_CLIENT</component>
-    </service>
-    <service name="HDFS">
-      <component>JOURNALNODE</component>
-      <component>NAMENODE</component>
-      <component>DATANODE</component>
-    </service>
+    <group name="ZOOKEEPER" title="Zookeeper">
+      <service name="ZOOKEEPER">
+        <component>ZOOKEEPER_SERVER</component>
+        <component>ZOOKEEPER_CLIENT</component>
+      </service>
+    </group>
+    <group name="CORE_MASTER" title="Core Masters">
+      <service name="HDFS">
+        <component>JOURNALNODE</component>
+        <component>NAMENODE</component>
+      </service>
+      <service name="YARN">
+        <component>RESOURCEMANAGER</component>
+      </service>
+    </group>
+    <group name="CORE_SLAVES" title="Core Slaves" xsi:type="colocated">
+      <service name="HDFS">
+        <component>DATANODE</component>
+      </service>
+      <service name="HBASE">
+        <component>REGIONSERVER</component>
+      </service>
+      <service name="YARN">
+        <component>NODEMANAGER</component>
+      </service>
+      
+      <batch>
+        <percent>20</percent>
+        <message>Please run additional tests</message>
+      </batch>
+    </group>
   </order>
+
   <processing>
     <service name="ZOOKEEPER">
       <component name="ZOOKEEPER_SERVER">
-        <batch xsi:type="count">
-          <count>2</count>
-        </batch>
         <pre-upgrade>
           <task xsi:type="manual">
             <message>this is pre</message>
           </task>
         </pre-upgrade>
         <upgrade>
-          <task xsi:type="execute">
-            <command>ls -l</command>
-          </task>
+          <task xsi:type="restart" />
         </upgrade>
         <post-upgrade>
-          <task xsi:type="manual">
-            <message>this is post</message>
+          <task xsi:type="execute">
+            <command>ls</command>
           </task>
         </post-upgrade>
       </component>
     </service>
     <service name="HDFS">
       <component name="NAMENODE">
-        <upgrade>
+        <pre-upgrade>
           <task xsi:type="execute">
             <command>su - {hdfs-user} -c 'dosomething'</command>
           </task>
@@ -65,18 +83,36 @@
           <task xsi:type="manual">
             <message>Update your database</message>
           </task>
+        </pre-upgrade>
+        <upgrade>
+          <task xsi:type="restart" />
         </upgrade>
+        <post-upgrade>
+          <task xsi:type="execute">
+            <command>ls</command>
+          </task>
+        </post-upgrade>
       </component>
       <component name="DATANODE">
-        <batch xsi:type="conditional">
-          <initial>15</initial>
-          <remaining>50</remaining>
-        </batch>
         <upgrade>
+          <task xsi:type="restart" />
+        </upgrade>
+      </component>
+    </service>
+    <service name="YARN">
+      <component name="RESOURCEMANAGER">
+        <pre-upgrade>
           <task xsi:type="execute">
             <command>ls</command>
           </task>
-        </upgrade>
+        </pre-upgrade>
+      </component>
+      <component name="NODEMANAGER">
+        <pre-upgrade>
+          <task xsi:type="execute">
+            <command>ls</command>
+          </task>
+        </pre-upgrade>
       </component>
     </service>
   </processing>