generic compute plan

Author: jmorel

substratest/factory.py

@@ -142,6 +142,7 @@ class Permissions:
 
 
 DEFAULT_PERMISSIONS = Permissions(public=True, authorized_ids=[])
+DEFAULT_OUT_MODEL_PERMISSIONS = Permissions(public=False, authorized_ids=[])
 
 
 @dataclasses.dataclass
@@ -243,12 +244,46 @@ class TesttupleSpec(_Spec):
 
 @dataclasses.dataclass
 class ComputePlanTraintupleSpec:
+    algo_key: str
     data_manager_key: str
     train_data_sample_keys: str
     traintuple_id: str
     in_models_ids: typing.List[str]
     tag: str
 
+    @property
+    def id(self):
+        return self.traintuple_id
+
+
+@dataclasses.dataclass
+class ComputePlanAggregatetupleSpec(_Spec):
+    aggregatetuple_id: str
+    algo_key: str
+    worker: str
+    in_models_ids: typing.List[str]
+    tag: str
+
+    @property
+    def id(self):
+        return self.aggregatetuple_id
+
+
+@dataclasses.dataclass
+class ComputePlanCompositeTraintupleSpec(_Spec):
+    composite_traintuple_id: str
+    algo_key: str
+    data_manager_key: str
+    train_data_sample_keys: str
+    in_head_model_id: str
+    in_trunk_model_id: str
+    tag: str
+    out_trunk_model_permissions: typing.Dict
+
+    @property
+    def id(self):
+        return self.composite_traintuple_id
+
 
 @dataclasses.dataclass
 class ComputePlanTesttupleSpec:
@@ -276,26 +311,69 @@ def _get_keys(obj, field='key'):
 
 @dataclasses.dataclass
 class ComputePlanSpec(_Spec):
-    algo_key: str
     objective_key: str
     traintuples: typing.List[ComputePlanTraintupleSpec]
+    composite_traintuples: typing.List[ComputePlanCompositeTraintupleSpec]
+    aggregatetuples: typing.List[ComputePlanAggregatetupleSpec]
     testtuples: typing.List[ComputePlanTesttupleSpec]
 
-    def add_traintuple(self, dataset, data_samples, traintuple_specs=None, tag=None):
-        traintuple_specs = traintuple_specs or []
+    def add_traintuple(self, algo, dataset, data_samples, in_models_tuples=None, tag=''):
+        in_models_tuples = in_models_tuples or []
         spec = ComputePlanTraintupleSpec(
+            algo_key=algo.key,
             traintuple_id=random_uuid(),
             data_manager_key=dataset.key,
             train_data_sample_keys=_get_keys(data_samples),
-            in_models_ids=[t.traintuple_id for t in traintuple_specs],
-            tag=tag or '',
+            in_models_ids=[t.id for t in in_models_tuples],
+            tag=tag,
         )
         self.traintuples.append(spec)
         return spec
 
+    def add_aggregatetuple(self, aggregate_algo, worker, in_models_tuples=None, tag=''):
+        in_models_tuples = in_models_tuples or []
+
+        for t in in_models_tuples:
+            assert isinstance(t, (ComputePlanTraintupleSpec, ComputePlanCompositeTraintupleSpec))
+
+        spec = ComputePlanAggregatetupleSpec(
+            aggregatetuple_id=random_uuid(),
+            algo_key=aggregate_algo.key,
+            worker=worker,
+            in_models_ids=[t.id for t in in_models_tuples],
+            tag=tag,
+        )
+        self.aggregatetuples.append(spec)
+        return spec
+
+    def add_composite_traintuple(self, composite_algo, dataset=None, data_samples=None,
+                                 in_head_model_tuple=None, in_trunk_model_tuple=None,
+                                 out_trunk_model_permissions=None, tag=''):
+        data_samples = data_samples or []
+
+        if in_head_model_tuple and in_trunk_model_tuple:
+            assert isinstance(in_head_model_tuple, ComputePlanCompositeTraintupleSpec)
+            assert isinstance(
+                in_trunk_model_tuple,
+                (ComputePlanCompositeTraintupleSpec, ComputePlanAggregatetupleSpec)
+            )
+
+        spec = ComputePlanCompositeTraintupleSpec(
+            composite_traintuple_id=random_uuid(),
+            algo_key=composite_algo.key,
+            data_manager_key=dataset.key if dataset else None,
+            train_data_sample_keys=_get_keys(data_samples),
+            in_head_model_id=in_head_model_tuple.id if in_head_model_tuple else None,
+            in_trunk_model_id=in_trunk_model_tuple.id if in_trunk_model_tuple else None,
+            out_trunk_model_permissions=out_trunk_model_permissions or DEFAULT_OUT_MODEL_PERMISSIONS,
+            tag=tag,
+        )
+        self.composite_traintuples.append(spec)
+        return spec
+
     def add_testtuple(self, traintuple_spec, tag=None):
         spec = ComputePlanTesttupleSpec(
-            traintuple_id=traintuple_spec.traintuple_id,
+            traintuple_id=traintuple_spec.id,
             tag=tag or '',
         )
         self.testtuples.append(spec)
