feat: introducing the experimental package and refactoring test str…

…ucture (#433) Signed-off-by: Terry Kong <[email protected]> Signed-off-by: NeMo-Aligner CI <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
NVIDIA · Jan 22, 2025 · 502ebde · 502ebde
1 parent 5f4f6d6
commit 502ebde
Show file tree

Hide file tree

Showing 20 changed files with 64 additions and 6 deletions.
diff --git a/tests/conftest.py → conftest.py b/tests/conftest.py → conftest.py
@@ -22,8 +22,8 @@
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
 from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
 from nemo_aligner.models.nlp.gpt.megatron_gpt_ppo_actor import MegatronGPTActorModel
+from nemo_aligner.testing.utils import Utils
 from nemo_aligner.utils.train_script_utils import init_distributed, resolve_and_create_trainer
-from tests.test_mcore_utilities import Utils
 
 dir_path = os.path.dirname(os.path.abspath(__file__))
 # TODO: This file exists because in cases where TRTLLM MPI communicators are involved,
@@ -67,7 +67,7 @@ def run_only_on_device_fixture(request, device):
 
 @pytest.fixture
 def init_model_parallel():
-    from tests.test_mcore_utilities import Utils
+    from nemo_aligner.testing.utils import Utils
 
     def initialize(*args, **kwargs):
         Utils.initialize_model_parallel(*args, **kwargs)

diff --git a/docs/user-guide-experimental/README.md b/docs/user-guide-experimental/README.md
@@ -0,0 +1,5 @@
+# Experimental Docs
+
+This directory contains documentation for features that are still experimental or under development and not yet ready for general use.
+
+More context can be found in the [experimental/README.md](../../nemo_aligner/experimental/README.md) file.
diff --git a/nemo_aligner/data/nlp/tests/__init__.py b/nemo_aligner/data/nlp/tests/__init__.py
diff --git a/tests/test_cai_utils.py → ..._aligner/data/nlp/tests/cai_utils_test.py b/tests/test_cai_utils.py → ..._aligner/data/nlp/tests/cai_utils_test.py
diff --git a/tests/test_datasets.py → nemo_aligner/data/nlp/tests/datasets_test.py b/tests/test_datasets.py → nemo_aligner/data/nlp/tests/datasets_test.py
diff --git a/nemo_aligner/experimental/README.md b/nemo_aligner/experimental/README.md
@@ -0,0 +1,50 @@
+# Experimental Package
+
+The `experimental` sub-package contains projects that are under active development and may not be fully stable.
+
+## Experimental Project Directory Structure:
+
+```
+NeMo-Aligner/
+├── docs/
+│   ├── user-guide/
+│   │   └── ppo.html
+│   └── user-guide-experimental/    <----- experimental docs
+│       └── new-thing.html
+├── nemo_aligner/
+│   ├── algorithms/
+│   ├── data/
+│   │   ├── datasets.py
+│   │   └── tests/
+│   │       └── datasets_test.py
+│   └── experimental/               <----- experimental sub-package
+│       ├── <proj-name>/
+│           ├── dataset.py          <----- experimental dataset
+│           ├── new_algo.py         <----- experimental algo
+│           ├── model.py            <----- experimental model
+│           └── tests/
+│               └── model_test.py   <----- experimental model test
+└── tests/
+    └── functional/
+        └── dpo.sh
+        └── test_cases/
+            └── dpo-llama3
+    └── functional_experimental/    <----- experimental functional tests (mirrors functional/ structure)
+        ├── new_algo.sh
+        └── test_cases/
+            └── new_algo-llama3
+```
+
+The directories below exist to organize experimental projects (source code), tests, and documentation.
+
+- [nemo_aligner/experimental/](../../nemo_aligner/experimental/): Main experimental sub-package containing projects under development
+- [tests/functional_experimental/](../../tests/functional_experimental/): Functional tests for experimental projects
+- [docs/user-guide-experimental/](../../docs/user-guide-experimental/): Documentation directory for experimental features and algorithms
+
+The `experimental` sub-package follows a modular structure where each project has its own directory (sub-package) containing implementation and tests.
+
+## Guidelines for "experimental/" Projects
+
+- **Scope**: Projects can include new model definitions, training loops, utilities, or unit tests.
+- **Independence**: Projects should ideally be independent. Dependence on other projects signals it might benefit from being added to core with tests (and documentation if applicable).
+- **Testing**: Must include at least one functional test [example](../../tests/functional/test_cases/dpo-llama3).
diff --git a/nemo_aligner/experimental/__init__.py b/nemo_aligner/experimental/__init__.py
diff --git a/tests/test_mcore_utilities.py → nemo_aligner/testing/utils.py b/tests/test_mcore_utilities.py → nemo_aligner/testing/utils.py
diff --git a/nemo_aligner/utils/tests/__init__.py b/nemo_aligner/utils/tests/__init__.py
diff --git a/tests/test_distributed.py → nemo_aligner/utils/tests/distributed_test.py b/tests/test_distributed.py → nemo_aligner/utils/tests/distributed_test.py
diff --git a/tests/test_ppo_utils.py → nemo_aligner/utils/tests/ppo_utils_test.py b/tests/test_ppo_utils.py → nemo_aligner/utils/tests/ppo_utils_test.py
diff --git a/tests/test_text_generation_utils.py → ...utils/tests/text_generation_utils_test.py b/tests/test_text_generation_utils.py → ...utils/tests/text_generation_utils_test.py
diff --git a/tests/test_trainer_utils.py → ...aligner/utils/tests/trainer_utils_test.py b/tests/test_trainer_utils.py → ...aligner/utils/tests/trainer_utils_test.py
diff --git a/tests/test_trt_llm.py → nemo_aligner/utils/tests/trt_llm_test.py b/tests/test_trt_llm.py → nemo_aligner/utils/tests/trt_llm_test.py
diff --git a/tests/test_utils.py → nemo_aligner/utils/tests/utils_test.py b/tests/test_utils.py → nemo_aligner/utils/tests/utils_test.py
diff --git a/tests/functional_experimental/README.md b/tests/functional_experimental/README.md
@@ -0,0 +1,3 @@
+# Experimental Functional Tests
+
+More context can be found in the [experimental/README.md](../../nemo_aligner/experimental/README.md) file.
diff --git a/tests/functional_experimental/test_cases/.gitkeep b/tests/functional_experimental/test_cases/.gitkeep
diff --git a/tests/functional_experimental/test_data/.gitkeep b/tests/functional_experimental/test_data/.gitkeep
diff --git a/tests/run_mpi_unit.sh b/tests/run_mpi_unit.sh
@@ -24,9 +24,9 @@ if [[ $NUM_GPUS_AVAILABLE -lt 2 ]]; then
 fi
 
 export PYTHONPATH=$(realpath ..):${PYTHONPATH:-}
-CUDA_VISIBLE_DEVICES=0,1 mpirun -np 2 --allow-run-as-root pytest .. -rA -s -x -vv --mpi $@ || true
+CUDA_VISIBLE_DEVICES=0,1 mpirun -np 2 --allow-run-as-root pytest ../nemo_aligner -rA -s -x -vv --mpi $@ || true
 
-if [[ -f PYTEST_SUCCESS ]]; then
+if [[ -f ../PYTEST_SUCCESS ]]; then
     echo SUCCESS
 else
     echo FAILURE

diff --git a/tests/run_unit.sh b/tests/run_unit.sh
@@ -24,9 +24,9 @@ if [[ $NUM_GPUS_AVAILABLE -lt 2 ]]; then
 fi
 
 export PYTHONPATH=$(realpath ..):${PYTHONPATH:-}
-CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node 2 -m pytest .. -rA -s -x -vv $@ || true
+CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node 2 -m pytest ../nemo_aligner -rA -s -x -vv $@ || true
 
-if [[ -f PYTEST_SUCCESS ]]; then
+if [[ -f ../PYTEST_SUCCESS ]]; then
     echo SUCCESS
 else
     echo FAILURE
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		# Experimental Functional Tests

		More context can be found in the [experimental/README.md](../../nemo_aligner/experimental/README.md) file.