Merge pull request #419 from JuliaParallel/jps/storage-tags

Add File and tofile helpers

Author: jpsamaroo

Project.toml

@@ -24,7 +24,7 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 ContextVariablesX = "0.1"
 DataStructures = "0.18"
 MacroTools = "0.5"
-MemPool = "0.4.3"
+MemPool = "0.4.4"
 Requires = "1"
 StatsBase = "0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34"
 TimespanLogging = "0.1"

docs/src/index.md

@@ -110,7 +110,30 @@ or by specifying the `worker` argument to `@mutable`:
 A = Dagger.@mutable worker=2 rand(1000, 1000)
 ```
 
-### Parallel reduction
+### Operate on distributed data
+
+Often we want to work with more than one piece of data; the common case of
+wanting one piece of data per worker is easy to do by using `@shard`:
+
+```julia
+X = Dagger.@shard myid()
+```
+
+This will execute `myid()` independently on every worker in your Julia
+cluster, and place references to each within a `Shard` object called `X`. We
+can then use `X` in task spawning, but we'll only get the result of
+`myid()` that corresponds to the worker that the task is running on:
+
+```julia
+for w in workers()
+    @show fetch(Dagger.@spawn scope=Dagger.scope(worker=w) identity(X))
+end
+```
+
+The above should print the result of `myid()` for each worker in `worker()`, as
+`identity(X)` receives only the value of `X` specific to that worker.
+
+### Reducing over distributed data
 
 Reductions are often parallelized by reducing a set of partitions on each
 worker, and then reducing those intermediate reductions on a single worker.
@@ -131,14 +154,74 @@ points to a set of histogram bins on each worker. When we `@spawn hist!`,
 Dagger passes in the random array and bins for only the specific worker that
 the task is run on; i.e. a call to `hist!` that runs on worker 2 will get a
 different `A` and `temp_bins` from a call to `hist!` on worker 3. All of the
-calls to `hist!` may run in parallel
+calls to `hist!` may run in parallel.
 
 By using `map` on `temp_bins`, we then make a copy of each worker's bins that
 we can safely return back to our current worker, and sum them together to get
 our total histogram.
 
 -----
 
+## Quickstart: File IO
+
+Dagger has support for loading and saving files that integrates seamlessly with
+its task system, in the form of `Dagger.File` and `Dagger.tofile`.
+
+!!! warn
+    These functions are not yet fully tested, so please make sure to take backups of any files that you load with them.
+
+### Loading files from disk
+
+In order to load one or more files from disk, Dagger provides the `File`
+function, which creates a lazy reference to a file:
+
+```julia
+f = Dagger.File("myfile.jls")
+```
+
+`f` is now a lazy reference to `"myfile.jls"`, and its contents can be loaded
+automatically by just passing the object to a task:
+
+```julia
+wait(Dagger.@spawn println(f))
+# Prints the loaded contents of the file
+```
+
+By default, `File` assumes that the file uses Julia's Serialization format;
+this can be easily changed to assume Arrow format, for example:
+
+```julia
+using Arrow
+f = Dagger.File("myfile.arrow"; serialize=Arrow.write, deserialize=Arrow.Table)
+```
+
+### Writing data to disk
+
+Saving data to disk is as easy as loading it; `tofile` provides this capability
+in a similar manner to `File`:
+
+```julia
+A = rand(1000)
+f = Dagger.tofile(A, "mydata.jls")
+```
+
+Like `File`, `f` can still be used to reference the file's data in tasks. It is
+likely most useful to use `tofile` at the end of a task to save results:
+
+```julia
+function make_data()
+    A = rand(1000)
+    return Dagger.tofile(A, "mydata.jls")
+end
+fetch(Dagger.@spawn make_data())
+# Data was also written to "mydata.jls"
+```
+
+`tofile` takes the same keyword arguments as `File`, allowing the format of
+data on disk to be specified as desired.
+
+-----
+
 ## Quickstart: Distributed Arrays
 
 Dagger's `DArray` type represents a distributed array, where a single large

src/file-io.jl

@@ -1,5 +1,83 @@
 export save, load
 using SparseArrays
+import MemPool: GenericFileDevice, FileRef
+
+struct File
+    path::String
+    chunk::Chunk
+end
+
+"""
+    File(path::AbstractString;
+         serialize::Base.Callable, deserialize::Base.Callable,
+         use_io::Bool, mmap::Bool) -> Dagger.File
+
+References data in the file at `path`, using the derialization function `deserialize`.
+`use_io` specifies whether `deserialize` takes an `IO` (the default) or takes a
+`String` path. `mmap` is experimental, and specifies whether `deserialize` can
+use MemPool's `MMWrap` memory mapping functionality (default is `false`).
+
+The file at `path` is not yet loaded when the call to `File` returns; passing
+it to a task will cause reading to occur, and the result will be passed to the
+task.
+"""
+function File(path::AbstractString;
+              serialize::Base.Callable=serialize,
+              deserialize::Base.Callable=deserialize,
+              use_io::Bool=true,
+              mmap::Bool=false)
+    if !isfile(path)
+        # FIXME: Once MemPool better propagates errors, this check won't be necessary
+        throw(ArgumentError("`Dagger.File` expected file to exist at \"$path\""))
+    end
+    if isabspath(path)
+        dir, file = dirname(path), basename(path)
+    else
+        dir, file = pwd(), basename(path)
+    end
+    Tdevice = GenericFileDevice{serialize, deserialize, use_io, mmap}
+    device = Tdevice(dir)
+    leaf_tag = MemPool.Tag(Tdevice=>file)
+    chunk = Dagger.tochunk(FileRef(path), OSProc(), ProcessScope();
+                           restore=true, retain=true, device, leaf_tag)
+    return File(path, chunk)
+end
+
+"""
+    tofile(data, path::AbstractString;
+           serialize::Base.Callable, deserialize::Base.Callable,
+           use_io::Bool, mmap::Bool) -> Dagger.File
+
+Writes `data` to disk at `path`, using the serialization function `serialize`.
+`use_io` specifies whether `serialize` takes an `IO` (the default) or takes a
+`String` path. `mmap` is experimental, and specifies whether `serialize` can
+use MemPool's `MMWrap` memory mapping functionality (default is `false`).
+
+The returned `File` object can be passed to tasks as an argument, or returned
+from tasks as a result.
+"""
+function tofile(@nospecialize(data), path::AbstractString;
+                serialize::Base.Callable=serialize,
+                deserialize::Base.Callable=deserialize,
+                use_io::Bool=true,
+                mmap::Bool=false)
+    if isabspath(path)
+        dir, file = dirname(path), basename(path)
+    else
+        dir, file = pwd(), basename(path)
+    end
+    Tdevice = GenericFileDevice{serialize, deserialize, use_io, mmap}
+    device = Tdevice(dir)
+    chunk = Dagger.tochunk(data, OSProc(), ProcessScope();
+                           retain=true, device, tag=file)
+    return File(path, chunk)
+end
+Base.fetch(file::File) = fetch(file.chunk)
+Base.show(io::IO, file::File) = print(io, "Dagger.File(path=\"$(file.path)\")")
+
+function move(from_proc::Processor, to_proc::Processor, file::File)
+    return move(from_proc, to_proc, file.chunk)
+end
 
 """
     FileReader

src/sch/Sch.jl

@@ -219,6 +219,16 @@ error will be displayed.
 when constructing `Chunk`s (such as when constructing the return value). The
 device must support `MemPool.CPURAMResource`. When `nothing`, uses
 `MemPool.GLOBAL_DEVICE[]`.
+- `storage_root_tag::Any=nothing`: If not `nothing`,
+specifies the MemPool storage leaf tag to associate with the thunk's result.
+This tag can be used by MemPool's storage devices to manipulate their behavior,
+such as the file name used to store data on disk."
+- `storage_leaf_tag::MemPool.Tag,Nothing}=nothing`: If not `nothing`,
+specifies the MemPool storage leaf tag to associate with the thunk's result.
+This tag can be used by MemPool's storage devices to manipulate their behavior,
+such as the file name used to store data on disk."
+- `storage_retain::Bool=false`: The value of `retain` to pass to
+`MemPool.poolset` when constructing the result `Chunk`.
 """
 Base.@kwdef struct ThunkOptions
     single::Union{Int,Nothing} = nothing
@@ -230,6 +240,9 @@ Base.@kwdef struct ThunkOptions
     checkpoint = nothing
     restore = nothing
     storage::Union{Chunk,Nothing} = nothing
+    storage_root_tag = nothing
+    storage_leaf_tag::Union{MemPool.Tag,Nothing} = nothing
+    storage_retain::Bool = false
 end
 
 """
@@ -249,7 +262,10 @@ function Base.merge(sopts::SchedulerOptions, topts::ThunkOptions)
                  allow_errors,
                  topts.checkpoint,
                  topts.restore,
-                 topts.storage)
+                 topts.storage,
+                 topts.storage_root_tag,
+                 topts.storage_leaf_tag,
+                 topts.storage_retain)
 end
 Base.merge(sopts::SchedulerOptions, ::Nothing) =
     ThunkOptions(sopts.single,
@@ -283,6 +299,9 @@ function populate_defaults(opts::ThunkOptions, Tf, Targs)
         maybe_default(:checkpoint),
         maybe_default(:restore),
         maybe_default(:storage),
+        maybe_default(:storage_root_tag),
+        maybe_default(:storage_leaf_tag),
+        maybe_default(:storage_retain),
     )
 end
 
@@ -530,9 +549,9 @@ function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
             node = unwrap_weak_checked(state.thunk_dict[thunk_id])
             if metadata !== nothing
                 state.worker_time_pressure[pid][proc] = metadata.time_pressure
-                to_storage = node.options.storage
-                state.worker_storage_pressure[pid][to_storage] = metadata.storage_pressure
-                state.worker_storage_capacity[pid][to_storage] = metadata.storage_capacity
+                #to_storage = fetch(node.options.storage)
+                #state.worker_storage_pressure[pid][to_storage] = metadata.storage_pressure
+                #state.worker_storage_capacity[pid][to_storage] = metadata.storage_capacity
                 state.worker_loadavg[pid] = metadata.loadavg
                 sig = signature(node, state)
                 state.signature_time_cost[sig] = (metadata.threadtime + get(state.signature_time_cost, sig, 0)) ÷ 2
@@ -1546,7 +1565,10 @@ function do_task(to_proc, task_desc)
 
         # Construct result
         # TODO: We should cache this locally
-        send_result || meta ? res : tochunk(res, to_proc; device, persist, cache=persist ? true : cache)
+        send_result || meta ? res : tochunk(res, to_proc; device, persist, cache=persist ? true : cache,
+                                            tag=options.storage_root_tag,
+                                            leaf_tag=something(options.storage_leaf_tag, MemPool.Tag()),
+                                            retain=options.storage_retain)
     catch ex
         bt = catch_backtrace()
         RemoteException(myid(), CapturedException(ex, bt))

test/file-io.jl

@@ -0,0 +1,33 @@
+using Serialization
+
+@testset "File IO" begin
+    @testset "File" begin
+        data = [1,2,3]
+        data_path = joinpath(tempdir(), "jl_" * join(rand('a':'z', 8)) * ".jls")
+        atexit() do
+            @assert isfile(data_path)
+            rm(data_path)
+        end
+        serialize(data_path, data)
+
+        data_c = Dagger.File(data_path)::Dagger.File
+        @test fetch(data_c) == data
+        @test fetch(Dagger.@spawn identity(data_c)) == data
+
+        @test isfile(data_path)
+    end
+    @testset "tofile" begin
+        data = [4,5,6]
+        data_path = joinpath(tempdir(), "jl_" * join(rand('a':'z', 8)) * ".jls")
+        atexit() do
+            @assert isfile(data_path)
+            rm(data_path)
+        end
+
+        data_c = Dagger.tofile(data, data_path)::Dagger.File
+        @test fetch(data_c) == data
+        @test fetch(Dagger.@spawn identity(data_c)) == data
+
+        @test isfile(data_path)
+    end
+end

test/runtests.jl

@@ -38,6 +38,7 @@ include("domain.jl")
 include("array.jl")
 include("cache.jl")
 include("diskcaching.jl")
+include("file-io.jl")
 
 try # TODO: Fault tolerance is sometimes unreliable
 #include("fault-tolerance.jl")