DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition #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

unroll load balancer docs

* observability docs * comments

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

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

skypilot-org#4084) * fix narrow window and show log path during exception * format * format

* refactor * lint * refactor, dataclass * refactor, dataclass * refactor * lint

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

* [Performance] Use new custom image to create GCP GPU VMs * update image tags for both CPU and GPU * always generate .sky/python_path --------- Co-authored-by: Yika Luo <yikaluo@Yikas-MacBook-Pro.local>

* Add H100 mega support on GCP * fix for some other regions * format * fix resource type * fix catalog fetching

sky/utils/common_utils.py

Co-authored-by: Tian Xia <cblmemo@gmail.com>

…or output (skypilot-org#4111) * add `ux_utils.print_exception_no_traceback()` for cleaner error output * Empty commit * remove unnecessary with block

…ceback()` wrappers (skypilot-org#4130) fix unnecessary with block for returning

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

deepspeed kubernetes fixes

* Bug fix for sky config file path resolution. * format * [OCI] Bug fix for image_id in Task YAML * [OCI]: Support more OS types (esp. oraclelinux) in addition to ubuntu. * format * Disable system firewall * Bug fix for validation of the Marketplace images * Update sky/clouds/oci.py Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com> * Update sky/clouds/oci.py Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com> * variable/function naming * address review comments: not to change the service_catalog api. call oci_catalog directly for get os type for a image. * Update sky/clouds/oci.py Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com> * Update sky/clouds/oci.py Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com> * Update sky/clouds/oci.py Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com> * address review comments --------- Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com>

cblmemo · 2024-10-21T17:39:07Z

@andylizf Why did we change the base branch back?

cblmemo

Thanks @andylizf ! Final nits. Please update the base branch and it should be ready to go!

sky/dag.py

sky/utils/common_utils.py

sky/utils/dag_utils.py

andylizf · 2024-10-21T18:17:13Z

@andylizf Why did we change the base branch back?

Because only the master branch has GitHub Actions. 😂I'll switch it back now.

Co-authored-by: Tian Xia <cblmemo@gmail.com>

cblmemo · 2024-10-21T18:54:41Z

@andylizf Why did we change the base branch back?

Because only the master branch has GitHub Actions. 😂I'll switch it back now.

Actually, considering this, should we create a branch on your repo and run those actions?

andylizf · 2024-10-21T19:56:02Z

@andylizf Why did we change the base branch back?

Because only the master branch has GitHub Actions. 😂I'll switch it back now.

Actually, considering this, should we create a branch on your repo and run those actions?

Good idea. You mean creating a branch like advanced-dag in my repo?

cblmemo · 2024-10-21T19:57:18Z

@andylizf Why did we change the base branch back?

Because only the master branch has GitHub Actions. 😂I'll switch it back now.

Actually, considering this, should we create a branch on your repo and run those actions?

Good idea. You mean creating a branch like advanced-dag in my repo?

yes, and deploy any necessary actions in the branch.

andylizf · 2024-10-21T20:32:01Z

@andylizf Why did we change the base branch back?

Because only the master branch has GitHub Actions. 😂I'll switch it back now.

Actually, considering this, should we create a branch on your repo and run those actions?

Good idea. You mean creating a branch like advanced-dag in my repo?

yes, and deploy any necessary actions in the branch.

But what reminds me is that GitHub handles actions in a bit of a weird way. Even if I add the GitHub Actions config directly to my PR, they can still run as expected, just like they're running now. So there's no need to create a PR in my repo. I can simply update the GitHub Actions in this PR to run when merging into advanced-dag. Once the PR detects it's targeting advanced-dag, the actions will trigger automatically.

cblmemo · 2024-10-21T20:34:47Z

@andylizf Why did we change the base branch back?

Because only the master branch has GitHub Actions. 😂I'll switch it back now.

Actually, considering this, should we create a branch on your repo and run those actions?

Good idea. You mean creating a branch like advanced-dag in my repo?

yes, and deploy any necessary actions in the branch.

But what reminds me is that GitHub handles actions in a bit of a weird way. Even if I add the GitHub Actions config directly to my PR, they can still run as expected, just like they're running now. So there's no need to create a PR in my repo. I can simply update the GitHub Actions in this PR to run when merging into advanced-dag. Once the PR detects it's targeting advanced-dag, the actions will trigger automatically.

Sounds great. Lets fix the error in the CI and it should be ready to go ;)

* 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 * Add dependencies feature for task dependency management (#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 * fix(jobs): type errors * refactor: `_update_failed_task_state` for unified error handling * refactor: separate finally block for a meaningful name * feat: simple parallel execution support * Apply suggestions from code review Co-authored-by: Tian Xia <cblmemo@gmail.com> * change wording all to up/downstream style * Add unique suffix to task names, fallback to timestamp if unnamed * Unify handling of single and multiple tasks without dependencies * Refactor tasks initialization: use list comprehension and fail fast * Fix remove task dependency description: upstream, not downstream Co-authored-by: Tian Xia <cblmemo@gmail.com> * Remove duplicated `self.edges`, use nx api instead * Revert "Add unique suffix to task names, fallback to timestamp if unnamed" Otherwise, users can not refer to the task by name in the DAG. This reverts commit 8486352. * comment the checking used as upstream logic * remove is_chain restriction in jobs launch * Change Static layered parallelism to dynamic queue parallelism * fix: add blocked_tasks set * refactor and update canceled tasks in database * format: some types are unsubscriptable * add some logging * Fooled again by the silly finally block, mistakenly thought it only runs on errors! * feat: redirect logging for each thread to a separate file to prevent interleaving output * fix due to reviwer's suggestions and some nits * chore: remove some debugging info * partially revert "update canceled tasks in database", view a task as a whole * add some comments to inform the future of thread-level redirector * cancell all tasks when a task failed * combine 3 sets to 1 dict * add some comments to illustrate `_try_add_successors_to_queue` only queue each node once * Cancel all non-running tasks when cancelling a job. Left a TODO for future cancellation policy discussion. * make those logging files and dir * refactor: stream_logs_by_id * refactor and format * provide a cli to print run.log of a certain subtask * format * add some comments and checks * clearer log * hint users to see the log * add some comments explaining why skip convetional code path * add with exception_no_traceback * fix os.makedirs, add expandusr first * fix: diamond named diamond * fix: use managed_job_id as dirname instead of runtimestamp * fix: make strategy_executor local to prevent race conditions * feat: implement stream logs for a task command * fix: forgot to add parentheses around the addition * chore: log indent * fix: early checking and better comments * refactor: reuse `stream_logs_by_id` for 2nd part run log tailing * fix: deal with tasks already finished * refactor: a generalized `follow_logs` and implement output checking on its basis * format * add return type annotations Co-authored-by: Tian Xia <cblmemo@gmail.com> * add all annotations and rename `task_queue` to `ready_queue` * Apply suggestions from code review Co-authored-by: Tian Xia <cblmemo@gmail.com> * revert one-liner conditional for better readibility * use tuple unpacking instead of `[1]` Co-authored-by: Tian Xia <cblmemo@gmail.com> * Apply suggestions from code review Co-authored-by: Tian Xia <cblmemo@gmail.com> * chore: fix more type annotations * refactor: clearer responsibility split between `_ThreadAwareOutput` and `RedirectOutputForThread` * feat: ensures `open` returns a `TextIO` * fix: unbounded `is_dag_chain` * fix: typing * revert type subscription, stuck by pylint * fix: use default false `cancel_all` argument to avoid breaking existing calls * fix: add explicit write/flush methods to _ThreadAwareOutput for logging redirection --------- Co-authored-by: Tian Xia <cblmemo@gmail.com>

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

romilbhardwaj and others added 2 commits October 14, 2024 17:57

[docs] Unroll k8s internal load balancer docs (skypilot-org#4083)

1ff843f

unroll load balancer docs

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

[docs] sky status --kubernetes docs (skypilot-org#4064)

a0243e5

* observability docs * comments

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

Michaelvll and others added 3 commits October 14, 2024 22:48

[UX] Show log after failure and fix the color issue with narrow window (

9243113

skypilot-org#4084) * fix narrow window and show log path during exception * format * format

[k8s] sky status --k8s refactor (skypilot-org#4079)

a4e2fcd

* refactor * lint * refactor, dataclass * refactor, dataclass * refactor * lint

add comments for add_edge

ff528a5

cblmemo reviewed Oct 15, 2024

View reviewed changes

andylizf and others added 6 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

[Performance] Use new GCP custom images (skypilot-org#4027)

53380e2

* [Performance] Use new custom image to create GCP GPU VMs * update image tags for both CPU and GPU * always generate .sky/python_path --------- Co-authored-by: Yika Luo <yikaluo@Yikas-MacBook-Pro.local>

[GCP] Add H100 mega (skypilot-org#4099)

724e97e

* Add H100 mega support on GCP * fix for some other regions * format * fix resource type * fix catalog fetching

cblmemo reviewed Oct 20, 2024

View reviewed changes

sky/utils/common_utils.py Outdated Show resolved Hide resolved

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

Fix remove task dependency description: upstream, not downstream

65d0bdd

Co-authored-by: Tian Xia <cblmemo@gmail.com>

Remove duplicated self.edges, use nx api instead

28b6482

[Serve] Add ux_utils.print_exception_no_traceback() for cleaner err…

03c2adb

…or output (skypilot-org#4111) * add `ux_utils.print_exception_no_traceback()` for cleaner error output * Empty commit * remove unnecessary with block

Partially revert: Remove unnecessary `ux_utils.print_exception_no_tra…

985df83

…ceback()` wrappers (skypilot-org#4130) fix unnecessary with block for returning

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

romilbhardwaj and others added 3 commits October 21, 2024 01:26

[examples] Deepspeed fixes + k8s support (skypilot-org#4124)

067a0a3

deepspeed kubernetes fixes

Empty commit

635c564

cblmemo reviewed Oct 21, 2024

View reviewed changes

sky/dag.py Outdated Show resolved Hide resolved

sky/dag.py Outdated Show resolved Hide resolved

sky/dag.py Outdated Show resolved Hide resolved

sky/utils/common_utils.py Outdated Show resolved Hide resolved

sky/utils/dag_utils.py Show resolved Hide resolved

andylizf and others added 3 commits October 21, 2024 11:19

Apply suggestions from code review

b63eb09

Co-authored-by: Tian Xia <cblmemo@gmail.com>

fix: typing.cast

95b0936

add TODOs for future function migration

32a7cc6

andylizf changed the base branch from master to advanced-dag October 21, 2024 18:49

andylizf changed the title ~~Add dependencies feature for task dependency management~~ DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition Oct 21, 2024

andylizf mentioned this pull request Oct 21, 2024

DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition andylizf/skypilot#1

Closed

andylizf added 2 commits October 21, 2024 13:23

remove dependencies wording to reduce ambiguity

b7a8948

temporarily add github actions

686bba0

Merge remote-tracking branch 'upstream/master' into job-dag

8992f23

cblmemo merged commit 7d93b75 into skypilot-org:advanced-dag Oct 21, 2024
18 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition #4067

DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition #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 commented Oct 21, 2024

cblmemo left a comment

andylizf commented Oct 21, 2024

cblmemo commented Oct 21, 2024

andylizf commented Oct 21, 2024

cblmemo commented Oct 21, 2024

andylizf commented Oct 21, 2024

cblmemo commented Oct 21, 2024

	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

DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition #4067

DAG API Enhancements: Introducing Downstream Task Parsing and Explicit Flow Definition #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

cblmemo Oct 15, 2024

Choose a reason for hiding this comment

andylizf Oct 15, 2024

Choose a reason for hiding this comment

cblmemo Oct 16, 2024

Choose a reason for hiding this comment

andylizf Oct 17, 2024

Choose a reason for hiding this comment

cblmemo Oct 17, 2024

Choose a reason for hiding this comment

andylizf Oct 17, 2024 • edited Loading

Choose a reason for hiding this comment

andylizf Oct 18, 2024

Choose a reason for hiding this comment

cblmemo Oct 19, 2024

Choose a reason for hiding this comment

cblmemo commented Oct 21, 2024

cblmemo left a comment

Choose a reason for hiding this comment

andylizf commented Oct 21, 2024

cblmemo commented Oct 21, 2024

andylizf commented Oct 21, 2024

cblmemo commented Oct 21, 2024

andylizf commented Oct 21, 2024

cblmemo commented Oct 21, 2024

andylizf commented Oct 10, 2024 •

edited

Loading

andylizf Oct 17, 2024 •

edited

Loading