fix: fix odps table raise device error (#15)

Author: SeaOfOcean

epl/cluster.py

@@ -80,6 +80,11 @@ def get_device(self, replica_idx, device_idx):
     """
     return self._slice_devices[replica_idx][device_idx]
 
+  @property
+  def slice_devices(self):
+    """Get device slices"""
+    return self._slice_devices
+
   def _device_in_local_worker(self, device):
     return common.get_task_index_from_device_str(device) == self._worker_index
 

epl/ir/graph.py

@@ -280,7 +280,7 @@ def check_and_set_cloned_dataset_need_clone(self):
       self.clone_dataset_related_ops = True
       return
     constructor_num_device = dict()
-    for devices in self._env.cluster.virtual_devices[0]:
+    for devices in self._env.cluster.virtual_devices[0].slice_devices:
       for device in devices:
         constructor_task_index = \
             common.get_task_index_from_device_str(device)
@@ -536,8 +536,6 @@ def add_operation(self, primitive_op):
     # Specific dataset api op.
     if op.type in constant.DATASET_API_OPS:
       self._dataset_api_op.append(op)
-      if op.type in constant.ODPS_TABLE_API_OPS:
-        self.check_and_set_cloned_dataset_need_clone()
 
     if current_model_phase == ModelPhase.ADD_FUNCTION:
       if not self._current_function_name:

epl/parallel/parallel.py

@@ -212,6 +212,9 @@ def do_parallelism(self):
     """Tranform original graph to parallel graph."""
     if self._graph.is_constructor:
       with ops.name_scope(constant.PARALLEL_STRATEGY):
+        for op in self._graph.dataset_api_op:
+          if op.type in constant.ODPS_TABLE_API_OPS:
+            self._graph.check_and_set_cloned_dataset_need_clone()
         self._fix_map_operations_to_taskgraph()
         if Env.get().config.offload.level == "v0":
           tf_logging.info("enable weight offload")

tests/graph_test.py

@@ -774,6 +774,27 @@ def test_multithread_prefetch(self):
         steps.append(step)
     self.assertEqual(steps, [0, 1, 2])
 
+  def test_check_and_set_cloned_dataset_need_clone(self):
+    epl.init()
+    epl.set_default_strategy(epl.replicate(1))
+    num_x = np.random.randint(0, 10, (500, 10)).astype(dtype=np.float32)
+    num_y = np.random.randint(0, 10, 500).astype(dtype=np.int32)
+    dataset = tf.data.Dataset.from_tensor_slices((num_x, num_y)) \
+        .batch(2).repeat(1)
+    iterator = dataset.make_initializable_iterator()
+    tf.add_to_collection(tf.GraphKeys.TABLE_INITIALIZERS, iterator.initializer)
+    source, target = iterator.get_next()
+    x = tf.layers.dense(inputs=source, units=16, activation=None)
+    x = tf.layers.dense(inputs=x, units=16, activation=None)
+    dense1 = tf.layers.dense(inputs=x, units=16, activation=None)
+    logits = tf.layers.dense(inputs=dense1, units=10, activation=None)
+    loss = tf.reduce_mean(logits)
+    global_step = tf.train.get_or_create_global_step()
+    optimizer = tf.train.MomentumOptimizer(learning_rate=0.001, momentum=0.9)
+    train_op = optimizer.minimize(loss, global_step=global_step)
+    epl.Graph.get().check_and_set_cloned_dataset_need_clone()
+
+
 
 # pylint: enable=missing-docstring,protected-access,unused-argument,
 # pylint: enable=line-too-long,bad-continuation,unused-variable