[Python] Log dependencies installed in submission environment

sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py

@@ -211,15 +211,13 @@ def test_environment_override_translation_legacy_worker_harness_image(self):
           p | ptransform.Create([1, 2, 3])
           | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
           | ptransform.GroupByKey())
-    self.assertEqual(
-        list(remote_runner.proto_pipeline.components.environments.values()),
-        [
-            beam_runner_api_pb2.Environment(
-                urn=common_urns.environments.DOCKER.urn,
-                payload=beam_runner_api_pb2.DockerPayload(
-                    container_image='LEGACY').SerializeToString(),
-                capabilities=environments.python_sdk_docker_capabilities())
-        ])
+    first = remote_runner.proto_pipeline.components.environments.values()
+    second = beam_runner_api_pb2.Environment(
+        urn=common_urns.environments.DOCKER.urn,
+        payload=beam_runner_api_pb2.DockerPayload(
+            container_image='LEGACY').SerializeToString(),
+        capabilities=environments.python_sdk_docker_capabilities())
+    self.assertTrue(first.__eq__(second))
 
   def test_environment_override_translation_sdk_container_image(self):
     self.default_properties.append('--experiments=beam_fn_api')
@@ -231,15 +229,13 @@ def test_environment_override_translation_sdk_container_image(self):
           p | ptransform.Create([1, 2, 3])
           | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
           | ptransform.GroupByKey())
-    self.assertEqual(
-        list(remote_runner.proto_pipeline.components.environments.values()),
-        [
-            beam_runner_api_pb2.Environment(
-                urn=common_urns.environments.DOCKER.urn,
-                payload=beam_runner_api_pb2.DockerPayload(
-                    container_image='FOO').SerializeToString(),
-                capabilities=environments.python_sdk_docker_capabilities())
-        ])
+    first = remote_runner.proto_pipeline.components.environments.values()
+    second = beam_runner_api_pb2.Environment(
+        urn=common_urns.environments.DOCKER.urn,
+        payload=beam_runner_api_pb2.DockerPayload(
+            container_image='FOO').SerializeToString(),
+        capabilities=environments.python_sdk_docker_capabilities())
+    self.assertTrue(first.__eq__(second))
 
   def test_remote_runner_translation(self):
     remote_runner = DataflowRunner()

sdks/python/apache_beam/runners/portability/stager.py

@@ -52,6 +52,7 @@
 import logging
 import os
 import shutil
+import subprocess
 import sys
 import tempfile
 from importlib.metadata import distribution
@@ -84,6 +85,8 @@
 WORKFLOW_TARBALL_FILE = 'workflow.tar.gz'
 REQUIREMENTS_FILE = 'requirements.txt'
 EXTRA_PACKAGES_FILE = 'extra_packages.txt'
+# Filename that stores the submission environment dependencies.
+SUBMISSION_ENV_DEPENDENCIES_FILE = 'submission_environment_dependencies.txt'
 # One of the choices for user to use for requirements cache during staging
 SKIP_REQUIREMENTS_CACHE = 'skip'
 
