feat: lake: record cache service in outputs (#12113)

This PR changes the alters the file format of outputs stored in the
local Lake cache to include an identifier indicating the service (if
any) the output came from. This will be used to enable lazily
downloading artifacts on-demand during builds.

Author: tydeu

src/lake/Lake/CLI/Help.lean

@@ -377,10 +377,11 @@ OPTIONS:
 Downloads artifacts for packages in the workspace from a remote cache service.
 The cache service used can be configured via the environment variables:
 
+  LAKE_CACHE_SERVICE            identifier recorded in ouptuts
   LAKE_CACHE_ARTIFACT_ENDPOINT  base URL for artifact downloads
   LAKE_CACHE_REVISION_ENDPOINT  base URL for the mapping download
 
-If neither of these are set, Lake will use Reservoir.
+If neither endpoint is set, Lake will use Reservoir.
 
 If an input-to-outputs mappings file, `--scope`, or `--repo` is provided,
 Lake will download artifacts for the root package. Otherwise, it will use

src/lake/Lake/CLI/Main.lean

@@ -63,6 +63,7 @@ public structure LakeOptions where
   offline : Bool := false
   outputsFile? : Option FilePath := none
   forceDownload : Bool := false
+  service? : Option String := none
   scope? : Option String := none
   /-- Was `scope?` set with `--repo` (and not `--scope`)? -/
   repoScope : Bool := false
@@ -248,6 +249,9 @@ def lakeLongOption : (opt : String) → CliM PUnit
 | "--wfail"       => modifyThe LakeOptions ({· with failLv := .warning})
 | "--iofail"      => modifyThe LakeOptions ({· with failLv := .info})
 | "--force-download" => modifyThe LakeOptions ({· with forceDownload := true})
+| "--service" => do
+  let service ← takeOptArg "--service" "service name"
+  modifyThe LakeOptions ({· with service? := some service})
 | "--scope" => do
   let scope ← takeOptArg "--scope" "cache scope"
   modifyThe LakeOptions ({· with scope? := some scope, repoScope := false})
@@ -396,7 +400,7 @@ protected def get : CliM PUnit := do
       | error "to use `cache get` with a mappings file, `--scope` or `--repo` must be set"
     let service : CacheService :=
       if let some artifactEndpoint := ws.lakeEnv.cacheArtifactEndpoint? then
-        .downloadArtsService artifactEndpoint
+        .downloadArtsService artifactEndpoint ws.lakeEnv.cacheService?
       else
         .reservoirService ws.lakeEnv.reservoirApiUrl
     let map ← CacheMap.load file
@@ -407,7 +411,7 @@ protected def get : CliM PUnit := do
     let service : CacheService ← id do
       match ws.lakeEnv.cacheArtifactEndpoint?, ws.lakeEnv.cacheRevisionEndpoint? with
       | some artifactEndpoint, some revisionEndpoint =>
-        return .downloadService artifactEndpoint revisionEndpoint
+        return .downloadService artifactEndpoint revisionEndpoint ws.lakeEnv.cacheService?
       | none, none =>
         return .reservoirService ws.lakeEnv.reservoirApiUrl
       | some artifactEndpoint, none =>
@@ -531,7 +535,7 @@ protected def add : CliM PUnit := do
     | _ => pure ws.root
   let scope := pkg.cacheScope
   let map ← CacheMap.load file
-  ws.lakeCache.writeMap scope map
+  ws.lakeCache.writeMap scope map opts.service?
 
 protected def help : CliM PUnit := do
   IO.println <| helpCache <| ← takeArgD ""

src/lake/Lake/Config/Cache.lean

@@ -211,6 +211,40 @@ public structure Cache where
   dir : FilePath
   deriving Inhabited
 
+/-- The current version of the output file format. -/
+def CacheOutput.schemaVersion : String := "2026-02-10"
+
+structure CacheOutput where
+  service? : Option String := none
+  data : Json
+  deriving ToJson
+
+namespace CacheOutput
+
+protected def toJson (out : CacheOutput) : Json :=
+  JsonObject.empty
+  |>.insert "schemaVersion" schemaVersion
+  |>.insert "service" out.service?
+  |>.insert "data" out.data
+
+instance : ToJson CacheOutput := ⟨CacheOutput.toJson⟩
+
+protected def fromJson? (json : Json) : Except String CacheOutput := do
+  if let .obj obj := json then
+    let obj := JsonObject.mk obj
+    if obj.contains "schemaVersion" then
+      -- presumably the new format
+      -- (the edge case of a custom output with a `schemaVersion` is not worth covering)
+      let service? ← obj.get? "service"
+      let data ← obj.get "data"
+      return {service?, data}
+  -- old format: just the data
+  return {data := json}
+
+instance : FromJson CacheOutput := ⟨CacheOutput.fromJson?⟩
+
+end CacheOutput
+
 namespace Cache
 
 /-- Returns the artifact directory for the Lake cache. -/
@@ -260,29 +294,34 @@ public def getArtifactPaths
   cache.outputsDir / scope / s!"{inputHash}.json"
 
 /-- Cache the outputs corresponding to the given input for the package.  -/
-public def writeOutputsCore
+def writeOutputsCore
   (cache : Cache) (scope : String) (inputHash : Hash) (outputs : Json)
+  (service? : Option String := none)
 : IO Unit := do
   let file := cache.outputsFile scope inputHash
   createParentDirs file
-  IO.FS.writeFile file outputs.compress
+  let out := {service?, data := outputs : CacheOutput}
+  IO.FS.writeFile file (toJson out).pretty
 
 /-- Cache the outputs corresponding to the given input for the package.  -/
 @[inline] public def writeOutputs
   [ToJson α] (cache : Cache) (scope : String) (inputHash : Hash) (outputs : α)
-: IO Unit := cache.writeOutputsCore scope inputHash (toJson outputs)
+  (service? : Option String := none)
+: IO Unit := cache.writeOutputsCore scope inputHash (toJson outputs) service?
 
 /-- Cache the input-to-outputs mappings from a `CacheMap`.  -/
-public def writeMap (cache : Cache) (scope : String) (map : CacheMap) : IO Unit :=
-  map.forM fun i o => cache.writeOutputsCore scope i o
+public def writeMap
+  (cache : Cache) (scope : String) (map : CacheMap) (service? : Option String := none)
+: IO Unit := map.forM fun i o => cache.writeOutputsCore scope i o service?
 
 /-- Retrieve the cached outputs corresponding to the given input for the package (if any). -/
 public def readOutputs? (cache : Cache) (scope : String) (inputHash : Hash) : LogIO (Option Json) := do
   let path := cache.outputsFile scope inputHash
   match (← IO.FS.readFile path |>.toBaseIO) with
   | .ok contents =>
-    match Json.parse contents with
-    | .ok outputs => return outputs
+    match Json.parse contents >>= fromJson? with
+    | .ok out =>
+      return CacheOutput.data out
     | .error e =>
       logWarning s!"{path}: invalid JSON: {e}"
       return none
@@ -336,6 +375,7 @@ the desired functions by using `CacheService`'s smart constructors
 -/
 public structure CacheService where
   private mk ::
+    private name? : Option String := none
     /- S3 Bucket -/
     private key : String := ""
     private artifactEndpoint : String := ""
@@ -353,19 +393,22 @@ namespace CacheService
 
 /-- Constructs a `CacheService` for a Reservoir endpoint. -/
 @[inline] public def reservoirService (apiEndpoint : String) (repoScope := false) : CacheService :=
-  {isReservoir := true, apiEndpoint, repoScope}
+  {name? := some "reservoir", isReservoir := true, apiEndpoint, repoScope}
 
 /-- Constructs a `CacheService` to upload artifacts and/or outputs to an S3 endpoint. -/
-@[inline] public def uploadService (key artifactEndpoint revisionEndpoint : String) : CacheService :=
-  {key, artifactEndpoint, revisionEndpoint}
+@[inline] public def uploadService
+  (key artifactEndpoint revisionEndpoint : String)
+: CacheService := {key, artifactEndpoint, revisionEndpoint}
 
-/-- Constructs a `CacheService` to download artifacts and/or outputs from to an S3 endpoint. -/
-@[inline] public def downloadService (artifactEndpoint revisionEndpoint : String) : CacheService :=
-  {artifactEndpoint, revisionEndpoint}
+/-- Constructs a `CacheService` to download artifacts and/or outputs from an S3 endpoint. -/
+@[inline] public def downloadService
+  (artifactEndpoint revisionEndpoint : String) (name? : Option String := none)
+: CacheService := {name?, artifactEndpoint, revisionEndpoint}
 
-/-- Constructs a `CacheService` to download just artifacts from to an S3 endpoint. -/
-@[inline] public def downloadArtsService (artifactEndpoint : String) : CacheService :=
-  {artifactEndpoint}
+/-- Constructs a `CacheService` to download just artifacts from an S3 endpoint. -/
+@[inline] public def downloadArtsService
+  (artifactEndpoint : String) (name? : Option String := none)
+: CacheService := {name?, artifactEndpoint}
 
 /--
 Reconfigures the cache service to interpret scopes as repositories (or not if `false`).
@@ -385,10 +428,10 @@ private def appendScope (endpoint : String) (scope : String) : String :=
   scope.split '/' |>.fold (init := endpoint) fun s component =>
     uriEncode component.copy s |>.push '/'
 
-private def s3ArtifactUrl (contentHash : Hash) (service : CacheService) (scope : String)  : String :=
+private def s3ArtifactUrl (contentHash : Hash) (service : CacheService) (scope : String) : String :=
   appendScope s!"{service.artifactEndpoint}/" scope ++ s!"{contentHash.hex}.art"
 
-public def artifactUrl (contentHash : Hash) (service : CacheService) (scope : String)  : String :=
+public def artifactUrl (contentHash : Hash) (service : CacheService) (scope : String) : String :=
   if service.isReservoir then
     let endpoint :=
       if service.repoScope then
@@ -435,7 +478,7 @@ public def downloadOutputArtifacts
   (map : CacheMap) (cache : Cache) (service : CacheService)
   (localScope remoteScope : String) (force := false)
 : LoggerIO Unit := do
-  cache.writeMap localScope map
+  cache.writeMap localScope map service.name?
   let descrs ← map.collectOutputDescrs
   service.downloadArtifacts descrs cache remoteScope force
 

src/lake/Lake/Config/Env.lean

@@ -64,6 +64,8 @@ public structure Env where
   cacheArtifactEndpoint? : Option String
   /-- The base URL for revision uploads and downloads from the cache (i.e., `LAKE_CACHE_REVISION_ENDPOINT`). -/
   cacheRevisionEndpoint? : Option String
+  /-- The name of the cache service (i.e., `LAKE_CACHE_SERVICE`). -/
+  cacheService? : Option String
   /-- The initial Lean library search path of the environment (i.e., `LEAN_PATH`). -/
   initLeanPath : SearchPath
   /-- The initial Lean source search path of the environment (i.e., `LEAN_SRC_PATH`). -/
@@ -163,6 +165,7 @@ public def compute
     cacheKey? := (← IO.getEnv "LAKE_CACHE_KEY").map (·.trimAscii.copy)
     cacheArtifactEndpoint? := (← IO.getEnv "LAKE_CACHE_ARTIFACT_ENDPOINT").map normalizeUrl
     cacheRevisionEndpoint? := (← IO.getEnv "LAKE_CACHE_REVISION_ENDPOINT").map normalizeUrl
+    cacheService? := (← IO.getEnv "LAKE_CACHE_SERVICE").map (·.trimAscii.copy)
     githashOverride := (← IO.getEnv "LEAN_GITHASH").getD ""
     toolchain
     initLeanPath := ← getSearchPath "LEAN_PATH",
@@ -273,6 +276,7 @@ public def baseVars (env : Env) : Array (String × Option String)  :=
     ("LAKE_CACHE_KEY", env.cacheKey?),
     ("LAKE_CACHE_ARTIFACT_ENDPOINT", env.cacheArtifactEndpoint?),
     ("LAKE_CACHE_REVISION_ENDPOINT", env.cacheRevisionEndpoint?),
+    ("LAKE_CACHE_SERVICE", env.cacheService?),
     ("LEAN", env.lean.lean.toString),
     ("LEAN_SYSROOT", env.lean.sysroot.toString),
     ("LEAN_AR", env.lean.ar.toString),

src/lake/Lake/Util/JsonObject.lean

@@ -26,6 +26,9 @@ namespace JsonObject
 @[inline] public def mk (val : Std.TreeMap.Raw String Json) : JsonObject :=
   val
 
+@[inline] public def empty : JsonObject :=
+  Std.TreeMap.Raw.empty
+
 @[inline] public protected def toJson (obj : JsonObject) : Json :=
   .obj obj
 
@@ -37,6 +40,9 @@ public instance : ToJson JsonObject := ⟨JsonObject.toJson⟩
 
 public instance : FromJson JsonObject := ⟨JsonObject.fromJson?⟩
 
+@[inline] public nonrec def contains (obj : JsonObject) (prop : String) : Bool :=
+  obj.contains prop
+
 public def insertJson (obj : JsonObject) (prop : String) (val : Json) : JsonObject :=
   obj.insert prop (toJson val) -- specializes `insert`
 

tests/lake/tests/env/test.sh

@@ -38,6 +38,7 @@ test_eq "foo" env env ELAN_TOOLCHAIN=foo $LAKE env printenv ELAN_TOOLCHAIN
 test_out "foo" env env -u LAKE_CACHE_DIR ELAN_HOME=/ ELAN_TOOLCHAIN=foo \
   $LAKE env printenv LAKE_CACHE_DIR
 LAKE_CACHE_DIR=foo test_eq "foo" env printenv LAKE_CACHE_DIR
+LAKE_CACHE_SERVICE=foo test_eq "foo" env printenv LAKE_CACHE_SERVICE
 LEAN_GITHASH=foo test_eq "foo" env printenv LEAN_GITHASH
 LEAN_AR=foo test_eq "foo" env printenv LEAN_AR
 LEAN_CC=foo test_eq "foo" env printenv LEAN_CC