Cleanup parallel_emerge a bit.

parallel_emerge

 #!/usr/bin/python2.6
 # Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Program to run emerge in parallel, for significant speedup.
 
 Usage:
  ./parallel_emerge --board=BOARD [emerge args] package
 
@@ skipping 13 matching lines @@
   Parallel Emerge unlocks two things during operation, here's what you
   must do to keep this safe:
     * Storage dir containing binary packages. - Don't emerge new
       packages while installing the existing ones.
     * Portage database - You must not examine deps while modifying the
       database. Therefore you may only parallelize "-p" read only access,
       or "--nodeps" write only access.
   Caveats:
     * Some ebuild packages have incorrectly specified deps, and running
       them in parallel is more likely to bring out these failures.
-    * Portage "world" is a record of explicitly installed packages. In
-      this parallel scheme, explicitly installed packages are installed
-      twice, once for the real install, and once for world file addition.
     * Some ebuilds (especially the build part) have complex dependencies
       that are not captured well by this script (it may be necessary to
       install an old package to build, but then install a newer version
-      of the same package for a runtime dep). This script is only
-      currently stable for binpkg installs.
+      of the same package for a runtime dep).
 """
 
 import os
 import re
 import shlex
 import subprocess
 import sys
 import tempfile
 import time
 
 
 def Usage():
   print "Usage:"
-  print " ./parallel_emerge --board=BOARD [emerge args] package"
+  print " ./parallel_emerge --board=BOARD --jobs=JOBS [emerge args] package"
   sys.exit(1)
 
 
 # These are dependencies that are not specified in the package,
 # but will prevent the package from installing.
 secret_deps = {}
 
 # Globals: package we are building, board we are targeting,
 # emerge args we are passing through.
 PACKAGE = None
 EMERGE_ARGS = ""
 BOARD = None
 
-# Runtime flags. TODO(): maybe make these commandline options or
-# environment veriables.
+# Runtime flags. TODO(): Maybe make these command-line options or
+# environment variables.
 VERBOSE = False
 AUTOCLEAN = False
 
 
 def ParseArgs(argv):
   """Set global vars based on command line.
 
   We need to be compatible with emerge arg format.
-  We scrape --board-XXX, and distinguish between args
+  We scrape --board=XXX and --jobs=XXX, and distinguish between args
   and package names.
-  TODO(): robustify argument processing, as it's possible to
+  TODO(): Robustify argument processing, as it's possible to
   pass in many two argument parameters that are difficult
-  to programmaitcally identify, although we don't currently
-  use any besides --bdeps <y|n>.
+  to programmatically identify, although we don't currently
+  use any besides --with-bdeps <y|n>.
   Args:
     argv: arguments list
   Returns:
     triplet of (package list, emerge argumens, board string)
   """
   if VERBOSE:
     print argv
   board_arg = None
+  jobs_arg = 0
   package_args = []
   emerge_passthru_args = ""
-  re_board = re.compile(r"--board=(?P<board>.*)")
   for arg in argv[1:]:
-    # Check if the arg begins with '-'
-    if arg[0] == "-" or arg == "y" or arg == "n":
-      # Specifically match "--board="
-      m = re_board.match(arg)
-      if m:
-        board_arg = m.group("board")
-      else:
-        # Pass through to emerge.
-        emerge_passthru_args = emerge_passthru_args + " " + arg
+    # Specifically match "--board=" and "--jobs=".
+    if arg.startswith("--board="):
+      board_arg = arg.replace("--board=", "")
+    elif arg.startswith("--jobs="):
+      try:
+        jobs_arg = int(arg.replace("--jobs=", ""))
+      except ValueError:
+        print "Unrecognized argument:", arg
+        Usage()
+        sys.exit(1)
+    elif arg.startswith("-") or arg == "y" or arg == "n":
+      # Not a package name, so pass through to emerge.
+      emerge_passthru_args = emerge_passthru_args + " " + arg
     else:
-      # Only non-dashed arg should be the target package.
       package_args.append(arg)
 
-  if not package_args:
+  if not package_args and not emerge_passthru_args:
     Usage()
     sys.exit(1)
 
+  # Default to lots of jobs
+  if jobs_arg <= 0:
+    jobs_arg = 256
+
   # Set globals.
-  return " ".join(package_args), emerge_passthru_args, board_arg
+  return " ".join(package_args), emerge_passthru_args, board_arg, jobs_arg
 
 
 def EmergeCommand():
   """Helper function to return the base emerge commandline.
 
   This is configured for board type, and including pass thru args,
