Add metadata to EagerThunk

Author: jpsamaroo

Project.toml

@@ -11,6 +11,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"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"

src/dtask.jl

@@ -39,6 +39,16 @@ end
 Options(;options...) = Options((;options...))
 Options(options...) = Options((;options...))
 
+"""
+    DTaskMetadata
+
+Represents some useful metadata pertaining to a `DTask`:
+- `return_type::Type` - The inferred return type of the task
+"""
+mutable struct DTaskMetadata
+    return_type::Type
+end
+
 """
     DTask
 
@@ -50,9 +60,11 @@ more details.
 mutable struct DTask
     uid::UInt
     future::ThunkFuture
+    metadata::DTaskMetadata
     finalizer_ref::DRef
     thunk_ref::DRef
-    DTask(uid, future, finalizer_ref) = new(uid, future, finalizer_ref)
+
+    DTask(uid, future, metadata, finalizer_ref) = new(uid, future, metadata, finalizer_ref)
 end
 
 const EagerThunk = DTask

src/submission.jl

@@ -218,15 +218,27 @@ function eager_process_options_submission_to_local(id_map, options::NamedTuple)
         return options
     end
 end
+
+function DTaskMetadata(spec::DTaskSpec)
+    f = chunktype(spec.f).instance
+    arg_types = ntuple(i->chunktype(spec.args[i][2]), length(spec.args))
+    return_type = Base.promote_op(f, arg_types...)
+    return DTaskMetadata(return_type)
+end
+
 function eager_spawn(spec::DTaskSpec)
     # Generate new DTask
     uid = eager_next_id()
     future = ThunkFuture()
+    metadata = DTaskMetadata(spec)
     finalizer_ref = poolset(DTaskFinalizer(uid); device=MemPool.CPURAMDevice())
 
     # Create unlaunched DTask
-    return DTask(uid, future, finalizer_ref)
+    return DTask(uid, future, metadata, finalizer_ref)
 end
+
+chunktype(t::DTask) = t.metadata.return_type
+
 function eager_launch!((spec, task)::Pair{DTaskSpec,DTask})
     # Assign a name, if specified
     eager_assign_name!(spec, task)