diff --git a/cros_update.py b/cros_update.py
index 0bf6a4f..5d4475a 100644
--- a/cros_update.py
+++ b/cros_update.py
@@ -43,6 +43,9 @@
   remote_access = None
   timeout_util = None
 
+# The build channel for recovering host's stateful partition
+STABLE_BUILD_CHANNEL = 'stable-channel'
+
 # Timeout for CrOS auto-update process.
 CROS_UPDATE_TIMEOUT_MIN = 30
 
@@ -86,6 +89,10 @@
                              dest='full_update', default=False,
                              help=('force a rootfs update, skip stateful '
                                    'update'))
+    self.parser.add_argument('--original_build', action='store', type=str,
+                             dest='original_build', default='',
+                             help=('force stateful update with the same '
+                                   'version of previous rootfs partition'))
 
   def ParseArgs(self):
     """Parse and process command line arguments."""
@@ -110,7 +117,7 @@
   """
   def __init__(self, host_name, build_name, static_dir, progress_tracker=None,
                log_file=None, au_tempdir=None, force_update=False,
-               full_update=False):
+               full_update=False, original_build=None):
     self.host_name = host_name
     self.build_name = build_name
     self.static_dir = static_dir
@@ -119,6 +126,7 @@
     self.au_tempdir = au_tempdir
     self.force_update = force_update
     self.full_update = full_update
+    self.original_build = original_build
 
   def _WriteAUStatus(self, content):
     if self.progress_tracker:
@@ -152,6 +160,19 @@
     self._WriteAUStatus('post-check rootfs update')
     cros_updater.PostCheckRootfsUpdate()
 
+  def _GetOriginalPayloadDir(self):
+    """Get the directory of original payload.
+
+    Returns:
+      The directory of original payload, whose format is like:
+          'static/stable-channel/link/3428.210.0'
+    """
+    if self.original_build:
+      return os.path.join(self.static_dir, '%s/%s' % (STABLE_BUILD_CHANNEL,
+                                                      self.original_build))
+    else:
+      return None
+
   def TriggerAU(self):
     """Execute auto update for cros_host.
 
@@ -171,11 +192,14 @@
 
         logging.debug('Remote device %s is connected', self.host_name)
         payload_dir = os.path.join(self.static_dir, self.build_name)
+        original_payload_dir = self._GetOriginalPayloadDir()
+
         chromeos_AU = auto_updater.ChromiumOSUpdater(
             device, self.build_name, payload_dir,
             dev_dir=os.path.abspath(os.path.dirname(__file__)),
             tempdir=self.au_tempdir,
             log_file=self.log_file,
+            original_payload_dir=original_payload_dir,
             yes=True)
         chromeos_AU.CheckPayloads()
 
@@ -248,11 +272,7 @@
     AU_parser.parser.print_help()
     sys.exit(1)
 
-  host_name = AU_parser.options.host_name
-  build_name = AU_parser.options.build_name
-  static_dir = AU_parser.options.static_dir
-  force_update = AU_parser.options.force_update
-  full_update = AU_parser.options.full_update
+  options = AU_parser.options
 
   # Use process group id as the unique id in track and log files, since
   # os.setsid is executed before the current process is run.
@@ -261,23 +281,25 @@
 
   # Setting log files for CrOS auto-update process.
   # Log file:  file to record every details of CrOS auto-update process.
-  log_file = cros_update_progress.GetExecuteLogFile(host_name, pgid)
+  log_file = cros_update_progress.GetExecuteLogFile(options.host_name, pgid)
   logging.info('Writing executing logs into file: %s', log_file)
   logConfig.SetFileHandler(log_file)
 
   # Create a progress_tracker for tracking CrOS auto-update progress.
-  progress_tracker = cros_update_progress.AUProgress(host_name, pgid)
+  progress_tracker = cros_update_progress.AUProgress(options.host_name, pgid)
 
   # Create a dir for temporarily storing devserver codes and logs.
-  au_tempdir = cros_update_progress.GetAUTempDirectory(host_name, pgid)
+  au_tempdir = cros_update_progress.GetAUTempDirectory(options.host_name, pgid)
 
   # Create cros_update instance to run CrOS auto-update.
-  cros_updater_trigger = CrOSUpdateTrigger(host_name, build_name, static_dir,
-                                           progress_tracker=progress_tracker,
-                                           log_file=log_file,
-                                           au_tempdir=au_tempdir,
-                                           force_update=force_update,
-                                           full_update=full_update)
+  cros_updater_trigger = CrOSUpdateTrigger(
+      options.host_name, options.build_name, options.static_dir,
+      progress_tracker=progress_tracker,
+      log_file=log_file,
+      au_tempdir=au_tempdir,
+      force_update=options.force_update,
+      full_update=options.full_update,
+      original_build=options.original_build)
 
   # Set timeout the cros-update process.
   try:
diff --git a/devserver.py b/devserver.py
index 7aa81c6..8ab9d68 100755
--- a/devserver.py
+++ b/devserver.py
@@ -64,8 +64,6 @@
 import artifact_info
 import build_artifact
 import cherrypy_ext
