Skip to content

fix: unskip dpo trainer integ tests #5835

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mujtaba1747 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

fix: unskip dpo trainer integ tests #5835

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4cd4a58
fix: unskip dpo trainer integ tests
May 7, 2026
adef22f
fix: unskip sft rlaif rlvr integ tests
May 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions sagemaker-train/tests/integ/train/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
8 changes: 3 additions & 5 deletions sagemaker-train/tests/integ/train/test_dpo_trainer_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,15 +22,14 @@
import pytest


@pytest.mark.skip(reason="Skipping GPU resource intensive test")
def test_dpo_trainer_lora_complete_workflow(sagemaker_session):
"""Test complete DPO training workflow with LORA."""
# Create DPOTrainer instance with comprehensive configuration
trainer = DPOTrainer(
model="meta-textgeneration-llama-3-2-1b-instruct",
training_type=TrainingType.LORA,
model_package_group="sdk-test-finetuned-models",
training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/dpo-oss-test-data/0.0.1",
training_dataset="s3://mc-flows-sdk-testing/input_data/dpo/preference_dataset_train_256.jsonl",
s3_output_path="s3://mc-flows-sdk-testing/output/",
accept_eula=True
)
Expand Down Expand Up @@ -61,16 +60,15 @@ def test_dpo_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_dpo_trainer_with_validation_dataset(sagemaker_session):
"""Test DPO trainer with both training and validation datasets."""

dpo_trainer = DPOTrainer(
model="meta-textgeneration-llama-3-2-1b-instruct",
training_type=TrainingType.LORA,
model_package_group="sdk-test-finetuned-models",
training_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/dpo-oss-test-data/0.0.1",
validation_dataset="arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/DataSet/dpo-oss-test-data/0.0.1",
training_dataset="s3://mc-flows-sdk-testing/input_data/dpo/preference_dataset_train_256.jsonl",
validation_dataset="s3://mc-flows-sdk-testing/input_data/dpo/preference_dataset_train_256.jsonl",
s3_output_path="s3://mc-flows-sdk-testing/output/",
accept_eula=True
)
Expand Down
9 changes: 3 additions & 6 deletions sagemaker-train/tests/integ/train/test_rlaif_trainer_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand All @@ -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
)
Expand Down Expand Up @@ -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."""

Expand All @@ -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
)
Expand All @@ -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."""

Expand All @@ -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
)
Expand Down
18 changes: 6 additions & 12 deletions sagemaker-train/tests/integ/train/test_rlvr_trainer_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
"""Integration tests for RLVR trainer"""
from __future__ import absolute_import

import os
import time
import pytest
import boto3
Expand All @@ -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."""

Expand All @@ -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
)
Expand Down Expand Up @@ -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."""

Expand All @@ -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
Expand Down Expand Up @@ -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",
Expand All @@ -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'

Expand Down
22 changes: 8 additions & 14 deletions sagemaker-train/tests/integ/train/test_sft_trainer_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
"""Integration tests for SFT trainer"""
from __future__ import absolute_import

import os
import time
import pytest
import boto3
Expand All @@ -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
)
Expand Down Expand Up @@ -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
)

Expand All @@ -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
Expand Down
Loading