feat: add head server deps

Signed-off-by: Sugam Devare <sdevare@nvidia.com>

Author: sdevare-nv

nemo_gym/cli.py

@@ -14,6 +14,7 @@
 import asyncio
 import json
 import shlex
+import subprocess
 import tomllib
 from glob import glob
 from os import environ, makedirs
@@ -27,13 +28,14 @@
 import rich
 import uvicorn
 from devtools import pprint
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import DictConfig, OmegaConf, open_dict
 from pydantic import BaseModel, Field
 from tqdm.auto import tqdm
 
 from nemo_gym import PARENT_DIR
 from nemo_gym.config_types import BaseNeMoGymCLIConfig
 from nemo_gym.global_config import (
+    HEAD_SERVER_DEPS_KEY_NAME,
     NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME,
     NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME,
     NEMO_GYM_RESERVED_TOP_LEVEL_KEYS,
@@ -48,18 +50,38 @@
 )
 
 
-def _setup_env_command(dir_path: Path) -> str:  # pragma: no cover
+def _capture_head_server_dependencies(global_config_dict: DictConfig) -> None:  # pragma: no cover
+    try:
+        result = subprocess.run(
+            ["uv", "pip", "freeze", "--exclude-editable"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        head_server_deps = result.stdout
+    except Exception as e:
+        print(f"Warning: Could not capture head server dependencies: {e}")
+        head_server_deps = None
+
+    with open_dict(global_config_dict):
+        global_config_dict[HEAD_SERVER_DEPS_KEY_NAME] = head_server_deps
+
+
+def _setup_env_command(dir_path: Path, head_server_deps: Optional[str] = None) -> str:  # pragma: no cover
+    install_cmd = "uv pip install -r requirements.txt"
+    if head_server_deps:
+        install_cmd += f" --constraint <(cat << 'EOF'\n{head_server_deps}\nEOF\n)"
+
     return f"""cd {dir_path} \\
     && uv venv --allow-existing \\
     && source .venv/bin/activate \\
-    && uv pip install -r requirements.txt \\
+    && {install_cmd} \\
    """
 
 
 def _run_command(command: str, working_directory: Path) -> Popen:  # pragma: no cover
     custom_env = environ.copy()
     custom_env["PYTHONPATH"] = f"{working_directory.absolute()}:{custom_env.get('PYTHONPATH', '')}"
-    print(f"Executing command:\n{command}\n")
     return Popen(command, executable="/bin/bash", shell=True, env=custom_env)
 
 
@@ -114,6 +136,9 @@ class RunHelper:  # pragma: no cover
     def start(self, global_config_dict_parser_config: GlobalConfigDictParserConfig) -> None:
         global_config_dict = get_global_config_dict(global_config_dict_parser_config=global_config_dict_parser_config)
 
+        # Capture head server dependencies and store in global config dict
+        _capture_head_server_dependencies(global_config_dict)
+
         # Assume Nemo Gym Run is for a single agent.
         escaped_config_dict_yaml_str = shlex.quote(OmegaConf.to_yaml(global_config_dict))
 
@@ -149,7 +174,9 @@ def start(self, global_config_dict_parser_config: GlobalConfigDictParserConfig)
 
             dir_path = PARENT_DIR / Path(first_key, second_key)
 
-            command = f"""{_setup_env_command(dir_path)} \\
+            head_server_deps = global_config_dict.get(HEAD_SERVER_DEPS_KEY_NAME)
+
+            command = f"""{_setup_env_command(dir_path, head_server_deps)} \\
     && {NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME}={escaped_config_dict_yaml_str} \\
     {NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME}={shlex.quote(top_level_path)} \\
     python {str(entrypoint_fpath)}"""

nemo_gym/global_config.py

@@ -35,11 +35,13 @@
 ENTRYPOINT_KEY_NAME = "entrypoint"
 DEFAULT_HOST_KEY_NAME = "default_host"
 HEAD_SERVER_KEY_NAME = "head_server"
+HEAD_SERVER_DEPS_KEY_NAME = "head_server_deps"
 NEMO_GYM_RESERVED_TOP_LEVEL_KEYS = [
     CONFIG_PATHS_KEY_NAME,
     ENTRYPOINT_KEY_NAME,
     DEFAULT_HOST_KEY_NAME,
     HEAD_SERVER_KEY_NAME,
+    HEAD_SERVER_DEPS_KEY_NAME,
 ]
 
 POLICY_BASE_URL_KEY_NAME = "policy_base_url"