Add ClimaArtifacts

.buildkite/pipeline.yml

@@ -19,6 +19,8 @@ steps:
       - "julia --project -e 'using CUDA; CUDA.precompile_runtime()'"
       - "julia --project -e 'using Pkg; Pkg.status()'"
 
+      - "julia --project=test/test_artifacts -e 'using Pkg; Pkg.develop(;path=\".\"); Pkg.instantiate()'"
+
     agents:
       slurm_cpus_per_task: 8
     env:
@@ -55,3 +57,28 @@ steps:
       slurm_nodes: 1
       slurm_ntasks_per_node: 2
       slurm_gpus_per_task: 1
+
+  # The artifacts jobs cannot be concurrent
+  - label: ":amphora: artifacts"
+    key: "artifacts"
+    command:
+      - julia --project=test/test_artifacts test/test_artifacts/test_clima_artifacts.jl
+    env:
+      CLIMACOMMS_TEST_DEVICE: CPU
+    agents:
+      slurm_nodes: 1
+      slurm_ntasks_per_node: 1
+    concurrency: 1
+    concurrency_group: 'artifacts'
+
+  - label: ":amphora: artifacts (MPI)"
+    key: "mpi_artifacts"
+    command:
+      - julia --project=test/test_artifacts test/test_artifacts/test_clima_artifacts.jl
+    env:
+      CLIMACOMMS_TEST_DEVICE: CPU
+    agents:
+      slurm_nodes: 1
+      slurm_ntasks_per_node: 2
+    concurrency: 1
+    concurrency_group: 'artifacts'

Project.toml

@@ -1,16 +1,10 @@
 name = "ClimaComms"
 uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
-authors = [
-    "Kiran Pamnany <[email protected]>",
-    "Simon Byrne <[email protected]>",
-    "Charles Kawczynski <[email protected]>",
-    "Sriharsha Kandala <[email protected]>",
-    "Jake Bolewski <[email protected]>",
-    "Gabriele Bozzola <[email protected]>",
-]
-version = "0.5.7"
+authors = ["Kiran Pamnany <[email protected]>", "Simon Byrne <[email protected]>", "Charles Kawczynski <[email protected]>", "Sriharsha Kandala <[email protected]>", "Jake Bolewski <[email protected]>", "Gabriele Bozzola <[email protected]>"]
+version = "0.5.8"
 
 [deps]
+Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 

src/ClimaArtifacts.jl

