fix(trainer): fix data path visitor (#1045)

* fix(trainer): fix data path visitor * fix(trainer): fix data path visitor * fix(trainer): fix data path visitor Co-authored-by: 杭卫强 <[email protected]>
bytedance · Sep 22, 2022 · 1cf0b69 · 1cf0b69
1 parent 9318944
commit 1cf0b69
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/deploy/scripts/trainer/run_trainer_master.sh b/deploy/scripts/trainer/run_trainer_master.sh
@@ -35,6 +35,7 @@ export_model=$(normalize_env_to_args "--export-model" $EXPORT_MODEL)
 shuffle=$(normalize_env_to_args "--shuffle" $SUFFLE_DATA_BLOCK)
 shuffle_in_day=$(normalize_env_to_args "--shuffle-in-day" $SHUFFLE_IN_DAY)
 local_data_source=$(normalize_env_to_args "--local-data-source" $LOCAL_DATA_SOURCE)
+local_data_path=$(normalize_env_to_args "--local-data-path" $LOCAL_DATA_PATH)
 local_start_date=$(normalize_env_to_args "--local-start-date" $LOCAL_START_DATE)
 local_end_date=$(normalize_env_to_args "--local-end-date" $LOCAL_END_DATE)
 
@@ -109,6 +110,6 @@ python main.py --master \
     $mode $sparse_estimator \
     $save_checkpoint_steps $save_checkpoint_secs \
     $summary_save_steps $summary_save_secs \
-    $local_data_source $local_start_date $local_end_date \
-    $epoch_num $start_date $end_date $shuffle $shuffle_in_day \
-    $extra_params $export_model
+    $local_data_source $local_data_path $local_start_date \
+    $local_end_date $epoch_num $start_date $end_date \
+    $shuffle $shuffle_in_day $extra_params $export_model
diff --git a/fedlearner/trainer/data_visitor.py b/fedlearner/trainer/data_visitor.py
@@ -349,12 +349,11 @@ def __init__(self,
                     continue
                 subdirname = os.path.relpath(dirname, data_path)
                 block_id = os.path.join(subdirname, filename)
-                datablock = _RawDataBlock(block_id,
-                                          os.path.join(dirname, filename),
-                                          None, None,
-                                          tm_pb.JOINED)
+                datablock = _RawDataBlock(
+                    id=block_id, data_path=os.path.join(dirname, filename),
+                    start_time=None, end_time=None, type=tm_pb.JOINED)
                 datablocks.append(datablock)
-        datablocks.sort(key=lambda x: x.end_time)
+        datablocks.sort(key=lambda x: x.id)
 
         fl_logging.info("create DataVisitor by local_data_path: %s",
                         local_data_path)
@@ -366,12 +365,11 @@ def __init__(self,
                         continue
                     subdirname = os.path.relpath(dirname, local_data_path)
                     block_id = os.path.join(subdirname, filename)
-                    datablock = _RawDataBlock(block_id,
-                                              os.path.join(dirname, filename),
-                                              None, None,
-                                              tm_pb.LOCAL)
+                    datablock = _RawDataBlock(
+                        id=block_id, data_path=os.path.join(dirname, filename),
+                        start_time=None, end_time=None, type=tm_pb.LOCAL)
                     local_datablocks.append(datablock)
-        local_datablocks.sort(key=lambda x: x.end_time)
+        local_datablocks.sort(key=lambda x: x.id)
 
         super(DataPathVisitor, self).__init__(datablocks, local_datablocks,
                                               epoch_num, shuffle_type)