@@ -486,8 +564,6 @@ def create_composite_traintuple(self, algo=None, objective=None, dataset=None,
                                     permissions=None):
         data_samples = data_samples or []
 
-        kwargs = {}
-
         if head_traintuple and trunk_traintuple:
             assert isinstance(head_traintuple, assets.CompositeTraintuple)
             assert isinstance(
@@ -511,7 +587,6 @@ def create_composite_traintuple(self, algo=None, objective=None, dataset=None,
             compute_plan_id=compute_plan_id,
             rank=rank,
             out_trunk_model_permissions=permissions or DEFAULT_PERMISSIONS,
-            **kwargs,
         )
 
     def create_testtuple(self, traintuple=None, tag=None):
@@ -520,10 +595,11 @@ def create_testtuple(self, traintuple=None, tag=None):
             tag=tag,
         )
 
-    def create_compute_plan(self, algo=None, objective=None):
+    def create_compute_plan(self, objective=None):
         return ComputePlanSpec(
-            algo_key=algo.key if algo else None,
             objective_key=objective.key if objective else None,
             traintuples=[],
+            composite_traintuples=[],
+            aggregatetuples=[],
             testtuples=[],
         )

tests/test_execution_compute_plan.py

@@ -22,24 +22,27 @@ def test_compute_plan(global_execution_env):
     algo_2 = session_2.add_algo(spec)
 
     # create compute plan
-    cp_spec = factory.create_compute_plan(algo=algo_2, objective=objective_1)
+    cp_spec = factory.create_compute_plan(objective=objective_1)
 
     # TODO add a testtuple in the compute plan
 
     traintuple_spec_1 = cp_spec.add_traintuple(
+        algo=algo_2,
         dataset=dataset_1,
         data_samples=dataset_1.train_data_sample_keys,
     )
 
     traintuple_spec_2 = cp_spec.add_traintuple(
+        algo=algo_2,
         dataset=dataset_2,
         data_samples=dataset_2.train_data_sample_keys,
     )
 
