feat: Prefilter Shard Records on Python Minor Version (#4214)

Signed-off-by: Julien Jerphanion <git@jjerphan.xyz>
Co-authored-by: Johan Mabille <johan.mabille@gmail.com>

Author: jjerphan

libmamba/CMakeLists.txt

@@ -247,6 +247,8 @@ set(
     ${LIBMAMBA_SOURCE_DIR}/core/query.cpp
     ${LIBMAMBA_SOURCE_DIR}/core/repo_checker_store.cpp
     ${LIBMAMBA_SOURCE_DIR}/core/run.cpp
+    ${LIBMAMBA_SOURCE_DIR}/core/shard_python_minor_prefilter.hpp
+    ${LIBMAMBA_SOURCE_DIR}/core/shard_python_minor_prefilter.cpp
     ${LIBMAMBA_SOURCE_DIR}/core/shell_init.cpp
     ${LIBMAMBA_SOURCE_DIR}/core/shards.cpp
     ${LIBMAMBA_SOURCE_DIR}/core/shard_index_loader.cpp

libmamba/include/mamba/api/channel_loader.hpp

@@ -7,11 +7,13 @@
 #ifndef MAMBA_API_CHANNEL_LOADER_HPP
 #define MAMBA_API_CHANNEL_LOADER_HPP
 
+#include <optional>
 #include <set>
 #include <string>
 #include <vector>
 
 #include "mamba/core/error_handling.hpp"
+#include "mamba/specs/version.hpp"
 
 namespace mamba
 {
@@ -41,6 +43,8 @@ namespace mamba
      * @param subdir_idx Index of the subdir to load in \p subdirs.
      * @param loaded_subdirs_with_shards Set of subdir names already loaded via shards (updated).
      * @param priorities Repo priorities aligned with \p subdirs.
+     * @param python_minor_version_for_prefilter Optional python minor for shard record prefiltering
+     * (from \c prepare_solver_context).
      * @return The repo for the requested subdir, or unexpected mamba_error on failure.
      */
     auto load_subdir_with_shards(
@@ -50,7 +54,8 @@ namespace mamba
         std::vector<SubdirIndexLoader>& subdirs,
         std::size_t subdir_idx,
         std::set<std::string>& loaded_subdirs_with_shards,
-        const std::vector<solver::libsolv::Priorities>& priorities
+        const std::vector<solver::libsolv::Priorities>& priorities,
+        std::optional<specs::Version> python_minor_version_for_prefilter = std::nullopt
     ) -> expected_t<solver::libsolv::RepoInfo>;
 
     class ChannelContext;
@@ -86,7 +91,8 @@ namespace mamba
         ChannelContext& channel_context,
         solver::libsolv::Database& database,
         MultiPackageCache& package_caches,
-        const std::vector<std::string>& root_packages = {}
+        const std::vector<std::string>& root_packages = {},
+        std::optional<specs::Version> python_minor_version_for_prefilter = std::nullopt
     ) -> expected_t<void, mamba_aggregated_error>;
 
     /* Brief Creates channels and mirrors objects,

libmamba/include/mamba/core/shards.hpp

@@ -22,6 +22,7 @@
 #include "mamba/fs/filesystem.hpp"
 #include "mamba/specs/authentication_info.hpp"
 #include "mamba/specs/channel.hpp"
+#include "mamba/specs/version.hpp"
 
 namespace mamba
 {
@@ -30,6 +31,12 @@ namespace mamba
      *
      * This class manages fetching and caching of individual shards from
      * a sharded repodata index.
+     *
+     * **Python minor prefilter:** When constructed with ``python_minor_version_for_prefilter``
+     * (e.g. 3.12), parsing a shard msgpack drops package records whose ``depends`` list constrains
+     * ``python`` to a range that does not contain that minor, reducing work for the solver.
+     * When that optional is unset, no such filtering is applied and all records in the shard
+     * are parsed (python compatibility is left to the solver).
      */
     class Shards
     {
@@ -47,6 +54,9 @@ namespace mamba
          * @param mirrors Optional base mirrors for channel-based downloads. When provided,
          *        extend_mirrors in fetch_shards will be initialized from these before adding
          *        absolute-URL mirrors.
+         * @param python_minor_version_for_prefilter If set, shard parsing filters out records whose
+         *        ``depends`` python constraints are incompatible with this minor; if unset,
+         *        no python-minor-based record filtering is performed.
          */
         Shards(
             ShardsIndexDict shards_index,
@@ -56,7 +66,8 @@ namespace mamba
             download::RemoteFetchParams remote_fetch_params,
             // 0 means: auto; value is normalized with normalize_to_affinity_concurrency().
             std::size_t download_threads = 0,
-            std::optional<std::reference_wrapper<const download::mirror_map>> mirrors = std::nullopt
+            std::optional<std::reference_wrapper<const download::mirror_map>> mirrors = std::nullopt,
+            std::optional<specs::Version> python_minor_version_for_prefilter = std::nullopt
         );
 
         /** Return the names of all packages available in this shard collection. */
@@ -119,6 +130,13 @@ namespace mamba
         /** Optional base mirrors for channel-based downloads. */
         std::optional<std::reference_wrapper<const download::mirror_map>> m_mirrors;
 
+        /**
+         * Environment python minor used when parsing shards to prefilter package records
+         * (see ``record_depends_on_python_minor_version_for_prefilter`` in shards.cpp).
+         * Empty means the prefilter is disabled.
+         */
+        std::optional<specs::Version> m_python_minor_version_for_prefilter;
+
         /** Visited shards, keyed by package name. */
         std::map<std::string, ShardDict> m_visited;
 

libmamba/src/api/channel_loader.cpp

@@ -24,6 +24,9 @@
 #include "mamba/solver/libsolv/repo_info.hpp"
 #include "mamba/specs/error.hpp"
 #include "mamba/specs/package_info.hpp"
+#include "mamba/specs/version.hpp"
+
+#include "utils.hpp"
 
 namespace mamba
 {
@@ -235,7 +238,8 @@ namespace mamba
             std::size_t subdir_idx,
             std::set<std::string>& loaded_subdirs_with_shards,
             const SubdirDownloadParams& subdir_params,
-            const std::vector<solver::libsolv::Priorities>& priorities
+            const std::vector<solver::libsolv::Priorities>& priorities,
+            std::optional<specs::Version> python_minor_version_for_prefilter
         )
         {
             auto& subdir = subdirs[subdir_idx];
@@ -253,7 +257,8 @@ namespace mamba
                     subdirs,
                     subdir_idx,
                     loaded_subdirs_with_shards,
-                    priorities
+                    priorities,
+                    python_minor_version_for_prefilter
                 );
 
                 if (!res)
@@ -434,7 +439,8 @@ namespace mamba
             const std::vector<solver::libsolv::Priorities>& priorities,
             const SubdirDownloadParams& subdir_params,
             bool is_retry,
-            std::vector<mamba_error>& error_list
+            std::vector<mamba_error>& error_list,
+            std::optional<specs::Version> python_minor_version_for_prefilter
         )
         {
             std::set<std::string> loaded_subdirs_with_shards;
@@ -475,7 +481,8 @@ namespace mamba
                     i,
                     loaded_subdirs_with_shards,
                     subdir_params,
-                    priorities
+                    priorities,
+                    python_minor_version_for_prefilter
                 );
 
                 if (result)
@@ -556,7 +563,11 @@ namespace mamba
                     continue;
                 }
                 SubdirIndexLoader subdir_index_loader = std::move(subdir_index_loader_result).value();
-                if (subdir_index_loader.valid_cache_found() && Console::can_report_status())
+
+                // Only show flat repodata cache status if we're not using shards and we have a
+                // valid cache
+                if (!ctx.repodata_use_shards && subdir_index_loader.valid_cache_found()
+                    && Console::can_report_status())
                 {
                     Console::stream()
                         << fmt::format("{:<50} {:>20}", subdir_index_loader.name(), "Using cache");
@@ -641,7 +652,8 @@ namespace mamba
         std::vector<SubdirIndexLoader>& subdirs,
         std::size_t subdir_idx,
         std::set<std::string>& loaded_subdirs_with_shards,
-        const std::vector<solver::libsolv::Priorities>& priorities
+        const std::vector<solver::libsolv::Priorities>& priorities,
+        std::optional<specs::Version> python_minor_version_for_prefilter
     ) -> expected_t<solver::libsolv::RepoInfo>
     {
         auto& subdir = subdirs[subdir_idx];
@@ -670,6 +682,15 @@ namespace mamba
         LOG_DEBUG << "Shard index fetched for " << subdir.name();
         const auto& channel = subdir.channel();
         std::string current_repodata_url = subdir.repodata_url().str();
+        if (python_minor_version_for_prefilter.has_value())
+        {
+            LOG_DEBUG << "Shard prefilter on python minor version enabled with "
+                      << python_minor_version_for_prefilter.value().to_string();
+        }
+        else
+        {
+            LOG_DEBUG << "Shard prefilter on python minor version disabled.";
+        }
 
         // For all subdirs sharing the same channel URL, fetch their shard indices and build
         //    a Shards instance per subdir; collect them into a RepodataSubset.
@@ -702,7 +723,8 @@ namespace mamba
                     ctx.authentication_info(),
                     ctx.remote_fetch_params,
                     normalize_to_affinity_concurrency(static_cast<int>(ctx.repodata_shards_threads)),
-                    std::cref(ctx.mirrors)
+                    std::cref(ctx.mirrors),
+                    python_minor_version_for_prefilter
                 );
                 url_to_subdir_idx[sdir_url] = j;
             }
