[dask] Make vanilla-pad work properlly with dask

In this MR, I'm.

1. Wrapping the pipeline with DaskWrapper, because it wasn't. I think that's why things were not working for you with SGE @ydayer. vanilla-pad was creating SGE jobs AND running everything locally.
2. Added the options "--dask-partition-size" and "--dask-n-workers", so the user has some freedom to set these parameters if the heuristics in place are unpleasant.

I'll merge this one right away because of the workshop.

Feel free to open issues if any

bob/pad/base/script/vanilla_pad.py

@@ -6,6 +6,7 @@ import click
 from bob.extension.scripts.click_helper import ConfigCommand
 from bob.extension.scripts.click_helper import ResourceOption
 from bob.extension.scripts.click_helper import verbosity_option
+from bob.pipelines.distributed import dask_get_partition_size
 
 
 @click.command(
@@ -67,10 +68,39 @@ from bob.extension.scripts.click_helper import verbosity_option
     help="If set, it will checkpoint all steps of the pipeline. Checkpoints will be saved in `--output`.",
     cls=ResourceOption,
 )
+@click.option(
+    "--dask-partition-size",
+    "-s",
+    help="If using Dask, this option defines the size of each dask.bag.partition."
+    "Use this option if the current heuristic that sets this value doesn't suit your experiment."
+    "(https://docs.dask.org/en/latest/bag-api.html?highlight=partition_size#dask.bag.from_sequence).",
+    default=None,
+    type=click.INT,
+    cls=ResourceOption,
+)
+@click.option(
+    "--dask-n-workers",
+    "-n",
+    help="If using Dask, this option defines the number of workers to start your experiment."
+    "Dask automatically scales up/down the number of workers due to the current load of tasks to be solved."
+    "Use this option if the current amount of workers set to start an experiment doesn't suit you.",
+    default=None,
+    type=click.INT,
+    cls=ResourceOption,
+)
 @verbosity_option(cls=ResourceOption)
 @click.pass_context
 def vanilla_pad(
-    ctx, pipeline, database, dask_client, groups, output, checkpoint, **kwargs
+    ctx,
+    pipeline,
+    database,
+    dask_client,
+    groups,
+    output,
+    checkpoint,
+    dask_partition_size,
+    dask_n_workers,
+    **kwargs,
 ):
     """Runs the simplest PAD pipeline."""
 
@@ -84,6 +114,8 @@ def vanilla_pad(
     import dask.bag
     from bob.extension.scripts.click_helper import log_parameters
     from bob.pipelines.distributed.sge import get_resource_requirements
+    from bob.pipelines.utils import isinstance_nested
+    from bob.pipelines.wrappers import DaskWrapper
 
     logger = logging.getLogger(__name__)
     log_parameters(logger)
@@ -95,8 +127,33 @@ def vanilla_pad(
             ["checkpoint"], pipeline, features_dir=output, model_path=output
         )
 
-    if dask_client is None:
-        logger.warning("`dask_client` not set. Your pipeline will run locally")
+    # Fetching samples
+    fit_samples = database.fit_samples()
+    total_samples = len(fit_samples)
+    predict_samples = dict()
+    for group in groups:
+        predict_samples[group] = database.predict_samples(group=group)
+        total_samples += len(predict_samples[group])
+
+    # Checking if the pipieline is dask-wrapped
+    first_step = pipeline[0]
+    if not isinstance_nested(first_step, "estimator", DaskWrapper):
+
+        # Scaling up if necessary
+        if dask_n_workers is not None and not isinstance(dask_client, str):
+            dask_client.cluster.scale(dask_n_workers)
+
+        # Defining the partition size
+        partition_size = None
+        if not isinstance(dask_client, str):
+            lower_bound = 25  # lower bound of 25 videos per chunk
+            partition_size = dask_get_partition_size(
+                dask_client.cluster, total_samples, lower_bound=lower_bound
+            )
+        if dask_partition_size is not None:
+            partition_size = dask_partition_size
+
+        pipeline = mario.wrap(["dask"], pipeline, partition_size=partition_size)
 
     # create an experiment info file
     with open(os.path.join(output, "Experiment_info.txt"), "wt") as f:
@@ -107,14 +164,12 @@ def vanilla_pad(
             f.write(f"Step {i}: {name}\n{estimator!r}\n")
 
     # train the pipeline
-    fit_samples = database.fit_samples()
     pipeline.fit(fit_samples)
 
     for group in groups:
 
         logger.info(f"Running vanilla biometrics for group {group}")
-        predict_samples = database.predict_samples(group=group)
-        result = pipeline.decision_function(predict_samples)
+        result = pipeline.decision_function(predict_samples[group])
 
         scores_path = os.path.join(output, f"scores-{group}")