fix: unskip sft rlaif rlvr integ tests

Author: Syed Jafri

sagemaker-train/tests/integ/train/conftest.py

@@ -38,3 +38,13 @@ def sagemaker_session():
 
     if region_manual_set and "AWS_DEFAULT_REGION" in os.environ:
         del os.environ["AWS_DEFAULT_REGION"]
+
+
+NOVA_REGION = "us-east-1"
+
+
+@pytest.fixture(scope="module")
+def sagemaker_session_us_east_1():
+    """Create a SageMaker session in us-east-1 for Nova model tests."""
+    boto_session = boto3.Session(region_name=NOVA_REGION)
+    return Session(boto_session=boto_session)

sagemaker-train/tests/integ/train/test_rlaif_trainer_integration.py

@@ -21,7 +21,6 @@
 import pytest
 
 
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
 def test_rlaif_trainer_lora_complete_workflow(sagemaker_session):
     """Test complete RLAIF training workflow with LORA."""
 
@@ -33,7 +32,7 @@ def test_rlaif_trainer_lora_complete_workflow(sagemaker_session):
         reward_prompt='Builtin.Summarize',
         mlflow_experiment_name="test-rlaif-finetuned-models-exp",
         mlflow_run_name="test-rlaif-finetuned-models-run",
-        training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/rlvr-rlaif-oss-test-data/0.0.1",
+        training_dataset="s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl",
         s3_output_path="s3://mc-flows-sdk-testing/output/",
         accept_eula=True
     )
@@ -61,7 +60,6 @@ def test_rlaif_trainer_lora_complete_workflow(sagemaker_session):
     assert training_job.output_model_package_arn is not None
 
 
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
 def test_rlaif_trainer_with_custom_reward_settings(sagemaker_session):
     """Test RLAIF trainer with different reward model and prompt."""
 
@@ -73,7 +71,7 @@ def test_rlaif_trainer_with_custom_reward_settings(sagemaker_session):
         reward_prompt="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/JsonDoc/rlaif-test-prompt/0.0.1",
         mlflow_experiment_name="test-rlaif-finetuned-models-exp",
         mlflow_run_name="test-rlaif-finetuned-models-run",
-        training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/rlvr-rlaif-oss-test-data/0.0.1",
+        training_dataset="s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl",
         s3_output_path="s3://mc-flows-sdk-testing/output/",
         accept_eula=True
     )
@@ -100,7 +98,6 @@ def test_rlaif_trainer_with_custom_reward_settings(sagemaker_session):
     assert training_job.output_model_package_arn is not None
 
 
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
 def test_rlaif_trainer_continued_finetuning(sagemaker_session):
     """Test complete RLAIF training workflow with LORA."""
 
@@ -112,7 +109,7 @@ def test_rlaif_trainer_continued_finetuning(sagemaker_session):
         reward_prompt='Builtin.Summarize',
         mlflow_experiment_name="test-rlaif-finetuned-models-exp",
         mlflow_run_name="test-rlaif-finetuned-models-run",
-        training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/rlvr-rlaif-oss-test-data/0.0.1",
+        training_dataset="s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl",
         s3_output_path="s3://mc-flows-sdk-testing/output/",
         accept_eula=True
     )

sagemaker-train/tests/integ/train/test_rlvr_trainer_integration.py

@@ -13,7 +13,6 @@
 """Integration tests for RLVR trainer"""
 from __future__ import absolute_import
 
-import os
 import time
 import pytest
 import boto3
@@ -22,7 +21,6 @@
 from sagemaker.train.common import TrainingType
 
 
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
 def test_rlvr_trainer_lora_complete_workflow(sagemaker_session):
     """Test complete RLVR training workflow with LORA."""
 
@@ -32,7 +30,7 @@ def test_rlvr_trainer_lora_complete_workflow(sagemaker_session):
         model_package_group="sdk-test-finetuned-models",
         mlflow_experiment_name="test-rlvr-finetuned-models-exp",
         mlflow_run_name="test-rlvr-finetuned-models-run",
-        training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/rlvr-rlaif-oss-test-data/0.0.1",
+        training_dataset="s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl",
         s3_output_path="s3://mc-flows-sdk-testing/output/",
         accept_eula=True
     )
@@ -60,7 +58,6 @@ def test_rlvr_trainer_lora_complete_workflow(sagemaker_session):
     assert training_job.output_model_package_arn is not None
 
 
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
 def test_rlvr_trainer_with_custom_reward_function(sagemaker_session):
     """Test RLVR trainer with custom reward function."""
 
@@ -70,7 +67,7 @@ def test_rlvr_trainer_with_custom_reward_function(sagemaker_session):
         model_package_group="sdk-test-finetuned-models",
         mlflow_experiment_name="test-rlvr-finetuned-models-exp",
         mlflow_run_name="test-rlvr-finetuned-models-run",
-        training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/rlvr-rlaif-oss-test-data/0.0.1",
+        training_dataset="s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl",
         s3_output_path="s3://mc-flows-sdk-testing/output/",
         custom_reward_function="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/JsonDoc/rlvr-test-rf/0.0.1",
         accept_eula=True