-  using global variables. TODO(): unglobalfy.
+  using global variables. TODO(): Unglobalfy.
   Returns:
     string containing emerge command.
   """
   emerge = "emerge"
   if BOARD:
     emerge += "-" + BOARD
   return emerge + " " + EMERGE_ARGS
 
 
 def GetDepsFromPortage(package):
   """Get dependency tree info by running emerge.
 
   Run 'emerge -p --debug package', and get a text output of all deps.
-  TODO(): Put dep caclation in a library, as cros_extract_deps
+  TODO(): Put dep calculation in a library, as cros_extract_deps
   also uses this code.
   Args:
-    package: string containing the packages to build.
+    package: String containing the packages to build.
   Returns:
-    text output of emerge -p --debug, which can be processed elsewhere.
+    Text output of emerge -p --debug, which can be processed elsewhere.
   """
   print "Calculating deps for package %s" % package
   cmdline = EmergeCommand() + " -p --debug --color=n " + package
   print "+ %s" % cmdline
 
   # Store output in a temp file as it is too big for a unix pipe.
   stderr_buffer = tempfile.TemporaryFile()
   stdout_buffer = tempfile.TemporaryFile()
   # Launch the subprocess.
   start = time.time()
@@ skipping 18 matching lines @@
     print "Failed to generate deps"
     sys.exit(1)
 
   return lines
 
 
 def DepsToTree(lines):
   """Regex the output from 'emerge --debug' to generate a nested dict of deps.
 
   Args:
-    lines: output from 'emerge -p --debug package'
+    lines: Output from 'emerge -p --debug package'.
   Returns:
-    dep_tree: nested dict of dependencies, as specified by emerge.
-              there may be dupes, or circular deps.
+    dep_tree: Nested dict of dependencies, as specified by emerge.
+              There may be dupes, or circular deps.
 
   We need to regex lines as follows:
   hard-host-depends depends on
   ('ebuild', '/', 'dev-lang/swig-1.3.36', 'merge') depends on
     ('ebuild', '/', 'dev-lang/perl-5.8.8-r8', 'merge') (buildtime)
   ('binary', '/.../rootfs/', 'sys-auth/policykit-0.9-r1', 'merge') depends on
   ('binary', '/.../rootfs/', 'x11-misc/xbitmaps-1.1.0', 'merge') (no children)
   """
 
   re_deps = re.compile(r"(?P<indent>\W*)\(\'(?P<pkgtype>\w+)\', "
                        r"\'(?P<destination>[\w/\.-]+)\',"
                        r" \'(?P<pkgdir>[\w\+-]+)/(?P<pkgname>[\w\+-]+)-"
                        r"(?P<version>\d+[\w\.-]*)\', \'(?P<action>\w+)\'\) "
                        r"(?P<deptype>(depends on|\(.*\)))")
   re_origdeps = re.compile(r"(?P<pkgname>[\w\+/-]+) depends on")
   re_installed_package = re.compile(
       r"\[(?P<desc>[^\]]*)\] "
       r"(?P<pkgdir>[\w\+-]+)/"
       r"(?P<pkgname>[\w\+-]+)-"
       r"(?P<version>\d+[\w\.-]*)( \["
       r"(?P<oldversion>\d+[\w\.-]*)\])?"
   )
-  re_failed = re.compile(r".*depends on.*")
+  re_failed = re.compile(r".*\) depends on.*")
   deps_tree = {}
   deps_stack = []
   deps_info = {}
   for line in lines:
     m = re_deps.match(line)
     m_orig = re_origdeps.match(line)
     m_installed = re_installed_package.match(line)
     if m:
       pkgname = m.group("pkgname")
       pkgdir = m.group("pkgdir")
       pkgtype = m.group("pkgtype")
       indent = m.group("indent")
       doins = m.group("action")
       deptype = m.group("deptype")
       depth = 1
       if not indent:
         depth = 0
       version = m.group("version")
 
       # If we are indented, we should have
       # found a "depends on" previously.
       if len(deps_stack) < depth:
         print "FAIL: corrupt input at:"
         print line
         print "No Parent."
         sys.exit(1)
 
       # Go step by step through stack and tree
-      # until we find our parent. Generate
+      # until we find our parent.
       updatedep = deps_tree
       for i in range(0, depth):
         updatedep = updatedep[deps_stack[i]]["deps"]
 
       # Pretty print what we've captured.
       indent = "|" + "".ljust(depth, "_")
       fullpkg = "%s/%s-%s" % (pkgdir, pkgname, version)
       if VERBOSE:
         print ("" + indent + " " + pkgdir + "/" + pkgname + " - " +
                version + " (" + pkgtype + ", " +  doins +
@@ skipping 48 matching lines @@
       fullpkg = "%s/%s-%s" % (pkgdir, pkgname, version)
       deps_info[fullpkg] = {"idx": len(deps_info),
                             "pkgdir": pkgdir,
                             "pkgname": pkgname,
                             "oldversion": oldversion,
                             "uninstall": uninstall}
     else:
       # Is this a package that failed to match our huge regex?
       m = re_failed.match(line)
       if m:
+        print "\n".join(lines)
         print "FAIL: Couldn't understand line:"
         print line
         sys.exit(1)
 
   return deps_tree, deps_info
 
 
 def PrintTree(deps, depth=""):
   """Print the deps we have seen in the emerge output.
 
   Args:
-   deps: dependency tree structure.
-   depth: allows printing the tree recursively, with indentation.
+   deps: Dependency tree structure.
+   depth: Allows printing the tree recursively, with indentation.
   """
   for entry in deps:
     action = deps[entry]["action"]
     print "%s %s (%s)" % (depth, entry, action)
     PrintTree(deps[entry]["deps"], depth=depth + "  ")
 
 
 def GenDependencyGraph(deps_tree, deps_info):
   """Generate a doubly linked dependency graph.
 
   Args:
-    deps_tree: dependency tree structure.
-    deps_info: more info on the dependencies.
+    deps_tree: Dependency tree structure.
+    deps_info: More details on the dependencies.
   Returns:
     Deps graph in the form of a dict of packages, with each package
     specifying a "needs" list and "provides" list.
   """
   deps_map = {}
 
   def ReverseTree(packages):
     """Convert tree to digraph.
 
     Take the tree of package -> requirements and reverse it to a digraph of
-    buildable packages -> packages they unblock
+    buildable packages -> packages they unblock.
     Args:
-      packages: tree(s) of dependencies
+      packages: Tree(s) of dependencies.
     Returns:
-      unsanitized digraph
+      Unsanitized digraph.
     """
     for pkg in packages:
       action = packages[pkg]["action"]
       this_pkg = deps_map.setdefault(
           pkg, {"needs": set(), "provides": set(), "action": "nomerge"})
       if action != "nomerge":
         this_pkg["action"] = action
         this_pkg["deps_info"] = deps_info.get(pkg)
       ReverseTree(packages[pkg]["deps"])
       for dep, dep_item in packages[pkg]["deps"].items():
         dep_pkg = deps_map[dep]
         dep_type = dep_item["deptype"]
         if dep_type != "(runtime_post)":
           dep_pkg["provides"].add(pkg)
           this_pkg["needs"].add(dep)
 
   def RemoveInstalledPackages():
-    """Remove installed packages, propagating dependencies"""
+    """Remove installed packages, propagating dependencies."""
 
     rm_pkgs = set(deps_map.keys()) - set(deps_info.keys())
     for pkg in rm_pkgs:
       this_pkg = deps_map[pkg]
       needs = this_pkg["needs"]
       provides = this_pkg["provides"]
       for dep in needs:
         dep_provides = deps_map[dep]["provides"]
         dep_provides.update(provides)
         dep_provides.discard(pkg)
         dep_provides.discard(dep)
       for target in provides:
         target_needs = deps_map[target]["needs"]
         target_needs.update(needs)
         target_needs.discard(pkg)
         target_needs.discard(target)
       del deps_map[pkg]
 
-
   def SanitizeDep(basedep, currdep, oldstack, limit):
     """Search for circular deps between basedep and currdep, then recurse.
 
     Args:
-      basedep: original dependency, top of stack.
-      currdep: bottom of our current recursion, bottom of stack.
-      oldstack: current dependency chain.
-      limit: how many more levels of recusion to go through, max.
+      basedep: Original dependency, top of stack.
+      currdep: Bottom of our current recursion, bottom of stack.
+      oldstack: Current dependency chain.
+      limit: How many more levels of recusion to go through, max.
     TODO(): Break RDEPEND preferentially.
     Returns:
       True iff circular dependencies are found.
     """
     if limit == 0:
       return
     for dep in deps_map[currdep]["needs"]:
       stack = oldstack + [dep]
       if basedep in deps_map[dep]["needs"] or dep == basedep:
         if dep != basedep:
-           stack += [basedep]
+          stack += [basedep]
         print "Remove cyclic dependency from:"
         for i in xrange(0, len(stack) - 1):
           print "  %s -> %s " % (stack[i], stack[i+1])
         return True
       if dep not in oldstack and SanitizeDep(basedep, dep, stack, limit - 1):
         return True
     return
 
   def SanitizeTree():
     """Remove circular dependencies up to cycle length 32."""
@@ skipping 64 matching lines @@
     print "Pending %s, Ready %s, Running %s, Retrying %s, Total %s" % (
         len(self._deps_map), len(self._emerge_queue),
         len(self._jobs), len(self._retry_queue), self._total_jobs)
 
   def _LaunchOneEmerge(self, target):
     """Run emerge --nodeps to do a single package install.
 
     If this is a pseudopackage, that means we're done, and can select in in the
     world file.
     Args:
-      target: the full package name of the package to install.
+      target: The full package name of the package to install.
               eg. "sys-apps/portage-2.17"
     Returns:
