Create drive in setup.py for PerfZero (#321)

lindong28 · tfboyd · commit 5054769776af · 2019-02-21T19:21:28.000-08:00
diff --git a/perfzero/lib/benchmark.py b/perfzero/lib/benchmark.py
@@ -26,7 +26,6 @@
 import time
 import traceback
 
-import perfzero.device_utils as device_utils
 import perfzero.perfzero_config as perfzero_config
 import perfzero.report_utils as report_utils
 import perfzero.utils as utils
@@ -45,11 +44,6 @@ def __init__(self, config):
 
   def _setup(self):
     """Download data and checkout git repository."""
-    # Set up the raid array.
-    start_time = time.time()
-    device_utils.create_drive_from_devices(self.config.root_data_dir,
-                                           self.config.gce_nvme_raid)
-    self.benchmark_execution_time['create_drive'] = time.time() - start_time
 
     # Acticate gcloud service
     start_time = time.time()
diff --git a/perfzero/lib/perfzero/perfzero_config.py b/perfzero/lib/perfzero/perfzero_config.py
@@ -39,6 +39,18 @@ def add_setup_parser_arguments(parser):
       type=str,
       help='''The gcloud key file url. When specified, it will be downloaded to the directory specified by the flag --workspace.
       ''')
+  parser.add_argument(
+      '--root_data_dir',
+      default='/data',
+      type=str,
+      help='The directory which should contain the dataset required by the becnhmark method.'
+      )
+  parser.add_argument(
+      '--gce_nvme_raid',
+      default=None,
+      type=str,
+      help='If set to non-empty string, create raid 0 array with devices at the directory specified by the flag --root_data_dir'
+      )
 
 
 def add_benchmark_parser_arguments(parser):
@@ -48,12 +60,6 @@ def add_benchmark_parser_arguments(parser):
       action='store_true',
       help='If set, do git pull for dependent git repositories'
       )
-  parser.add_argument(
-      '--gce_nvme_raid',
-      default=None,
-      type=str,
-      help='If set to non-empty string, create raid 0 array with devices at the directory specified by the flag --root_data_dir'
-      )
   parser.add_argument(
       '--gcs_downloads',
       default=None,
@@ -185,7 +191,6 @@ def __init__(self, mode, flags=None):
     self.mode = mode
     self.flags = flags
     if mode == 'flags':
-      self.gce_nvme_raid = flags.gce_nvme_raid
       self.gcs_downloads_str = flags.gcs_downloads
       self.data_downloads_str = flags.data_downloads
       self.git_repos_str = flags.git_repos
diff --git a/perfzero/lib/setup.py b/perfzero/lib/setup.py
@@ -21,6 +21,7 @@
 import os
 import time
 
+import perfzero.device_utils as device_utils
 import perfzero.perfzero_config as perfzero_config
 import perfzero.utils as utils
 
@@ -46,6 +47,12 @@
   utils.active_gcloud_service(FLAGS.gcloud_key_file_url, workspace_dir, download_only=True)  # pylint: disable=line-too-long
   setup_execution_time['download_token'] = time.time() - start_time
 
+  # Set up the raid array.
+  start_time = time.time()
+  device_utils.create_drive_from_devices(FLAGS.root_data_dir, FLAGS.gce_nvme_raid)  # pylint: disable=line-too-long
+  setup_execution_time['create_drive'] = time.time() - start_time
+
+  # Create docker image
   start_time = time.time()
   dockerfile_path = os.path.join(project_dir, FLAGS.dockerfile_path)
   docker_tag = 'temp/tf-gpu'
