Add metadata to EagerThunk

Author: jpsamaroo

Project.toml

@@ -9,6 +9,7 @@ Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
+Mmap = "a63ad114-7e13-5084-954f-fe012c677804"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"

src/eager_thunk.jl

@@ -36,6 +36,16 @@ end
 Options(;options...) = Options((;options...))
 Options(options...) = Options((;options...))
 
+"""
+    EagerThunkMetadata
+
+Represents some useful metadata pertaining to an `EagerThunk`:
+- `return_type::Type` - The inferred return type of the task
+"""
+mutable struct EagerThunkMetadata
+    return_type::Type
+end
+
 """
     EagerThunk
 
@@ -46,9 +56,11 @@ be `fetch`'d or `wait`'d on at any time.
 mutable struct EagerThunk
     uid::UInt
     future::ThunkFuture
+    metadata::EagerThunkMetadata
     finalizer_ref::DRef
     thunk_ref::DRef
-    EagerThunk(uid, future, finalizer_ref) = new(uid, future, finalizer_ref)
+    EagerThunk(uid, future, metadata, finalizer_ref) =
+        new(uid, future, metadata, finalizer_ref)
 end
 
 Base.isready(t::EagerThunk) = isready(t.future)

src/submission.jl

@@ -217,14 +217,22 @@ function eager_process_options_submission_to_local(id_map, options::NamedTuple)
         return options
     end
 end
+function EagerThunkMetadata(spec::EagerTaskSpec)
+    f = chunktype(spec.f).instance
+    arg_types = ntuple(i->chunktype(spec.args[i][2]), length(spec.args))
+    return_type = Base._return_type(f, Base.to_tuple_type(arg_types))
+    return EagerThunkMetadata(return_type)
+end
+chunktype(t::EagerThunk) = t.metadata.return_type
 function eager_spawn(spec::EagerTaskSpec)
     # Generate new EagerThunk
     uid = eager_next_id()
     future = ThunkFuture()
+    metadata = EagerThunkMetadata(spec)
     finalizer_ref = poolset(EagerThunkFinalizer(uid); device=MemPool.CPURAMDevice())
 
     # Return unlaunched EagerThunk
-    return EagerThunk(uid, future, finalizer_ref)
+    return EagerThunk(uid, future, metadata, finalizer_ref)
 end
 function eager_launch!((spec, task)::Pair{EagerTaskSpec,EagerThunk})
     # Lookup EagerThunk -> ThunkID