@@ -98,14 +95,10 @@ def test_rlvr_trainer_with_custom_reward_function(sagemaker_session):
     assert training_job.output_model_package_arn is not None
 
 
-# @pytest.mark.skipif(os.environ.get('AWS_DEFAULT_REGION') != 'us-east-1', reason="Nova models only available in us-east-1")
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
-def test_rlvr_trainer_nova_workflow(sagemaker_session):
+def test_rlvr_trainer_nova_workflow(sagemaker_session_us_east_1):
     """Test RLVR training workflow with Nova model."""
-    import os
-    os.environ['SAGEMAKER_REGION'] = 'us-east-1'
+    # sagemaker_session_us_east_1 fixture is defined in conftest.py (us-east-1 region)
 
-    # For fine-tuning 
     rlvr_trainer = RLVRTrainer(
         model="nova-textgeneration-lite-v2",
         model_package_group="sdk-test-finetuned-models",
@@ -115,7 +108,8 @@ def test_rlvr_trainer_nova_workflow(sagemaker_session):
         validation_dataset="s3://mc-flows-sdk-testing-us-east-1/input_data/rlvr-nova/grpo-64-sample.jsonl",
         s3_output_path="s3://mc-flows-sdk-testing-us-east-1/output/",
         custom_reward_function="arn:aws:sagemaker:us-east-1:729646638167:hub-content/sdktest/JsonDoc/rlvr-nova-test-rf/0.0.1",
-        accept_eula=True
+        accept_eula=True,
+        sagemaker_session=sagemaker_session_us_east_1
     )
     rlvr_trainer.hyperparameters.data_s3_path = 's3://example-bucket'
 

sagemaker-train/tests/integ/train/test_sft_trainer_integration.py

@@ -13,7 +13,6 @@
 """Integration tests for SFT trainer"""
 from __future__ import absolute_import
 
-import os
 import time
 import pytest
 import boto3
@@ -22,15 +21,14 @@
 from sagemaker.train.common import TrainingType
 
 
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
 def test_sft_trainer_lora_complete_workflow(sagemaker_session):
     """Test complete SFT training workflow with LORA."""
 
     sft_trainer = SFTTrainer(
         model="meta-textgeneration-llama-3-2-1b-instruct",
         training_type=TrainingType.LORA,
         model_package_group="arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models",
-        training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/sft-oss-test-data/0.0.1",
+        training_dataset="s3://mc-flows-sdk-testing/input_data/sft/sample_data_256_final.jsonl",
         s3_output_path="s3://mc-flows-sdk-testing/output/",
         accept_eula=True
     )
@@ -58,16 +56,15 @@ def test_sft_trainer_lora_complete_workflow(sagemaker_session):
     assert training_job.output_model_package_arn is not None
 
 
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
 def test_sft_trainer_with_validation_dataset(sagemaker_session):
     """Test SFT trainer with both training and validation datasets."""
 
     sft_trainer = SFTTrainer(
         model="meta-textgeneration-llama-3-2-1b-instruct",
         training_type=TrainingType.LORA,
         model_package_group="arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models",
-        training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/sft-oss-test-data/0.0.1",
-        validation_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/sft-oss-test-data/0.0.1",
+        training_dataset="s3://mc-flows-sdk-testing/input_data/sft/sample_data_256_final.jsonl",
+        validation_dataset="s3://mc-flows-sdk-testing/input_data/sft/sample_data_256_final.jsonl",
         accept_eula=True
     )
 
@@ -92,22 +89,19 @@ def test_sft_trainer_with_validation_dataset(sagemaker_session):
     assert hasattr(training_job, 'output_model_package_arn')
 
 
-# @pytest.mark.skipif(os.environ.get('AWS_DEFAULT_REGION') != 'us-east-1', reason="Nova models only available in us-east-1")
-@pytest.mark.skip(reason="Skipping GPU resource intensive test")
-def test_sft_trainer_nova_workflow(sagemaker_session):
+def test_sft_trainer_nova_workflow(sagemaker_session_us_east_1):
     """Test SFT trainer with Nova model."""
-    import os
-    os.environ['SAGEMAKER_REGION'] = 'us-east-1'
+    # sagemaker_session_us_east_1 fixture is defined in conftest.py (us-east-1 region)
 
-    # For fine-tuning 
     sft_trainer_nova = SFTTrainer(
         model="nova-textgeneration-lite-v2",
         training_type=TrainingType.LORA, 
         model_package_group="sdk-test-finetuned-models",
         mlflow_experiment_name="test-nova-finetuned-models-exp",
         mlflow_run_name="test-nova-finetuned-models-run",
-        training_dataset="arn:aws:sagemaker:us-east-1:729646638167:hub-content/sdktest/DataSet/sft-nova-test-dataset/0.0.1",
-        s3_output_path="s3://mc-flows-sdk-testing-us-east-1/output/"
+        training_dataset="s3://mc-flows-sdk-testing-us-east-1/input_data/sft-nova/sft_8_samples.jsonl",
+        s3_output_path="s3://mc-flows-sdk-testing-us-east-1/output/",
+        sagemaker_session=sagemaker_session_us_east_1
     )
 
     # Create training job