tests: fix doctest in serve (#1584)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/flash/core/serve/dag/optimization.py

@@ -4,11 +4,11 @@
 
 from flash.core.serve.dag.task import flatten, get, get_dependencies, ishashable, istask, reverse_dict, subs, toposort
 from flash.core.serve.dag.utils import key_split
+from flash.core.utilities.imports import _TOPIC_SERVE_AVAILABLE
 
 # Skip doctests if requirements aren't available
-# if not _TOPIC_SERVE_AVAILABLE:
-# FixMe: all these test need to be updated
-__doctest_skip__ = ["*"]
+if not _TOPIC_SERVE_AVAILABLE:
+    __doctest_skip__ = ["*"]
 
 
 def cull(dsk, keys):
@@ -22,10 +22,10 @@ def cull(dsk, keys):
     >>> from flash.core.serve.dag.utils_test import add, inc
     >>> d = {'x': 1, 'y': (inc, 'x'), 'out': (add, 'x', 10)}
     >>> dsk, dependencies = cull(d, 'out')
-    >>> dsk
-    {'x': 1, 'out': (add, 'x', 10)}
+    >>> dsk  # doctest: +ELLIPSIS
+    {'out': (<function add at ...>, 'x', 10), 'x': 1}
     >>> dependencies
-    {'x': set(), 'out': set(['x'])}
+    {'out': ['x'], 'x': []}
 
     Returns
     -------
@@ -99,14 +99,14 @@ def fuse_linear(dsk, keys=None, dependencies=None, rename_keys=True):
     >>> from flash.core.serve.dag.utils_test import inc
     >>> d = {'a': 1, 'b': (inc, 'a'), 'c': (inc, 'b')}
     >>> dsk, dependencies = fuse(d)
-    >>> dsk
-    {'a-b-c': (inc, (inc, 1)), 'c': 'a-b-c'}
+    >>> dsk  # doctest: +ELLIPSIS
+    {'c': 'a-b-c', 'a-b-c': (<function inc at ...>, (<function inc at ...>, 1))}
     >>> dsk, dependencies = fuse(d, rename_keys=False)
-    >>> dsk
-    {'c': (inc, (inc, 1))}
+    >>> dsk  # doctest: +ELLIPSIS
+    {'c': (<function inc at ...>, (<function inc at ...>, 1))}
     >>> dsk, dependencies = fuse(d, keys=['b'], rename_keys=False)
-    >>> dsk
-    {'b': (inc, 1), 'c': (inc, 'b')}
+    >>> dsk  # doctest: +ELLIPSIS
+    {'b': (<function inc at ...>, 1), 'c': (<function inc at ...>, 'b')}
 
     Returns
     -------
@@ -230,12 +230,12 @@ def inline(dsk, keys=None, inline_constants=True, dependencies=None):
     --------
     >>> from flash.core.serve.dag.utils_test import add, inc
     >>> d = {'x': 1, 'y': (inc, 'x'), 'z': (add, 'x', 'y')}
-    >>> inline(d)
-    {'x': 1, 'y': (inc, 1), 'z': (add, 1, 'y')}
-    >>> inline(d, keys='y')
-    {'x': 1, 'y': (inc, 1), 'z': (add, 1, (inc, 1))}
-    >>> inline(d, keys='y', inline_constants=False)
-    {'x': 1, 'y': (inc, 1), 'z': (add, 'x', (inc, 'x'))}
+    >>> inline(d)  # doctest: +ELLIPSIS
+    {'x': 1, 'y': (<function inc at ...>, 1), 'z': (<function add at ...>, 1, 'y')}
+    >>> inline(d, keys='y')  # doctest: +ELLIPSIS
+    {'x': 1, 'y': (<function inc at ...>, 1), 'z': (<function add at ...>, 1, (<function inc at ...>, 1))}
+    >>> inline(d, keys='y', inline_constants=False)  # doctest: +ELLIPSIS
+    {'x': 1, 'y': (<function inc at ...>, 'x'), 'z': (<function add at ...>, 'x', (<function inc at ...>, 'x'))}
     """
     if dependencies and isinstance(next(iter(dependencies.values())), list):
         dependencies = {k: set(v) for k, v in dependencies.items()}
@@ -282,18 +282,20 @@ def inline_functions(dsk, output, fast_functions=None, inline_constants=False, d
     ...        'i': (inc, 'x'),
     ...        'd': (double, 'y'),
     ...        'x': 1, 'y': 1}
-    >>> inline_functions(dsk, [], [inc])
-    {'out': (add, (inc, 'x'), 'd'),
-     'd': (double, 'y'),
-     'x': 1, 'y': 1}
+    >>> inline_functions(dsk, [], [inc])  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
+    {'x': 1,
+     'out': (<function add at ...>, (<function inc at ...>, 'x'), 'd'),
+     'd': (<function <lambda> at ...>, 'y'),
+     'y': 1}
 
     Protect output keys.  In the example below ``i`` is not inlined because it
     is marked as an output key.
 
-    >>> inline_functions(dsk, ['i', 'out'], [inc, double])
-    {'out': (add, 'i', (double, 'y')),
-     'i': (inc, 'x'),
-     'x': 1, 'y': 1}
+    >>> inline_functions(dsk, ['i', 'out'], [inc, double])  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
+    {'y': 1,
+     'out': (<function add at ...>, 'i', (<function <lambda> at ...>, 'y')),
+     'i': (<function inc at ...>, 'x'),
+     'x': 1}
     """
     if not fast_functions:
         return dsk
@@ -334,8 +336,8 @@ def functions_of(task):
     --------
     >>> from flash.core.serve.dag.utils_test import add, inc, mul
     >>> task = (add, (mul, 1, 2), (inc, 3))
-    >>> functions_of(task)
-    set([add, mul, inc])
+    >>> sorted(functions_of(task), key=str)  # doctest: +ELLIPSIS
+    [<function add at ...>, <function inc at ...>, <function mul at ...>]
     """
     funcs = set()
 

src/flash/core/serve/dag/order.py

@@ -80,11 +80,11 @@
 from math import log
 
 from flash.core.serve.dag.task import get_dependencies, getcycle, reverse_dict
+from flash.core.utilities.imports import _TOPIC_SERVE_AVAILABLE
 
 # Skip doctests if requirements aren't available
-# if not _TOPIC_SERVE_AVAILABLE:
-# FixMe: all these test need to be updated
-__doctest_skip__ = ["*"]
+if not _TOPIC_SERVE_AVAILABLE:
+    __doctest_skip__ = ["*"]
 
 
 def order(dsk, dependencies=None):
@@ -541,6 +541,7 @@ def graph_metrics(dependencies, dependents, total_dependencies):
 
     Examples
     --------
+    >>> from flash.core.serve.dag.task import get_deps
     >>> from flash.core.serve.dag.utils_test import add, inc
     >>> dsk = {'a1': 1, 'b1': (inc, 'a1'), 'b2': (inc, 'a1'), 'c1': (inc, 'b1')}
     >>> dependencies, dependents = get_deps(dsk)
@@ -616,6 +617,7 @@ def ndependencies(dependencies, dependents):
 
     Examples
     --------
+    >>> from flash.core.serve.dag.task import get_deps
     >>> from flash.core.serve.dag.utils_test import add, inc
     >>> dsk = {'a': 1, 'b': (inc, 'a'), 'c': (inc, 'b')}
     >>> dependencies, dependents = get_deps(dsk)

src/flash/core/serve/dag/rewrite.py

@@ -1,11 +1,11 @@
 from collections import deque
 
 from flash.core.serve.dag.task import istask, subs
+from flash.core.utilities.imports import _TOPIC_SERVE_AVAILABLE
 
 # Skip doctests if requirements aren't available
-# if not _TOPIC_SERVE_AVAILABLE:
-# FixMe: all these test need to be updated
-__doctest_skip__ = ["*"]
+if not _TOPIC_SERVE_AVAILABLE:
+    __doctest_skip__ = ["*"]
 
 
 def head(task):
@@ -224,8 +224,8 @@ class RuleSet:
     >>> rs.rewrite((add, 2, 0))       # Apply ruleset to single task
     2
 
-    >>> rs.rewrite((f, (g, 'a', 3)))
-    (h, 'a', 3)
+    >>> rs.rewrite((f, (g, 'a', 3)))  # doctest: +ELLIPSIS
+    (<function h at ...>, 'a', 3)
 
     >>> dsk = {'a': (add, 2, 0),      # Apply ruleset to full dask graph
     ...        'b': (f, (g, 'a', 3))}
@@ -328,8 +328,10 @@ def rewrite(self, task, strategy="bottom_up"):
         Suppose there was a function `add` that returned the sum of 2 numbers,
         and another function `double` that returned twice its input:
 
-        >>> add = lambda x, y: x + y
-        >>> double = lambda x: 2*x
+        >>> def add(x, y):
+        ...     return x + y
+        >>> def double(x):
+        ...     return 2*x
 
         Now suppose `double` was *significantly* faster than `add`, so
         you'd like to replace all expressions `(add, x, x)` with `(double,
@@ -341,14 +343,14 @@ def rewrite(self, task, strategy="bottom_up"):
         This can then be applied to terms to perform the rewriting:
 
         >>> term = (add, (add, 2, 2), (add, 2, 2))
-        >>> rs.rewrite(term)
-        (double, (double, 2))
+        >>> rs.rewrite(term)  # doctest: +ELLIPSIS
+        (<function double at ...>, (<function double at ...>, 2))
 
         If we only wanted to apply this to the top level of the term, the
         `strategy` kwarg can be set to "top_level".
 
-        >>> rs.rewrite(term)
-        (double, (add, 2, 2))
+        >>> rs.rewrite(term)  # doctest: +ELLIPSIS
+        (<function double at ...>, (<function double at ...>, 2))
         """
         return strategies[strategy](self, task)
 

src/flash/core/serve/dag/task.py

@@ -1,10 +1,11 @@
 from collections import defaultdict
 from typing import List, Sequence
 
+from flash.core.utilities.imports import _TOPIC_SERVE_AVAILABLE
+
 # Skip doctests if requirements aren't available
-# if not _TOPIC_SERVE_AVAILABLE:
-# FixMe: all these test need to be updated
-__doctest_skip__ = ["*"]
+if not _TOPIC_SERVE_AVAILABLE:
+    __doctest_skip__ = ["*"]
 
 no_default = "__no_default__"
 
@@ -152,8 +153,8 @@ def get_dependencies(dsk, key=None, task=no_default, as_list=False):
     set()
     >>> get_dependencies(dsk, 'y')
     {'x'}
-    >>> get_dependencies(dsk, 'z')
-    {'x', 'y'}
+    >>> sorted(get_dependencies(dsk, 'z'))
+    ['x', 'y']
     >>> get_dependencies(dsk, 'w')  # Only direct dependencies
     {'z'}
     >>> get_dependencies(dsk, 'a')  # Ignore non-keys
@@ -238,7 +239,7 @@ def reverse_dict(d):
     >>> a, b, c = 'abc'
     >>> d = {a: [b, c], b: [c]}
     >>> reverse_dict(d)
-    {'a': set([]), 'b': set(['a']}, 'c': set(['a', 'b'])}
+    defaultdict(None, {'a': set(), 'b': {'a'}, 'c': {'b', 'a'}})
     """
     result = defaultdict(set)
     _add = set.add
@@ -256,8 +257,8 @@ def subs(task, key, val):
     Examples
     --------
     >>> from flash.core.serve.dag.utils_test import inc
-    >>> subs((inc, 'x'), 'x', 1)
-    (inc, 1)
+    >>> subs((inc, 'x'), 'x', 1)  # doctest: +ELLIPSIS
+    (<function inc at ...>, 1)
     """
     type_task = type(task)
     if not (type_task is tuple and task and callable(task[0])):  # istask(task):
@@ -373,6 +374,7 @@ def getcycle(d, keys):
 
     Examples
     --------
+    >>> from flash.core.serve.dag.utils_test import inc
     >>> d = {'x': (inc, 'z'), 'y': (inc, 'x'), 'z': (inc, 'y')}
     >>> getcycle(d, 'x')
     ['x', 'z', 'y', 'x']
@@ -426,6 +428,7 @@ def quote(x):
 
     Examples
     --------
+    >>> from flash.core.serve.dag.utils_test import add
     >>> quote((add, 1, 2))
     (literal<type=tuple>,)
     """