Use runtime env var for RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES

[praateekmahajan]

Currently its needed that an existing ray cluster is started with the following env varibale
RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=0

However if a user starts a ray cluster without this env variable, then the gpu allocation doesn't work out and error out here

cosmos-xenna/cosmos_xenna/ray_utils/stage_worker.py

Lines 376 to 381 in c62aa91

	if self._worker.allocation.gpus and gpu.get_num_gpus() == 0:
	raise RuntimeError(
	"Worker is a GPU worker, but no GPUs are available. This likely means that the ray cluster was not "
	"started with 'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=0'. Xenna needs this env variable to be set "
	"before cluster creation as it works around ray's gpu allocation mechanisms."
	)

I believe this PR should allow us to attach to existing cluster where that envvar wasn't set.

Earlier

RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=0 ray start --num-cpus 16 --num-gpus 4 --port 1234

RAY_ADDRESS=localhost:1234 python main.py

Now

ray start --num-cpus 16 --num-gpus 4 --port 1234

RAY_ADDRESS=localhost:1234 python main.py

TODO haven't tested really for edge cases

[pkuwangh]

I think you don't have to delete these?

[praateekmahajan]

We do! If we don't and this code path is executed first then these variables will "live" on the driver outside the ray session as well. So any subsequent (non xenna) ray pipeline if it does a ray.init(..) again it'll have these env variables.

run_xenna_pipeline()
run_ray_data_pipeline() # this will fail because `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=0`

[abhinavg4]

I think we need this even with Xenna. I'm not sure how but without this, the following case fails for me:

Reusing cluster
Start ray cluster using ray start --port 1234  (no env variables) and pass RAY_ADDRESS=localhost:1234 python xenna.py

Which means that even with this PR if we have these variables set, xenna would fail.

Start ray cluster using ray start --port 1234  (no env variables) and RAY_ADDRESS=localhost:1234 RAY_MAX_LIMIT_FROM_API_SERVER=40000  RAY_MAX_LIMIT_FROM_DATA_SOURCE=40000 python xenna.py

[abhinavg4]

Maybe we unset these variables here as a solution?

[pkuwangh]

For this

run_xenna_pipeline()
run_ray_data_pipeline() # this will fail because `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=0`

I think you better start a process for each call to make sure things are clean.
I had to do this otherwise ray metrics are quite messed up.

[praateekmahajan]

Thanks @abhinavg4 for catching that. So we need the RAY_MAX_LIMIT_* envvars when we start the cluster. The reason it wasn't caught before was because here we use a default value of 10k, which Ray Clusters anyway support

cosmos-xenna/cosmos_xenna/pipelines/private/monitoring.py

Line 137 in c62aa91

limit = int(os.environ.get("RAY_MAX_LIMIT_FROM_API_SERVER", "10000"))

However if we change that value to str(API_LIMIT) then my PR in the previous stage would've failed.

Essentially there are some cluster level env variables, driver level env variables and worker level env variables.
RAY_MAX_LIMIT needs to be a cluster level variable (and because of the monitoring.py also a driver level variable)
While NOSET_CUDA can be a worker level variable