fix integs

Author: benieric

Diff for: tests/integ/sagemaker/modules/train/test_model_trainer.py

@@ -35,21 +35,22 @@
     },
 }
 
-DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310"
+PARAM_SCRIPT_SOURCE_DIR = f"{DATA_DIR}/modules/params_script"
+PARAM_SCRIPT_SOURCE_CODE = SourceCode(
+    source_dir=PARAM_SCRIPT_SOURCE_DIR,
+    requirements="requirements.txt",
+    entry_script="train.py",
+)
 
+DEFAULT_CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py31"
 
-def test_hp_contract_basic_py_script(modules_sagemaker_session):
-    source_code = SourceCode(
-        source_dir=f"{DATA_DIR}/modules/params_script",
-        requirements="requirements.txt",
-        entry_script="train.py",
-    )
 
+def test_hp_contract_basic_py_script(modules_sagemaker_session):
     model_trainer = ModelTrainer(
         sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
         hyperparameters=EXPECTED_HYPERPARAMETERS,
-        source_code=source_code,
+        source_code=PARAM_SCRIPT_SOURCE_CODE,
         base_job_name="hp-contract-basic-py-script",
     )
 
@@ -59,6 +60,7 @@ def test_hp_contract_basic_py_script(modules_sagemaker_session):
 def test_hp_contract_basic_sh_script(modules_sagemaker_session):
     source_code = SourceCode(
         source_dir=f"{DATA_DIR}/modules/params_script",
+        requirements="requirements.txt",
         entry_script="train.sh",
     )
     model_trainer = ModelTrainer(
@@ -73,17 +75,13 @@ def test_hp_contract_basic_sh_script(modules_sagemaker_session):
 
 
 def test_hp_contract_mpi_script(modules_sagemaker_session):
-    source_code = SourceCode(
-        source_dir=f"{DATA_DIR}/modules/params_script",
-        entry_script="train.py",
-    )
     compute = Compute(instance_type="ml.m5.xlarge", instance_count=2)
     model_trainer = ModelTrainer(
         sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
         compute=compute,
         hyperparameters=EXPECTED_HYPERPARAMETERS,
-        source_code=source_code,
+        source_code=PARAM_SCRIPT_SOURCE_CODE,
         distributed=MPI(),
         base_job_name="hp-contract-mpi-script",
     )
@@ -92,17 +90,13 @@ def test_hp_contract_mpi_script(modules_sagemaker_session):
 
 
 def test_hp_contract_torchrun_script(modules_sagemaker_session):
-    source_code = SourceCode(
-        source_dir=f"{DATA_DIR}/modules/params_script",
-        entry_script="train.py",
-    )
     compute = Compute(instance_type="ml.m5.xlarge", instance_count=2)
     model_trainer = ModelTrainer(
         sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
         compute=compute,
         hyperparameters=EXPECTED_HYPERPARAMETERS,
-        source_code=source_code,
+        source_code=PARAM_SCRIPT_SOURCE_CODE,
         distributed=Torchrun(),
         base_job_name="hp-contract-torchrun-script",
     )
@@ -111,33 +105,23 @@ def test_hp_contract_torchrun_script(modules_sagemaker_session):
 
 
 def test_hp_contract_hyperparameter_json(modules_sagemaker_session):
-    source_dir = f"{DATA_DIR}/modules/params_script"
-    source_code = SourceCode(
-        source_dir=source_dir,
-        entry_script="train.py",
-    )
     model_trainer = ModelTrainer(
         sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
-        hyperparameters=f"{source_dir}/hyperparameters.json",
-        source_code=source_code,
+        hyperparameters=f"{PARAM_SCRIPT_SOURCE_DIR}/hyperparameters.json",
+        source_code=PARAM_SCRIPT_SOURCE_CODE,
         base_job_name="hp-contract-hyperparameter-json",
     )
     assert model_trainer.hyperparameters == EXPECTED_HYPERPARAMETERS
     model_trainer.train()
 
 
 def test_hp_contract_hyperparameter_yaml(modules_sagemaker_session):
-    source_dir = f"{DATA_DIR}/modules/params_script"
-    source_code = SourceCode(
-        source_dir=source_dir,
-        entry_script="train.py",
-    )
     model_trainer = ModelTrainer(
         sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
-        hyperparameters=f"{source_dir}/hyperparameters.yaml",
-        source_code=source_code,
+        hyperparameters=f"{PARAM_SCRIPT_SOURCE_DIR}/hyperparameters.yaml",
+        source_code=PARAM_SCRIPT_SOURCE_CODE,
         base_job_name="hp-contract-hyperparameter-yaml",
     )
     assert model_trainer.hyperparameters == EXPECTED_HYPERPARAMETERS