Implement parallel execution for DAG tasks

[andylizf]

Closes #4055

This PR implements parallel execution for DAG tasks in the jobs controller, addressing issue #4055. The changes allow for efficient execution of complex DAGs with independent tasks running concurrently, significantly improving performance for workflows with parallel components.

Changes

• Modified JobsController to identify and execute parallel task groups
• Implemented thread-safe task execution and monitoring
• Added concurrent resource management and cleanup

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
• All smoke tests: pytest tests/test_smoke.py
• Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
• Backward compatibility tests: conda deactivate; bash -i tests/backward_compatibility_tests.sh

[andylizf]

@cblmemo I'm currently working on implementing a cancellation mechanism for tasks that have already started or are queued for execution (similar to your setup with replicas preparing to launch). I'm currently using future.cancel, but it doesn't seem to fully address cancellation of tasks that are already in progress. I haven't switched to using thread.event yet, which might improve this.

That said, I noticed you used Process for managing the launch and termination of replicas. I don't see any clear advantages to using Process over Thread, especially since Thread should handle task cancellation just as well without the overhead of creating separate processes. Could you clarify your reasoning for choosing Process here? Is there a specific limitation you're addressing with this approach?

[cblmemo]

@cblmemo I'm currently working on implementing a cancellation mechanism for tasks that have already started or are queued for execution (similar to your setup with replicas preparing to launch). I'm currently using future.cancel, but it doesn't seem to fully address cancellation of tasks that are already in progress. I haven't switched to using thread.event yet, which might improve this.

That said, I noticed you used Process for managing the launch and termination of replicas. I don't see any clear advantages to using Process over Thread, especially since Thread should handle task cancellation just as well without the overhead of creating separate processes. Could you clarify your reasoning for choosing Process here? Is there a specific limitation you're addressing with this approach?

This is mainly due to logging. Threading will share a same sys.stdout which makes the following code not feasible.

I cannot find a way to do this kind of logging redirection back then. If you figured out a way, pls let me know ;)

skypilot/sky/utils/ux_utils.py

Lines 80 to 121 in 7971aa2

	class RedirectOutputForProcess:
	"""Redirects stdout and stderr to a file.
	This class enabled output redirect for multiprocessing.Process.
	Example usage:
	p = multiprocessing.Process(
	target=RedirectOutputForProcess(func, file_name).run, args=...)
	This is equal to:
	p = multiprocessing.Process(target=func, args=...)
	Plus redirect all stdout/stderr to file_name.
	"""
	def __init__(self, func: Callable, file: str, mode: str = 'w') -> None:
	self.func = func
	self.file = file
	self.mode = mode
	def run(self, *args, **kwargs):
	with open(self.file, self.mode, encoding='utf-8') as f:
	sys.stdout = f
	sys.stderr = f
	# reconfigure logger since the logger is initialized before
	# with previous stdout/stderr
	sky_logging.reload_logger()
	logger = sky_logging.init_logger(__name__)
	# The subprocess_util.run('sky status') inside
	# sky.execution::_execute cannot be redirect, since we cannot
	# directly operate on the stdout/stderr of the subprocess. This
	# is because some code in skypilot will specify the stdout/stderr
	# of the subprocess.
	try:
	self.func(*args, **kwargs)
	except Exception as e: # pylint: disable=broad-except
	logger.error(f'Failed to run {self.func.__name__}. '
	f'Details: {common_utils.format_exception(e)}')
	with ux_utils.enable_traceback():
	logger.error(f' Traceback:\n{traceback.format_exc()}')
	raise

[andylizf]

• Controller-side cluster's launch logs accessibility: Currently, logs like ~/sky_logs/sky-2024-10-28-04-14-45-846061/task_3_launch.log are only available on the controller machine and can only be accessed remotely.
• Cluster-side task run log persistence: As demonstrated in the example where task 0's logs are no longer available after completion, we need to implement proper log retention for completed tasks.

Hi @cblmemo, could you check the PR? All TODOs are done; log download isn't added due to complexity - run logs are accessible only if the cluster isn’t terminated.

[andylizf]

By the way, I refactored _follow_replica_logs, but the logic may differ slightly. Could you take a look?

[cblmemo]

Thanks for adding this feature @andylizf ! This is awesome. LGTM except for some nits. After testing it should be ready to go ;)

Before merging, please update all tests you've done in the PR description :))

[cblmemo]

Why change this?

[andylizf]

As mentiond in #4128 (comment), we should not cancelled job already started running.

[cblmemo]

Lets remove the else and revert the indents? We should reduce indent if possible.

[andylizf]

Thank you for the input! I agree that too much indentation isn’t ideal, but in this case, a couple of levels seem fine. Early exits work well for handling specific errors, while here, both cases are part of the function’s main flow.