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