Modify --requirements_file behavior to cache locally packages

The packages will be staged and later installed in the workers in a
manner that reduces to zero (potentially) PyPI downloads during worker
startup. To avoid any downloads the requirements file must contain the
transitive  closure of all dependencies which is the typical result of
running 'pip freeze'.

----Release Notes----
The --requirements_file option caches locally all the packages specified
in the requirements file and then stages the files during job submission.

Introduce --requirements_cache as an option to override the default location
for the cache of downloaded packages.

[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=120947879

Author: silviulica

google/cloud/dataflow/utils/dependency.py

@@ -198,8 +198,27 @@ def _stage_extra_packages(extra_packages,
   return resources
 
 
-def stage_job_resources(options, file_copy=_dependency_file_copy,
-                        build_setup_args=None, temp_dir=None):
+def _populate_requirements_cache(requirements_file, cache_dir):
+  # The 'pip download' command will not download again if it finds the
+  # tarball with the proper version already present.
+  # It will get the packages downloaded in the order they are presented in
+  # the requirements file and will not download package dependencies.
+  cmd_args = [
+      'pip', 'install', '--download', cache_dir,
+      '-r', requirements_file,
+      # Download from PyPI source distributions.
+      '--no-binary', ':all:']
+  logging.info('Executing command: %s', cmd_args)
+  result = processes.call(cmd_args)
+  if result != 0:
+    raise RuntimeError(
+        'Failed to execute command: %s. Exit code %d',
+        cmd_args, result)
+
+
+def stage_job_resources(
+    options, file_copy=_dependency_file_copy, build_setup_args=None,
+    temp_dir=None, populate_requirements_cache=_populate_requirements_cache):
   """Creates (if needed) and stages job resources to options.staging_location.
 
   Args:
@@ -214,6 +233,8 @@ def stage_job_resources(options, file_copy=_dependency_file_copy,
       testing.
     temp_dir: Temporary folder where the resource building can happen. If None
       then a unique temp directory will be created. Used only for testing.
+    populate_requirements_cache: Callable for populating the requirements cache.
+      Used only for testing.
 
   Returns:
     A list of file names (no paths) for the resources staged. All the files
@@ -247,6 +268,20 @@ def stage_job_resources(options, file_copy=_dependency_file_copy,
                                   REQUIREMENTS_FILE)
     file_copy(setup_options.requirements_file, staged_path)
     resources.append(REQUIREMENTS_FILE)
+    requirements_cache_path = (
+        os.path.join(tempfile.gettempdir(), 'dataflow-requirements-cache')
+        if setup_options.requirements_cache is None
+        else setup_options.requirements_cache)
+    # Populate cache with packages from requirements and stage the files
+    # in the cache.
+    if not os.path.exists(requirements_cache_path):
+      os.makedirs(requirements_cache_path)
+    populate_requirements_cache(
+        setup_options.requirements_file, requirements_cache_path)
+    for pkg in  glob.glob(os.path.join(requirements_cache_path, '*')):
+      file_copy(pkg, utils.path.join(google_cloud_options.staging_location,
+                                     os.path.basename(pkg)))
+      resources.append(os.path.basename(pkg))
 
   # Handle a setup file if present.
   # We will build the setup package locally and then copy it to the staging

google/cloud/dataflow/utils/dependency_test.py

@@ -43,6 +43,11 @@ def create_temp_file(self, path, contents):
       f.write(contents)
       return f.name
 
+  def populate_requirements_cache(self, requirements_file, cache_dir):
+    _ = requirements_file
+    self.create_temp_file(os.path.join(cache_dir, 'abc.txt'), 'nothing')
+    self.create_temp_file(os.path.join(cache_dir, 'def.txt'), 'nothing')
+
   def test_no_staging_location(self):
     with self.assertRaises(RuntimeError) as cm:
       dependency.stage_job_resources(PipelineOptions())
@@ -98,12 +103,16 @@ def test_with_requirements_file(self):
     self.create_temp_file(
         os.path.join(source_dir, dependency.REQUIREMENTS_FILE), 'nothing')
     self.assertEqual(
-        [dependency.REQUIREMENTS_FILE,
-         names.PICKLED_MAIN_SESSION_FILE],
-        dependency.stage_job_resources(options))
+        sorted([dependency.REQUIREMENTS_FILE, names.PICKLED_MAIN_SESSION_FILE,
+                'abc.txt', 'def.txt']),
+        sorted(dependency.stage_job_resources(
+            options,
+            populate_requirements_cache=self.populate_requirements_cache)))
     self.assertTrue(
         os.path.isfile(
             os.path.join(staging_dir, dependency.REQUIREMENTS_FILE)))
+    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt')))
+    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt')))
 
   def test_requirements_file_not_present(self):
     staging_dir = tempfile.mkdtemp()
@@ -112,12 +121,38 @@ def test_requirements_file_not_present(self):
       options.view_as(GoogleCloudOptions).staging_location = staging_dir
       self.update_options(options)
       options.view_as(SetupOptions).requirements_file = 'nosuchfile'
-      dependency.stage_job_resources(options)
+      dependency.stage_job_resources(
+          options, populate_requirements_cache=self.populate_requirements_cache)
     self.assertEqual(
         cm.exception.message,
         'The file %s cannot be found. It was specified in the '
         '--requirements_file command line option.' % 'nosuchfile')
 
+  def test_with_requirements_file_and_cache(self):
+    staging_dir = tempfile.mkdtemp()
+    source_dir = tempfile.mkdtemp()
+
+    options = PipelineOptions()
+    options.view_as(GoogleCloudOptions).staging_location = staging_dir
+    self.update_options(options)
+    options.view_as(SetupOptions).requirements_file = os.path.join(
+        source_dir, dependency.REQUIREMENTS_FILE)
+    options.view_as(SetupOptions).requirements_cache = os.path.join(
+        tempfile.gettempdir(), 'alternative-cache-dir')
+    self.create_temp_file(
+        os.path.join(source_dir, dependency.REQUIREMENTS_FILE), 'nothing')
+    self.assertEqual(
+        sorted([dependency.REQUIREMENTS_FILE, names.PICKLED_MAIN_SESSION_FILE,
+                'abc.txt', 'def.txt']),
+        sorted(dependency.stage_job_resources(
+            options,
+            populate_requirements_cache=self.populate_requirements_cache)))
+    self.assertTrue(
+        os.path.isfile(
+            os.path.join(staging_dir, dependency.REQUIREMENTS_FILE)))
+    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt')))
+    self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt')))
+
   def test_with_setup_file(self):
     staging_dir = tempfile.mkdtemp()
     source_dir = tempfile.mkdtemp()

google/cloud/dataflow/utils/options.py

@@ -364,10 +364,19 @@ def _add_argparse_args(cls, parser):
         help=
         ('Path to a requirements file containing package dependencies. '
          'Typically it is produced by a pip freeze command. More details: '
-         'https://pip.pypa.io/en/latest/reference/pip_freeze.html. If '
-         'specified, the worker will install the required dependenciesi before'
-         ' running any custom code. Typically the file is named '
-         'requirements.txt.'))
+         'https://pip.pypa.io/en/latest/reference/pip_freeze.html. '
+         'If used, all the packages specified will be downloaded, '
+         'cached (use --requirements_cache to change default location), '
+         'and then staged so that they can be automatically installed in '
+         'workers during startup. The cache is refreshed as needed '
+         'avoiding extra downloads for existing packages. Typically the '
+         'file is named requirements.txt.'))
+    parser.add_argument(
+        '--requirements_cache',
+        default=None,
+        help=
+        ('Path to a folder to cache the packages specified in '
+         'the requirements file using the --requirements_file option.'))
     parser.add_argument(
         '--setup_file',
         default=None,