Add Mock setup for mtrl integ tests

Author: Roja Reddy Sareddy

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

@@ -22,6 +22,13 @@
 DEFAULT_REGION = "us-west-2"
 
 
+@pytest.fixture(autouse=True, scope="session")
+def use_private_hub():
+    os.environ["SAGEMAKER_HUB_NAME"] = "sdktest"
+    yield
+    del os.environ["SAGEMAKER_HUB_NAME"]
+
+
 @pytest.fixture(scope="module")
 def sagemaker_session():
     region = os.environ.get("AWS_DEFAULT_REGION")

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

@@ -41,7 +41,7 @@ def _get_test_config():
     boto_session = boto3.Session(region_name=_REGION)
     account_id = boto_session.client("sts").get_caller_identity()["Account"]
     return {
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "agent_arn": f"arn:aws:bedrock-agentcore:{_REGION}:{account_id}:runtime/sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS",
         "dataset": f"s3://sagemaker-rft-{account_id}/prompts/gsm8k_small/prompts.parquet",
         "s3_output_path": f"s3://sagemaker-{_REGION}-{account_id}/model-evaluation/output-artifacts/",
@@ -113,7 +113,7 @@ def mtrl_trainer(sagemaker_session_mtrl, test_config):
 
     trainer = object.__new__(MultiTurnRLTrainer)
     trainer._model_name = test_config["base_model"]
-    trainer._model_arn = f"arn:aws:sagemaker:{_REGION}:aws:hub-content/SageMakerPublicHub/Model/{test_config['base_model']}/1.0.0"
+    trainer._model_arn = f"arn:aws:sagemaker:{_REGION}:{test_config['account_id']}:hub-content/sdktest/Model/{test_config['base_model']}/0.0.1"
     trainer.agent_env = test_config["agent_arn"]
     trainer.bedrock_agentcore_qualifier = "DEFAULT"
     trainer.output_model_package_group = test_config["model_package_group"]

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

@@ -148,7 +148,7 @@ def _get_3p_test_config():
     boto_session = boto3.Session(region_name=_REGION)
     account_id = boto_session.client("sts").get_caller_identity()["Account"]
     return {
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "dataset": os.environ.get(
             "MTRL_3P_DATASET",
             f"s3://sagemaker-rft-{account_id}/prompts/gsm8k_small/prompts.parquet",

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

@@ -53,7 +53,7 @@ def _get_account_id():
     "729646638167": {
         "env_name": "PROD",
         "existing_job_name": "openai-reasoning-gpt-oss-20b-mtrl-20260602215955",
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:729646638167:runtime/sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS",
         "dataset": "s3://sagemaker-rft-729646638167/prompts/gsm8k_small/prompts.parquet",
         "s3_output_path": "s3://sagemaker-us-west-2-729646638167/mtrl-integ/eval-output/",
@@ -65,7 +65,7 @@ def _get_account_id():
     "391266019386": {
         "env_name": "PREPROD",
         "existing_job_name": "mtrl-integ-gpt-oss-agentcore-1779143704358",
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:391266019386:runtime/mtrl_integ_gsm8k_streaming-bIz4H5Echk",
         "dataset": "s3://sagemaker-rft-beta-391266019386/prompts/gsm8k_small/prompts.parquet",
         "s3_output_path": "s3://sagemaker-us-west-2-391266019386/mtrl-integ/eval-output/",
@@ -77,7 +77,7 @@ def _get_account_id():
     "742774200982": {
         "env_name": "BETA",
         "existing_job_name": "openai-reasoning-gpt-oss-20b-mtrl-20260601114439",
-        "base_model": "openai-reasoning-gpt-oss-20b",
+        "base_model": "mock-oss-test",
         "agent_core_arn": "arn:aws:bedrock-agentcore:us-west-2:742774200982:runtime/sagemaker_rft_prod_gsm8k_streaming-UwSB6LEfEq",
         "dataset": "s3://sagemaker-rft-beta-742774200982/prompts/gsm8k_small/prompts.parquet",
         "s3_output_path": "s3://sagemaker-us-west-2-742774200982/mtrl-integ/eval-output/",

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

@@ -39,7 +39,8 @@ def _get_account_id():
     return _ACCOUNT_ID
 
 AGENT_RUNTIME_ID = "sagemaker_rft_prod_gsm8k_streaming-Yk6O377mUS"
-BASE_MODEL = "openai-reasoning-gpt-oss-20b"
+#BASE_MODEL = "openai-reasoning-gpt-oss-20b"
+BASE_MODEL = "mock-oss-test"
 EXISTING_JOB_NAME = "openai-reasoning-gpt-oss-20b-mtrl-20260602215955"
 
 
@@ -192,3 +193,6 @@ def test_list_bedrock_agentcore_runtimes(self, sagemaker_session):
             session=sagemaker_session.boto_session
         )
         assert isinstance(runtimes, list)
+
+
+