Description

This PRs goal is to ensure we are able to run different executor pipeline in the same python process e.g. Xenna pipeline followed by Ray Data pipeline (and vice versa). It is currently not possible to do so without shutdown because ray preserve envvars from the previous sessions.

If https://github.com/nvidia-cosmos/cosmos-xenna/pull/6/files does not merge, even then this PR works as expected as Xenna sets os.environ[RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES] which means if Ray Data runs after Xenna then Ray Data will not set CUDA_VISIBLE_DEVICES i.e. use gpu_id=0. We solve that by forcefully setting RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES="" inside envvars while doing ray.init inside Ray Data.

Workarounds needed

However because we introduce ray.shutdown() now

1. running two subsequent Ray Data pipelines will fail due to [data] Failed to submit task to actor after ray.shutdown() and re-ray.init() in data pipeline for an existing cluster ray-project/ray#54841
2. Whenever we create a ray cluster inside the process, it'll be killed by the first pipeline. To avoid that we can use subprocess.run to create our cluster

See the following for clearer understanding

1. You start a Ray Cluster outside the process and run pipelines in the same process

ray start --head --port 1234

os.environ["RAY_ADDRESS"] = "localhost:1234"
ray.init(ignore_reinit=True) # this will connect to 1234
run_pipeline("xenna")
ray.shutdown()

ray.init(ignore_reinit=True) # this will connect to 1234
run_pipeline("ray data")
ray.shutdown()

2. You start a ray cluster inside the process and run pipelines in the same process

ray.init(ignore_reinit=True) # this will start a new cluster
run_pipeline("xenna")
ray.shutdown() # this will kill the ray cluster


ray.init(ignore_reinit=True) # this will start a new cluster
run_pipeline("xenna")
ray.shutdown()

3. You start a ray cluster explicitly inside the new process

ray.init(ignore_reinit=True) # this will start a new cluster A

ray.init(ignore_reinit=True) # this will connect to cluster A
run_pipeline("xenna")
ray.shutdown() # this will kill the cluster A


ray.init(ignore_reinit=True) # this will start cluster B
run_pipeline("xenna")
ray.shutdown()

4. You use subprocess to start a ray cluster

subprocess.run(["ray", "start", "--head", "--port", "1234"]) # this will start a new cluster A

ray.init(ignore_reinit=True) # this will connect to cluster A @ 1234
run_pipeline("xenna")
ray.shutdown() # this will exit session

ray.init(ignore_reinit=True) # this will connect to cluster A @ 1234
run_pipeline("xenna")
ray.shutdown()

Usage

# Add snippet demonstrating usage

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.