Skip to content

Commit

Permalink
Speed up launch --async and fix jobs launch
Browse files Browse the repository at this point in the history
  • Loading branch information
Michaelvll committed Aug 12, 2024
1 parent 68c0548 commit c04ad87
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 14 deletions.
2 changes: 1 addition & 1 deletion sky/api/requests/payloads.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ class JobStatusBody(RequestBody):

class JobsLaunchBody(RequestBody):
task: str
name: str
name: Optional[str]
detach_run: bool
retry_until_up: bool

Expand Down
23 changes: 12 additions & 11 deletions sky/api/sdk.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,19 +167,20 @@ def launch(

dag = api_common.upload_mounts_to_api_server(task)

cluster_status = None
request_id = status([cluster_name])
clusters = get(request_id)
if not clusters:
# Show the optimize log before the prompt if the cluster does not exist.
request_id = optimize(dag)
stream_and_get(request_id)
else:
cluster_record = clusters[0]
cluster_status = cluster_record['status']

confirm_shown = False

if need_confirmation:
cluster_status = None
request_id = status([cluster_name])
clusters = get(request_id)
if not clusters:
# Show the optimize log before the prompt if the cluster does not exist.
request_id = optimize(dag)
stream_and_get(request_id)
else:
cluster_record = clusters[0]
cluster_status = cluster_record['status']

# Prompt if (1) --cluster is None, or (2) cluster doesn't exist, or (3)
# it exists but is STOPPED.
prompt = None
Expand Down
4 changes: 2 additions & 2 deletions sky/jobs/api/sdk.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,9 @@ def launch(
dag_utils.dump_chain_dag_to_yaml(dag, f.name)
dag_str = f.read()

request_id = sdk.optimize(dag)
sdk.stream_and_get(request_id)
if need_confirmation:
request_id = sdk.optimize(dag)
sdk.stream_and_get(request_id)
prompt = f'Launching a managed job {dag.name!r}. Proceed?'
if prompt is not None:
click.confirm(prompt, default=True, abort=True, show_default=True)
Expand Down

0 comments on commit c04ad87

Please sign in to comment.