compute plan future

Author: jmorel

substratest/assets.py

@@ -50,6 +50,50 @@ def get(self):
         return self._asset
 
 
+class ComputePlanFuture(Future):
+    _keys_properties = {
+        'ComputePlan': {
+            'traintuple_keys': 'traintuples',
+            'composite_traintuple_keys': 'composite_traintuples',
+            'aggregatetuple_keys': 'aggregatetuples',
+            'testtuple_keys': 'testtuples',
+        },
+        'ComputePlanCreated': {
+            'traintuple_keys': 'traintuple_keys',
+            'composite_traintuple_keys': 'composite_traintuple_keys',
+            'aggregatetuple_keys': 'aggregatetuple_keys',
+            'testtuple_keys': 'testtuple_keys',
+
+        },
+    }
+
+    def __init__(self, asset, session):
+        self._asset = asset
+        self._getter = session.get_compute_plan
+        for k, v in enumerate(self._keys_properties[asset.__class__.name]):
+            setattr(self, f'_{k}', getattr(asset, v))
+        self._get_traintuple = session.get_traintuple
+        self._get_composite_traintuple = session.get_composite_traintuple
+        self._get_aggregatetuple = session.get_aggregatetuple
+        self._get_testtuple = session.get_testtuple
+
+    def wait(self, timeout=FUTURE_TIMEOUT, raises=True):
+        """wait until all tuples are completed (done or failed)."""
+        for key in self._traintuple_keys:
+            self._get_traintuple(key).future().wait(timeout, raises)
+        for key in self._composite_traintuple_keys:
+            self._get_composite_traintuple(key).future().wait(timeout, raises)
+        for key in self._aggregatetuple_keys:
+            self._get_aggregatetuple(key).future().wait(timeout, raises)
+        for key in self._testtuple_keys:
+            self._get_testtuple(key).future().wait(timeout, raises)
+
+        return self.get()
+
+    def get(self):
+        return self._getter(self._asset.key)
+
+
 class _FutureMixin(abc.ABC):
     def attach(self, session):
         """Attach session to asset."""
@@ -60,7 +104,11 @@ def future(self):
         """Returns future from asset."""
         assert hasattr(self, 'status')
         assert hasattr(self, 'key')
-        return Future(self, self._session)
+
+        try:
+            return self.Meta.FutureCls(self, self._session)
+        except AttributeError:
+            return Future(self, self._session)
 
 
 def _convert(name):
@@ -321,24 +369,48 @@ class Meta:
 
 
 @dataclasses.dataclass
-class ComputePlanCreated(_Asset):
+class ComputePlanCreated(_Asset, _FutureMixin):
     compute_plan_id: str
     traintuple_keys: typing.List[str]
+    composite_traintuple_keys: typing.List[str]
+    aggregatetuple_keys: typing.List[str]
     testtuple_keys: typing.List[str]
 
+    class Meta:
+        FutureCls = ComputePlanFuture
+
 
 @dataclasses.dataclass
 class ComputePlan(_Asset):
     compute_plan_id: str
     algo_key: str
     objective_key: str
     traintuples: typing.List[str]
+    composite_traintuples: typing.List[str]
+    aggregatetuples: typing.List[str]
     testtuples: typing.List[str]
 
+    class Meta:
+        FutureCls = ComputePlanFuture
+
     def __post_init__(self):
         if self.testtuples is None:
             self.testtuples = []
 
+    def list_traintuples(self, session):
+        return session.list_traintuples(filters=[f'traintuple:computePlanId:{self.compute_plan_id}'])
+
+    def list_composite_traintuples(self, session):
+        return session.list_composite_traintuples(
+            filters=[f'composite_traintuple:computePlanId:{self.compute_plan_id}']
+        )
+
+    def list_aggregatetuples(self, session):
+        return session.list_aggregatetuples(filters=[f'aggregatetuple:computePlanId:{self.compute_plan_id}'])
+
+    def list_testtuples(self, session):
+        return session.list_testtuples(filters=[f'testtuple:computePlanId:{self.compute_plan_id}'])
+
 
 @dataclasses.dataclass(frozen=True)
 class Node(_Asset):

substratest/client.py

@@ -117,7 +117,7 @@ def add_testtuple(self, spec):
 
     def add_compute_plan(self, spec):
         res = self._client.add_compute_plan(spec.to_dict())
-        compute_plan = assets.ComputePlanCreated.load(res)
+        compute_plan = assets.ComputePlanCreated.load(res).attach(self)
         return compute_plan
 
     def list_compute_plan(self, *args, **kwargs):

tests/test_execution_compute_plan.py

@@ -182,26 +182,19 @@ def test_compute_plan_single_session_failure(global_execution_env):
     cp_spec.add_testtuple(traintuple_spec_3)
 
     # Submit compute plan and wait for it to complete
-    cp = session.add_compute_plan(cp_spec)
+    cp_created = session.add_compute_plan(cp_spec)
+    cp = cp_created.future().wait()
 
-    traintuples = [
-        session.get_traintuple(key).future().wait(raises=False)
-        for key in cp.traintuple_keys
-    ]
-
-    testtuples = [
-        session.get_testtuple(key).future().wait(raises=False)
-        for key in cp.testtuple_keys
-    ]
+    traintuples = cp.list_traintuples(session)
+    testtuples = cp.list_testtuples(session)
 
     # All the train/test tuples should be marked as failed
     for t in traintuples + testtuples:
         assert t.status == 'failed'
 
-    compute_plan = session.get_compute_plan(cp.compute_plan_id)
-    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)
+    assert cp_created.compute_plan_id == cp.compute_plan_id
+    assert set(cp_created.traintuple_keys) == set(cp.traintuples)
+    assert set(cp_created.testtuple_keys) == set(cp.testtuples)
 
 
 def test_compute_plan_aggregate_composite_traintuples(factory, session_1, session_2):
@@ -286,7 +279,13 @@ def test_compute_plan_aggregate_composite_traintuples(factory, session_1, sessio
             traintuple_spec=composite_traintuple_spec,
         )
 
-    session_1.add_compute_plan(cp_spec).future().wait()
+    cp = session_1.add_compute_plan(cp_spec).future().wait()
+    tuples = (cp.list_traintuple(session_1) +
+              cp.list_composite_traintuples(session_1) +
+              cp.list_aggregate_tuples(session_1) +
+              cp.list_testtuples(session_1))
+    for t in tuples:
+        assert t.status == 'done'
 
 
 def test_compute_plan_circular_dependency_failure(factory, session):