-
Notifications
You must be signed in to change notification settings - Fork 5.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Core] Merge Driver/Job's runtime environment when it conflicts (#39208)
When the job and driver both specifies the runtime env, driver's runtime env is ignored. This is a confusing behavior and we'd like to merge instead. This PR merges the runtime envs, and if there's a conflict raises an exception. User can override the job's runtime env by driver's runtime env by specifying the env var.
- Loading branch information
Showing
10 changed files
with
426 additions
and
44 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
49 changes: 49 additions & 0 deletions
49
dashboard/modules/job/tests/subprocess_driver_scripts/driver_runtime_env_inheritance.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,49 @@ | ||
import argparse | ||
import sys | ||
import time | ||
|
||
import ray | ||
|
||
if __name__ == "__main__": | ||
parser = argparse.ArgumentParser(description="Dashboard agent.") | ||
parser.add_argument( | ||
"--conflict", | ||
type=str, | ||
) | ||
parser.add_argument( | ||
"--worker-process-setup-hook", | ||
type=str, | ||
) | ||
|
||
args = parser.parse_args() | ||
|
||
if args.worker_process_setup_hook: | ||
ray.init( | ||
runtime_env={ | ||
"worker_process_setup_hook": lambda: print( | ||
args.worker_process_setup_hook | ||
) | ||
} | ||
) | ||
|
||
@ray.remote | ||
def f(): | ||
pass | ||
|
||
ray.get(f.remote()) | ||
time.sleep(5) | ||
sys.exit(0) | ||
|
||
if args.conflict == "pip": | ||
ray.init(runtime_env={"pip": ["numpy"]}) | ||
print(ray._private.worker.global_worker.runtime_env) | ||
elif args.conflict == "env_vars": | ||
ray.init(runtime_env={"env_vars": {"A": "1"}}) | ||
print(ray._private.worker.global_worker.runtime_env) | ||
else: | ||
ray.init( | ||
runtime_env={ | ||
"env_vars": {"C": "1"}, | ||
} | ||
) | ||
print(ray._private.worker.global_worker.runtime_env) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,162 @@ | ||
import time | ||
import json | ||
import sys | ||
|
||
import pytest | ||
|
||
from ray.job_submission import JobSubmissionClient, JobStatus | ||
from ray.cluster_utils import Cluster | ||
from ray.dashboard.modules.job.tests.conftest import ( | ||
_driver_script_path, | ||
) | ||
|
||
|
||
def wait_until_status(client, job_id, status_to_wait_for, timeout_seconds=20): | ||
start = time.time() | ||
while time.time() - start <= timeout_seconds: | ||
status = client.get_job_status(job_id) | ||
print(f"status: {status}") | ||
if status in status_to_wait_for: | ||
return | ||
time.sleep(1) | ||
raise Exception | ||
|
||
|
||
def test_job_driver_inheritance(): | ||
try: | ||
c = Cluster() | ||
c.add_node(num_cpus=1) | ||
# If using a remote cluster, replace 127.0.0.1 with the head node's IP address. | ||
client = JobSubmissionClient("http://127.0.0.1:8265") | ||
driver_script_path = _driver_script_path("driver_runtime_env_inheritance.py") | ||
job_id = client.submit_job( | ||
entrypoint=f"python {driver_script_path}", | ||
runtime_env={ | ||
"env_vars": {"A": "1", "B": "2"}, | ||
"pip": ["requests"], | ||
}, | ||
) | ||
|
||
def wait(job_id): | ||
wait_until_status( | ||
client, | ||
job_id, | ||
{JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED}, | ||
timeout_seconds=60, | ||
) | ||
|
||
def get_runtime_env_from_logs(client, job_id): | ||
wait(job_id) | ||
logs = client.get_job_logs(job_id) | ||
print(logs) | ||
assert client.get_job_status(job_id) == JobStatus.SUCCEEDED | ||
return json.loads(logs.strip().split("\n")[-1]) | ||
|
||
# Test key is merged | ||
print("Test key merged") | ||
runtime_env = get_runtime_env_from_logs(client, job_id) | ||
assert runtime_env["env_vars"] == {"A": "1", "B": "2", "C": "1"} | ||
assert runtime_env["pip"] == {"packages": ["requests"], "pip_check": False} | ||
|
||
# Test worker process setuphook works. | ||
print("Test key setup hook") | ||
expected_str = "HELLOWORLD" | ||
job_id = client.submit_job( | ||
entrypoint=( | ||
f"python {driver_script_path} " | ||
f"--worker-process-setup-hook {expected_str}" | ||
), | ||
runtime_env={ | ||
"env_vars": {"A": "1", "B": "2"}, | ||
}, | ||
) | ||
wait(job_id) | ||
logs = client.get_job_logs(job_id) | ||
assert expected_str in logs | ||
|
||
# Test raise an exception upon key conflict | ||
print("Test conflicting pip") | ||
job_id = client.submit_job( | ||
entrypoint=f"python {driver_script_path} --conflict=pip", | ||
runtime_env={"pip": ["numpy"]}, | ||
) | ||
wait(job_id) | ||
status = client.get_job_status(job_id) | ||
logs = client.get_job_logs(job_id) | ||
assert status == JobStatus.FAILED | ||
assert "Failed to merge the Job's runtime env" in logs | ||
|
||
# Test raise an exception upon env var conflict | ||
print("Test conflicting env vars") | ||
job_id = client.submit_job( | ||
entrypoint=f"python {driver_script_path} --conflict=env_vars", | ||
runtime_env={ | ||
"env_vars": {"A": "1"}, | ||
}, | ||
) | ||
wait(job_id) | ||
status = client.get_job_status(job_id) | ||
logs = client.get_job_logs(job_id) | ||
assert status == JobStatus.FAILED | ||
assert "Failed to merge the Job's runtime env" in logs | ||
finally: | ||
c.shutdown() | ||
|
||
|
||
def test_job_driver_inheritance_override(monkeypatch): | ||
monkeypatch.setenv("RAY_OVERRIDE_JOB_RUNTIME_ENV", "1") | ||
|
||
try: | ||
c = Cluster() | ||
c.add_node(num_cpus=1) | ||
# If using a remote cluster, replace 127.0.0.1 with the head node's IP address. | ||
client = JobSubmissionClient("http://127.0.0.1:8265") | ||
driver_script_path = _driver_script_path("driver_runtime_env_inheritance.py") | ||
job_id = client.submit_job( | ||
entrypoint=f"python {driver_script_path}", | ||
runtime_env={ | ||
"env_vars": {"A": "1", "B": "2"}, | ||
"pip": ["requests"], | ||
}, | ||
) | ||
|
||
def wait(job_id): | ||
wait_until_status( | ||
client, | ||
job_id, | ||
{JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED}, | ||
timeout_seconds=60, | ||
) | ||
|
||
def get_runtime_env_from_logs(client, job_id): | ||
wait(job_id) | ||
logs = client.get_job_logs(job_id) | ||
print(logs) | ||
assert client.get_job_status(job_id) == JobStatus.SUCCEEDED | ||
return json.loads(logs.strip().split("\n")[-1]) | ||
|
||
# Test conflict resolution regular field | ||
job_id = client.submit_job( | ||
entrypoint=f"python {driver_script_path} --conflict=pip", | ||
runtime_env={"pip": ["torch"]}, | ||
) | ||
runtime_env = get_runtime_env_from_logs(client, job_id) | ||
print(runtime_env) | ||
assert runtime_env["pip"] == {"packages": ["numpy"], "pip_check": False} | ||
|
||
# Test raise an exception upon env var conflict | ||
job_id = client.submit_job( | ||
entrypoint=f"python {driver_script_path} --conflict=env_vars", | ||
runtime_env={ | ||
"env_vars": {"A": "2"}, | ||
}, | ||
) | ||
runtime_env = get_runtime_env_from_logs(client, job_id) | ||
print(runtime_env) | ||
assert runtime_env["env_vars"]["A"] == "1" | ||
finally: | ||
c.shutdown() | ||
|
||
|
||
if __name__ == "__main__": | ||
sys.exit(pytest.main(["-v", __file__])) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.