-import cros_update
-import cros_update_progress
 import common_util
 import devserver_constants
 import downloader
@@ -95,6 +93,24 @@
        'not be collected.', e)
   psutil = None
 
+# Use try-except to skip unneccesary import for simple use case, eg. running
+# devserver on host.
+try:
+  import cros_update
+  import cros_update_progress
+except ImportError as e:
+  _Log('cros_update cannot be imported: %r', e)
+  cros_update = None
+  cros_update_progress = None
+
+# only import setup_chromite before chromite import.
+import setup_chromite # pylint: disable=unused-import
+try:
+  from chromite.lib.paygen import gspaths
+except ImportError as e:
+  _Log('chromite cannot be imported: %r', e)
+  gspaths = None
+
 try:
   import android_build
 except ImportError as e:
@@ -478,6 +494,14 @@
 
 
 def _check_base_args_for_auto_update(kwargs):
+  """Check basic args required for auto-update.
+
+  Args:
+    kwargs: the parameters to be checked.
+
+  Raises:
+    DevServerHTTPError if required parameters don't exist in kwargs.
+  """
   if 'host_name' not in kwargs:
     raise common_util.DevServerHTTPError(KEY_ERROR_MSG % 'host_name')
 
@@ -486,6 +510,18 @@
 
 
 def _parse_boolean_arg(kwargs, key):
+  """Parse boolean arg from kwargs.
+
+  Args:
+    kwargs: the parameters to be checked.
+    key: the key to be parsed.
+
+  Returns:
+    The boolean value of kwargs[key], or False if key doesn't exist in kwargs.
+
+  Raises:
+    DevServerHTTPError if kwargs[key] is not a boolean variable.
+  """
   if key in kwargs:
     if kwargs[key] == 'True':
       return True
@@ -497,6 +533,34 @@
   else:
     return False
 
+def _parse_string_arg(kwargs, key):
+  """Parse string arg from kwargs.
+
+  Args:
+    kwargs: the parameters to be checked.
+    key: the key to be parsed.
+
+  Returns:
+    The string value of kwargs[key], or None if key doesn't exist in kwargs.
+  """
+  if key in kwargs:
+    return kwargs[key]
+  else:
+    return None
+
+def _build_uri_from_build_name(build_name):
+  """Get build url from a given build name.
+
+  Args:
+    build_name: the build name to be parsed, whose format is
+        'board/release_version'.
+
+  Returns:
+    The release_archive_url on Google Storage for this build name.
+  """
+  return gspaths.ChromeosReleases.BuildUri(
+      cros_update.STABLE_BUILD_CHANNEL, build_name.split('/')[0],
+      build_name.split('/')[1])
 
 class ApiRoot(object):
   """RESTful API for Dev Server information."""
@@ -822,12 +886,24 @@
     force_update = _parse_boolean_arg(kwargs, 'force_update')
     full_update = _parse_boolean_arg(kwargs, 'full_update')
     async = _parse_boolean_arg(kwargs, 'async')
+    original_build = _parse_string_arg(kwargs, 'original_build')
 
     if async:
       path = os.path.dirname(os.path.abspath(__file__))
       execute_file = os.path.join(path, 'cros_update.py')
       args = (AUTO_UPDATE_CMD % (execute_file, host_name, build_name,
                                  updater.static_dir))
+
+      # The original_build's format is like: link/3428.210.0
+      # The corresponding release_archive_url's format is like:
+      #    gs://chromeos-releases/stable-channel/link/3428.210.0
+      if original_build:
+        release_archive_url = _build_uri_from_build_name(original_build)
+        # First staging the stateful.tgz synchronousely.
+        self.stage(files='stateful.tgz', async=False,
+                   archive_url=release_archive_url)
+        args = ('%s --original_build %s' % (args, original_build))
+
       if force_update:
         args = ('%s --force_update' % args)
 
@@ -845,7 +921,8 @@
       return json.dumps((True, pid))
     else:
       cros_update_trigger = cros_update.CrOSUpdateTrigger(
-          host_name, build_name, updater.static_dir)
+          host_name, build_name, updater.static_dir, force_update=force_update,
+          full_update=full_update, original_build=original_build)
       cros_update_trigger.TriggerAU()
 
   @cherrypy.expose
@@ -1494,8 +1571,18 @@
       The count of processes that match the given command pattern.
     """
     try:
-      return int(subprocess.check_output(
-          'pgrep -fc "%s"' % process_cmd_pattern, shell=True))
+      # Use Popen instead of check_output since the latter cannot run with old
+      # python version (less than 2.7)
+      proc = subprocess.Popen(
+          'pgrep -fc "%s"' % process_cmd_pattern,
+          stdout=subprocess.PIPE,
+          stderr=subprocess.PIPE,
+          shell=True)
+      cmd_output, cmd_error = proc.communicate()
+      if cmd_error:
+        _Log('Error happened when getting process count: %s' % cmd_error)
+
+      return int(cmd_output)
     except subprocess.CalledProcessError:
       return 0