tf-encrypted / moose

Secure distributed dataflow framework for encrypted machine learning and data processing
Apache License 2.0
56 stars 15 forks source link

Imputing new host placements in an AbstractComputation #1112

Closed jvmncs closed 2 years ago

jvmncs commented 2 years ago

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.])