@@ -761,7 +783,8 @@ namespace mamba
             solver::libsolv::Database& database,
             MultiPackageCache& package_caches,
             const std::vector<std::string>& root_packages,
-            bool is_retry
+            bool is_retry,
+            std::optional<specs::Version> python_minor_version_for_prefilter
         )
         {
             std::vector<SubdirIndexLoader> subdirs;
@@ -808,7 +831,8 @@ namespace mamba
                 priorities,
                 subdir_params,
                 is_retry,
-                error_list
+                error_list,
+                python_minor_version_for_prefilter
             );
 
             if (loading_failed)
@@ -824,7 +848,8 @@ namespace mamba
                         database,
                         package_caches,
                         root_packages,
-                        retry
+                        retry,
+                        python_minor_version_for_prefilter
                     );
                 }
                 error_list.emplace_back(
@@ -843,11 +868,20 @@ namespace mamba
         ChannelContext& channel_context,
         solver::libsolv::Database& database,
         MultiPackageCache& package_caches,
-        const std::vector<std::string>& root_packages
+        const std::vector<std::string>& root_packages,
+        std::optional<specs::Version> python_minor_version_for_prefilter
     ) -> expected_t<void, mamba_aggregated_error>
     {
         bool retry = false;
-        return load_channels_impl(ctx, channel_context, database, package_caches, root_packages, retry);
+        return load_channels_impl(
+            ctx,
+            channel_context,
+            database,
+            package_caches,
+            root_packages,
+            retry,
+            std::move(python_minor_version_for_prefilter)
+        );
     }
 
     void init_channels(Context& context, ChannelContext& channel_context)

