Add http/2 with httpx strategy for shards

conda_libmamba_solver/index.py

@@ -108,7 +108,7 @@
     PackageInfo,
 )
 
-from conda_libmamba_solver.shards_subset import build_repodata_subset
+from conda_libmamba_solver.shards_subset import RepodataSubset, build_repodata_subset
 
 from .mamba_utils import logger_callback
 
@@ -156,6 +156,21 @@ def _is_sharded_repodata_enabled():
     return context.plugins.use_sharded_repodata is True  # type: ignore
 
 
+def _sharded_repodata_strategy():
+    """
+    Which algorithm should we use to collect sharded repodata?
+    """
+    strategy = context.plugins.sharded_repodata_strategy.lower()  # type: ignore
+    if RepodataSubset.has_strategy(strategy):
+        return strategy
+    log.warning(
+        "Unknown sharded_repodata_strategy '%s', falling back to '%s'.",
+        strategy,
+        RepodataSubset.DEFAULT_STRATEGY,
+    )
+    return RepodataSubset.DEFAULT_STRATEGY
+
+
 _SUPPORTS_PYTHON_SITE_PACKAGES = hasattr(PackageInfo, "python_site_packages_path")
 
 
@@ -377,18 +392,26 @@ def _load_channel_repo_info_shards(
         self, urls_to_channel: dict[str, Channel]
     ) -> list[_ChannelRepoInfo] | None:
         """
-        Load repository information from sharded repodata cache.
+        Load repository information by fetching and processing repodata shards.
         """
         # make a subset of possible dependencies
         root_packages = (*self.in_state.installed.keys(), *self.in_state.requested)
         channel_data = build_repodata_subset(root_packages, urls_to_channel)
         if channel_data is None:
             return  # caller should fall back to repodata.json
 
-        channel_repo_infos = self._load_repo_info_from_repodata_dict(channel_data)
+        channel_repo_infos = self._load_repo_info_from_repodata_shards(channel_data)
 
         return channel_repo_infos
 
+    def _build_repodata_subset(
+        self, root_packages: tuple[str, ...], urls_to_channel: dict[str, Channel]
+    ) -> dict[str, ShardBase]:
+        # split into a separate method for tests
+        return build_repodata_subset(
+            root_packages, urls_to_channel, algorithm=_sharded_repodata_strategy()
+        )
+
     def _load_channel_repo_info_json(
         self, urls_to_channel: dict[str, Channel], try_solv: bool
     ) -> list[_ChannelRepoInfo]:
@@ -615,12 +638,12 @@ def _load_pkgs_cache(self, pkgs_dirs: PathsType) -> list[RepoInfo]:
         return repos
 
     @time_recorder(module_name=__name__)
-    def _load_repo_info_from_repodata_dict(
+    def _load_repo_info_from_repodata_shards(
         self, repodata_subset: dict[str, ShardBase]
     ) -> list[_ChannelRepoInfo]:
         """
-        Load repository information from deserialized repodata.json-like
-        structures.
+        Load repository information from already-fetched ShardBase objects, that
+        produce in-memory repodata.json-like dicts.
         """
         repos = []
         for channel_url, shardlike in repodata_subset.items():

conda_libmamba_solver/plugin.py

@@ -44,3 +44,9 @@ def conda_settings():
         description="Enable use of sharded repodata when available.",
         parameter=PrimitiveParameter(False, element_type=bool),
     )
+
+    yield CondaSetting(
+        name="sharded_repodata_strategy",
+        description="Strategy to collect sharded repodata.",
+        parameter=PrimitiveParameter("auto", element_type=str),
+    )

conda_libmamba_solver/shards.py

@@ -247,6 +247,12 @@ def build_repodata(self) -> RepodataDict:
                 repodata[package_group].update(shard[package_group])
         return repodata
 
