CLI help and command help; misc improvements (#229)

Signed-off-by: Brian Yu <bxyu@nvidia.com>

Author: bxyu-nvidia

nemo_gym/cli.py

@@ -14,6 +14,7 @@
 import asyncio
 import json
 import shlex
+import tomllib
 from glob import glob
 from os import environ, makedirs
 from os.path import exists
@@ -23,13 +24,15 @@
 from time import sleep
 from typing import Dict, List, Optional
 
+import rich
 import uvicorn
 from devtools import pprint
 from omegaconf import DictConfig, OmegaConf
-from pydantic import BaseModel
+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 (
     NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME,
     NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME,
@@ -60,23 +63,32 @@ def _run_command(command: str, working_directory: Path) -> Popen:  # pragma: no
     return Popen(command, executable="/bin/bash", shell=True, env=custom_env)
 
 
-class RunConfig(BaseModel):
-    entrypoint: str
+class RunConfig(BaseNeMoGymCLIConfig):
+    entrypoint: str = Field(
+        description="Entrypoint for this command. This must be a relative path with 2 parts. Should look something like `responses_api_agents/simple_agent`."
+    )
 
 
 class TestConfig(RunConfig):
-    should_validate_data: bool = False
+    should_validate_data: bool = Field(
+        default=False,
+        description="Whether or not to validate the example data (examples, metrics, rollouts, etc) for this server.",
+    )
 
-    dir_path: Path = None  # initialized in model_post_init
+    _dir_path: Path  # initialized in model_post_init
 
     def model_post_init(self, context):  # pragma: no cover
         # TODO: This currently only handles relative entrypoints. Later on we can resolve the absolute path.
-        self.dir_path = Path(self.entrypoint)
+        self._dir_path = Path(self.entrypoint)
         assert not self.dir_path.is_absolute()
         assert len(self.dir_path.parts) == 2
 
         return super().model_post_init(context)
 
+    @property
+    def dir_path(self) -> Path:
+        return self._dir_path
+
 
 class ServerInstanceDisplayConfig(BaseModel):
     process_name: str
@@ -274,6 +286,10 @@ def check_http_server_statuses(self) -> List[ServerStatus]:
 def run(
     global_config_dict_parser_config: Optional[GlobalConfigDictParserConfig] = None,
 ):  # pragma: no cover
+    global_config_dict = get_global_config_dict(global_config_dict_parser_config=global_config_dict_parser_config)
+    # Just here for help
+    BaseNeMoGymCLIConfig.model_validate(global_config_dict)
+
     rh = RunHelper()
     rh.start(global_config_dict_parser_config)
     rh.run_forever()
@@ -386,8 +402,11 @@ def _format_pct(count: int, total: int) -> str:  # pragma: no cover
     return f"{count} / {total} ({100 * count / total:.2f}%)"
 
 
-class TestAllConfig(BaseModel):
-    fail_on_total_and_test_mismatch: bool = False
+class TestAllConfig(BaseNeMoGymCLIConfig):
+    fail_on_total_and_test_mismatch: bool = Field(
+        default=False,
+        description="There may be situations where there are an un-equal number of servers that exist vs have tests. This flag will fail the test job if this mismatch exists.",
+    )
 
 
 def test_all():  # pragma: no cover
@@ -474,6 +493,10 @@ def test_all():  # pragma: no cover
 
 
 def dev_test():  # pragma: no cover
+    global_config_dict = get_global_config_dict()
+    # Just here for help
+    BaseNeMoGymCLIConfig.model_validate(global_config_dict)
+
     proc = Popen("pytest --cov=. --durations=10", shell=True)
     exit(proc.wait())
 
@@ -592,4 +615,28 @@ def init_resources_server():  # pragma: no cover
 
 def dump_config():  # pragma: no cover
     global_config_dict = get_global_config_dict()
+    # Just here for help
+    BaseNeMoGymCLIConfig.model_validate(global_config_dict)
+
     print(OmegaConf.to_yaml(global_config_dict, resolve=True))
+
+
+def display_help():  # pragma: no cover
+    global_config_dict = get_global_config_dict()
+    # Just here for help
+    BaseNeMoGymCLIConfig.model_validate(global_config_dict)
+
+    pyproject_path = Path(PARENT_DIR) / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        pyproject_data = tomllib.load(f)
+
+    project_scripts = pyproject_data["project"]["scripts"]
+    rich.print("""Run a command with `+h=true` or `+help=true` to see more detailed information!
+
+[bold]Available CLI scripts[/bold]
+-----------------""")
+    for script in project_scripts:
+        if not script.startswith("ng_"):
+            continue
+
+        print(script)

nemo_gym/config_types.py

@@ -14,6 +14,7 @@
 from enum import Enum
 from typing import Any, ClassVar, Dict, List, Literal, Optional, Union
 
+import rich
 from omegaconf import DictConfig, OmegaConf
 from pydantic import (
     BaseModel,
@@ -23,6 +24,60 @@
     ValidationError,
     model_validator,
 )
+from rich.text import Text
+
+
+########################################
+# Base CLI configs
+########################################
+
+
+class BaseNeMoGymCLIConfig(BaseModel):
+    @model_validator(mode="before")
+    @classmethod
+    def pre_process(cls, data):
+        if not (data.get("h") or data.get("help")):
+            return data
+
+        rich.print(f"""Displaying help for [bold]{cls.__name__}[/bold]
+""")
+        # We use __doc__ directly here since inspect.getdoc will inherit the doc from parent classes.
+        class_doc = cls.__doc__
+        if class_doc:
+            rich.print(f"""[bold]Description[/bold]
+-----------
+{class_doc.strip()}
+""")
+
+        fields = cls.model_fields.items()
+        if fields:
+            rich.print("""[bold]Parameters[/bold]
+----------""")
+
+            prefixes: List[Text] = []
+            suffixes: List[Text] = []
+            for field_name, field in fields:
+                description_str = field.description if field.description else ""
+
+                # Not sure if there is a better way to get this annotation_str, e.g. using typing.get_args or typing.get_origin
+                annotation_str = (
+                    field.annotation.__name__ if isinstance(field.annotation, type) else str(field.annotation)
+                )
+                annotation_str = annotation_str.replace("typing.", "")
+
+                prefixes.append(Text.from_markup(f"- [blue]{field_name}[/blue] [yellow]({annotation_str})[/yellow]"))
+                suffixes.append(description_str)
+
+            max_prefix_length = max(map(len, prefixes))
+            ljust_length = max_prefix_length + 3
+            for prefix, suffix in zip(prefixes, suffixes):
+                prefix.align("left", ljust_length)
+                rich.print(prefix + suffix)
+        else:
+            print("There are no arguments to this CLI command!")
+
+        # Exit after help is printed.
+        exit()
 
 
 ########################################
@@ -63,10 +118,14 @@ def is_server_ref(config_dict: DictConfig) -> Optional[ServerRef]:
 ########################################
 
 
-class UploadJsonlDatasetGitlabConfig(BaseModel):
-    dataset_name: str
-    version: str  # Must be x.x.x
-    input_jsonl_fpath: str
+class UploadJsonlDatasetGitlabConfig(BaseNeMoGymCLIConfig):
+    """
+    Upload a local jsonl dataset artifact to Gitlab.
+    """
+
+    dataset_name: str = Field(description="The dataset name.")
+    version: str = Field(description="The version of this dataset. Must be in the format `x.x.x`.")
+    input_jsonl_fpath: str = Field(description="Path to the jsonl file to upload.")
 
 
 class JsonlDatasetGitlabIdentifer(BaseModel):
@@ -75,8 +134,11 @@ class JsonlDatasetGitlabIdentifer(BaseModel):
     artifact_fpath: str
 
 
-class DownloadJsonlDatasetGitlabConfig(JsonlDatasetGitlabIdentifer):
-    output_fpath: str
+class DownloadJsonlDatasetGitlabConfig(JsonlDatasetGitlabIdentifer, BaseNeMoGymCLIConfig):
+    dataset_name: str = Field(description="The dataset name.")
+    version: str = Field(description="The version of this dataset. Must be in the format `x.x.x`.")
+    artifact_fpath: str = Field(description="The filepath to the artifact to download.")
+    output_fpath: str = Field(description="Where to save the downloaded dataset.")
 
 
 DatasetType = Union[Literal["train"], Literal["validation"], Literal["example"]]

nemo_gym/dataset_viewer.py

@@ -23,10 +23,11 @@
     ResponseInputItemParam,
     ResponseReasoningItemParam,
 )
-from pydantic import BaseModel, ConfigDict
+from pydantic import ConfigDict, Field
 from tqdm.auto import tqdm
 
 from nemo_gym.base_resources_server import BaseVerifyResponse
+from nemo_gym.config_types import BaseNeMoGymCLIConfig
 from nemo_gym.server_utils import get_global_config_dict
 from nemo_gym.train_data_utils import (
     DatasetMetrics,
@@ -203,8 +204,8 @@ def extra_info_to_messages(d: DatasetViewerVerifyResponse) -> List[ChatMessage]:
     return messages
 
 
-class JsonlDatasetViewerConfig(BaseModel):
-    jsonl_fpath: str
+class JsonlDatasetViewerConfig(BaseNeMoGymCLIConfig):
+    jsonl_fpath: str = Field(description="Filepath to a local jsonl file to view.")
 
 
 def get_aggregate_metrics(data: List[DatasetViewerVerifyResponse]) -> Dict[str, Any]:

nemo_gym/openai_utils.py

@@ -422,7 +422,7 @@ class NeMoGymChatCompletionCreateParamsNonStreaming(BaseModel):
 ########################################
 
 
-class NeMoGymAsyncOpenAI(BaseModel):
+class NeMoGymAsyncOpenAI(BaseModel):  # pragma: no cover
     """This is just a stub class that wraps around aiohttp"""
 
     base_url: str

nemo_gym/rollout_collection.py

@@ -22,7 +22,7 @@
 from pydantic import BaseModel, Field
 from tqdm.asyncio import tqdm
 
-from nemo_gym.config_types import BaseServerConfig
+from nemo_gym.config_types import BaseNeMoGymCLIConfig, BaseServerConfig
 from nemo_gym.server_utils import (
     GlobalAIOHTTPAsyncClientConfig,
     ServerClient,
@@ -32,14 +32,30 @@
 )
 
 
-class RolloutCollectionConfig(BaseModel):
-    agent_name: str
-    input_jsonl_fpath: str
-    output_jsonl_fpath: str
-    limit: Optional[int] = None
-    num_repeats: Optional[int] = None
-    num_samples_in_parallel: Optional[int] = None
-    responses_create_params: Dict[str, Any] = Field(default_factory=dict)
+class RolloutCollectionConfig(BaseNeMoGymCLIConfig):
+    """
+    Perform a batch of rollout collection.
+    """
+
+    agent_name: str = Field(description="The agent to collect rollouts from.")
+    input_jsonl_fpath: str = Field(
+        description="The input data source to use to collect rollouts, in the form of a file path to a jsonl file."
+    )
+    output_jsonl_fpath: str = Field(description="The output data jsonl file path.")
+    limit: Optional[int] = Field(
+        default=None, description="Maximum number of examples to load and take from the input dataset."
+    )
+    num_repeats: Optional[int] = Field(
+        default=None,
+        description="The number of times to repeat each example to run. Useful if you want to calculate mean@k e.g. mean@4 or mean@16.",
+    )
+    num_samples_in_parallel: Optional[int] = Field(
+        default=None, description="Limit the number of concurrent samples running at once."
+    )
+    responses_create_params: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Overrides for the responses_create_params e.g. temperature, max_output_tokens, etc.",
+    )
 
 
 class RolloutCollectionHelper(BaseModel):  # pragma: no cover

nemo_gym/train_data_utils.py

@@ -29,6 +29,7 @@
 from nemo_gym.config_types import (
     AGENT_REF_KEY,
     AgentServerRef,
+    BaseNeMoGymCLIConfig,
     DatasetConfig,
     DatasetType,
     DownloadJsonlDatasetGitlabConfig,
@@ -42,10 +43,15 @@
 )
 
 
-class TrainDataProcessorConfig(BaseModel):
-    output_dirpath: str
-    mode: Union[Literal["train_preparation"], Literal["example_validation"]]
-    should_download: bool = False
+class TrainDataProcessorConfig(BaseNeMoGymCLIConfig):
+    output_dirpath: str = Field(description="Path to the directory to save the outputs.")
+    mode: Union[Literal["train_preparation"], Literal["example_validation"]] = Field(
+        description="Whether to do train_preparation or example_validation."
+    )
+    should_download: bool = Field(
+        default=False,
+        description="Whether or not to download missing datasets. By default, no datasets will be downloaded.",
+    )
 
     @property
     def in_scope_dataset_types(self) -> List[DatasetType]:

pyproject.toml

@@ -246,6 +246,10 @@ ng_prepare_data = "nemo_gym.train_data_utils:prepare_data"
 nemo_gym_dump_config = "nemo_gym.cli:dump_config"
 ng_dump_config = "nemo_gym.cli:dump_config"
 
+# Display help
+nemo_gym_help = "nemo_gym.cli:display_help"
+ng_help = "nemo_gym.cli:display_help"
+
 
 [tool.setuptools.packages.find]
 where = ["."]

tests/unit_tests/test_cli.py

@@ -11,10 +11,48 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import sys
+import tomllib
+from importlib import import_module
+from io import StringIO
+from pathlib import Path
+
+from omegaconf import OmegaConf
+from pytest import MonkeyPatch, raises
+
+import nemo_gym.global_config
+from nemo_gym import PARENT_DIR
 from nemo_gym.cli import RunConfig
 
 
 # TODO: Eventually we want to add more tests to ensure that the CLI flows do not break
 class TestCLI:
     def test_sanity(self) -> None:
         RunConfig(entrypoint="", name="")
+
+    def test_pyproject_scripts(self) -> None:
+        pyproject_path = Path(PARENT_DIR) / "pyproject.toml"
+        with pyproject_path.open("rb") as f:
+            pyproject_data = tomllib.load(f)
+
+        project_scripts = pyproject_data["project"]["scripts"]
+
+        for script_name, import_path in project_scripts.items():
+            # Dedupe `nemo_gym_*` from `ng_*` commands
+            if not script_name.startswith("ng_"):
+                continue
+
+            # We only test `+h=true` and not `+help=true`
+            print(f"Running `{script_name} +h=true`")
+
+            module, fn = import_path.split(":")
+            fn = getattr(import_module(module), fn)
+
+            with MonkeyPatch.context() as mp:
+                mp.setattr(nemo_gym.global_config, "_GLOBAL_CONFIG_DICT", OmegaConf.create({"h": True}))
+
+                text_trap = StringIO()
+                mp.setattr(sys, "stdout", text_trap)
+
+                with raises(SystemExit):
+                    fn()