From 938fbcdb74f0713567b93f24af841cf0ce4af30c Mon Sep 17 00:00:00 2001
From: Jordi Chauzi <jchauzi@upmem.com>
Date: Thu, 16 Jul 2020 11:16:06 +0200
Subject: [PATCH] checksum: handle DPU_ALLOCATE_ALL as a value of the nr_dpus
 parameter

It is notably used in the README example
---
 checksum/host/host.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/checksum/host/host.py b/checksum/host/host.py
index 995fdb2..1fbe941 100644
--- a/checksum/host/host.py
+++ b/checksum/host/host.py
@@ -16,6 +16,7 @@
 
 
 from dpu import DpuSet
+from dpu import ALLOCATE_ALL
 import argparse
 import os
 import random
@@ -90,8 +91,17 @@ def create_test_file():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('nr_dpus', type=int)
+    parser.add_argument('nr_dpus')
     parser.add_argument('nr_tasklets', type=int)
     args = parser.parse_args()
 
-    main(args.nr_dpus, args.nr_tasklets)
+    if args.nr_dpus == 'DPU_ALLOCATE_ALL':
+        nr_dpus = ALLOCATE_ALL
+    else:
+        try:
+            nr_dpus = int(args.nr_dpus)
+        except ValueError:
+            parser.error("argument nr_dpus: invalid value: '{}'".format(args.nr_dpus))
+            sys.exit(os.EX_USAGE)
+
+    main(nr_dpus, args.nr_tasklets)