update docstrings

dask4dvc/cli/__init__.py

@@ -93,9 +93,6 @@ def run(
 
     Run 'dvc exp run --run-all' with full parallel execution using Dask distributed.
     """
-    # TODO rename to exp run
-    cleanup = cleanup
-
     if name is None:
         tmp_dirs = dask4dvc.dvc_handling.load_all_exp_to_tmp_dir()
 
@@ -111,7 +108,7 @@ def run(
                 for tmp_dir in tmp_dirs.values():
                     shutil.rmtree(tmp_dir)
                 answer = input(
-                    "Press Enter to load experiments and close dask client (yes/no)"
+                    "Press Enter to load experiments and close dask client (yes/no) "
                 )
 
         if answer == "yes":
@@ -133,7 +130,7 @@ def run(
 
 
 def version_callback(value: bool):
-    """Get the installed zntrack version"""
+    """Get the installed dask4dvc version"""
     if value:
         typer.echo(f"dask4dvc {dask4dvc.__version__}")
         raise typer.Exit()
@@ -145,7 +142,8 @@ def main(
         None, "--version", callback=version_callback, is_eager=True
     ),
 ):
-    """
-    <update me>
+    """Dask4DVC
+
+    Run the DVC graph or DVC experiments in parallel using dask.
     """
     _ = version  # this would be greyed out otherwise

dask4dvc/dvc_handling.py

@@ -12,7 +12,10 @@
 
 
 def get_dvc_graph(cwd=None) -> nx.DiGraph:
-    """Use the dvc dag command to get the graph into networkx format"""
+    """Use the dvc dag command to get the graph into networkx format
+
+    Running 'dvc dag' can be slow on bigger graphs.
+    """
     dot_string = subprocess.run(
         ["dvc", "dag", "--dot"], capture_output=True, check=True, cwd=cwd
     )
@@ -31,6 +34,9 @@ def run_dvc_repro_in_cwd(node_name: str, cwd=None, deps=None) -> None:
     cwd: str
         working directory
 
+    deps: list[dask.distributed.Future]|Future
+        Any dependencies for the dask graph
+
     Raises
     -------
     subprocess.CalledProcessError: if dvc cmd fails
@@ -96,7 +102,9 @@ def prepare_dvc_workspace(
     return tmp_dir
 
 
-def submit_dvc_stage(name, deps=None, cwd: pathlib.Path = None, cleanup: bool = True):
+def submit_dvc_stage(
+    name: str, deps=None, cwd: pathlib.Path = None, cleanup: bool = True
+):
     """Run a DVC stage
 
     1. prepare a temporary workspace
@@ -107,7 +115,8 @@ def submit_dvc_stage(name, deps=None, cwd: pathlib.Path = None, cleanup: bool =
     ----------
     name: str
         Name of the Node
-    deps: any dependencies for Dask to build the graph
+    deps: list[dask.distributed.Future]|Future
+        any dependencies for Dask to build the graph
     cwd: pathlib.Path
         The working directory for the repro command.
         Will be None for 'dvc repro' and set to a custom directory for e.g. 'dvc exp run'
@@ -130,12 +139,11 @@ def submit_dvc_stage(name, deps=None, cwd: pathlib.Path = None, cleanup: bool =
 
 
 def load_all_exp_to_tmp_dir() -> typing.Dict[str, pathlib.Path]:
-    """Load all queued expments into temporary directories each.
+    """Load all queued experiments into temporary directories each.
 
     Returns
     -------
     A dictionary of the created directories with the experiment names as keys.
-
     """
     queued_exp = get_queued_exp_names()
     tmp_dirs = {}

dask4dvc/utils/graph.py

@@ -59,8 +59,11 @@ def iterate_over_nodes(
 
 
 def submit_to_dask(
-    client: dask.distributed.Client, node_pairs, cmd: typing.Callable, **kwargs
-):
+    client: dask.distributed.Client,
+    node_pairs: typing.List[tuple],
+    cmd: typing.Callable,
+    **kwargs,
+) -> typing.Dict[str, dask.distributed.Future]:
     """Use the node_pairs to map the DVC graph onto a dask graph.
 
     Parameters
@@ -78,6 +81,10 @@ def submit_to_dask(
     Returns
     -------
 
+    dict: A dictionary with a unique submission id as key and the respective dask node
+            as value. The unique key is added for parallel execution of experiments with
+            the same Node name.
+
     """
 
     task_name_map = {}