Add integ test

Author: benieric

tests/data/modules/custom_drivers/driver.py

@@ -0,0 +1,34 @@
+import json
+import os
+import subprocess
+import sys
+
+
+def main():
+    driver_config = json.loads(os.environ["SM_DISTRIBUTED_CONFIG"])
+    process_count_per_node = driver_config["process_count_per_node"]
+    assert process_count_per_node != None
+
+    hps = json.loads(os.environ["SM_HP"])
+    assert hps != None
+    assert isinstance(hps, dict)
+
+    source_dir = os.environ["SM_SOURCE_DIR"]
+    assert source_dir == "/opt/ml/input/data/code"
+    sm_drivers_dir = os.environ["SM_DRIVER_DIR"]
+    assert sm_drivers_dir == "/opt/ml/input/data/sm_drivers/drivers"
+
+    entry_script = os.environ["SM_ENTRY_SCRIPT"]
+    assert entry_script != None
+
+    python = sys.executable
+
+    command = [python, entry_script]
+    print(f"Running command: {command}")
+    subprocess.run(command, check=True)
+
+
+if __name__ == "__main__":
+    print("Running custom driver script")
+    main()
+    print("Finished running custom driver script")

tests/data/modules/scripts/entry_script.py

@@ -0,0 +1,19 @@
+import json
+import os
+import time
+
+
+def main():
+    hps = json.loads(os.environ["SM_HPS"])
+    assert hps != None
+    print(f"Hyperparameters: {hps}")
+
+    print("Running pseudo training script")
+    for epochs in range(hps["epochs"]):
+        print(f"Epoch: {epochs}")
+        time.sleep(1)
+    print("Finished running pseudo training script")
+
+
+if __name__ == "__main__":
+    main()

tests/integ/sagemaker/modules/train/test_model_trainer.py

@@ -17,7 +17,7 @@
 
 from sagemaker.modules.train import ModelTrainer
 from sagemaker.modules.configs import SourceCode, Compute
-from sagemaker.modules.distributed import MPI, Torchrun
+from sagemaker.modules.distributed import MPI, Torchrun, DistributedConfig
 
 EXPECTED_HYPERPARAMETERS = {
     "integer": 1,
@@ -106,3 +106,33 @@ def test_hp_contract_torchrun_script(modules_sagemaker_session):
     )
 
     model_trainer.train()
+
+
+def test_custom_distributed_driver(modules_sagemaker_session):
+    class CustomDriver(DistributedConfig):
+        process_count_per_node: int = 2
+
+        @property
+        def driver_dir(self) -> str:
+            return f"{DATA_DIR}/modules/custom_drivers"
+
+        @property
+        def driver_script(self) -> str:
+            return "driver.py"
+
+    source_code = SourceCode(
+        source_dir=f"{DATA_DIR}/modules/scripts",
+        entry_script="entry_script.py",
+    )
+
+    hyperparameters = {"epochs": 10}
+
+    model_trainer = ModelTrainer(
+        sagemaker_session=modules_sagemaker_session,
+        training_image=DEFAULT_CPU_IMAGE,
+        hyperparameters=hyperparameters,
+        source_code=source_code,
+        distributed=CustomDriver(),
+        base_job_name="custom-distributed-driver",
+    )
+    model_trainer.train()