Imputing new host placements in an AbstractComputation
jvmncs opened this issue · comments
It's common for us to define some common computations with standardized, placeholder host placements, e.g. alice
/bob
/carole
, like so:
alice = pm.host_placement("alice")
bob = pm.host_placement("bob")
carole = pm.host_placement("carole")
@pm.computation
def my_comp():
with alice:
x = pm.constant(1, dtype=pm.float64)
with bob:
y = pm.constant(2, dtype=pm.float64)
with carole:
z = pm.add(x, y)
res = pm.save("sum", z)
return res
The user of this function should not be pigeonholed into using a runtime with host identities alice
, bob
, and carole
. Previously the runtimes had a role_assignment
parameter that allowed for more flexibility in this, mapping "roles" in the computation to "identities" found in the runtime. But since this role assignment led to a clunky API for the Runtime abstraction, we removed it in #1099. We now expect that each host placement in a Computation matches at least one identity in the Runtime by name.
Thus, we should make it possible to swap out host placements after a user has already defined their AbstractComputation, e.g.
runtime = pm.LocalMooseRuntime(["hospital-1", "hospital-2", "data-scientist"])
runtime.set_default()
my_comp.with_role_map({"alice": "hospital-1", "bob": "hospital-2", "carole": "data-scientist"})()
z = runtime.read_value_from_storage("data-scientist", "sum")
print(z)
# array([3.])