Change dir path to distributed_drivers

Author: benieric

src/sagemaker/modules/distributed.py

@@ -133,7 +133,7 @@ def driver_dir(self) -> str:
         Returns:
             str: Path to directory containing the driver script
         """
-        return os.path.join(SM_DRIVERS_LOCAL_PATH, "drivers")
+        return os.path.join(SM_DRIVERS_LOCAL_PATH, "distributed_drivers")
 
     @property
     def driver_script(self) -> str:
@@ -169,7 +169,7 @@ def driver_dir(self) -> str:
         Returns:
             str: Path to directory containing the driver script
         """
-        return os.path.join(SM_DRIVERS_LOCAL_PATH, "drivers")
+        return os.path.join(SM_DRIVERS_LOCAL_PATH, "distributed_drivers")
 
     @property
     def driver_script(self) -> str:

src/sagemaker/modules/templates.py

@@ -21,12 +21,12 @@
 
 EXECUTE_BASIC_SCRIPT_DRIVER = """
 echo "Running Basic Script driver"
-$SM_PYTHON_CMD /opt/ml/input/data/sm_drivers/drivers/basic_script_driver.py
+$SM_PYTHON_CMD /opt/ml/input/data/sm_drivers/distributed_drivers/basic_script_driver.py
 """
 
 EXEUCTE_DISTRIBUTED_DRIVER = """
 echo "Running {driver_name} Driver"
-$SM_PYTHON_CMD /opt/ml/input/data/sm_drivers/drivers/{driver_script}
+$SM_PYTHON_CMD /opt/ml/input/data/sm_drivers/distributed_drivers/{driver_script}
 """
 
 TRAIN_SCRIPT_TEMPLATE = """

src/sagemaker/modules/train/container_drivers/drivers/__init__.py renamed to src/sagemaker/modules/train/container_drivers/distributed_drivers/__init__.py

@@ -48,7 +48,7 @@
 SM_OUTPUT_FAILURE = "/opt/ml/output/failure"
 SM_OUTPUT_DATA_DIR = "/opt/ml/output/data"
 SM_SOURCE_DIR_PATH = "/opt/ml/input/data/code"
-SM_DRIVER_DIR_PATH = "/opt/ml/input/data/sm_drivers/drivers"
+SM_DISTRIBUTED_DRIVER_DIR_PATH = "/opt/ml/input/data/sm_drivers/distributed_drivers"
 
 SM_MASTER_ADDR = "algo-1"
 SM_MASTER_PORT = 7777
@@ -173,7 +173,7 @@ def set_env(
 
     distributed = read_distributed_json()
     if distributed:
-        env_vars["SM_DRIVER_DIR"] = SM_DRIVER_DIR_PATH
+        env_vars["SM_DISTRIBUTED_DRIVER_DIR"] = SM_DISTRIBUTED_DRIVER_DIR_PATH
         env_vars["SM_DISTRIBUTED_CONFIG"] = distributed
 
     # Data Channels

src/sagemaker/modules/train/container_drivers/drivers/basic_script_driver.py renamed to src/sagemaker/modules/train/container_drivers/distributed_drivers/basic_script_driver.py

@@ -569,7 +569,7 @@ def train(
             # If distributed is provided, overwrite code under <root>/drivers
             if self.distributed:
                 distributed_driver_dir = self.distributed.driver_dir
-                driver_dir = os.path.join(tmp_dir.name, "drivers")
+                driver_dir = os.path.join(tmp_dir.name, "distributed_drivers")
                 shutil.copytree(distributed_driver_dir, driver_dir, dirs_exist_ok=True)
 
             # If source code is provided, create a channel for the source code

src/sagemaker/modules/train/container_drivers/drivers/mpi_driver.py renamed to src/sagemaker/modules/train/container_drivers/distributed_drivers/mpi_driver.py

@@ -15,8 +15,8 @@ def main():
 
     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"
+    sm_drivers_dir = os.environ["SM_DISTRIBUTED_DRIVER_DIR"]
+    assert sm_drivers_dir == "/opt/ml/input/data/sm_drivers/distributed_drivers"
 
     entry_script = os.environ["SM_ENTRY_SCRIPT"]
     assert entry_script != None

src/sagemaker/modules/train/container_drivers/drivers/mpi_utils.py renamed to src/sagemaker/modules/train/container_drivers/distributed_drivers/mpi_utils.py

@@ -98,7 +98,7 @@
 export SM_MASTER_PORT='7777'
 export SM_SOURCE_DIR='/opt/ml/input/data/code'
 export SM_ENTRY_SCRIPT='train.py'
-export SM_DRIVER_DIR='/opt/ml/input/data/sm_drivers/drivers'
+export SM_DISTRIBUTED_DRIVER_DIR='/opt/ml/input/data/sm_drivers/distributed_drivers'
 export SM_DISTRIBUTED_CONFIG='{"process_count_per_node": 2}'
 export SM_CHANNEL_TRAIN='/opt/ml/input/data/train'
 export SM_CHANNEL_VALIDATION='/opt/ml/input/data/validation'

src/sagemaker/modules/train/container_drivers/drivers/torchrun_driver.py renamed to src/sagemaker/modules/train/container_drivers/distributed_drivers/torchrun_driver.py

@@ -22,7 +22,7 @@
 sys.modules["utils"] = MagicMock()
 sys.modules["mpi_utils"] = MagicMock()
 
-from sagemaker.modules.train.container_drivers.drivers import mpi_driver  # noqa: E402
+from sagemaker.modules.train.container_drivers.distributed_drivers import mpi_driver  # noqa: E402
 
 
 DUMMY_MPI_COMMAND = [

src/sagemaker/modules/train/container_drivers/scripts/environment.py

@@ -27,7 +27,7 @@
 mock_utils.get_python_executable = Mock(return_value="/usr/bin/python")
 
 with patch.dict("sys.modules", {"utils": mock_utils}):
-    from sagemaker.modules.train.container_drivers.drivers.mpi_utils import (
+    from sagemaker.modules.train.container_drivers.distributed_drivers.mpi_utils import (
         CustomHostKeyPolicy,
         _can_connect,
         write_status_file_to_workers,

src/sagemaker/modules/train/model_trainer.py

@@ -21,7 +21,9 @@
 
 sys.modules["utils"] = MagicMock()
 
-from sagemaker.modules.train.container_drivers.drivers import torchrun_driver  # noqa: E402
+from sagemaker.modules.train.container_drivers.distributed_drivers import (  # noqa: E402
+    torchrun_driver,
+)
 
 DUMMY_DISTRIBUTED = {"process_count_per_node": 2}