Skip to content

Parallelize initial Rust download in bootstrap #110427

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 30, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 86 additions & 15 deletions src/bootstrap/bootstrap.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import tempfile

from time import time
from multiprocessing import Pool, cpu_count

try:
import lzma
Expand Down Expand Up @@ -392,6 +393,48 @@ def channel(self):
return self.version + "-" + self.date


class DownloadInfo:
"""A helper class that can be pickled into a parallel subprocess"""

def __init__(
self,
base_download_url,
download_path,
bin_root,
tarball_path,
tarball_suffix,
checksums_sha256,
pattern,
verbose,
):
self.base_download_url = base_download_url
self.download_path = download_path
self.bin_root = bin_root
self.tarball_path = tarball_path
self.tarball_suffix = tarball_suffix
self.checksums_sha256 = checksums_sha256
self.pattern = pattern
self.verbose = verbose

def download_component(download_info):
if not os.path.exists(download_info.tarball_path):
get(
download_info.base_download_url,
download_info.download_path,
download_info.tarball_path,
download_info.checksums_sha256,
verbose=download_info.verbose,
)

def unpack_component(download_info):
unpack(
download_info.tarball_path,
download_info.tarball_suffix,
download_info.bin_root,
match=download_info.pattern,
verbose=download_info.verbose,
)

class RustBuild(object):
"""Provide all the methods required to build Rust"""
def __init__(self):
Expand Down Expand Up @@ -429,17 +472,49 @@ def download_toolchain(self):
self.program_out_of_date(self.rustc_stamp(), key)):
if os.path.exists(bin_root):
shutil.rmtree(bin_root)

key = self.stage0_compiler.date
cache_dst = os.path.join(self.build_dir, "cache")
rustc_cache = os.path.join(cache_dst, key)
if not os.path.exists(rustc_cache):
os.makedirs(rustc_cache)

tarball_suffix = '.tar.gz' if lzma is None else '.tar.xz'
filename = "rust-std-{}-{}{}".format(
rustc_channel, self.build, tarball_suffix)
pattern = "rust-std-{}".format(self.build)
self._download_component_helper(filename, pattern, tarball_suffix)
filename = "rustc-{}-{}{}".format(rustc_channel, self.build,
tarball_suffix)
self._download_component_helper(filename, "rustc", tarball_suffix)
filename = "cargo-{}-{}{}".format(rustc_channel, self.build,
tarball_suffix)
self._download_component_helper(filename, "cargo", tarball_suffix)

toolchain_suffix = "{}-{}{}".format(rustc_channel, self.build, tarball_suffix)

tarballs_to_download = [
("rust-std-{}".format(toolchain_suffix), "rust-std-{}".format(self.build)),
("rustc-{}".format(toolchain_suffix), "rustc"),
("cargo-{}".format(toolchain_suffix), "cargo"),
]

tarballs_download_info = [
DownloadInfo(
base_download_url=self.download_url,
download_path="dist/{}/{}".format(self.stage0_compiler.date, filename),
bin_root=self.bin_root(),
tarball_path=os.path.join(rustc_cache, filename),
tarball_suffix=tarball_suffix,
checksums_sha256=self.checksums_sha256,
pattern=pattern,
verbose=self.verbose,
)
for filename, pattern in tarballs_to_download
]

# Download the components serially to show the progress bars properly.
for download_info in tarballs_download_info:
download_component(download_info)

# Unpack the tarballs in parallle.
# In Python 2.7, Pool cannot be used as a context manager.
p = Pool(min(len(tarballs_download_info), cpu_count()))
try:
p.map(unpack_component, tarballs_download_info)
finally:
p.close()

if self.should_fix_bins_and_dylibs():
self.fix_bin_or_dylib("{}/bin/cargo".format(bin_root))

Expand All @@ -455,13 +530,9 @@ def download_toolchain(self):
rust_stamp.write(key)

def _download_component_helper(
self, filename, pattern, tarball_suffix,
self, filename, pattern, tarball_suffix, rustc_cache,
):
key = self.stage0_compiler.date
cache_dst = os.path.join(self.build_dir, "cache")
rustc_cache = os.path.join(cache_dst, key)
if not os.path.exists(rustc_cache):
os.makedirs(rustc_cache)

tarball = os.path.join(rustc_cache, filename)
if not os.path.exists(tarball):
Expand Down