+    def reset(self):
+        """
+        Clear visited shards.
+        """
+        self.visited.clear()
+
 
 class ShardLike(ShardBase):
     """

conda_libmamba_solver/shards_cache.py

@@ -134,8 +134,6 @@ def insert(self, raw_shard: AnnotatedRawShard):
             package: package name
             raw_shard: msgpack.zst compressed shard data
         """
-        # decompress and return shard for convenience, also to validate? unless
-        # caller would rather retrieve the shard from another thread.
         with self.conn as c:
             c.execute(
                 "INSERT OR IGNORE INTO SHARDS (url, package, shard) VALUES (?, ?, ?)",
@@ -153,11 +151,21 @@ def retrieve(self, url) -> ShardDict | None:
                 else None
             )  # type: ignore
 
-    def retrieve_multiple(self, urls: list[str]) -> dict[str, ShardDict | None]:
+    def retrieve_multiple(self, urls: list[str]) -> dict[str, ShardDict]:
         """
         Query database for cached shard urls.
 
-        Return a dict of urls in cache mapping to the Shard or None if not present.
+        Return a dict of urls in cache mapping to the Shard. If not in cache,
+        that url will not appear in the result.
+        """
+        return {k: v[0] for k, v in self.retrieve_multiple_size(urls).items()}
+
+    def retrieve_multiple_size(self, urls: list[str]) -> dict[str, tuple[ShardDict, int]]:
+        """
+        Query database for cached shard urls; include compressed size.
+
+        Return a dict of urls in cache mapping to the Shard, omit URLs that are
+        not found.
         """
         if not urls:
             return {}  # this optimization does not save a noticeable amount of time.
@@ -168,12 +176,13 @@ def retrieve_multiple(self, urls: list[str]) -> dict[str, ShardDict | None]:
 
         query = f"SELECT url, shard FROM shards WHERE url IN ({','.join(('?',) * len(urls))}) ORDER BY url"
         with self.conn as c:
-            result: dict[str, ShardDict | None] = {
-                row["url"]: msgpack.loads(
-                    dctx.decompress(row["shard"], max_output_size=ZSTD_MAX_SHARD_SIZE)
+            result: dict[str, tuple[ShardDict, int]] = {
+                row["url"]: (
+                    msgpack.loads(
+                        dctx.decompress(row["shard"], max_output_size=ZSTD_MAX_SHARD_SIZE)
+                    ),
+                    len(row["shard"]),
                 )
-                if row
-                else None
                 for row in c.execute(query, urls)  # type: ignore
             }
             return result

conda_libmamba_solver/shards_subset.py

@@ -289,7 +289,30 @@ def _reachable_bfs(self, root_packages):
                     if not next_node.visited:
                         node_queue.append(next_node)
 
-    def reachable_pipelined(self, root_packages):
+    def reachable_httpx(self, root_packages):
+        from .shards_subset_http2 import network_fetch_thread_httpx
+
+        return self.reachable_pipelined(root_packages, network_fetch_thread_httpx)
+
+    def reachable_pycurl(self, root_packages):
+        """
+        httpx + pycurl transport for httpx (curl sockets in event loop async implementation)
+        """
+        from .shards_subset_http2 import network_fetch_thread_httpx_pycurl
+
+        return self.reachable_pipelined(root_packages, network_fetch_thread_httpx_pycurl)
+
+    def reachable_pycurl2(self, root_packages):
+        """
+        httpx + pycurl transport for httpx (delegate to thread async implementation)
+        """
+        from .shards_subset_http2 import network_fetch_thread_httpx_pycurl2
+
+        return self.reachable_pipelined(root_packages, network_fetch_thread_httpx_pycurl2)
+
+    # could try sync + threads http/2 also...
+
+    def reachable_pipelined(self, root_packages, network_worker=None):
         """
         Fetch all packages reachable from `root_packages`' by following
         dependencies.
@@ -303,7 +326,8 @@ def reachable_pipelined(self, root_packages):
         # empty shards.
         if context.offline:
             network_worker = offline_nofetch_thread
-        else:
+
+        if network_worker is None:
             network_worker = network_fetch_thread
 
         # Ignore cache on shards object, use our own. Necessary if there are no
@@ -342,6 +366,9 @@ def _reachable_pipelined(
             daemon=True,  # may have to set to False if we ever want to run in a subinterpreter
         )
 
+        if network_worker is None:
+            network_worker = network_fetch_thread
+
         network_thread = threading.Thread(
             target=network_worker,
             args=(cache_miss_queue, shard_out_queue, cache, self.shardlikes),
@@ -493,7 +520,7 @@ def drain_pending(
 def build_repodata_subset(
     root_packages: Iterable[str],
     channels: dict[str, Channel],
-    algorithm: Literal["bfs", "pipelined"] = RepodataSubset.DEFAULT_STRATEGY,
+    algorithm: Literal["bfs", "pipelined", "httpx"] = RepodataSubset.DEFAULT_STRATEGY,
 ) -> dict[str, ShardBase] | None:
     """
     Retrieve all necessary information to build a repodata subset.

conda_libmamba_solver/shards_subset_http2.py

@@ -0,0 +1,170 @@
+# Copyright (C) 2022 Anaconda, Inc
+# Copyright (C) 2023 conda
+# SPDX-License-Identifier: BSD-3-Clause
+"""
+http/2 network fetching of shards, proof of concept.
+
+The greater concurrency of http/2 should be very beneficial for sharded
+repodata.
+"""
+
+import asyncio
+import concurrent.futures
+import logging
+from queue import Queue
+
+import httpx
+import msgpack
+import pycurltx
+import zstandard
+
+from .shards import ShardBase, Shards
+from .shards_cache import ZSTD_MAX_SHARD_SIZE, AnnotatedRawShard, ShardCache
+from .shards_subset import NodeId, exception_to_queue
+from .shards_typing import ShardDict
+
+log = logging.getLogger(__name__)
+
+
+async def _network_fetch_loop_httpx(
+    in_queue: Queue[list[NodeId] | None],
+    shard_out_queue: Queue[list[tuple[NodeId, ShardDict] | Exception] | None],
+    cache: ShardCache,
+    shardlikes: list[ShardBase],
+    transport=None,
+):
+    cache = cache.copy()
+    dctx = zstandard.ZstdDecompressor()
+    shardlikes_by_url = {s.url: s for s in shardlikes}
+
+    async def fetch(client: httpx.AsyncClient, url: str, node_id: NodeId):
+        response = await client.get(url)
+        response.raise_for_status()
+        data = response.content
+        return (url, node_id, data)
+
+    async def submit(client, node_id):
+        # this worker should only recieve network node_id's:
+        shardlike = shardlikes_by_url[node_id.channel]
+        if not isinstance(shardlike, Shards):
+            log.warning("network_fetch_thread got non-network shardlike")
+            return
+        url = shardlikes_by_url[node_id.channel].shard_url(node_id.package)
+        return await fetch(client, url, node_id)
+
+    async def get_work():
+        with concurrent.futures.ThreadPoolExecutor() as thread_pool:
+            # this may be faster than asyncio.to_thread()
+            while True:
+                batch = await asyncio.get_running_loop().run_in_executor(thread_pool, in_queue.get)
+                if batch is None:
+                    break
+                for node_id in batch:
+                    yield node_id
+
+    def handle_result(task: asyncio.Task):
+        tasks.remove(task)
+        url, node_id, data = task.result()  # can raise exceptions from fetch
+        log.debug("Fetch %s (%s bytes)", url, len(data))
+        # Decompress and parse. If it decodes as
+        # msgpack.zst, insert into cache. Then put "known
+        # good" shard into out queue.
+        shard: ShardDict = msgpack.loads(
+            dctx.decompress(data, max_output_size=ZSTD_MAX_SHARD_SIZE)
+        )  # type: ignore[assign]
+        cache.insert(AnnotatedRawShard(url, node_id.package, data))
+        shard_out_queue.put([(node_id, shard)])
+
+    # if transport is not None, http2 comes from pycurl
+    async with httpx.AsyncClient(http2=(transport is None), transport=transport) as client:
+        results: list[tuple[NodeId, ShardDict] | Exception] = []
+
+        tasks = set()
+        async for node_id in get_work():
+            task = asyncio.create_task(submit(client, node_id))
+            task.add_done_callback(handle_result)
+            tasks.add(task)
+
+        # this loop is unnecessary as handle_result sends results to shard_out_queue
+        for task in asyncio.as_completed(tasks):
+            try:
+                result = await task
+                results.append(result)
+            except Exception as e:
+                results.append(e)
+
+
+@exception_to_queue
+def network_fetch_thread_httpx(
+    in_queue: Queue[list[NodeId] | None],
+    shard_out_queue: Queue[list[tuple[NodeId, ShardDict] | Exception] | None],
+    cache: ShardCache,
+    shardlikes: list[ShardBase],
+):
+    """
+    in_queue contains urls to fetch over the network.
+    While the in_queue has not received a sentinel None, empty everything from
+    the queue. Fetch all of them over the network. Fetched shards go to
+    shard_out_queue.
+    Unhandled exceptions also go to shard_out_queue, and exit this thread.
+    """
+    for shard in shardlikes:
+        if not shard.url.startswith("http"):
+            raise ValueError(f"Unsupported shard in http/2 fetch thread: {shard.url}")
+    asyncio.run(_network_fetch_loop_httpx(in_queue, shard_out_queue, cache, shardlikes))
+
+
+@exception_to_queue
+def network_fetch_thread_httpx_pycurl(
+    in_queue: Queue[list[NodeId] | None],
+    shard_out_queue: Queue[list[tuple[NodeId, ShardDict] | Exception] | None],
+    cache: ShardCache,
+    shardlikes: list[ShardBase],
+):
+    """
+    in_queue contains urls to fetch over the network.
+    While the in_queue has not received a sentinel None, empty everything from
+    the queue. Fetch all of them over the network. Fetched shards go to
+    shard_out_queue.
+    Unhandled exceptions also go to shard_out_queue, and exit this thread.
+    """
+    for shard in shardlikes:
+        if not shard.url.startswith("http"):
+            raise ValueError(f"Unsupported shard in http/2 fetch thread: {shard.url}")
+    asyncio.run(
+        _network_fetch_loop_httpx(
+            in_queue,
+            shard_out_queue,
+            cache,
+            shardlikes,
+            transport=pycurltx.PyCurlAsyncMultiSocketTransport(max_connections=100),
+        )
+    )
+
+
+@exception_to_queue
+def network_fetch_thread_httpx_pycurl2(
+    in_queue: Queue[list[NodeId] | None],
+    shard_out_queue: Queue[list[tuple[NodeId, ShardDict] | Exception] | None],
+    cache: ShardCache,
+    shardlikes: list[ShardBase],
+):
+    """
+    in_queue contains urls to fetch over the network.
+    While the in_queue has not received a sentinel None, empty everything from
+    the queue. Fetch all of them over the network. Fetched shards go to
+    shard_out_queue.
+    Unhandled exceptions also go to shard_out_queue, and exit this thread.
+    """
+    for shard in shardlikes:
+        if not shard.url.startswith("http"):
+            raise ValueError(f"Unsupported shard in http/2 fetch thread: {shard.url}")
+    asyncio.run(
+        _network_fetch_loop_httpx(
+            in_queue,
+            shard_out_queue,
+            cache,
+            shardlikes,
+            transport=pycurltx.PyCurlAsyncTransport(timeout=10.0),
+        )
+    )

tests/requirements.txt

@@ -7,3 +7,7 @@ pytest-codspeed >=4
 conda-forge::pytest-mock
 # needed for many conda tests
 flask
+httpx
+h2
+hpack
+hyperframe