DeepSeek/smallpond
1
0
Fork 0
mirror of https://github.com/deepseek-ai/smallpond synced 2025-06-26 18:27:45 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

JobManager: move all arguments to __init__

Browse Source 
Signed-off-by: Runji Wang <runji@deepseek.com>
This commit is contained in:
Runji Wang
2025-02-28 20:03:55 +08:00
parent ed112db42a
commit 5a56a052bf
4 changed files with 79 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
smallpond/execution/driver.py
View File

@@ -372,10 +372,10 @@ class Driver(object):
elif args.mode == "scheduler": elif args.mode == "scheduler":
assert plan is not None assert plan is not None
jobmgr = JobManager( jobmgr = JobManager(
args.data_root, args.python_venv, args.task_image, args.platform data_root=args.data_root,
) python_venv=args.python_venv,
exec_plan = jobmgr.run( task_image=args.task_image,
plan, platform=args.platform,
job_id=args.job_id, job_id=args.job_id,
job_time=args.job_time, job_time=args.job_time,
job_name=args.job_name, job_name=args.job_name,
@@ -412,6 +412,7 @@ class Driver(object):
disable_log_rotation=args.disable_log_rotation, disable_log_rotation=args.disable_log_rotation,
output_path=args.output_path, output_path=args.output_path,
) )
exec_plan = jobmgr.run(plan)
retval = exec_plan if exec_plan.successful else None retval = exec_plan if exec_plan.successful else None
if stop_process_on_done: if stop_process_on_done:

97
smallpond/execution/manager.py
View File