libmamba/src/api/install.cpp

@@ -376,6 +376,33 @@ namespace mamba
     {
         using Request = solver::Request;
 
+        // When the user explicitly asks for ``python`` in the requested specs, also inject a
+        // plain ``pip`` request unless it is already present. This complements
+        // ``add_pip_as_python_dependency`` at the repo level and makes sure that the Request
+        // is in phase with the root packages including both ``python`` and ``pip`` when requested.
+        bool wants_python = false;
+        bool wants_pip = false;
+        for (const auto& s : specs)
+        {
+            const auto maybe_name = specs::MatchSpec::extract_name(s);
+            if (!maybe_name.has_value())
+            {
+                continue;
+            }
+            if (maybe_name.value() == "python")
+            {
+                wants_python = true;
+            }
+            else if (maybe_name.value() == "pip")
+            {
+                wants_pip = true;
+            }
+        }
+        if (wants_python && !wants_pip)
+        {
+            specs.emplace_back("pip");
+        }
+
         const auto& prefix_pkgs = prefix_data.records();
 
         auto request = Request();
@@ -555,7 +582,13 @@ namespace mamba
             auto& no_env = config.at("no_env").value<bool>();
 
             validate_target_prefix_and_channels(ctx, create_env);
-            auto [db, package_caches] = prepare_solver_context(ctx, channel_context, raw_specs);
+            auto [db, package_caches] = prepare_solver_context(
+                ctx,
+                channel_context,
+                raw_specs,
+                is_retry,
+                no_py_pin
+            );
 
             auto prefix_data = load_prefix_data_and_installed(ctx, channel_context, db);
 

libmamba/src/api/update.cpp

@@ -4,6 +4,11 @@
 //
 // The full license is in the file LICENSE, distributed with this software.
 
+#include <cctype>
+
+#include <fmt/format.h>
+
+#include "mamba/api/channel_loader.hpp"
 #include "mamba/api/configuration.hpp"
 #include "mamba/api/install.hpp"
 #include "mamba/api/update.hpp"
@@ -147,7 +152,13 @@ namespace mamba
             auto& retry_clean_cache = config.at("retry_clean_cache").value<bool>();
 
             validate_target_prefix_and_channels(ctx, /* create_env= */ false);
-            auto [db, package_caches] = prepare_solver_context(ctx, channel_context, raw_update_specs);
+            auto [db, package_caches] = prepare_solver_context(
+                ctx,
+                channel_context,
+                raw_update_specs,
+                is_retry,
+                no_py_pin
+            );
 
             auto prefix_data = load_prefix_data_and_installed(ctx, channel_context, db);