-    _ = cp_spec.add_traintuple(
+    cp_spec.add_traintuple(
+        algo=algo_2,
         dataset=dataset_1,
         data_samples=dataset_1.train_data_sample_keys,
-        traintuple_specs=[traintuple_spec_1, traintuple_spec_2],
+        in_models_tuples=[traintuple_spec_1, traintuple_spec_2],
     )
 
     # submit compute plan and wait for it to complete
@@ -83,25 +86,28 @@ def test_compute_plan_single_session_success(global_execution_env):
     spec = factory.create_algo()
     algo = session.add_algo(spec)
 
-    cp_spec = factory.create_compute_plan(algo=algo, objective=objective)
+    cp_spec = factory.create_compute_plan(objective=objective)
 
     traintuple_spec_1 = cp_spec.add_traintuple(
+        algo=algo,
         dataset=dataset,
         data_samples=[data_sample_1]
     )
     cp_spec.add_testtuple(traintuple_spec_1)
 
     traintuple_spec_2 = cp_spec.add_traintuple(
+        algo=algo,
         dataset=dataset,
         data_samples=[data_sample_2],
-        traintuple_specs=[traintuple_spec_1]
+        in_models_tuples=[traintuple_spec_1]
     )
     cp_spec.add_testtuple(traintuple_spec_2)
 
     traintuple_spec_3 = cp_spec.add_traintuple(
+        algo=algo,
         dataset=dataset,
         data_samples=[data_sample_3],
-        traintuple_specs=[traintuple_spec_2]
+        in_models_tuples=[traintuple_spec_2]
     )
     cp_spec.add_testtuple(traintuple_spec_3)
 
@@ -150,25 +156,28 @@ def test_compute_plan_single_session_failure(global_execution_env):
     spec = factory.create_algo(py_script=sbt.factory.INVALID_ALGO_SCRIPT)
     algo = session.add_algo(spec)
 
-    cp_spec = factory.create_compute_plan(algo=algo, objective=objective)
+    cp_spec = factory.create_compute_plan(objective=objective)
 
     traintuple_spec_1 = cp_spec.add_traintuple(
+        algo=algo,
         dataset=dataset,
         data_samples=[data_sample_1]
     )
     cp_spec.add_testtuple(traintuple_spec_1)
 
     traintuple_spec_2 = cp_spec.add_traintuple(
+        algo=algo,
         dataset=dataset,
         data_samples=[data_sample_2],
-        traintuple_specs=[traintuple_spec_1]
+        in_models_tuples=[traintuple_spec_1]
     )
     cp_spec.add_testtuple(traintuple_spec_2)
 
     traintuple_spec_3 = cp_spec.add_traintuple(
+        algo=algo,
         dataset=dataset,
         data_samples=[data_sample_3],
-        traintuple_specs=[traintuple_spec_2]
+        in_models_tuples=[traintuple_spec_2]
     )
     cp_spec.add_testtuple(traintuple_spec_3)
 
@@ -193,3 +202,123 @@ def test_compute_plan_single_session_failure(global_execution_env):
     assert cp.compute_plan_id == compute_plan.compute_plan_id
     assert set(cp.traintuple_keys) == set(compute_plan.traintuples)
     assert set(cp.testtuple_keys) == set(compute_plan.testtuples)
+
+
+def test_compute_plan_aggregate_composite_traintuples(factory, session_1, session_2):
+    """
+    Compute plan version of the `test_aggregate_composite_traintuples` method from `test_execution.py`
+    """
+    aggregate_worker = session_1.node_id
+    sessions = [session_1, session_2]
+    number_of_rounds = 2
+
+    # register objectives, datasets, and data samples
+    datasets = []
+    for s in sessions:
+        # register one dataset per node
+        spec = factory.create_dataset()
+        dataset = s.add_dataset(spec)
+        datasets.append(dataset)
+
+        # register one data sample per dataset per round of aggregation
+        for _ in range(number_of_rounds):
+            spec = factory.create_data_sample(test_only=False, datasets=[dataset])
+            s.add_data_sample(spec)
+    # reload datasets (to ensure they are properly linked with the created data samples)
+    datasets = [
+        sessions[i].get_dataset(d.key)
+        for i, d in enumerate(list(datasets))
+    ]
+    # register test data on first node
+    spec = factory.create_data_sample(test_only=True, datasets=[datasets[0]])
+    test_data_sample = sessions[0].add_data_sample(spec)
+    # register objective on first node
+    spec = factory.create_objective(
+        dataset=datasets[0],
+        data_samples=[test_data_sample],
+    )
+    objective = sessions[0].add_objective(spec)
+
+    # register algos on first node
+    spec = factory.create_composite_algo()
+    composite_algo = sessions[0].add_composite_algo(spec)
+    spec = factory.create_aggregate_algo()
+    aggregate_algo = sessions[0].add_aggregate_algo(spec)
+
+    # launch execution
+    previous_aggregatetuple_spec = None
+    previous_composite_traintuple_specs = []
+
+    cp_spec = factory.create_compute_plan(objective=objective)
+
+    for round_ in range(number_of_rounds):
+        # create composite traintuple on each node
+        composite_traintuple_specs = []
+        for index, dataset in enumerate(datasets):
+            kwargs = {}
+            if previous_aggregatetuple_spec:
+                kwargs = {
+                    'in_head_model_tuple': previous_composite_traintuple_specs[index],
+                    'in_trunk_model_tuple': previous_aggregatetuple_spec,
+                }
+            spec = cp_spec.add_composite_traintuple(
+                composite_algo=composite_algo,
+                dataset=dataset,
+                data_samples=[dataset.train_data_sample_keys[0 + round_]],
+                **kwargs,
+            )
+            composite_traintuple_specs.append(spec)
+
+        # create aggregate on its node
+        spec = cp_spec.add_aggregatetuple(
+            aggregate_algo=aggregate_algo,
+            worker=aggregate_worker,
+            in_models_tuples=composite_traintuple_specs,
+        )
+
+        # save state of round
+        previous_aggregatetuple_spec = spec
+        previous_composite_traintuple_specs = composite_traintuple_specs
+
+    # last round: create associated testtuple
+    for composite_traintuple_spec in previous_composite_traintuple_specs:
+        cp_spec.add_testtuple(
+            traintuple_spec=composite_traintuple_spec,
+        )
+
+    session_1.add_compute_plan(cp_spec).future().wait()
+
+
+def test_compute_plan_circular_dependency_failure(factory, session):
+    spec = factory.create_dataset()
+    dataset = session.add_dataset(spec)
+
+    spec = factory.create_algo()
+    algo = session.add_algo(spec)
+
+    spec = factory.create_data_sample(test_only=False, datasets=[dataset])
+    data_sample = session.add_data_sample(spec)
+
+    spec = factory.create_objective(dataset=dataset)
+    objective = session.add_objective(spec)
+
+    cp_spec = factory.create_compute_plan(objective=objective)
+
+    traintuple_spec_1 = cp_spec.add_traintuple(
+        dataset=dataset,
+        algo=algo,
+        data_samples=[data_sample]
+    )
+
+    traintuple_spec_2 = cp_spec.add_traintuple(
+        dataset=dataset,
+        algo=algo,
+        data_samples=[data_sample]
+    )
+
+    traintuple_spec_1.in_models_ids.append(traintuple_spec_2.id)
+    traintuple_spec_2.in_models_ids.append(traintuple_spec_1.id)
+
+    # TODO make sur the creation is rejected
+    cp = session.add_compute_plan(cp_spec)
+    assert False