@@ -52,20 +52,11 @@ class JobManager(object):
def __init__( def __init__(
self, self,
*,
data_root: Optional[str] = None, data_root: Optional[str] = None,
python_venv: Optional[str] = None, python_venv: Optional[str] = None,
task_image: str = "default", task_image: str = "default",
platform: Optional[str] = None, platform: Optional[str] = None,
) -> None:
self.platform = get_platform(platform)
self.data_root = os.path.abspath(data_root or self.platform.default_data_root())
self.python_venv = python_venv
self.task_image = task_image
@logger.catch(reraise=True, message="job manager terminated unexpectedly")
def run(
self,
plan: LogicalPlan,
job_id: Optional[str] = None, job_id: Optional[str] = None,
job_time: Optional[float] = None, job_time: Optional[float] = None,
job_name: str = "smallpond", job_name: str = "smallpond",
@@ -107,9 +98,15 @@ class JobManager(object):
sched_state_observers: Optional[List[Scheduler.StateObserver]] = None, sched_state_observers: Optional[List[Scheduler.StateObserver]] = None,
output_path: Optional[str] = None, output_path: Optional[str] = None,
**kwargs, **kwargs,
) -> ExecutionPlan: ) -> None:
self.platform = get_platform(platform)
logger.info(f"using platform: {self.platform}") logger.info(f"using platform: {self.platform}")
self.data_root = os.path.abspath(data_root or self.platform.default_data_root())
self.python_venv = python_venv
self.task_image = task_image
self.manifest_only_final_results = manifest_only_final_results
job_id = JobId(hex=job_id or self.platform.default_job_id()) job_id = JobId(hex=job_id or self.platform.default_job_id())
job_time = ( job_time = (
datetime.fromtimestamp(job_time) datetime.fromtimestamp(job_time)
@@ -145,7 +142,7 @@ class JobManager(object):
self.platform.shared_log_root() if share_log_analytics else None self.platform.shared_log_root() if share_log_analytics else None
) )
runtime_ctx = RuntimeContext( self.runtime_ctx = RuntimeContext(
job_id, job_id,
job_time, job_time,
self.data_root, self.data_root,
@@ -166,27 +163,13 @@ class JobManager(object):
output_path=output_path, output_path=output_path,
**kwargs, **kwargs,
) )
runtime_ctx.initialize(socket.gethostname(), root_exist_ok=True) self.runtime_ctx.initialize(socket.gethostname(), root_exist_ok=True)
logger.info( logger.info(
f"command-line arguments: {' '.join([os.path.basename(sys.argv[0]), *sys.argv[1:]])}" f"command-line arguments: {' '.join([os.path.basename(sys.argv[0]), *sys.argv[1:]])}"
) )
dump(runtime_ctx, runtime_ctx.runtime_ctx_path, atomic_write=True) dump(self.runtime_ctx, self.runtime_ctx.runtime_ctx_path, atomic_write=True)
logger.info(f"saved runtime context at {runtime_ctx.runtime_ctx_path}") logger.info(f"saved runtime context at {self.runtime_ctx.runtime_ctx_path}")
dump(plan, runtime_ctx.logcial_plan_path, atomic_write=True)
logger.info(f"saved logcial plan at {runtime_ctx.logcial_plan_path}")
plan.graph().render(runtime_ctx.logcial_plan_graph_path, format="png")
logger.info(
f"saved logcial plan graph at {runtime_ctx.logcial_plan_graph_path}.png"
)
exec_plan = Planner(runtime_ctx).create_exec_plan(
plan, manifest_only_final_results
)
dump(exec_plan, runtime_ctx.exec_plan_path, atomic_write=True)
logger.info(f"saved execution plan at {runtime_ctx.exec_plan_path}")
sidecar_list = sidecars or [] sidecar_list = sidecars or []
@@ -225,27 +208,26 @@ class JobManager(object):
sched_state_observers = sched_state_observers or [] sched_state_observers = sched_state_observers or []
if enable_sched_state_dump: if enable_sched_state_dump:
sched_state_exporter = SchedStateExporter(runtime_ctx.sched_state_path) sched_state_exporter = SchedStateExporter(self.runtime_ctx.sched_state_path)
sched_state_observers.insert(0, sched_state_exporter) sched_state_observers.insert(0, sched_state_exporter)
if os.path.exists(runtime_ctx.sched_state_path): if os.path.exists(self.runtime_ctx.sched_state_path):
logger.warning( logger.warning(
f"loading scheduler state from: {runtime_ctx.sched_state_path}" f"loading scheduler state from: {self.runtime_ctx.sched_state_path}"
) )
scheduler: Scheduler = load(runtime_ctx.sched_state_path) self.scheduler: Scheduler = load(self.runtime_ctx.sched_state_path)
scheduler.sched_epoch += 1 self.scheduler.sched_epoch += 1
scheduler.sched_state_observers = sched_state_observers self.scheduler.sched_state_observers = sched_state_observers
else: else:
scheduler = Scheduler( self.scheduler = Scheduler(
exec_plan, max_retry_count=max_retry_count,
max_retry_count, max_fail_count=max_fail_count,
max_fail_count, prioritize_retry=prioritize_retry,
prioritize_retry, speculative_exec=speculative_exec,
speculative_exec, stop_executor_on_failure=stop_executor_on_failure,
stop_executor_on_failure, nonzero_exitcode_as_oom=nonzero_exitcode_as_oom,
nonzero_exitcode_as_oom, remove_output_root=remove_output_root,
remove_output_root, sched_state_observers=sched_state_observers,
sched_state_observers,
) )
# start executors # start executors
self.platform.start_job( self.platform.start_job(
@@ -257,7 +239,7 @@ class JobManager(object):
"--data_root", "--data_root",
self.data_root, self.data_root,
"--runtime_ctx_path", "--runtime_ctx_path",
runtime_ctx.runtime_ctx_path, self.runtime_ctx.runtime_ctx_path,
"executor", "executor",
], ],
envs={ envs={
@@ -278,6 +260,25 @@ class JobManager(object):
), ),
) )
@logger.catch(reraise=True, message="job manager terminated unexpectedly")
def run(
self,
plan: LogicalPlan,
) -> ExecutionPlan:
dump(plan, self.runtime_ctx.logcial_plan_path, atomic_write=True)
logger.info(f"saved logcial plan at {self.runtime_ctx.logcial_plan_path}")
plan.graph().render(self.runtime_ctx.logcial_plan_graph_path, format="png")
logger.info(
f"saved logcial plan graph at {self.runtime_ctx.logcial_plan_graph_path}.png"
)
exec_plan = Planner(self.runtime_ctx).create_exec_plan(
plan, self.manifest_only_final_results
)
dump(exec_plan, self.runtime_ctx.exec_plan_path, atomic_write=True)
logger.info(f"saved execution plan at {self.runtime_ctx.exec_plan_path}")
# run scheduler # run scheduler
scheduler.run() self.scheduler.run(exec_plan)
return scheduler.exec_plan return exec_plan

