Skip to content

Commit

Permalink
parallelize artifact downloads
Browse files Browse the repository at this point in the history
  • Loading branch information
@IanButterworth
IanButterworth committed Aug 7, 2024
1 parent 7aef1f0 commit 9a5ce60
Show file tree
Hide file tree
Showing 3 changed files with 104 additions and 60 deletions.
125 changes: 73 additions & 52 deletions src/Artifacts.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,7 @@ function download_artifact(
verbose::Bool = false,
quiet_download::Bool = false,
io::IO=stderr_f(),
progress::Union{Function, Nothing} = nothing,
)
if artifact_exists(tree_hash)
return true
Expand All @@ -323,8 +324,8 @@ function download_artifact(
temp_dir = mktempdir(artifacts_dir)

try
download_verify_unpack(tarball_url, tarball_hash, temp_dir, ignore_existence=true, verbose=verbose,
quiet_download=quiet_download, io=io)
download_verify_unpack(tarball_url, tarball_hash, temp_dir;
ignore_existence=true, verbose, quiet_download, io, progress)
calc_hash = SHA1(GitTools.tree_hash(temp_dir))

# Did we get what we expected? If not, freak out.
Expand Down Expand Up @@ -408,68 +409,88 @@ function ensure_artifact_installed(name::String, meta::Dict, artifacts_toml::Str
platform::AbstractPlatform = HostPlatform(),
verbose::Bool = false,
quiet_download::Bool = false,
io::IO=stderr_f())
hash = SHA1(meta["git-tree-sha1"])
progress_callback::Union{Function,Nothing} = nothing,
io::IO=stderr_f(),
progress::Union{Function,Nothing}=nothing)

hash = SHA1(meta["git-tree-sha1"])
if !artifact_exists(hash)
errors = Any[]
# first try downloading from Pkg server
# TODO: only do this if Pkg server knows about this package
if (server = pkg_server()) !== nothing
url = "$server/artifact/$hash"
download_success = let url=url
@debug "Downloading artifact from Pkg server" name artifacts_toml platform url
with_show_download_info(io, name, quiet_download) do
download_artifact(hash, url; verbose=verbose, quiet_download=quiet_download, io=io)
end
end
# download_success is either `true` or an error object
if download_success === true
return artifact_path(hash)
else
@debug "Failed to download artifact from Pkg server" download_success
push!(errors, (url, download_success))
if isnothing(progress_callback) || verbose == true
return try_artifact_download_sources(name, hash, meta, artifacts_toml; platform, verbose, quiet_download, io)
else
return Threads.@spawn begin
try_artifact_download_sources(name, meta, artifacts_toml; platform, quiet_download=true, io, progress)
end
end
else
return artifact_path(hash)
end
end

# If this artifact does not exist on-disk already, ensure it has download
# information, then download it!
if !haskey(meta, "download")
error("Cannot automatically install '$(name)'; no download section in '$(artifacts_toml)'")
function try_artifact_download_sources(
name::String, hash::SHA1, meta::Dict, artifacts_toml::String;
platform::AbstractPlatform=HostPlatform(),
verbose::Bool=false,
quiet_download::Bool=false,
io::IO=stderr_f(),
progress::Union{Function,Nothing}=nothing)

errors = Any[]
# first try downloading from Pkg server
# TODO: only do this if Pkg server knows about this package
if (server = pkg_server()) !== nothing
url = "$server/artifact/$hash"
download_success = let url = url
@debug "Downloading artifact from Pkg server" name artifacts_toml platform url
with_show_download_info(io, name, quiet_download) do
download_artifact(hash, url; verbose, quiet_download, io, progress)
end
end
# download_success is either `true` or an error object
if download_success === true
return artifact_path(hash)
else
@debug "Failed to download artifact from Pkg server" download_success
push!(errors, (url, download_success))
end
end

# Attempt to download from all sources
for entry in meta["download"]
url = entry["url"]
tarball_hash = entry["sha256"]
download_success = let url=url
@debug "Downloading artifact" name artifacts_toml platform url
with_show_download_info(io, name, quiet_download) do
download_artifact(hash, url, tarball_hash; verbose=verbose, quiet_download=quiet_download, io=io)
end
end
# download_success is either `true` or an error object
if download_success === true
return artifact_path(hash)
else
@debug "Failed to download artifact" download_success
push!(errors, (url, download_success))
# If this artifact does not exist on-disk already, ensure it has download
# information, then download it!
if !haskey(meta, "download")
error("Cannot automatically install '$(name)'; no download section in '$(artifacts_toml)'")
end

# Attempt to download from all sources
for entry in meta["download"]
url = entry["url"]
tarball_hash = entry["sha256"]
download_success = let url = url
@debug "Downloading artifact" name artifacts_toml platform url
with_show_download_info(io, name, quiet_download) do
download_artifact(hash, url, tarball_hash; verbose, quiet_download, io, progress)
end
end
errmsg = """
Unable to automatically download/install artifact '$(name)' from sources listed in '$(artifacts_toml)'.
Sources attempted:
"""
for (url, err) in errors
errmsg *= "- $(url)\n"
errmsg *= " Error: $(sprint(showerror, err))\n"
# download_success is either `true` or an error object
if download_success === true
return artifact_path(hash)
else
@debug "Failed to download artifact" download_success
push!(errors, (url, download_success))
end
error(errmsg)
else
return artifact_path(hash)
end
errmsg = """
Unable to automatically download/install artifact '$(name)' from sources listed in '$(artifacts_toml)'.
Sources attempted:
"""
for (url, err) in errors
errmsg *= "- $(url)\n"
errmsg *= " Error: $(sprint(showerror, err))\n"
end
error(errmsg)
end


