Julia.withPackages improvements (#433332)

This commit is contained in:
Nick Cao
2025-09-03 17:36:40 -04:00
committed by GitHub
12 changed files with 282 additions and 106 deletions
+33 -9
View File
@@ -11,6 +11,7 @@
# Artifacts dependencies
fetchurl,
gcc,
glibc,
pkgs,
stdenv,
@@ -79,7 +80,9 @@ let
PythonCall = [ "PyCall" ];
};
# Invoke Julia resolution logic to determine the full dependency closure
# Invoke Julia resolution logic to determine the full dependency closure. Also
# gather information on the Julia standard libraries, which we'll need to
# generate a Manifest.toml.
packageOverridesRepoified = lib.mapAttrs util.repoifySimple packageOverrides;
closureYaml = callPackage ./package-closure.nix {
inherit
@@ -90,6 +93,9 @@ let
;
packageOverrides = packageOverridesRepoified;
};
stdlibInfos = callPackage ./stdlib-infos.nix {
inherit julia;
};
# Generate a Nix file consisting of a map from dependency UUID --> package info with fetchgit call:
# {
@@ -181,6 +187,27 @@ let
"${dependencyUuidToRepoYaml}" \
"$out"
'';
project =
runCommand "julia-project"
{
buildInputs = [
(python3.withPackages (
ps: with ps; [
toml
pyyaml
]
))
git
];
}
''
python ${./python}/project.py \
"${closureYaml}" \
"${stdlibInfos}" \
'${lib.generators.toJSON { } overridesOnly}' \
"${dependencyUuidToRepoYaml}" \
"$out"
'';
# Next, deal with artifacts. Scan each artifacts file individually and generate a Nix file that
# produces the desired Overrides.toml.
@@ -220,7 +247,7 @@ let
;
}
// lib.optionalAttrs (!stdenv.targetPlatform.isDarwin) {
inherit glibc;
inherit gcc glibc;
}
);
overridesJson = writeTextFile {
@@ -235,8 +262,7 @@ let
"$out"
'';
# Build a Julia project and depot. The project contains Project.toml/Manifest.toml, while the
# depot contains package build products (including the precompiled libraries, if precompile=true)
# Build a Julia project and depot under $out/project and $out/depot respectively
projectAndDepot = callPackage ./depot.nix {
inherit
closureYaml
@@ -247,12 +273,8 @@ let
precompile
;
julia = juliaWrapped;
inherit project;
registry = minimalRegistry;
packageNames =
if makeTransitiveDependenciesImportable then
lib.mapAttrsToList (uuid: info: info.name) dependencyUuidToInfo
else
packageNames;
};
in
@@ -276,7 +298,9 @@ runCommand "julia-${julia.version}-env"
inherit artifactsNix;
inherit overridesJson;
inherit overridesToml;
inherit project;
inherit projectAndDepot;
inherit stdlibInfos;
};
}
(
+18 -21
View File
@@ -14,7 +14,7 @@
juliaCpuTarget,
overridesToml,
packageImplications,
packageNames,
project,
precompile,
registry,
}:
@@ -44,7 +44,7 @@ runCommand "julia-depot"
(python3.withPackages (ps: with ps; [ pyyaml ]))
]
++ extraLibs;
inherit precompile registry;
inherit precompile project registry;
}
(
''
@@ -52,19 +52,21 @@ runCommand "julia-depot"
echo "Building Julia depot and project with the following inputs"
echo "Julia: ${julia}"
echo "Project: $project"
echo "Registry: $registry"
echo "Overrides ${overridesToml}"
mkdir -p $out/project
export JULIA_PROJECT="$out/project"
cp "$project/Manifest.toml" "$JULIA_PROJECT/Manifest.toml"
cp "$project/Project.toml" "$JULIA_PROJECT/Project.toml"
mkdir -p $out/depot/artifacts
export JULIA_DEPOT_PATH="$out/depot"
cp ${overridesToml} $out/depot/artifacts/Overrides.toml
# These can be useful to debug problems
# export JULIA_DEBUG=Pkg
# export JULIA_DEBUG=loading
# export JULIA_DEBUG=Pkg,loading
${setJuliaSslCaRootsPath}
@@ -104,26 +106,21 @@ runCommand "julia-depot"
Pkg.Registry.add(Pkg.RegistrySpec(path="${registry}"))
input = ${lib.generators.toJSON { } packageNames} ::Vector{String}
# No need to Pkg.activate() since we set JULIA_PROJECT above
println("Running Pkg.instantiate()")
Pkg.instantiate()
if isfile("extra_package_names.txt")
append!(input, readlines("extra_package_names.txt"))
end
# Build is a separate step from instantiate.
# Needed for packages like Conda.jl to set themselves up.
println("Running Pkg.build()")
Pkg.build()
input = unique(input)
if !isempty(input)
println("Adding packages: " * join(input, " "))
Pkg.add(input; preserve=PRESERVE_NONE)
Pkg.instantiate()
if "precompile" in keys(ENV) && ENV["precompile"] != "0" && ENV["precompile"] != ""
if isdefined(Sys, :CPU_NAME)
println("Precompiling with CPU_NAME = " * Sys.CPU_NAME)
end
Pkg.precompile()
if "precompile" in keys(ENV) && ENV["precompile"] != "0" && ENV["precompile"] != ""
if isdefined(Sys, :CPU_NAME)
println("Precompiling with CPU_NAME = " * Sys.CPU_NAME)
end
Pkg.precompile()
end
# Remove the registry to save space
@@ -43,12 +43,21 @@ let
println(io, "- name: " * spec.name)
println(io, " uuid: " * string(spec.uuid))
println(io, " version: " * string(spec.version))
println(io, " tree_hash: " * string(spec.tree_hash))
if endswith(spec.name, "_jll") && haskey(deps_map, spec.uuid)
println(io, " depends_on: ")
for (dep_name, dep_uuid) in pairs(deps_map[spec.uuid])
println(io, " \"$(dep_name)\": \"$(dep_uuid)\"")
end
end
println(io, " deps: ")
for (dep_name, dep_uuid) in pairs(deps_map[spec.uuid])
println(io, " - name: \"$(dep_name)\"")
println(io, " uuid: \"$(dep_uuid)\"")
end
if spec.name in input
println(io, " is_input: true")
end
end
end
'';
@@ -47,14 +47,17 @@ def get_archive_derivation(uuid, artifact_name, url, sha256, closure_dependencie
''"""
else:
# We provide gcc.cc.lib by default in order to get some common libraries
# like libquadmath.so. A number of packages expect this to be available and
# will give linker errors if it isn't.
fixup = f"""fixupPhase = let
libs = lib.concatMap (lib.mapAttrsToList (k: v: v.path))
[{" ".join(["uuid-" + x for x in depends_on])}];
in ''
find $out -type f -executable -exec \
patchelf --set-rpath \$ORIGIN:\$ORIGIN/../lib:${{lib.makeLibraryPath (["$out" glibc] ++ libs ++ (with pkgs; [{" ".join(other_libs)}]))}} {{}} \;
patchelf --set-rpath \\$ORIGIN:\\$ORIGIN/../lib:${{lib.makeLibraryPath (["$out" glibc gcc.cc.lib] ++ libs ++ (with pkgs; [{" ".join(other_libs)}]))}} {{}} \\;
find $out -type f -executable -exec \
patchelf --set-interpreter ${{glibc}}/lib/ld-linux-x86-64.so.2 {{}} \;
patchelf --set-interpreter ${{glibc}}/lib/ld-linux-x86-64.so.2 {{}} \\;
''"""
return f"""stdenv.mkDerivation {{
@@ -145,7 +148,7 @@ def main():
if is_darwin:
f.write("{ lib, fetchurl, pkgs, stdenv }:\n\n")
else:
f.write("{ lib, fetchurl, glibc, pkgs, stdenv }:\n\n")
f.write("{ lib, fetchurl, gcc, glibc, pkgs, stdenv }:\n\n")
f.write("rec {\n")
@@ -24,14 +24,15 @@ with open(desired_packages_path, "r") as f:
uuid_to_versions = defaultdict(list)
for pkg in desired_packages:
uuid_to_versions[pkg["uuid"]].append(pkg["version"])
uuid_to_versions[pkg["uuid"]].append(pkg["version"])
with open(dependencies_path, "r") as f:
uuid_to_store_path = yaml.safe_load(f)
os.makedirs(out_path)
registry = toml.load(registry_path / "Registry.toml")
full_registry = toml.load(registry_path / "Registry.toml")
registry = full_registry.copy()
registry["packages"] = {k: v for k, v in registry["packages"].items() if k in uuid_to_versions}
for (uuid, versions) in uuid_to_versions.items():
@@ -80,20 +81,48 @@ for (uuid, versions) in uuid_to_versions.items():
if (registry_path / path / f).exists():
shutil.copy2(registry_path / path / f, out_path / path)
# Copy the Versions.toml file, trimming down to the versions we care about
# Copy the Versions.toml file, trimming down to the versions we care about.
# In the case where versions=None, this is a weak dep, and we keep all versions.
all_versions = toml.load(registry_path / path / "Versions.toml")
versions_to_keep = {k: v for k, v in all_versions.items() if k in versions}
versions_to_keep = {k: v for k, v in all_versions.items() if k in versions} if versions != None else all_versions
for k, v in versions_to_keep.items():
del v["nix-sha256"]
with open(out_path / path / "Versions.toml", "w") as f:
toml.dump(versions_to_keep, f)
# Fill in the local store path for the repo
if not uuid in uuid_to_store_path: continue
package_toml = toml.load(registry_path / path / "Package.toml")
package_toml["repo"] = "file://" + uuid_to_store_path[uuid]
with open(out_path / path / "Package.toml", "w") as f:
toml.dump(package_toml, f)
if versions is None:
# This is a weak dep; just grab the whole Package.toml
shutil.copy2(registry_path / path / "Package.toml", out_path / path / "Package.toml")
elif uuid in uuid_to_store_path:
# Fill in the local store path for the repo
package_toml = toml.load(registry_path / path / "Package.toml")
package_toml["repo"] = "file://" + uuid_to_store_path[uuid]
with open(out_path / path / "Package.toml", "w") as f:
toml.dump(package_toml, f)
# Look for missing weak deps and include them. This can happen when our initial
# resolve step finds dependencies, but we fail to resolve them at the project.py
# stage. Usually this happens because the package that depends on them does so
# as a weak dep, but doesn't have a Package.toml in its repo making this clear.
for pkg in desired_packages:
for dep in (pkg.get("deps", []) or []):
uuid = dep["uuid"]
if not uuid in uuid_to_versions:
entry = full_registry["packages"].get(uuid)
if not entry:
print(f"""WARNING: found missing UUID but couldn't resolve it: {uuid}""")
continue
# Add this entry back to the minimal Registry.toml
registry["packages"][uuid] = entry
# Bring over the Package.toml
path = Path(entry["path"])
if (out_path / path / "Package.toml").exists():
continue
Path(out_path / path).mkdir(parents=True, exist_ok=True)
shutil.copy2(registry_path / path / "Package.toml", out_path / path / "Package.toml")
# Finally, dump the Registry.toml
with open(out_path / "Registry.toml", "w") as f:
toml.dump(registry, f)
+104
View File
@@ -0,0 +1,104 @@
from collections import defaultdict
import json
import os
from pathlib import Path
import sys
import toml
import yaml
desired_packages_path = Path(sys.argv[1])
stdlib_infos_path = Path(sys.argv[2])
package_overrides = json.loads(sys.argv[3])
dependencies_path = Path(sys.argv[4])
out_path = Path(sys.argv[5])
with open(desired_packages_path, "r") as f:
desired_packages = yaml.safe_load(f) or []
with open(stdlib_infos_path, "r") as f:
stdlib_infos = yaml.safe_load(f) or []
with open(dependencies_path, "r") as f:
uuid_to_store_path = yaml.safe_load(f)
result = {
"deps": defaultdict(list)
}
for pkg in desired_packages:
if pkg["uuid"] in package_overrides:
info = package_overrides[pkg["uuid"]]
result["deps"][info["name"]].append({
"uuid": pkg["uuid"],
"path": info["src"],
})
continue
path = uuid_to_store_path.get(pkg["uuid"], None)
isStdLib = False
if pkg["uuid"] in stdlib_infos["stdlibs"]:
path = stdlib_infos["stdlib_root"] + "/" + stdlib_infos["stdlibs"][pkg["uuid"]]["name"]
isStdLib = True
if path:
if (Path(path) / "Project.toml").exists():
project_toml = toml.load(Path(path) / "Project.toml")
deps = []
weak_deps = project_toml.get("weakdeps", {})
extensions = project_toml.get("extensions", {})
if "deps" in project_toml:
# Build up deps for the manifest, excluding weak deps
weak_deps_uuids = weak_deps.values()
for (dep_name, dep_uuid) in project_toml["deps"].items():
if not (dep_uuid in weak_deps_uuids):
deps.append(dep_name)
else:
# Not all projects have a Project.toml. In this case, use the deps we
# calculated from the package resolve step. This isn't perfect since it
# will fail to properly split out weak deps, but it's better than nothing.
print(f"""WARNING: package {pkg["name"]} didn't have a Project.toml in {path}""")
deps = [x["name"] for x in pkg.get("deps", [])]
weak_deps = {}
extensions = {}
tree_hash = pkg.get("tree_hash", "")
result["deps"][pkg["name"]].append({
"version": pkg["version"],
"uuid": pkg["uuid"],
"git-tree-sha1": (tree_hash if tree_hash != "nothing" else None) or None,
"deps": deps or None,
"weakdeps": weak_deps or None,
"extensions": extensions or None,
# We *don't* set "path" here, because then Julia will try to use the
# read-only Nix store path instead of cloning to the depot. This will
# cause packages like Conda.jl to fail during the Pkg.build() step.
#
# "path": None if isStdLib else path ,
})
else:
print("WARNING: adding a package that we didn't have a path for, and it doesn't seem to be a stdlib", pkg)
result["deps"][pkg["name"]].append({
"version": pkg["version"],
"uuid": pkg["uuid"],
"deps": [x["name"] for x in pkg["deps"]]
})
os.makedirs(out_path)
with open(out_path / "Manifest.toml", "w") as f:
f.write(f'julia_version = "{stdlib_infos["julia_version"]}"\n')
f.write('manifest_format = "2.0"\n\n')
toml.dump(result, f)
with open(out_path / "Project.toml", "w") as f:
f.write('[deps]\n')
for pkg in desired_packages:
if pkg.get("is_input", False):
f.write(f'''{pkg["name"]} = "{pkg["uuid"]}"\n''')
@@ -24,7 +24,7 @@ def ensure_version_valid(version):
Ensure a version string is a valid Julia-parsable version.
It doesn't really matter what it looks like as it's just used for overrides.
"""
return re.sub('[^0-9\.]','', version)
return re.sub('[^0-9.]','', version)
with open(out_path, "w") as f:
f.write("{fetchgit}:\n")
@@ -41,6 +41,9 @@ with open(out_path, "w") as f:
treehash = "{treehash}";
}};\n""")
elif uuid in registry["packages"]:
# The treehash is missing for stdlib packages. Don't bother downloading these.
if (not ("tree_hash" in pkg)) or pkg["tree_hash"] == "nothing": continue
registry_info = registry["packages"][uuid]
path = registry_info["path"]
packageToml = toml.load(registry_path / path / "Package.toml")
@@ -65,7 +68,8 @@ with open(out_path, "w") as f:
treehash = "{version_to_use["git-tree-sha1"]}";
}};\n""")
else:
# print("Warning: couldn't figure out what to do with pkg in sources_nix.py", pkg)
# This is probably a stdlib
# print("WARNING: couldn't figure out what to do with pkg in sources_nix.py", pkg)
pass
f.write("}")
+3 -3
View File
@@ -3,7 +3,7 @@
fetchFromGitHub {
owner = "CodeDownIO";
repo = "General";
rev = "998c6da1553dc0776dfff314d2f1bd5af488ed71";
sha256 = "sha256-57RiIPTu9895mdk3oSfo7I3PYw7G0BfJG1u+mYkJeLk=";
# date = "2024-07-01T12:22:35+00:00";
rev = "4b19a1dc55d2877e85a5d0e98702b75872210e9d";
sha256 = "sha256-mVeBTpEQnW3fvJu1+4T8z+earMjEgtdy0tnZnAxz/pk=";
# date = "2025-08-12T05:20:40+00:00";
}
@@ -46,54 +46,6 @@ end
foreach(pkg -> ctx.env.project.deps[pkg.name] = pkg.uuid, pkgs)
# Save the original pkgs for later. We might need to augment it with the weak dependencies
orig_pkgs = pkgs
orig_pkgs = deepcopy(pkgs)
pkgs, deps_map = _resolve(ctx.io, ctx.env, ctx.registries, pkgs, PRESERVE_NONE, ctx.julia_version)
if VERSION >= VersionNumber("1.9")
while true
# Check for weak dependencies, which appear on the RHS of the deps_map but not in pkgs.
# Build up weak_name_to_uuid
uuid_to_name = Dict()
for pkg in pkgs
uuid_to_name[pkg.uuid] = pkg.name
end
weak_name_to_uuid = Dict()
for (uuid, deps) in pairs(deps_map)
for (dep_name, dep_uuid) in pairs(deps)
if !haskey(uuid_to_name, dep_uuid)
weak_name_to_uuid[dep_name] = dep_uuid
end
end
end
if isempty(weak_name_to_uuid)
break
end
# We have nontrivial weak dependencies, so add each one to the initial pkgs and then re-run _resolve
println("Found weak dependencies: $(keys(weak_name_to_uuid))")
orig_uuids = Set([pkg.uuid for pkg in orig_pkgs])
for (name, uuid) in pairs(weak_name_to_uuid)
if uuid in orig_uuids
continue
end
pkg = PackageSpec(name, uuid)
push!(orig_uuids, uuid)
push!(orig_pkgs, pkg)
ctx.env.project.deps[name] = uuid
entry = Pkg.Types.manifest_info(ctx.env.manifest, uuid)
if VERSION >= VersionNumber("1.11")
orig_pkgs[length(orig_pkgs)] = update_package_add(ctx, pkg, entry, nothing, nothing, false)
else
orig_pkgs[length(orig_pkgs)] = update_package_add(ctx, pkg, entry, false)
end
end
global pkgs, deps_map = _resolve(ctx.io, ctx.env, ctx.registries, orig_pkgs, PRESERVE_NONE, ctx.julia_version)
end
end
@@ -0,0 +1,36 @@
{
julia,
runCommand,
}:
let
juliaExpression = ''
using Pkg
open(ENV["out"], "w") do io
println(io, "stdlib_root: \"$(Sys.STDLIB)\"")
println(io, "julia_version: \"$(string(VERSION))\"")
stdlibs = Pkg.Types.stdlibs()
println(io, "stdlibs:")
for (uuid, (name, version)) in stdlibs
println(io, " \"$(uuid)\": ")
println(io, " name: $name")
println(io, " version: $version")
end
end
'';
in
runCommand "julia-stdlib-infos.yml"
{
buildInputs = [
julia
];
}
''
# Prevent a warning where Julia tries to download package server info
export JULIA_PKG_SERVER=""
julia -e '${juliaExpression}';
''
@@ -12,8 +12,8 @@
module Main (main) where
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Data.Aeson as A hiding (Options, defaultOptions)
import qualified Data.Aeson.Key as A
import qualified Data.Aeson.KeyMap as HM
@@ -24,12 +24,14 @@ import Data.Text as T hiding (count)
import qualified Data.Vector as V
import qualified Data.Yaml as Yaml
import GHC.Generics
import Options.Applicative
import Options.Applicative hiding (info)
import System.Exit
import System.FilePath
import Test.Sandwich hiding (info)
import Test.Sandwich
import UnliftIO.Exception
import UnliftIO.MVar
import UnliftIO.Process
import UnliftIO.QSem
data Args = Args {
@@ -67,13 +69,15 @@ main = do
Left err -> throwIO $ userError ("Couldn't decode names and counts YAML file: " <> show err)
Right x -> pure x
runSandwichWithCommandLineArgs' defaultOptions argsParser $ do
runSandwichWithCommandLineArgs' defaultOptions argsParser $ parallel $ do
miscTests args
describe ("Building environments for top " <> show topN <> " Julia packages") $
parallelN parallelism $
forM_ (L.take topN namesAndCounts) $ \(NameAndCount {..}) ->
testExpr args name [i|#{juliaAttr}.withPackages ["#{name}"]|]
introduce "Introduce parallel semaphore" parallelSemaphore (liftIO $ newQSem parallelism) (const $ return ()) $
parallel $
forM_ (L.take topN namesAndCounts) $ \(NameAndCount {..}) ->
around "Claim semaphore" claimRunSlot $
testExpr args name [i|#{juliaAttr}.withPackages ["#{name}"]|]
miscTests :: Args -> SpecFree ctx IO ()
miscTests args@(Args {..}) = describe "Misc tests" $ do
@@ -89,6 +93,9 @@ miscTests args@(Args {..}) = describe "Misc tests" $ do
};
}) [ "HelloWorld" ]|]
describe "misc cases" $ do
testExpr args "Optimization" [iii|(#{juliaAttr}.withPackages) [ "Optimization" "OptimizationOptimJL" ]|]
-- * Low-level
testExpr :: Args -> Text -> String -> SpecFree ctx IO ()
@@ -98,7 +105,9 @@ testExpr _args name expr = do
let cp = proc "nix" ["build", "--impure", "--no-link", "--json", "--expr", [i|with import ../../../../. {}; #{expr}|]]
output <- readCreateProcessWithLogging cp ""
juliaPath <- case A.eitherDecode (BL8.pack output) of
Right (A.Array ((V.!? 0) -> Just (A.Object (aesonLookup "outputs" -> Just (A.Object (aesonLookup "out" -> Just (A.String t))))))) -> pure (JuliaPath ((T.unpack t) </> "bin" </> "julia"))
Right (A.Array ((V.!? 0) -> Just (A.Object (aesonLookup "outputs" -> Just (A.Object (aesonLookup "out" -> Just (A.String t))))))) -> do
info [i|built: #{t}|]
pure (JuliaPath ((T.unpack t) </> "bin" </> "julia"))
x -> expectationFailure ("Couldn't parse output: " <> show x)
getContext julia >>= flip modifyMVar_ (const $ return (Just juliaPath))
@@ -113,3 +122,8 @@ testExpr _args name expr = do
where
aesonLookup :: Text -> HM.KeyMap v -> Maybe v
aesonLookup = HM.lookup . A.fromText
claimRunSlot :: (HasParallelSemaphore ctx) => ExampleT ctx IO a -> ExampleT ctx IO ()
claimRunSlot f = do
s <- getContext parallelSemaphore
bracket_ (liftIO $ waitQSem s) (liftIO $ signalQSem s) (void f)
+5 -1
View File
@@ -1,5 +1,6 @@
{
gitMinimal,
lib,
runCommand,
}:
@@ -10,7 +11,10 @@
# See https://github.com/NixOS/nixpkgs/pull/97467#issuecomment-689315186
addPackagesToPython =
python: packages:
if python ? "env" then
# TODO: this stopped working because "env" ended up being a key of the base
# derivation like "python3" as well. Is there a robust way to determine if
# this Python is already wrapped?
if python ? "env" && lib.isDerivation python.env then
python.override (old: {
extraLibs = old.extraLibs ++ packages;
})