Add `dependencies` feature for task dependency management #4067

andylizf · 2024-10-10T23:31:30Z

Resolve #4054

This PR introduces the dependencies feature to SkyPilot, enabling users to specify task dependencies in their YAML configurations. This enhancement allows for more complex workflow definitions and ensures proper execution order of interdependent tasks.

Key changes:

Extended YAML parsing to support dependencies section
Implemented task graph construction based on dependencies
Provide much powerful interface for graph construction
Updated YAML generation to include dependency information

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 · 2024-10-11T23:43:57Z

@cblmemo Hi, could you kindly review this PR when you have a moment? Thanks!

cblmemo

Thanks for adding this feature @andylizf ! This would be really helpful for specifying complex workflows. Left some comments. It mostly looks good to me ;)

I created a branch for our dag project and lets merge all features to this branch first. We might want to do some final checks before we merge this to our master branch and this advanced-dag branch can serve as a buffer :))

examples/dag/rhombic.yml

sky/dag.py

cblmemo · 2024-10-14T22:10:41Z

Also, as mentioned here, lets change the target branch to merge to advanced-dag.

cblmemo

Thanks @andylizf for the awesome work! Left some discussion here. I tested and the basic functionality works fine!

sky/dag.py

cblmemo · 2024-10-15T17:59:12Z

sky/utils/dag_utils.py

+ if multi_tasks:
+ raise ValueError('Multiple task definitions without a valid'
+ 'header.')


Lets default to use chain yaml for backward compatibility?

Is it also a chain yaml only if there is a name section?

Could you elaborate?

I see your initial point now. We can default to chain YAML without header, add a generate_dag_name function, and use the YAML filename as the default name. Sound good?

Mostly sounds good to me. for the name of the dag, lets keep align with current master branch?

skypilot/sky/utils/dag_utils.py

Lines 84 to 94 in 3e98afe

configs = common_utils.read_yaml_all(path)

dag_name = None

if set(configs[0].keys()) == {'name'}:

dag_name = configs[0]['name']

configs = configs[1:]

elif len(configs) == 1:

dag_name = configs[0].get('name')

if len(configs) == 0:

# YAML has only `name: xxx`. Still instantiate a task.

configs = [{'name': dag_name}]

IIUC, the current master's behavior is that if the header is missing and it is a multitasking pipeline, then dag_name will remain None. Is this the expected behavior?

Allowing dag.name to be None creates a problematic cycle. If a user loads a DAG without a name, dag.name becomes None. When dumping this DAG to YAML, it results in name: null. Later, when the controller loads this YAML, it parses null back to None. This leads to an AssertionError in _get_dag_and_name function, which assumes dag.name is not None. This cycle of None -> null -> None causes unexpected runtime errors in the existing codebase and should be addressed at once to ensure system stability.

skypilot/sky/jobs/controller.py

Lines 39 to 43 in 71a95f4

def _get_dag_and_name(dag_yaml: str) -> Tuple['sky.Dag', str]:

dag = dag_utils.load_chain_dag_from_yaml(dag_yaml)

dag_name = dag.name

assert dag_name is not None, dag

return dag, dag_name

Pls check this

skypilot/sky/jobs/core.py

Line 67 in c6ae536

dag_utils.maybe_infer_and_fill_dag_and_task_names(dag)

sky/dag.py

examples/dag/diamond.yml

sky/utils/dag_utils.py

…#4067) provide an example, edited from pipeline.yml more focus on dependencies for user dag lib more powerful user interface load and dump new yaml format fix fix: reversed logic in add_edge rename refactor due to reviewer's comments generate task.name if not given add comments for add_edge add `print_exception_no_traceback` when raise make `Dag.tasks` a property print dependencies for `__repr__` move `get_unique_task_name` to common_utils rename methods to use downstream/edge terminology

cblmemo

Thanks for the prompt fix @andylizf ! Left some discussion. Most looks good to me!

sky/dag.py

sky/serve/serve_utils.py

cblmemo · 2024-10-19T23:08:26Z

sky/dag.py