@@ -0,0 +1,150 @@
+module ClimaArtifacts
+
+import Base.BinaryPlatforms: HostPlatform
+import Artifacts as JuliaArtifacts
+
+import ..SingletonCommsContext, ..MPICommsContext, ..iamroot, ..barrier
+
+const ACCESSED_ARTIFACTS::Set{String} = Set(String[])
+
+root_or_singleton(::SingletonCommsContext) = true
+root_or_singleton(ctx::MPICommsContext) = iamroot(ctx)
+
+maybe_wait(::SingletonCommsContext) = nothing
+maybe_wait(ctx::MPICommsContext) = barrier(ctx)
+
+# This code is largely a re-implementation of Artifacts.artifact_str extended to add
+# instrumentation and control MPI
+macro clima_artifact(name, context)
+    # Find Artifacts.toml file we're going to load from
+    srcfile = string(__source__.file)
+    if (
+        (isinteractive() && startswith(srcfile, "REPL[")) ||
+        (!isinteractive() && srcfile == "none")
+    ) && !isfile(srcfile)
+        srcfile = pwd()
+    end
+    local artifacts_toml = JuliaArtifacts.find_artifacts_toml(srcfile)
+    if isnothing(artifacts_toml)
+        error(
+            string(
+                "Cannot locate '(Julia)Artifacts.toml' file when attempting to use artifact '",
+                name,
+                "' in '",
+                __module__,
+                "'",
+            ),
+        )
+    end
+
+    # Load Artifacts.toml at compile time, so that we don't have to use `__source__.file`
+    # at runtime, which gets stale if the `.ji` file is relocated.
+    local artifact_dict = JuliaArtifacts.load_artifacts_toml(artifacts_toml)
+
+    # Invalidate calling .ji file if Artifacts.toml file changes
+    Base.include_dependency(artifacts_toml)
+
+    # Check if the user has provided `LazyArtifacts`, and thus supports lazy artifacts
+    # If not, check to see if `Pkg` or `Pkg.Artifacts` has been imported.
+    lazyartifacts = nothing
+    for module_name in (:LazyArtifacts, :Pkg, :Artifacts)
+        if isdefined(__module__, module_name)
+            lazyartifacts = GlobalRef(__module__, module_name)
+            break
+        end
+    end
+
+    # Artifacts.artifact_str deals with platforms, but we do not need to support that
+    # feature
+    platform = HostPlatform()
+
+    # If `name` is a constant, we can actually load and parse the `Artifacts.toml` file now,
+    # saving the work from runtime.
+    if isa(name, AbstractString)
+        # To support slash-indexing, we need to split the artifact name from the path tail:
+        artifact_name, artifact_path_tail, hash =
+            JuliaArtifacts.artifact_slash_lookup(
+                name,
+                artifact_dict,
+                artifacts_toml,
+                platform,
+            )
+        return quote
+            # We call JuliaArtifacts._artifact_str twice, the first time only with the root
+            # process (to avoid race conditions), the second time to ensure that all the
+            # processes have the artifact string
+            if Base.invokelatest(root_or_singleton, $(esc(context)))
+                artifact_path = Base.invokelatest(
+                    JuliaArtifacts._artifact_str,
+                    $(__module__),
+                    $(artifacts_toml),
+                    $(artifact_name),
+                    $(artifact_path_tail),
+                    $(artifact_dict),
+                    $(hash),
+                    $(platform),
+                    $(lazyartifacts),
+                )::String
+                push!(ACCESSED_ARTIFACTS, artifact_path)
+            end
+            Base.invokelatest(maybe_wait, $(esc(context)))
+            # When we call _artifact_str again, we can now assume that the artifact is
+            # available
+            return Base.invokelatest(
+                JuliaArtifacts._artifact_str,
+                $(__module__),
+                $(artifacts_toml),
+                $(artifact_name),
+                $(artifact_path_tail),
+                $(artifact_dict),
+                $(hash),
+                $(platform),
+                $(lazyartifacts),
+            )::String
+        end
+    else
+        # If artifact_name is not a string (e.g., it is a variable), we have to do all the
+        # work at runtime
+        return quote
+            local platform = $(esc(platform))
+            local artifact_name, artifact_path_tail, hash =
+                JuliaArtifacts.artifact_slash_lookup(
+                    $(esc(name)),
+                    $(artifact_dict),
+                    $(artifacts_toml),
+                    platform,
+                )
+            # We call JuliaArtifacts._artifact_str twice, the first time only with the root
+            # process (to avoid race conditions), the second time to ensure that all the
+            # processes have the artifact string
+            if Base.invokelatest(root_or_singleton, $(esc(context)))
+                artifact_path = Base.invokelatest(
+                    JuliaArtifacts._artifact_str,
+                    $(__module__),
+                    $(artifacts_toml),
+                    artifact_name,
+                    artifact_path_tail,
+                    $(artifact_dict),
+                    hash,
+                    platform,
+                    $(lazyartifacts),
+                )::String
+                push!(ACCESSED_ARTIFACTS, artifact_path)
+            end
+            Base.invokelatest(maybe_wait, $(esc(context)))
+            return Base.invokelatest(
+                JuliaArtifacts._artifact_str,
+                $(__module__),
+                $(artifacts_toml),
+                artifact_name,
+                artifact_path_tail,
+                $(artifact_dict),
+                hash,
+                platform,
+                $(lazyartifacts),
+            )::String
+        end
+    end
+end
+
+end

src/ClimaComms.jl

@@ -16,4 +16,8 @@ include("context.jl")
 include("singleton.jl")
 include("mpi.jl")
 
+include("ClimaArtifacts.jl")
+import .ClimaArtifacts: @clima_artifact
+export @clima_artifact
+
 end # module

test/test_artifacts/Artifacts.toml

@@ -0,0 +1,7 @@
+[socrates]
+git-tree-sha1 = "43563e7631a7eafae1f9f8d9d332e3de44ad7239"
+lazy = true
+
+    [[socrates.download]]
+    url = "https://github.com/staticfloat/small_bin/raw/master/socrates.tar.gz"
+    sha256 = "e65d2f13f2085f2c279830e863292312a72930fee5ba3c792b14c33ce5c5cc58"

test/test_artifacts/Project.toml

@@ -0,0 +1,9 @@
+name = "ClimaTestArtifact"
+
+[deps]
+Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[compat]
+ClimaComms = "0.5.8"

test/test_artifacts/test_clima_artifacts.jl

@@ -0,0 +1,14 @@
+using Artifacts
+using Test
+
+import ClimaComms
+
+const context = ClimaComms.context()
+
+@show expected_path = artifact"socrates"
+
+# Remove the artifact, so that we test that we are downloading it
+Base.Filesystem.rm(dirname(expected_path), recursive = true)
+@info "Removed artifact"
+
+@test ClimaComms.@clima_artifact("socrates", context) == expected_path