function with_show_download_info(f, io, name, quiet_download)
fancyprint = can_fancyprint(io)
if !quiet_download
Expand All @@ -485,7 +506,7 @@ function with_show_download_info(f, io, name, quiet_download)
if !quiet_download
fancyprint && print(io, "\033[1A") # move cursor up one line
fancyprint && print(io, "\033[2K") # clear line
if success
if success
fancyprint && printpkgstyle(io, :Downloaded, "artifact: $name")
else
printpkgstyle(io, :Failure, "artifact: $name", color = :red)
Expand Down
23 changes: 21 additions & 2 deletions src/Operations.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -828,16 +828,35 @@ function download_artifacts(env::EnvCache;
pkg_root === nothing || push!(pkg_roots, pkg_root)
end
push!(pkg_roots, dirname(env.project_file))
used_artifact_tomls = Set{String}()
download_jobs = Channel{Function}(Inf)
# download_progresses = # Dict for each live job of (total, now) -> nothing
for pkg_root in pkg_roots
for (artifacts_toml, artifacts) in collect_artifacts(pkg_root; platform)
# For each Artifacts.toml, install each artifact we've collected from it
for name in keys(artifacts)
ensure_artifact_installed(name, artifacts[name], artifacts_toml;
# @info name
push!(download_jobs,
() -> ensure_artifact_installed(name, artifacts[name], artifacts_toml;
verbose, quiet_download=!(usable_io(io)), io=io)
)
end
write_env_usage(artifacts_toml, "artifact_usage.toml")
push!(used_artifact_tomls, artifacts_toml)
end
end
close(download_jobs)

sema = Base.Semaphore(20)
# TODO: figure out printing & error handling
@sync for f in download_jobs
Threads.@spawn Base.acquire(sema) do
f()
end
end

for f in used_artifact_tomls
write_env_usage(f, "artifact_usage.toml")
end
end

function check_artifacts_downloaded(pkg_root::String; platform::AbstractPlatform=HostPlatform())
Expand Down
16 changes: 10 additions & 6 deletions src/PlatformEngines.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -255,7 +255,8 @@ function download(
verbose::Bool = false,
headers::Vector{Pair{String,String}} = Pair{String,String}[],
auth_header::Union{Pair{String,String}, Nothing} = nothing,
io::IO=stderr_f()
io::IO=stderr_f(),
progress::Union{Nothing,Function} = nothing, # (total, now) -> nothing
)
if auth_header === nothing
auth_header = get_auth_header(url, verbose=verbose)
Expand All @@ -268,7 +269,9 @@ function download(
end

do_fancy = verbose && can_fancyprint(io)
progress = if do_fancy
progress = if !isnothing(progress)
progress
elseif do_fancy
bar = MiniProgressBar(header="Downloading", color=Base.info_color())
start_progress(io, bar)
let bar=bar
Expand All @@ -286,7 +289,7 @@ function download(
nothing
end
try
Downloads.download(url, dest; headers, progress)
Downloads.download(url, dest; headers, progress, downloader=Downloads.Downloader(; grace=0))
finally
do_fancy && end_progress(io, bar)
end
Expand Down Expand Up @@ -326,6 +329,7 @@ function download_verify(
verbose::Bool = false,
force::Bool = false,
quiet_download::Bool = false,
progress::Union{Nothing,Function} = nothing, # (total, now) -> nothing
)
# Whether the file existed in the first place
file_existed = false
Expand All @@ -352,7 +356,7 @@ function download_verify(
attempts = 3
for i in 1:attempts
try
download(url, dest; verbose=verbose || !quiet_download)
download(url, dest; verbose=verbose || !quiet_download, progress)
break
catch err
@debug "download and verify failed on attempt $i/$attempts" url dest err
Expand Down Expand Up @@ -466,6 +470,7 @@ function download_verify_unpack(
verbose::Bool = false,
quiet_download::Bool = false,
io::IO=stderr_f(),
progress::Union{Nothing,Function} = nothing, # (total, now) -> nothing
)
# First, determine whether we should keep this tarball around
remove_tarball = false
Expand Down Expand Up @@ -510,8 +515,7 @@ function download_verify_unpack(

# Download the tarball; if it already existed and we needed to remove it
# then we should remove the unpacked path as well
should_delete = !download_verify(url, hash, tarball_path;
force=force, verbose=verbose, quiet_download=quiet_download)
should_delete = !download_verify(url, hash, tarball_path; force, verbose, quiet_download, progress)
if should_delete
if verbose
@info("Removing dest directory $(dest) as source tarball changed")
Expand Down

0 comments on commit 9a5ce60

Please sign in to comment.