Repository: intel/dffml

Issue:
df: memory: async iterator operations
Functions which are async iterators (`async def` with `yield`) should be supported by `@op`, as well as `dffml/df/memory.py` needs to be able to run one of these kinds of operations. Everytime the operation `yield`s, whatever it `yield`'d should be added to the input network context.

I think we'll likely want to combine `run` and `run_dispatch` so that we only have the `run` function, it's looking redundant (at first glance) *whoever submits a PR please figure this out and consolidate if appropriate*

https://github.com/intel/dffml/blob/7dfec5050c0e2f205fdd58524a420b861dc7fd9e/dffml/df/memory.py#L861-L976

Discussion:
pdxjohnny: @aghinsa I added GetMulti (https://github.com/intel/dffml/pull/525) if you could use it to grab all the counts and confirm that they're all there then we can merge this. Thanks!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d6506363d..7c6fd5cef 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Python prompts `>>>` can now be enabled or disabled for easy copying of code into interactive sessions.
 - `GetMulti` operation which gets all Inputs of a given definition
+- Support for async generator operations
 ### Changed
 - Renamed `"arg"` to `"plugin"`.
 
@@ -21,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `validate` parameter in `Input` takes `Operation.instance_name`
 - New db source can utilize any database that inherits from `BaseDatabase`
 - Logistic Regression with SAG optimizer
+- Test tensorflow DNNEstimator documentation examples in CI
 - shouldi got an operation to run cargo-audit on rust code.
 - Moved all the downloads to tests/downloads to speed the CI test.
 - Test tensorflow DNNEstimator documentation exaples in CI
diff --git a/configloader/png/setup.py b/configloader/png/setup.py
index 061ada3f6..ca1e88b37 100644
--- a/configloader/png/setup.py
+++ b/configloader/png/setup.py
@@ -9,10 +9,7 @@ spec = importlib.util.spec_from_file_location(
 common = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(common)
 
-common.KWARGS["install_requires"] += [
-    "Pillow>=7.0.0",
-    "numpy>=1.16.2",
-]
+common.KWARGS["install_requires"] += ["Pillow>=7.0.0", "numpy>=1.16.2"]
 common.KWARGS["entry_points"] = {
     "dffml.configloader": [
         f"mnistpng = {common.IMPORT_NAME}.configloader:PNGConfigLoader"
diff --git a/dffml/df/memory.py b/dffml/df/memory.py
index 5b9b13d7e..74e1ab4e8 100644
--- a/dffml/df/memory.py
+++ b/dffml/df/memory.py
@@ -599,10 +599,7 @@ class MemoryOperationNetwork(BaseOperationNetwork, BaseMemoryDataFlowObject):
     @classmethod
     def args(cls, args, *above) -> Dict[str, Arg]:
         cls.config_set(
-            args,
-            above,
-            "ops",
-            Arg(type=Operation.load, nargs="+", default=[]),
+            args, above, "ops", Arg(type=Operation.load, nargs="+", default=[])
         )
         return args
 
@@ -973,45 +970,51 @@ class MemoryOperationImplementationNetworkContext(
                 await parameter_set._asdict(),
             )
             if outputs is None:
-                return []
-        # Create a list of inputs from the outputs using the definition mapping
-        try:
-            inputs = []
-            if operation.expand:
-                expand = operation.expand
-            else:
-                expand = []
-            parents = [
-                item.origin async for item in parameter_set.parameters()
-            ]
-            for key, output in outputs.items():
-                if not key in expand:
-                    output = [output]
-                for value in output:
-                    new_input = Input(
-                        value=value,
-                        definition=operation.outputs[key],
-                        parents=parents,
-                        origin=(operation.instance_name, key),
+                return
+            if not inspect.isasyncgen(outputs):
+
+                async def to_async_gen(x):
+                    yield x
+
+                outputs = to_async_gen(outputs)
+        async for an_output in outputs:
+            # Create a list of inputs from the outputs using the definition mapping
+            try:
+                inputs = []
+                if operation.expand:
+                    expand = operation.expand
+                else:
+                    expand = []
+                parents = [
+                    item.origin async for item in parameter_set.parameters()
+                ]
+                for key, output in an_output.items():
+                    if not key in expand:
+                        output = [output]
+                    for value in output:
+                        new_input = Input(
+                            value=value,
+                            definition=operation.outputs[key],
+                            parents=parents,
+                            origin=(operation.instance_name, key),
+                        )
+                        new_input.validated = set_valid
+                        inputs.append(new_input)
+            except KeyError as error:
+                raise KeyError(
+                    "Value %s missing from output:definition mapping %s(%s)"
+                    % (
+                        str(error),
+                        operation.instance_name,
+                        ", ".join(operation.outputs.keys()),
                     )
-                    new_input.validated = set_valid
-                    inputs.append(new_input)
-        except KeyError as error:
-            raise KeyError(
-                "Value %s missing from output:definition mapping %s(%s)"
-                % (
-                    str(error),
-                    operation.instance_name,
-                    ", ".join(operation.outputs.keys()),
+                ) from error
+            # Add the input set made from the outputs to the input set network
+            await octx.ictx.add(
+                MemoryInputSet(
+                    MemoryInputSetConfig(ctx=parameter_set.ctx, inputs=inputs)
                 )
-            ) from error
-        # Add the input set made from the outputs to the input set network
-        await octx.ictx.add(
-            MemoryInputSet(
-                MemoryInputSetConfig(ctx=parameter_set.ctx, inputs=inputs)
             )
-        )
-        return inputs
 
     async def dispatch(
         self,
@@ -1432,14 +1435,13 @@ class MemoryOrchestratorContext(BaseOrchestratorContext):
         tasks = set()
         # String representing the context we are executing operations for
         ctx_str = (await ctx.handle()).as_string()
+        # schedule running of operations with no inputs
+        async for task in self.nctx.dispatch_auto_starts(self, ctx):
+            tasks.add(task)
         # Create initial events to wait on
         # TODO(dfass) Make ictx.added(ctx) specific to dataflow
         input_set_enters_network = asyncio.create_task(self.ictx.added(ctx))
         tasks.add(input_set_enters_network)
-        # schedule running of operations with no inputs
-        async for task in self.nctx.dispatch_auto_starts(self, ctx):
-            tasks.add(task)
-
         try:
             # Return when outstanding operations reaches zero
             while tasks:
@@ -1492,9 +1494,7 @@ class MemoryOrchestratorContext(BaseOrchestratorContext):
                                 self.config.dataflow,
                                 unvalidated_input_set,
                             ):
-                                await self.rctx.add(
-                                    operation, parameter_set
-                                )  # is this required here?
+                                await self.rctx.add(operation, parameter_set)
                                 dispatch_operation = await self.nctx.dispatch(
                                     self, operation, parameter_set
                                 )
diff --git a/dffml/operation/io.py b/dffml/operation/io.py
index ec5c55e2a..8ab34d40c 100644
--- a/dffml/operation/io.py
+++ b/dffml/operation/io.py
@@ -91,9 +91,7 @@ class AcceptUserInput(OperationImplementation):
         self.loop = None
 
 
-@op(
-    inputs={"data": DataToPrint}, outputs={}, conditions=[],
-)
+@op(inputs={"data": DataToPrint}, outputs={}, conditions=[])
 async def print_output(data: str):
     """
     Print the output on stdout using python print()
diff --git a/model/scratch/setup.py b/model/scratch/setup.py
index 5545f3b54..d749b58b4 100644
--- a/model/scratch/setup.py
+++ b/model/scratch/setup.py
@@ -67,7 +67,7 @@ setup(
     packages=find_packages(),
     entry_points={
         "dffml.model": [
-            "scratchlgrsag = dffml_model_scratch.logisticregression:LogisticRegression",
+            "scratchlgrsag = dffml_model_scratch.logisticregression:LogisticRegression"
         ]
     },
 )
diff --git a/model/transformers/dffml_model_transformers/ner/ner_model.py b/model/transformers/dffml_model_transformers/ner/ner_model.py
index 97e98cf57..0bcd06012 100644
--- a/model/transformers/dffml_model_transformers/ner/ner_model.py
+++ b/model/transformers/dffml_model_transformers/ner/ner_model.py
@@ -84,9 +84,7 @@ class NERModelConfig:
     output_dir: str = field(
         "The output directory where the model checkpoints will be written",
         default=str(
-            pathlib.Path(
-                "~", ".cache", "dffml", "transformers", "checkpoints",
-            )
+            pathlib.Path("~", ".cache", "dffml", "transformers", "checkpoints")
         ),
     )
     config_name: str = field(
@@ -222,7 +220,7 @@ class NERModelContext(ModelContext):
         x_cols: Dict[str, Any] = {
             feature: []
             for feature in (
-                [self.parent.config.sid.NAME, self.parent.config.words.NAME,]
+                [self.parent.config.sid.NAME, self.parent.config.words.NAME]
             )
         }
         y_cols = []
@@ -242,7 +240,7 @@ class NERModelContext(ModelContext):
                 all_records.append(record)
         for record in all_records:
             for feature, results in record.features(
-                [self.parent.config.sid.NAME, self.parent.config.words.NAME,]
+                [self.parent.config.sid.NAME, self.parent.config.words.NAME]
             ).items():
                 x_cols[feature].append(np.array(results))
             y_cols.append(record.feature(self.parent.config.predict.NAME))
diff --git a/model/transformers/examples/ner/ner_model.py b/model/transformers/examples/ner/ner_model.py
index 8f9a775a3..cc32cadd1 100644
--- a/model/transformers/examples/ner/ner_model.py
+++ b/model/transformers/examples/ner/ner_model.py
@@ -21,8 +21,8 @@ print("Accuracy:", accuracy(model, CSVSource(filename="train.csv")))
 # Make prediction
 for i, features, prediction in predict(
     model,
-    {"SentenceID": 1, "Words": "DFFML models can do NER",},
-    {"SentenceID": 2, "Words": "DFFML models can do regression",},
+    {"SentenceID": 1, "Words": "DFFML models can do NER"},
+    {"SentenceID": 2, "Words": "DFFML models can do regression"},
 ):
     features["Tag"] = prediction["Tag"]["value"]
     print(features)


Test patch:
diff --git a/tests/df/test_async_iter.py b/tests/df/test_async_iter.py
new file mode 100644
index 000000000..f22e1ef28
--- /dev/null
+++ b/tests/df/test_async_iter.py
@@ -0,0 +1,74 @@
+from dffml.df.types import DataFlow, Input, Definition
+from dffml.df.base import op
+from dffml.df.memory import MemoryOrchestrator
+from dffml.operation.output import GetMulti
+from dffml.util.asynctestcase import AsyncTestCase
+from dffml.operation.mapping import MAPPING
+
+
+CountStart = Definition(name="count_start", primitive="int")
+Count = Definition(name="count", primitive="int")
+Number = Definition(name="number", primitive="int")
+
+
+@op(inputs={"count_start": CountStart}, outputs={"count": Count})
+async def counter(count_start):
+    for i in range(count_start, count_start + 5):
+        yield {"count": i}
+
+
+@op(inputs={}, outputs={"count": Count})
+async def counter_auto_start():
+    for i in range(5):
+        yield {"count": i}
+
+
+@op(inputs={"number_in": Count}, outputs={"number_out": Number})
+def echo_num(number_in: int):
+    return {"number_out": number_in}
+
+
+class TestAsyncIter(AsyncTestCase):
+    async def test_gen_with_input(self):
+        test_dataflow = DataFlow.auto(GetMulti, counter, echo_num)
+        test_dataflow.seed.append(
+            Input(
+                value=[echo_num.op.outputs["number_out"].name],
+                definition=GetMulti.op.inputs["spec"],
+            )
+        )
+        test_dataflow.implementations[counter.op.name] = counter.imp
+        test_dataflow.implementations[echo_num.op.name] = echo_num.imp
+
+        test_inputs = {"TestCount": [Input(value=1, definition=CountStart)]}
+        async with MemoryOrchestrator.withconfig({}) as orchestrator:
+            async with orchestrator(test_dataflow) as octx:
+                async for ctx_str, results in octx.run(test_inputs):
+                    self.assertIn("number", results)
+                    self.assertEqual(
+                        set([1, 2, 3, 4, 5]), set(results["number"])
+                    )
+
+    async def test_gen_auto_start(self):
+        test_dataflow = DataFlow.auto(GetMulti, counter, echo_num)
+        test_dataflow.seed.append(
+            Input(
+                value=[echo_num.op.outputs["number_out"].name],
+                definition=GetMulti.op.inputs["spec"],
+            )
+        )
+        test_dataflow.implementations[
+            counter_auto_start.op.name
+        ] = counter_auto_start.imp
+        test_dataflow.implementations[echo_num.op.name] = echo_num.imp
+
+        test_inputs = {
+            "TestCountAutoStart": [Input(value=1, definition=CountStart)]
+        }
+        async with MemoryOrchestrator.withconfig({}) as orchestrator:
+            async with orchestrator(test_dataflow) as octx:
+                async for ctx_str, results in octx.run(test_inputs):
+                    self.assertIn("number", results)
+                    self.assertEqual(
+                        set([1, 2, 3, 4, 5]), set(results["number"])
+                    )
diff --git a/tests/df/test_auto_start.py b/tests/df/test_auto_start.py
index cf132a89f..e4ffea8e7 100644
--- a/tests/df/test_auto_start.py
+++ b/tests/df/test_auto_start.py
@@ -25,7 +25,7 @@ class TestAutoStart(AsyncTestCase):
                     definition=GetSingle.op.inputs["spec"],
                 )
             ],
-            implementations={announce.op.name: announce.imp,},
+            implementations={announce.op.name: announce.imp},
         )
 
         self.dataflow = dataflow