+ """
+ self.name = name
+ self._task_name_lookup: Dict[str, 'task.Task'] = {}
+ self.edges: Dict['task.Task', Set['task.Task']] = {}


I'm a little bit confused why we need to maintain our own edges. doesnt nx.DiGraph provide this ?

Good catch. Ditched self.edges and are now fully using nx.DiGraph for edge stuff. Makes sense since we got rid of self.tasks too. These help keeping things cleaner.

sky/task.py

cblmemo · 2024-10-19T23:09:39Z

sky/utils/common_utils.py

+def get_unique_task_name(_: 'task_lib.Task') -> str:
+ timestamp = int(time.time())
+ unique_suffix = uuid.uuid4().hex[:6]
+ name = f'task_{timestamp}_{unique_suffix}'
+ return name


Lets respect the original task name and append uuid suffix to it.

Well, this function is only called when task name is None. However, to keep each task name unique, I would do what you said and add the uuid suffix to every task name while keeping this logic when the task name is null.

Co-authored-by: Tian Xia <[email protected]>

sky/dag.py

cblmemo · 2024-10-20T21:54:31Z

sky/utils/common_utils.py

-def get_unique_task_name(_: 'task_lib.Task') -> str:
- timestamp = int(time.time())
+def get_unique_task_name(task: 'task_lib.Task') -> str:
+ name = task.name or f'task_{int(time.time())}'


Suggested change

name = task.name or f'task_{int(time.time())}'

name = task.name if task.name is not None else f'task_{int(time.time())}'

nit: Always use if foo is None.
https://google.github.io/styleguide/pyguide.html#2144-decision

It's intentional. I think we should keep it as is. Catching empty strings is also what we want. If the task name is an empty string, we still want to fall back to the timestamp.

Co-authored-by: Tian Xia <[email protected]>

…amed" Otherwise, users can not refer to the task by name in the DAG. This reverts commit 8486352.

andylizf added 6 commits October 10, 2024 16:29

provide an example, edited from pipeline.yml

7c0965a

more focus on dependencies for user dag lib

6949ac1

more powerful user interface

55db40b

load and dump new yaml format

db7ff9f

fix

054cc26

fix: reversed logic in add_edge

24ef94e

andylizf force-pushed the job-dag branch from c70a8bd to 24ef94e Compare October 11, 2024 06:56

andylizf marked this pull request as ready for review October 11, 2024 07:34

cblmemo reviewed Oct 14, 2024

View reviewed changes

rename

129bdbf

andylizf changed the title ~~Add depends_on feature for task dependency management~~ Add dependencies feature for task dependency management Oct 14, 2024

refactor due to reviewer's comments

12ec5a4

andylizf force-pushed the job-dag branch from 3cb6bc5 to 12ec5a4 Compare October 14, 2024 18:45

generate task.name if not given

9497a3e

andylizf force-pushed the job-dag branch from 8117d09 to 9497a3e Compare October 14, 2024 19:17

andylizf changed the base branch from master to advanced-dag October 14, 2024 22:25

add comments for add_edge

ff528a5

cblmemo reviewed Oct 15, 2024

View reviewed changes

andylizf added 4 commits October 15, 2024 12:20

add print_exception_no_traceback when raise

04c6f9d

make Dag.tasks a property

4985813

print dependencies for __repr__

48a2826

move get_unique_task_name to common_utils

78d826d

andylizf mentioned this pull request Oct 17, 2024

Refactor optimizer: Ensure proper removal of sink/dummy nodes in DAG #4109

Open

rename methods to use downstream/edge terminology

e88acc1

cblmemo reviewed Oct 19, 2024

View reviewed changes

andylizf and others added 2 commits October 20, 2024 14:03

Apply suggestions from code review

e4fbb28

Co-authored-by: Tian Xia <[email protected]>

change wording all to up/downstream style

a27969b

This comment was marked as duplicate.

Sign in to view

andylizf added 3 commits October 20, 2024 14:28

Add unique suffix to task names, fallback to timestamp if unnamed

8486352

Unify handling of single and multiple tasks without dependencies

c14980e

Refactor tasks initialization: use list comprehension and fail fast

66fc864

cblmemo reviewed Oct 20, 2024

View reviewed changes

sky/dag.py Outdated Show resolved Hide resolved

cblmemo reviewed Oct 20, 2024

View reviewed changes

andylizf and others added 4 commits October 20, 2024 15:05

Fix remove task dependency description: upstream, not downstream

65d0bdd

Co-authored-by: Tian Xia <[email protected]>

Remove duplicated self.edges, use nx api instead

28b6482

Revert "Add unique suffix to task names, fallback to timestamp if unn…

1792ba6

…amed" Otherwise, users can not refer to the task by name in the DAG. This reverts commit 8486352.

comment the checking used as upstream logic

f600d16

andylizf changed the base branch from advanced-dag to master October 20, 2024 23:56

Empty commit

635c564

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `dependencies` feature for task dependency management #4067

Add `dependencies` feature for task dependency management #4067

andylizf commented Oct 10, 2024 •

edited

Loading

andylizf commented Oct 11, 2024

cblmemo left a comment

cblmemo commented Oct 14, 2024

cblmemo left a comment

cblmemo Oct 15, 2024

andylizf Oct 15, 2024

cblmemo Oct 16, 2024

andylizf Oct 17, 2024

cblmemo Oct 17, 2024

andylizf Oct 17, 2024 •

edited

Loading

andylizf Oct 18, 2024

cblmemo Oct 19, 2024

cblmemo left a comment

cblmemo Oct 19, 2024

andylizf Oct 20, 2024

cblmemo Oct 19, 2024

andylizf Oct 20, 2024

This comment was marked as duplicate.

cblmemo Oct 20, 2024

andylizf Oct 20, 2024 •

edited

Loading

	configs = common_utils.read_yaml_all(path)
	dag_name = None
	if set(configs[0].keys()) == {'name'}:
	dag_name = configs[0]['name']
	configs = configs[1:]
	elif len(configs) == 1:
	dag_name = configs[0].get('name')

	if len(configs) == 0:
	# YAML has only `name: xxx`. Still instantiate a task.
	configs = [{'name': dag_name}]

	def _get_dag_and_name(dag_yaml: str) -> Tuple['sky.Dag', str]:
	dag = dag_utils.load_chain_dag_from_yaml(dag_yaml)
	dag_name = dag.name
	assert dag_name is not None, dag
	return dag, dag_name

	name = task.name or f'task_{int(time.time())}'
	name = task.name if task.name is not None else f'task_{int(time.time())}'

Add dependencies feature for task dependency management #4067

Are you sure you want to change the base?

Add dependencies feature for task dependency management #4067

Conversation

andylizf commented Oct 10, 2024 • edited Loading

andylizf commented Oct 11, 2024

cblmemo left a comment

Choose a reason for hiding this comment

cblmemo commented Oct 14, 2024

cblmemo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andylizf Oct 17, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cblmemo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

This comment was marked as duplicate.

Choose a reason for hiding this comment

andylizf Oct 20, 2024 • edited Loading

Choose a reason for hiding this comment

Add `dependencies` feature for task dependency management #4067

Add `dependencies` feature for task dependency management #4067

andylizf commented Oct 10, 2024 •

edited

Loading

andylizf Oct 17, 2024 •

edited

Loading

andylizf Oct 20, 2024 •

edited

Loading