Repository: sql-machine-learning/elasticdl

Issue:
Ensure records per task to be integral multiple of mini-batch size.
Currently, EDL train command contains two parameters: minibatch_size and records_per_task. Because there is no constraint between the value of these two parameters(such as records_per_task = 10, minibatch_size = 8), the records in one mini-batch may come from two tasks.
In order to make the process logic simple and clear, we make sure that records_per_task to be integral multiple of minibatch_size. But we can't assume that user will set these two parameter according to the constraints. Change the parameter name from records_per_task to minibatch_count_per_task.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/args.py b/elasticdl/python/common/args.py
--- a/elasticdl/python/common/args.py
+++ b/elasticdl/python/common/args.py
@@ -173,7 +173,12 @@ def add_common_params(parser):
         help="The name of the Kubernetes namespace where ElasticDL "
         "pods will be created",
     )
-    parser.add_argument("--records_per_task", type=int, required=True)
+    parser.add_argument(
+        "--num_minibatches_per_task",
+        type=int,
+        help="The number of minibatches per task",
+        required=True,
+    )
     parser.add_argument(
         "--cluster_spec",
         help="The file that contains user-defined cluster specification",
diff --git a/elasticdl/python/master/main.py b/elasticdl/python/master/main.py
--- a/elasticdl/python/master/main.py
+++ b/elasticdl/python/master/main.py
@@ -94,11 +94,13 @@ def main():
         args.evaluation_data_dir,
         args.prediction_data_dir,
     )
+
+    records_per_task = args.minibatch_size * args.num_minibatches_per_task
     task_d = _make_task_dispatcher(
         args.training_data_dir,
         args.evaluation_data_dir,
         args.prediction_data_dir,
-        args.records_per_task,
+        records_per_task,
         args.num_epochs,
     )
     model_module = load_module(


Test patch:
diff --git a/scripts/client_test.sh b/scripts/client_test.sh
--- a/scripts/client_test.sh
+++ b/scripts/client_test.sh
@@ -15,7 +15,7 @@ if [[ "$JOB_TYPE" == "train" ]]; then
       --worker_resource_request="cpu=400m,memory=2048Mi" \
       --worker_resource_limit="cpu=1,memory=3072Mi" \
       --minibatch_size=64 \
-      --records_per_task=100 \
+      --num_minibatches_per_task=2 \
       --num_workers=2 \
       --checkpoint_steps=10 \
       --evaluation_steps=15 \
@@ -37,7 +37,7 @@ elif [[ "$JOB_TYPE" == "evaluate" ]]; then
       --worker_resource_request="cpu=400m,memory=2048Mi" \
       --worker_resource_limit="cpu=1,memory=3072Mi" \
       --minibatch_size=64 \
-      --records_per_task=100 \
+      --num_minibatches_per_task=2 \
       --num_workers=2 \
       --evaluation_steps=15 \
       --job_name=test-mnist-evaluate \
@@ -55,7 +55,7 @@ elif [[ "$JOB_TYPE" == "predict" ]]; then
       --worker_resource_request="cpu=400m,memory=2048Mi" \
       --worker_resource_limit="cpu=1,memory=3072Mi" \
       --minibatch_size=64 \
-      --records_per_task=100 \
+      --num_minibatches_per_task=2 \
       --num_workers=2 \
       --job_name=test-mnist-predict \
       --log_level=INFO \