@@ -162,6 +165,7 @@ def create_job_resources(options,  # type: PipelineOptions
                            pypi_requirements=None, # type: Optional[List[str]]
                            populate_requirements_cache=None,  # type: Optional[Callable[[str, str, bool], None]]
                            skip_prestaged_dependencies=False, # type: Optional[bool]
+                           log_submission_env_dependencies=True, # type: Optional[bool]
                            ):
     """For internal use only; no backwards-compatibility guarantees.
 
@@ -365,6 +369,11 @@ def create_job_resources(options,  # type: PipelineOptions
             Stager._create_file_stage_to_artifact(
                 pickled_session_file, names.PICKLED_MAIN_SESSION_FILE))
 
+    # stage the submission environment dependencies
+    if log_submission_env_dependencies:
+      resources.extend(
+          Stager._create_stage_submission_env_dependencies(temp_dir))
+
     return resources
 
   def stage_job_resources(self,
@@ -837,3 +846,20 @@ def _create_beam_sdk(sdk_remote_location, temp_dir):
     return [
         Stager._create_file_stage_to_artifact(local_download_file, staged_name)
     ]
+
+  @staticmethod
+  def _create_stage_submission_env_dependencies(temp_dir):
+    try:
+      local_dependency_file_path = os.path.join(
+          temp_dir, SUBMISSION_ENV_DEPENDENCIES_FILE)
+      dependencies = subprocess.check_output(
+          [sys.executable, '-m', 'pip', 'freeze'])
+      with open(local_dependency_file_path, 'w') as f:
+        f.write(str(dependencies))
+      return [
+          Stager._create_file_stage_to_artifact(
+              local_dependency_file_path, SUBMISSION_ENV_DEPENDENCIES_FILE)
+      ]
+    except Exception:
+      _LOGGER.debug("couldn't stage dependencies from submission environment")
+      return []

sdks/python/apache_beam/runners/portability/stager_test.py

@@ -162,7 +162,7 @@ def test_no_main_session(self):
     options.view_as(SetupOptions).save_main_session = False
     self.update_options(options)
 
-    self.assertEqual([],
+    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
 
@@ -180,7 +180,10 @@ def test_with_main_session(self):
     options.view_as(SetupOptions).pickle_library = pickler.USE_DILL
     self.update_options(options)
 
-    self.assertEqual([names.PICKLED_MAIN_SESSION_FILE],
+    self.assertEqual([
+        names.PICKLED_MAIN_SESSION_FILE,
+        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+    ],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
     self.assertTrue(
@@ -199,7 +202,7 @@ def test_main_session_not_staged_when_using_cloudpickle(self):
     # session is saved when pickle_library==cloudpickle.
     options.view_as(SetupOptions).pickle_library = pickler.USE_CLOUDPICKLE
     self.update_options(options)
-    self.assertEqual([],
+    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
 
@@ -208,10 +211,19 @@ def test_default_resources(self):
     options = PipelineOptions()
     self.update_options(options)
 
-    self.assertEqual([],
+    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
 
+  def test_no_submission_env_staging(self):
+    staging_dir = self.make_temp_dir()
+    options = PipelineOptions()
+    self.update_options(options)
+
+    resources = self.stager.create_job_resources(
+        options, staging_dir, log_submission_env_dependencies=False)
+    self.assertEqual([], resources)
+
   def test_with_requirements_file(self):
     staging_dir = self.make_temp_dir()
     requirements_cache_dir = self.make_temp_dir()
@@ -225,7 +237,12 @@ def test_with_requirements_file(self):
     self.create_temp_file(
         os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing')
     self.assertEqual(
-        sorted([stager.REQUIREMENTS_FILE, 'abc.txt', 'def.txt']),
+        sorted([
+            stager.REQUIREMENTS_FILE,
+            'abc.txt',
+            'def.txt',
+            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+        ]),
         sorted(
             self.stager.create_and_stage_job_resources(
                 options,
@@ -246,9 +263,11 @@ def test_with_pypi_requirements(self):
         pypi_requirements=['nothing>=1.0,<2.0'],
         populate_requirements_cache=self.populate_requirements_cache,
         staging_location=staging_dir)[1]
-    self.assertEqual(3, len(resources))
+    self.assertEqual(4, len(resources))
     self.assertTrue({'abc.txt', 'def.txt'} <= set(resources))
-    generated_requirements = (set(resources) - {'abc.txt', 'def.txt'}).pop()
+    generated_requirements = (
+        set(resources) -
+        {'abc.txt', 'def.txt', stager.SUBMISSION_ENV_DEPENDENCIES_FILE}).pop()
     with open(os.path.join(staging_dir, generated_requirements)) as f:
       data = f.read()
     self.assertEqual('nothing>=1.0,<2.0', data)
@@ -282,7 +301,12 @@ def test_with_requirements_file_and_cache(self):
     self.create_temp_file(
         os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing')
     self.assertEqual(
-        sorted([stager.REQUIREMENTS_FILE, 'abc.txt', 'def.txt']),
+        sorted([
+            stager.REQUIREMENTS_FILE,
+            'abc.txt',
+            'def.txt',
+            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+        ]),
         sorted(
             self.stager.create_and_stage_job_resources(
                 options,
@@ -313,7 +337,9 @@ def test_requirements_cache_not_populated_when_cache_disabled(self):
           populate_requirements_cache=self.populate_requirements_cache,
           staging_location=staging_dir)[1]
       assert not populate_requirements_cache.called
-      self.assertEqual([stager.REQUIREMENTS_FILE], resources)
+      self.assertEqual(
+          [stager.REQUIREMENTS_FILE, stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
+          resources)
       self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'abc.txt')))
       self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'def.txt')))
 
@@ -378,7 +404,8 @@ def test_sdk_location_default(self):
     _, staged_resources = self.stager.create_and_stage_job_resources(
         options, temp_dir=self.make_temp_dir(), staging_location=staging_dir)
 
-    self.assertEqual([], staged_resources)
+    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
+                     staged_resources)
 
   def test_sdk_location_local_directory(self):
     staging_dir = self.make_temp_dir()
@@ -391,7 +418,10 @@ def test_sdk_location_local_directory(self):
     self.update_options(options)
     options.view_as(SetupOptions).sdk_location = sdk_location
 
-    self.assertEqual([names.STAGED_SDK_SOURCES_FILENAME],
+    self.assertEqual([
+        names.STAGED_SDK_SOURCES_FILENAME,
+        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+    ],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
     tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME)
@@ -409,7 +439,10 @@ def test_sdk_location_local_source_file(self):
     self.update_options(options)
     options.view_as(SetupOptions).sdk_location = sdk_location
 
-    self.assertEqual([names.STAGED_SDK_SOURCES_FILENAME],
+    self.assertEqual([
+        names.STAGED_SDK_SOURCES_FILENAME,
+        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+    ],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
     tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME)
@@ -427,7 +460,7 @@ def test_sdk_location_local_wheel_file(self):
     self.update_options(options)
     options.view_as(SetupOptions).sdk_location = sdk_location
 
-    self.assertEqual([sdk_filename],
+    self.assertEqual([sdk_filename, stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
     tarball_path = os.path.join(staging_dir, sdk_filename)
@@ -463,7 +496,10 @@ def test_sdk_location_remote_source_file(self, *unused_mocks):
     self.update_options(options)
     options.view_as(SetupOptions).sdk_location = sdk_location
 
-    self.assertEqual([names.STAGED_SDK_SOURCES_FILENAME],
+    self.assertEqual([
+        names.STAGED_SDK_SOURCES_FILENAME,
+        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+    ],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
 
@@ -485,7 +521,7 @@ def file_download(_, to_path):
     with mock.patch('apache_beam.runners.portability.stager_test'
                     '.stager.Stager._download_file',
                     staticmethod(file_download)):
-      self.assertEqual([sdk_filename],
+      self.assertEqual([sdk_filename, stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
                        self.stager.create_and_stage_job_resources(
                            options, staging_location=staging_dir)[1])
 
@@ -509,7 +545,10 @@ def file_download(_, to_path):
     with mock.patch('apache_beam.runners.portability.stager_test'
                     '.stager.Stager._download_file',
                     staticmethod(file_download)):
-      self.assertEqual([names.STAGED_SDK_SOURCES_FILENAME],
+      self.assertEqual([
+          names.STAGED_SDK_SOURCES_FILENAME,
+          stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+      ],
                        self.stager.create_and_stage_job_resources(
                            options, staging_location=staging_dir)[1])
 
@@ -551,7 +590,8 @@ def test_with_extra_packages(self):
             'xyz2.tar',
             'whl.whl',
             'remote_file.tar.gz',
-            stager.EXTRA_PACKAGES_FILE
+            stager.EXTRA_PACKAGES_FILE,
+            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
         ],
                          self.stager.create_and_stage_job_resources(
                              options, staging_location=staging_dir)[1])
@@ -659,7 +699,13 @@ def test_with_jar_packages(self):
       with mock.patch('apache_beam.runners.portability.stager_test'
                       '.stager.Stager._is_remote_path',
                       staticmethod(self.is_remote_path)):
-        self.assertEqual(['abc.jar', 'xyz.jar', 'ijk.jar', 'remote.jar'],
+        self.assertEqual([
+            'abc.jar',
+            'xyz.jar',
+            'ijk.jar',
+            'remote.jar',
+            stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+        ],
                          self.stager.create_and_stage_job_resources(
                              options, staging_location=staging_dir)[1])
     self.assertEqual(['/tmp/remote/remote.jar'], self.remote_copied_files)
@@ -719,7 +765,8 @@ def test_populate_requirements_cache_with_bdist(self):
       resources = self.stager.create_and_stage_job_resources(
           options, staging_location=staging_dir)[1]
       for f in resources:
-        if f != stager.REQUIREMENTS_FILE:
+        if (f != stager.REQUIREMENTS_FILE and
+            f != stager.SUBMISSION_ENV_DEPENDENCIES_FILE):
           self.assertTrue(('.tar.gz' in f) or ('.whl' in f))
 
   # requirements cache will populated only with sdists/sources
@@ -744,7 +791,8 @@ def test_populate_requirements_cache_with_sdist(self):
           options, staging_location=staging_dir)[1]
 
       for f in resources:
-        if f != stager.REQUIREMENTS_FILE:
+        if (f != stager.REQUIREMENTS_FILE and
+            f != stager.SUBMISSION_ENV_DEPENDENCIES_FILE):
           self.assertTrue('.tar.gz' in f)
           self.assertTrue('.whl' not in f)
 
@@ -777,7 +825,8 @@ def test_populate_requirements_cache_with_local_files(self):
               stager.REQUIREMENTS_FILE,
               stager.EXTRA_PACKAGES_FILE,
               'nothing.tar.gz',
-              'local_package.tar.gz'
+              'local_package.tar.gz',
+              stager.SUBMISSION_ENV_DEPENDENCIES_FILE
           ]),
           sorted(resources))
 

sdks/python/apache_beam/transforms/environments.py

@@ -124,9 +124,9 @@ def __init__(self,
         dict(resource_hints) if resource_hints else {})
 
   def __eq__(self, other):
+    # don't compare artifacts since they have different hashes in their names.
     return (
-        self.__class__ == other.__class__ and
-        self._artifacts == other._artifacts
+        self.__class__ == other.__class__
         # Assuming that we don't have instances of the same Environment subclass
         # with different set of capabilities.
         and self._resource_hints == other._resource_hints)
@@ -580,7 +580,7 @@ def from_options(cls, options):
         url,
         params=params,
         capabilities=python_sdk_capabilities(),
-        artifacts=python_sdk_dependencies(options),
+        artifacts=python_sdk_dependencies(options, logDependencies=False),
         resource_hints=resource_hints_from_options(options))
 
 
@@ -605,7 +605,7 @@ def from_options(cls, options):
     # type: (PortableOptions) -> EmbeddedPythonEnvironment
     return cls(
         capabilities=python_sdk_capabilities(),
-        artifacts=python_sdk_dependencies(options),
+        artifacts=python_sdk_dependencies(options, logDependencies=False),
         resource_hints=resource_hints_from_options(options),
     )
 
@@ -692,11 +692,11 @@ def from_options(cls, options):
           state_cache_size=config.get('state_cache_size'),
           data_buffer_time_limit_ms=config.get('data_buffer_time_limit_ms'),
           capabilities=python_sdk_capabilities(),
-          artifacts=python_sdk_dependencies(options))
+          artifacts=python_sdk_dependencies(options, logDependencies=False))
     else:
       return cls(
           capabilities=python_sdk_capabilities(),
-          artifacts=python_sdk_dependencies(options),
+          artifacts=python_sdk_dependencies(options, logDependencies=False),
           resource_hints=resource_hints_from_options(options))
 
   @staticmethod
@@ -838,7 +838,7 @@ def _python_sdk_capabilities_iter():
   yield common_urns.protocols.DATA_SAMPLING.urn
 
 
-def python_sdk_dependencies(options, tmp_dir=None):
+def python_sdk_dependencies(options, tmp_dir=None, logDependencies=True):
   if tmp_dir is None:
     tmp_dir = tempfile.mkdtemp()
   skip_prestaged_dependencies = options.view_as(
@@ -850,4 +850,5 @@ def python_sdk_dependencies(options, tmp_dir=None):
           artifact[0] + artifact[1]
           for artifact in PyPIArtifactRegistry.get_artifacts()
       ],
-      skip_prestaged_dependencies=skip_prestaged_dependencies)
+      skip_prestaged_dependencies=skip_prestaged_dependencies,
+      log_submission_env_dependencies=logDependencies)