chore(test): migrate unit tests from unittest to pytest nvidia test p… (#2792)

This PR replaces unittest with pytest.

Part of https://github.com/meta-llama/llama-stack/issues/2680

cc @leseb

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>
This commit is contained in:
Mustafa Elbehery 2025-07-18 11:49:45 +02:00 committed by GitHub
parent d7cc38e934
commit 55713abe7d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -5,7 +5,6 @@
# the root directory of this source tree. # the root directory of this source tree.
import os import os
import unittest
import warnings import warnings
from unittest.mock import patch from unittest.mock import patch
@ -27,14 +26,13 @@ from llama_stack.providers.remote.post_training.nvidia.post_training import (
) )
class TestNvidiaParameters(unittest.TestCase): class TestNvidiaParameters:
def setUp(self): @pytest.fixture(autouse=True)
os.environ["NVIDIA_BASE_URL"] = "http://nemo.test" def setup_and_teardown(self):
"""Setup and teardown for each test method."""
os.environ["NVIDIA_CUSTOMIZER_URL"] = "http://nemo.test" os.environ["NVIDIA_CUSTOMIZER_URL"] = "http://nemo.test"
config = NvidiaPostTrainingConfig( config = NvidiaPostTrainingConfig(customizer_url=os.environ["NVIDIA_CUSTOMIZER_URL"], api_key=None)
base_url=os.environ["NVIDIA_BASE_URL"], customizer_url=os.environ["NVIDIA_CUSTOMIZER_URL"], api_key=None
)
self.adapter = NvidiaPostTrainingAdapter(config) self.adapter = NvidiaPostTrainingAdapter(config)
self.make_request_patcher = patch( self.make_request_patcher = patch(
@ -48,7 +46,8 @@ class TestNvidiaParameters(unittest.TestCase):
"updated_at": "2025-03-04T13:07:47.543605", "updated_at": "2025-03-04T13:07:47.543605",
} }
def tearDown(self): yield
self.make_request_patcher.stop() self.make_request_patcher.stop()
def _assert_request_params(self, expected_json): def _assert_request_params(self, expected_json):
@ -166,8 +165,8 @@ class TestNvidiaParameters(unittest.TestCase):
self.run_async( self.run_async(
self.adapter.supervised_fine_tune( self.adapter.supervised_fine_tune(
job_uuid=required_job_uuid, # Required parameter job_uuid=required_job_uuid,
model=required_model, # Required parameter model=required_model,
checkpoint_dir="", checkpoint_dir="",
algorithm_config=algorithm_config, algorithm_config=algorithm_config,
training_config=convert_pydantic_to_json_value(training_config), training_config=convert_pydantic_to_json_value(training_config),
@ -198,7 +197,6 @@ class TestNvidiaParameters(unittest.TestCase):
data_config = DataConfig( data_config = DataConfig(
dataset_id="test-dataset", dataset_id="test-dataset",
batch_size=8, batch_size=8,
# Unsupported parameters
shuffle=True, shuffle=True,
data_format=DatasetFormat.instruct, data_format=DatasetFormat.instruct,
validation_dataset_id="val-dataset", validation_dataset_id="val-dataset",
@ -207,20 +205,16 @@ class TestNvidiaParameters(unittest.TestCase):
optimizer_config = OptimizerConfig( optimizer_config = OptimizerConfig(
lr=0.0001, lr=0.0001,
weight_decay=0.01, weight_decay=0.01,
# Unsupported parameters
optimizer_type=OptimizerType.adam, optimizer_type=OptimizerType.adam,
num_warmup_steps=100, num_warmup_steps=100,
) )
efficiency_config = EfficiencyConfig( efficiency_config = EfficiencyConfig(enable_activation_checkpointing=True)
enable_activation_checkpointing=True # Unsupported parameter
)
training_config = TrainingConfig( training_config = TrainingConfig(
n_epochs=1, n_epochs=1,
data_config=data_config, data_config=data_config,
optimizer_config=optimizer_config, optimizer_config=optimizer_config,
# Unsupported parameters
efficiency_config=efficiency_config, efficiency_config=efficiency_config,
max_steps_per_epoch=1000, max_steps_per_epoch=1000,
gradient_accumulation_steps=4, gradient_accumulation_steps=4,
@ -228,7 +222,6 @@ class TestNvidiaParameters(unittest.TestCase):
dtype="bf16", dtype="bf16",
) )
# Capture warnings
with warnings.catch_warnings(record=True) as w: with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always") warnings.simplefilter("always")
@ -236,7 +229,7 @@ class TestNvidiaParameters(unittest.TestCase):
self.adapter.supervised_fine_tune( self.adapter.supervised_fine_tune(
job_uuid="test-job", job_uuid="test-job",
model="meta-llama/Llama-3.1-8B-Instruct", model="meta-llama/Llama-3.1-8B-Instruct",
checkpoint_dir="test-dir", # Unsupported parameter checkpoint_dir="test-dir",
algorithm_config=LoraFinetuningConfig( algorithm_config=LoraFinetuningConfig(
type="LoRA", type="LoRA",
apply_lora_to_mlp=True, apply_lora_to_mlp=True,
@ -246,8 +239,8 @@ class TestNvidiaParameters(unittest.TestCase):
lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"], lora_attn_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
), ),
training_config=convert_pydantic_to_json_value(training_config), training_config=convert_pydantic_to_json_value(training_config),
logger_config={"test": "value"}, # Unsupported parameter logger_config={"test": "value"},
hyperparam_search_config={"test": "value"}, # Unsupported parameter hyperparam_search_config={"test": "value"},
) )
) )
@ -265,7 +258,6 @@ class TestNvidiaParameters(unittest.TestCase):
"gradient_accumulation_steps", "gradient_accumulation_steps",
"max_validation_steps", "max_validation_steps",
"dtype", "dtype",
# required unsupported parameters
"rank", "rank",
"apply_lora_to_output", "apply_lora_to_output",
"lora_attn_modules", "lora_attn_modules",
@ -273,7 +265,3 @@ class TestNvidiaParameters(unittest.TestCase):
] ]
for field in fields: for field in fields:
assert any(field in text for text in warning_texts) assert any(field in text for text in warning_texts)
if __name__ == "__main__":
unittest.main()