fix accidentally double-registering setuptools for ipex (#9341)

### Problem

When creating `.ipex` files as per #8793, we don't add resolved distributions to `self._distributions` in `PexBuilderWrapper`. This leads to the annoying `_sanitize_requirements()` hack introduced into `ipex_launcher.py` in that PR (having to remove duplicate `setuptools` requirements).

### Solution

- Add all resolved distributions to `self._distributions` in `PexBuilderWrapper`, even when `generate_ipex=True`.
- Remove the separate list of transitively-resolved requirements passed through several methods in `PexBuilderWrapper`.

src/python/pants/backend/python/subsystems/ipex/ipex_launcher.py

@@ -16,7 +16,6 @@
 from pex.interpreter import PythonInterpreter
 from pex.pex_builder import PEXBuilder
 from pex.pex_info import PexInfo
-from pkg_resources import Requirement
 
 APP_CODE_PREFIX = "user_files/"
 
@@ -35,24 +34,6 @@ def _log(message):
     sys.stderr.write(message + "\n")
 
 
-def _sanitize_requirements(requirements):
-    """Remove duplicate keys such as setuptools or pex which may be injected multiple times into the
-    resulting ipex when first executed."""
-    project_names = []
-    new_requirements = {}
-
-    for r in requirements:
-        r = Requirement(r)
-        if r.marker and not r.marker.evaluate():
-            continue
-        if r.name not in new_requirements:
-            project_names.append(r.name)
-            new_requirements[r.name] = str(r)
-    sanitized_requirements = [new_requirements[n] for n in project_names]
-
-    return sanitized_requirements
-
-
 def modify_pex_info(pex_info, **kwargs):
     new_info = json.loads(pex_info.dump())
     new_info.update(kwargs)
@@ -86,10 +67,6 @@ def _hydrate_pex_file(self, hydrated_pex_file):
     # Perform a fully pinned intransitive resolve to hydrate the install cache.
     resolver_settings = ipex_info["resolver_settings"]
 
-    sanitized_requirements = _sanitize_requirements(bootstrap_info.requirements)
-    bootstrap_info = modify_pex_info(bootstrap_info, requirements=sanitized_requirements)
-    bootstrap_builder.info = bootstrap_info
-
     resolved_distributions = resolver.resolve(
         requirements=bootstrap_info.requirements,
         cache=bootstrap_info.pex_root,

src/python/pants/python/pex_build_util.py

@@ -214,17 +214,10 @@ def extract_single_dist_for_current_platform(self, reqs, dist_key) -> Distributi
         :raises: :class:`self.SingleDistExtractionError` if no dists or multiple dists matched the
                  given `dist_key`.
         """
-        distributions, _transitive_requirements = self.resolve_distributions(
-            reqs, platforms=["current"]
-        )
+        distributions = self.resolve_distributions(reqs, platforms=["current"])
         try:
             matched_dist = assert_single_element(
-                list(
-                    dist
-                    for _, dists in distributions.items()
-                    for dist in dists
-                    if dist.key == dist_key
-                )
+                dist for dists in distributions.values() for dist in dists if dist.key == dist_key
             )
         except (StopIteration, ValueError) as e:
             raise self.SingleDistExtractionError(
@@ -234,7 +227,7 @@ def extract_single_dist_for_current_platform(self, reqs, dist_key) -> Distributi
 
     def resolve_distributions(
         self, reqs: List[PythonRequirement], platforms: Optional[List[Platform]] = None,
-    ) -> Tuple[Dict[str, List[Distribution]], List[PythonRequirement]]:
+    ) -> Dict[str, List[Distribution]]:
         """Multi-platform dependency resolution.
 
         :param reqs: A list of :class:`PythonRequirement` to resolve.
@@ -254,10 +247,10 @@ def resolve_distributions(
                 find_links.add(req.repository)
 
         # Resolve the requirements into distributions.
-        distributions, transitive_requirements = self._resolve_multi(
+        distributions = self._resolve_multi(
             self._builder.interpreter, list(deduped_reqs), platforms, list(find_links),
         )
-        return (distributions, transitive_requirements)
+        return distributions
 
     def add_resolved_requirements(
         self,
@@ -279,9 +272,7 @@ def add_resolved_requirements(
                                                                         pex dependency to the output ipex file, and therefore needs to
                                                                         override the default behavior of this method.
         """
-        distributions, transitive_requirements = self.resolve_distributions(
-            reqs, platforms=platforms
-        )
+        distributions = self.resolve_distributions(reqs, platforms=platforms)
         locations: Set[str] = set()
         for platform, dists in distributions.items():
             for dist in dists:
@@ -290,25 +281,21 @@ def add_resolved_requirements(
                         self._log.debug(
                             f"  *AVOIDING* dumping distribution into ipex: .../{os.path.basename(dist.location)}"
                         )
+                        self._register_distribution(dist)
                     else:
                         self._log.debug(
                             f"  Dumping distribution: .../{os.path.basename(dist.location)}"
                         )
                         self.add_distribution(dist)
                 locations.add(dist.location)
-        # In addition to the top-level requirements, we add all the requirements matching the resolved
-        # distributions to the resulting pex. If `generate_ipex=True` is set, we need to have all the
-        # transitive requirements resolved in order to hydrate the .ipex with an intransitive resolve.
-        if self._generate_ipex and not override_ipex_build_do_actually_add_distribution:
-            self.add_direct_requirements(transitive_requirements)
 
     def _resolve_multi(
         self,
         interpreter: PythonInterpreter,
         requirements: List[PythonRequirement],
         platforms: Optional[List[Platform]],
         find_links: Optional[List[str]],
-    ) -> Tuple[Dict[str, List[Distribution]], List[PythonRequirement]]:
+    ) -> Dict[str, List[Distribution]]:
         """Multi-platform dependency resolution for PEX files.
 
         Returns a tuple containing a list of distributions that must be included in order to satisfy a
@@ -336,9 +323,6 @@ def _resolve_multi(
         self._all_find_links.update(OrderedSet(find_links))
 
         distributions: Dict[str, List[Distribution]] = defaultdict(list)
-        transitive_requirements: List[PythonRequirement] = []
-
-        all_find_links = [*python_repos.repos, *find_links]
 
         for platform in platforms:
             requirements_cache_dir = os.path.join(
@@ -349,17 +333,15 @@ def _resolve_multi(
                 interpreter=interpreter,
                 platform=platform,
                 indexes=python_repos.indexes,
-                find_links=all_find_links,
+                find_links=find_links,
                 cache=requirements_cache_dir,
                 allow_prereleases=python_setup.resolver_allow_prereleases,
                 manylinux=python_setup.manylinux,
             )
             for resolved_dist in resolved_dists:
-                dist = resolved_dist.distribution
-                transitive_requirements.append(dist.as_requirement())
-                distributions[platform].append(dist)
+                distributions[platform].append(resolved_dist.distribution)
 
-        return (distributions, transitive_requirements)
+        return distributions
 
     def _create_source_dumper(self, tgt: Target) -> Callable[[str], None]:
         buildroot = get_buildroot()
@@ -446,6 +428,11 @@ def _shuffle_underlying_pex_builder(self) -> Tuple[PexInfo, Path]:
             self._builder.info, self._builder.interpreter.identity
         )
 
+        # Remove all the original top-level requirements in favor of the transitive == requirements.
+        self._builder.info = ipex_launcher.modify_pex_info(self._builder.info, requirements=[])
+        transitive_reqs = [dist.as_requirement() for dist in self._distributions.values()]
+        self.add_direct_requirements(transitive_reqs)
+
         orig_info = self._builder.info.copy()
 
         orig_chroot = self._builder.chroot()
@@ -456,6 +443,8 @@ def _shuffle_underlying_pex_builder(self) -> Tuple[PexInfo, Path]:
             self._builder.info, self._builder.interpreter.identity
         )
 
+        self._distributions = {}
+
         return (orig_info, Path(orig_chroot.path()))
 
     def _shuffle_original_build_info_into_ipex(self):