Format

Author: zhongchun

mars/deploy/oscar/tests/test_ray_error_recovery.py

@@ -1,3 +1,17 @@
+# Copyright 1999-2021 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
 import numpy as np
 import pandas as pd

mars/services/task/execution/mars/executor.py

@@ -37,7 +37,6 @@
 from .....utils import Timer
 from ....context import ThreadedServiceContext
 from ....cluster.api import ClusterAPI
-from ....cluster.core import NodeStatus
 from ....lifecycle.api import LifecycleAPI
 from ....meta.api import MetaAPI
 from ....scheduling import SchedulingAPI
@@ -95,7 +94,8 @@ def __init__(
         self._meta_api = meta_api
 
         self._stage_processors = []
-		self._stage_id_to_processor = weakref.WeakValueDictionary()
+        self._stage_tile_progresses = []
+        self._stage_id_to_processor = weakref.WeakValueDictionary()
         self._cur_stage_processor = None
         self._result_tileables_lifecycle = None
         self._subtask_decref_events = dict()

mars/services/task/execution/mars/stage.py

@@ -279,8 +279,8 @@ async def set_subtask_result(self, result: SubtaskResult, band: BandType = None)
                 )
                 await self.cancel()
         else:
-            await async_call(
-                self._scheduling_api.finish_subtasks([result.subtask_id], bands=[band])
+            await self._scheduling_api.finish_subtasks(
+                [result.subtask_id], bands=[band]
             )
             logger.debug(
                 "Continue to schedule subtasks after subtask %s finished.",
@@ -306,7 +306,7 @@ async def set_subtask_result(self, result: SubtaskResult, band: BandType = None)
                         self.task.task_id,
                         self.stage_id,
                     )
-                    await async_call(self._schedule_subtasks([to_schedule_subtask]))
+                    await self._schedule_subtasks([to_schedule_subtask])
                 except KeyError:
                     logger.exception("Got KeyError.")
 

mars/services/task/execution/ray/executor.py

@@ -1074,8 +1074,8 @@ def gc():
                     last_check_slow_time = curr_time
             # Fast to next loop and give it a chance to update object_ref_to_subtask.
             await asyncio.sleep(interval_seconds if len(ready_objects) == 0 else 0)
-	
-	def get_stage_generation_order(self, stage_id: str):
+
+    def get_stage_generation_order(self, stage_id: str):
         raise NotImplementedError(
             "RayTaskExecutor doesn't support stage generation order."
         )

mars/services/task/supervisor/processor.py

@@ -30,7 +30,7 @@
 from ....typing import TileableType, ChunkType
 from ....utils import Timer
 from ...context import FailOverContext
-from ...subtask import SubtaskResult, Subtask
+from ...subtask import SubtaskResult, SubtaskGraph, Subtask
 from ..core import Task, TaskResult, TaskStatus, new_task_id
 from ..execution.api import TaskExecutor, ExecutionChunkResult
 from .preprocessor import TaskPreprocessor
@@ -481,8 +481,9 @@ def _finish(self):
 
     def is_done(self) -> bool:
         return self.done.is_set()
+
     def get_generation_order(self, stage_id: str):
         return self._executor.get_stage_generation_order(stage_id)
 
     def get_subtask(self, chunk_data_key: str):
-        return self._chunk_data_key_to_subtask.get(chunk_data_key)
+        return self._chunk_data_key_to_subtask.get(chunk_data_key)

mars/services/task/supervisor/tests/task_preprocessor.py

@@ -177,7 +177,7 @@ def analyze(
             self._config,
             chunk_to_subtasks,
             shuffle_fetch_type=shuffle_fetch_type,
-			stage_id=stage_id,
+            stage_id=stage_id,
         )
         subtask_graph = analyzer.gen_subtask_graph()
         results = set(

mars/storage/ray.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import inspect
+import logging
 
 from typing import Any, Dict, List, Tuple
 from ..lib import sparse
@@ -28,6 +29,7 @@
 from .core import BufferWrappedFileObject, StorageFileObject
 
 ray = lazy_import("ray")
+logger = logging.getLogger(__name__)
 
 
 # TODO(fyrestone): make the SparseMatrix pickleable.