-      triplet containing (target name, subprocess object, output buffer object)
+      Triplet containing (target name, subprocess object, output buffer object).
     """
     if target.startswith("original-"):
       # "original-" signifies one of the packages we originally requested.
       # Since we have explicitly installed the versioned package as a dep of
       # this, we only need to tag in "world" that we are done with this
       # install request. "--select -n" indicates an addition to "world"
       # without an actual install.
       newtarget = target.replace("original-", "")
       cmdline = EmergeCommand() + " --nodeps --select --noreplace " + newtarget
     else:
       # This package is a dependency of something we specifically
       # requested. Therefore we should install it but not allow it
       # in the "world" file, which represents explicit intalls.
       # "--oneshot" here will prevent it from being tagged in world.
       cmdline = EmergeCommand() + " --nodeps --oneshot =" + target
       deps_info = self._deps_map[target]["deps_info"]
       if deps_info["uninstall"]:
         package = "%(pkgdir)s/%(pkgname)s-%(oldversion)s" % deps_info
         cmdline += " && %s -1C =%s" % (EmergeCommand(), package)
 
     print "+ %s" % cmdline
 
     # Store output in a temp file as it is too big for a unix pipe.
     stdout_buffer = tempfile.TemporaryFile()
     # Modify the environment to disable locking.
     portage_env = os.environ.copy()
     portage_env["PORTAGE_LOCKS"] = "false"
     portage_env["UNMERGE_DELAY"] = "0"
     # Autoclean rummages around in the portage database and uninstalls
-    # old packages. Definitely not necessary for build_image. However
-    # it may be necessary for incremental build_packages. It may also
-    # not be parallel safe.
+    # old packages. It's not parallel safe, so we skip it. Instead, we
+    # handle the cleaning ourselves by uninstalling old versions of any
+    # new packages we install.
     if not AUTOCLEAN:
       portage_env["AUTOCLEAN"] = "no"
     # Launch the subprocess.
     emerge_proc = subprocess.Popen(
         cmdline, shell=True, stdout=stdout_buffer,
         stderr=subprocess.STDOUT, bufsize=64*1024, env=portage_env)
 
     return (target, emerge_proc, stdout_buffer)
 
   def _Finish(self, target):
@@ skipping 13 matching lines @@
       print "Retrying emerge of %s." % target
 
   def Run(self):
     """Run through the scheduled ebuilds.
 
     Keep running so long as we have uninstalled packages in the
     dependency graph to merge.
     """
     while self._deps_map:
       # If we have packages that are ready, kick them off.
-      if self._emerge_queue:
+      if self._emerge_queue and len(self._jobs) < JOBS:
         target = self._emerge_queue.pop(0)
         action = self._deps_map[target]["action"]
         # We maintain a tree of all deps, if this doesn't need
         # to be installed just free up it's children and continue.
         # It is possible to reinstall deps of deps, without reinstalling
         # first level deps, like so:
         # chromeos (merge) -> eselect (nomerge) -> python (merge)
         if action == "nomerge":
           self._Finish(target)
         else:
           # Kick off the build if it's marked to be built.
           print "Emerging %s (%s)" % (target, action)
           job = self._LaunchOneEmerge(target)
           # Append it to the active jobs list.
           self._jobs.append(job)
         continue
       # Wait a bit to see if maybe some jobs finish. You can't
       # wait on a set of jobs in python, so we'll just poll.
       time.sleep(1)
 
       # Check here that we are actually waiting for something.
       if (not self._emerge_queue and
           not self._jobs and
           self._deps_map):
-        # If we have failed on a package retry it now.
+        # If we have failed on a package, retry it now.
         if self._retry_queue:
           self._Retry()
         # If we have failed a package twice, just give up.
         elif self._failed:
           for failure, output in self._failed.items():
             print "Package failed: %s" % failure
             print output
           PrintDepsMap(self._deps_map)
           print "Packages failed: %s" % " ,".join(self._failed.keys())
           sys.exit(1)
@@ skipping 10 matching lines @@
           # Clean up the subprocess.
           job.wait()
           # Get the output if we want to print it.
           stdout.seek(0)
           output = stdout.read()
 
           # Remove from active jobs list, we are done with this process.
           self._jobs.remove((target, job, stdout))
 
           # Print if necessary.
-          if VERBOSE:
+          if VERBOSE or job.returncode != 0:
             print output
           if job.returncode != 0:
             # Handle job failure.
             if target in self._failed:
               # If this job has failed previously, give up.
               print "Failed %s. Your build has failed." % target
             else:
               # Queue up this build to try again after a long while.
               self._retry_queue.append(target)
               self._failed[target] = output
               print "Failed %s, retrying later." % target
           else:
             if target in self._failed and self._retry_queue:
               # If we have successfully retried a failed package, and there
               # are more failed packages, try the next one. We will only have
               # one retrying package actively running at a time.
               self._Retry()
 
             print "Completed %s" % target
             # Mark as completed and unblock waiting ebuilds.
             self._Finish(target)
 
             # Print an update.
             self._Status()
 
 
 # Main control code.
+PACKAGE, EMERGE_ARGS, BOARD, JOBS = ParseArgs(sys.argv)
+
+if not PACKAGE:
+  # No packages. Pass straight through to emerge.
+  # Allows users to just type ./parallel_emerge --depclean
+  sys.exit(os.system(EmergeCommand()))
+
 print "Starting fast-emerge."
-PACKAGE, EMERGE_ARGS, BOARD = ParseArgs(sys.argv)
 print " Building package %s on %s (%s)" % (PACKAGE, EMERGE_ARGS, BOARD)
-
 print "Running emerge to generate deps"
 deps_output = GetDepsFromPortage(PACKAGE)
 print "Processing emerge output"
 dependency_tree, dependency_info = DepsToTree(deps_output)
 if VERBOSE:
   print "Print tree"
   PrintTree(dependency_tree)
 
 print "Generate dependency graph."
 dependency_graph = GenDependencyGraph(dependency_tree, dependency_info)
 
 if VERBOSE:
   PrintDepsMap(dependency_graph)
 
 # Run the queued emerges.
 scheduler = EmergeQueue(dependency_graph)
 scheduler.Run()
 
 print "Done"