42
smallpond/execution/scheduler.py
View File

@@ -12,7 +12,7 @@ from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from enum import Enum from enum import Enum
from functools import cached_property from functools import cached_property
from typing import Any, Callable, Dict, Iterable, List, Literal, Set, Tuple, Union from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Set, Tuple, Union
import numpy as np import numpy as np
from loguru import logger from loguru import logger
@@ -380,20 +380,16 @@ class Scheduler(object):
def __init__( def __init__(
self, self,
exec_plan: ExecutionPlan, *,
max_retry_count: int = DEFAULT_MAX_RETRY_COUNT, max_retry_count: int = DEFAULT_MAX_RETRY_COUNT,
max_fail_count: int = DEFAULT_MAX_FAIL_COUNT, max_fail_count: int = DEFAULT_MAX_FAIL_COUNT,
prioritize_retry=False, prioritize_retry: bool = False,
speculative_exec: Literal["disable", "enable", "aggressive"] = "enable", speculative_exec: Literal["disable", "enable", "aggressive"] = "enable",
stop_executor_on_failure=False, stop_executor_on_failure: bool = False,
nonzero_exitcode_as_oom=False, nonzero_exitcode_as_oom: bool = False,
remove_output_root=False, remove_output_root: bool = False,
sched_state_observers=None, sched_state_observers: Optional[List[StateObserver]] = None,
) -> None: ) -> None:
self.ctx = exec_plan.ctx
self.exec_plan = exec_plan
self.logical_plan: LogicalPlan = self.exec_plan.logical_plan
self.logical_nodes = self.logical_plan.nodes
self.max_retry_count = max_retry_count self.max_retry_count = max_retry_count
self.max_fail_count = max_fail_count self.max_fail_count = max_fail_count
self.standalone_mode = self.ctx.num_executors == 0 self.standalone_mode = self.ctx.num_executors == 0
@@ -410,16 +406,12 @@ class Scheduler(object):
# task states # task states
self.local_queue: List[Task] = [] self.local_queue: List[Task] = []
self.sched_queue: List[Task] = [] self.sched_queue: List[Task] = []
self.tasks: Dict[str, Task] = self.exec_plan.tasks self.tasks: Dict[str, Task] = {}
self.scheduled_tasks: Dict[TaskRuntimeId, Task] = OrderedDict() self.scheduled_tasks: Dict[TaskRuntimeId, Task] = {}
self.finished_tasks: Dict[TaskRuntimeId, Task] = OrderedDict() self.finished_tasks: Dict[TaskRuntimeId, Task] = {}
self.succeeded_tasks: Dict[str, Task] = OrderedDict() self.succeeded_tasks: Dict[str, Task] = {}
self.nontrivial_tasks = dict( self.nontrivial_tasks: Dict[str, Task] = {}
(key, task) self.succeeded_nontrivial_tasks: Dict[str, Task] = {}
for (key, task) in self.tasks.items()
if not task.exec_on_scheduler
)
self.succeeded_nontrivial_tasks: Dict[str, Task] = OrderedDict()
# executor pool # executor pool
self.local_executor = LocalExecutor.create(self.ctx, "localhost") self.local_executor = LocalExecutor.create(self.ctx, "localhost")
self.available_executors = {self.local_executor.id: self.local_executor} self.available_executors = {self.local_executor.id: self.local_executor}
@@ -431,9 +423,6 @@ class Scheduler(object):
self.probe_epoch = 0 self.probe_epoch = 0
self.sched_epoch = 0 self.sched_epoch = 0
def __call__(self, *args: Any, **kwds: Any) -> Any:
return self.run()
def __getstate__(self): def __getstate__(self):
state = self.__dict__.copy() state = self.__dict__.copy()
del state["sched_state_observers"] del state["sched_state_observers"]
@@ -982,7 +971,10 @@ class Scheduler(object):
status = "failure" status = "failure"
fout.write(f"{status}@{int(time.time())}") fout.write(f"{status}@{int(time.time())}")
def run(self) -> bool: def run(self, exec_plan: ExecutionPlan) -> bool:
"""
Run the execution plan.
"""
mp.current_process().name = f"SchedulerMainProcess#{self.sched_epoch}" mp.current_process().name = f"SchedulerMainProcess#{self.sched_epoch}"
logger.info( logger.info(
f"start to run scheduler #{self.sched_epoch} on {socket.gethostname()}" f"start to run scheduler #{self.sched_epoch} on {socket.gethostname()}"

13
tests/test_fabric.py
View File

@@ -28,11 +28,14 @@ generate_data()
def run_scheduler( def run_scheduler(
runtime_ctx: RuntimeContext, scheduler: Scheduler, queue: queue.Queue runtime_ctx: RuntimeContext,
scheduler: Scheduler,
queue: queue.Queue,
exec_plan: ExecutionPlan,
): ):
runtime_ctx.initialize("scheduler") runtime_ctx.initialize("scheduler")
scheduler.add_state_observer(Scheduler.StateObserver(SaveSchedState(queue))) scheduler.add_state_observer(Scheduler.StateObserver(SaveSchedState(queue)))
retval = scheduler.run() retval = scheduler.run(exec_plan)
print(f"scheduler exited with value {retval}", file=sys.stderr) print(f"scheduler exited with value {retval}", file=sys.stderr)
@@ -206,9 +209,7 @@ class TestFabric(unittest.TestCase):
self.queue_manager = Manager() self.queue_manager = Manager()
self.sched_states = self.queue_manager.Queue() self.sched_states = self.queue_manager.Queue()
exec_plan = Planner(runtime_ctx).create_exec_plan(plan)
scheduler = Scheduler( scheduler = Scheduler(
exec_plan,
max_retry_count=max_retry_count, max_retry_count=max_retry_count,
max_fail_count=max_fail_count, max_fail_count=max_fail_count,
prioritize_retry=prioritize_retry, prioritize_retry=prioritize_retry,
@@ -216,6 +217,8 @@ class TestFabric(unittest.TestCase):
stop_executor_on_failure=stop_executor_on_failure, stop_executor_on_failure=stop_executor_on_failure,
nonzero_exitcode_as_oom=nonzero_exitcode_as_oom, nonzero_exitcode_as_oom=nonzero_exitcode_as_oom,
) )
exec_plan = Planner(runtime_ctx).create_exec_plan(plan)
self.latest_state = scheduler self.latest_state = scheduler
self.executors = [ self.executors = [
Executor.create(runtime_ctx, f"executor-{i}") for i in range(num_executors) Executor.create(runtime_ctx, f"executor-{i}") for i in range(num_executors)
@@ -225,7 +228,7 @@ class TestFabric(unittest.TestCase):
target=run_scheduler, target=run_scheduler,
# XXX: on macOS, scheduler state observer will be cleared when cross-process # XXX: on macOS, scheduler state observer will be cleared when cross-process
# so we pass the queue and add the observer in the new process # so we pass the queue and add the observer in the new process
args=(runtime_ctx, scheduler, self.sched_states), args=(runtime_ctx, scheduler, self.sched_states, exec_plan),
name="scheduler", name="scheduler",
) )
] ]