diff --git a/nemo_rl/distributed/worker_groups.py b/nemo_rl/distributed/worker_groups.py index b94fbb37510..f1944295122 100644 --- a/nemo_rl/distributed/worker_groups.py +++ b/nemo_rl/distributed/worker_groups.py @@ -27,7 +27,10 @@ ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.distributed.worker_group_utils import recursive_merge_options -from nemo_rl.utils.venvs import create_local_venv_on_each_node +from nemo_rl.utils.venvs import ( + create_local_venv_on_each_node, + patch_transformers_module_dir, +) @dataclass @@ -529,6 +532,7 @@ def _create_workers_from_bundle_indices( } runtime_env["env_vars"]["VIRTUAL_ENV"] = py_executable runtime_env["env_vars"]["UV_PROJECT_ENVIRONMENT"] = py_executable + patch_transformers_module_dir(runtime_env["env_vars"]) extra_options = {"runtime_env": runtime_env, "name": name} diff --git a/nemo_rl/utils/venvs.py b/nemo_rl/utils/venvs.py index 2d9c8018e20..117a54409df 100644 --- a/nemo_rl/utils/venvs.py +++ b/nemo_rl/utils/venvs.py @@ -187,3 +187,18 @@ def create_local_venv_on_each_node(py_executable: str, venv_name: str): ray.util.remove_placement_group(pg) # Return mapping from node IP to venv python path return paths[0] + + +# Need to set PYTHONPATH to include transformers downloaded modules. +# Assuming the cache directory is the same cross venvs. +def patch_transformers_module_dir(env_vars: dict[str, str]): + from transformers.utils.hub import TRANSFORMERS_CACHE + + module_dir = os.path.join(TRANSFORMERS_CACHE, "..", "modules") + assert module_dir is not None, "TRANSFORMERS_CACHE should exist." + if "PYTHONPATH" not in env_vars: + env_vars["PYTHONPATH"] = module_dir + else: + env_vars["PYTHONPATH"] = f"{module_dir}:{env_vars['PYTHONPATH']}" + + return env_vars