Include none types when creating config jsons for safer reference

Author: benieric

src/sagemaker/modules/distributed.py

@@ -128,10 +128,20 @@ class Torchrun(DistributedConfig):
 
     @property
     def driver_dir(self) -> str:
+        """Directory containing the driver script.
+
+        Returns:
+            str: Path to directory containing the driver script
+        """
         return os.path.join(SM_DRIVERS_LOCAL_PATH, "drivers")
 
     @property
     def driver_script(self) -> str:
+        """Name of the driver script.
+
+        Returns:
+            str: Name of the driver script file
+        """
         return "torchrun_driver.py"
 
 
@@ -154,8 +164,18 @@ class MPI(DistributedConfig):
 
     @property
     def driver_dir(self) -> str:
+        """Directory containing the driver script.
+
+        Returns:
+            str: Path to directory containing the driver script
+        """
         return os.path.join(SM_DRIVERS_LOCAL_PATH, "drivers")
 
     @property
     def driver_script(self) -> str:
+        """Name of the driver script.
+
+        Returns:
+            str: Name of the driver script
+        """
         return "mpi_driver.py"

src/sagemaker/modules/train/container_drivers/drivers/mpi_driver.py

@@ -77,7 +77,7 @@ def main():
 
         host_list = json.loads(os.environ["SM_HOSTS"])
         host_count = int(os.environ["SM_HOST_COUNT"])
-        process_count = int(distributed_config.get("process_count_per_node", 0))
+        process_count = int(distributed_config["process_count_per_node"] or 0)
         process_count = get_process_count(process_count)
 
         if process_count > 1:
@@ -87,7 +87,7 @@ def main():
             host_count=host_count,
             host_list=host_list,
             num_processes=process_count,
-            additional_options=distributed_config.get("mpi_additional_options", []),
+            additional_options=distributed_config["mpi_additional_options"] or [],
             entry_script_path=entry_script,
         )
 

src/sagemaker/modules/train/container_drivers/drivers/torchrun_driver.py

@@ -69,7 +69,7 @@ def create_commands():
     distributed_config = json.loads(os.environ["SM_DISTRIBUTED_CONFIG"])
     hyperparameters = json.loads(os.environ["SM_HPS"])
 
-    process_count = int(distributed_config.get("process_count_per_node", 0))
+    process_count = int(distributed_config["process_count_per_node"] or 0)
     process_count = get_process_count(process_count)
     host_count = int(os.environ["SM_HOST_COUNT"])
 

src/sagemaker/modules/train/model_trainer.py

@@ -773,7 +773,7 @@ def _write_source_code_json(self, tmp_dir: TemporaryDirectory, source_code: Sour
         """Write the source code configuration to a JSON file."""
         file_path = os.path.join(tmp_dir.name, SOURCE_CODE_JSON)
         with open(file_path, "w") as f:
-            dump = source_code.model_dump(exclude_none=True) if source_code else {}
+            dump = source_code.model_dump() if source_code else {}
             f.write(json.dumps(dump))
 
     def _write_distributed_json(
@@ -784,7 +784,7 @@ def _write_distributed_json(
         """Write the distributed runner configuration to a JSON file."""
         file_path = os.path.join(tmp_dir.name, DISTRIBUTED_JSON)
         with open(file_path, "w") as f:
-            dump = distributed.model_dump(exclude_none=True) if distributed else {}
+            dump = distributed.model_dump() if distributed else {}
             f.write(json.dumps(dump))
 
     def _prepare_train_script(

tests/unit/sagemaker/modules/train/container_drivers/scripts/test_enviornment.py

@@ -73,6 +73,15 @@
     },
 }
 
+SOURCE_CODE = {
+    "source_dir": "code",
+    "entry_script": "train.py",
+}
+
+DISTRIBUTED_CONFIG = {
+    "process_count_per_node": 2,
+}
+
 OUTPUT_FILE = os.path.join(os.path.dirname(__file__), "sm_training.env")
 
 # flake8: noqa
@@ -87,6 +96,10 @@
 export SM_LOG_LEVEL='20'
 export SM_MASTER_ADDR='algo-1'
 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_CONFIG='{"process_count_per_node": 2}'
 export SM_CHANNEL_TRAIN='/opt/ml/input/data/train'
 export SM_CHANNEL_VALIDATION='/opt/ml/input/data/validation'
 export SM_CHANNELS='["train", "validation"]'
@@ -110,6 +123,14 @@
 """
 
 
+@patch(
+    "sagemaker.modules.train.container_drivers.scripts.environment.read_source_code_json",
+    return_value=SOURCE_CODE,
+)
+@patch(
+    "sagemaker.modules.train.container_drivers.scripts.environment.read_distributed_json",
+    return_value=DISTRIBUTED_CONFIG,
+)
 @patch("sagemaker.modules.train.container_drivers.scripts.environment.num_cpus", return_value=8)
 @patch("sagemaker.modules.train.container_drivers.scripts.environment.num_gpus", return_value=0)
 @patch("sagemaker.modules.train.container_drivers.scripts.environment.num_neurons", return_value=0)
@@ -122,7 +143,13 @@
     side_effect=safe_deserialize,
 )
 def test_set_env(
-    mock_safe_deserialize, mock_safe_serialize, mock_num_cpus, mock_num_gpus, mock_num_neurons
+    mock_safe_deserialize,
+    mock_safe_serialize,
+    mock_num_neurons,
+    mock_num_gpus,
+    mock_num_cpus,
+    mock_read_distributed_json,
+    mock_read_source_code_json,
 ):
     with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}):
         set_env(
@@ -135,6 +162,8 @@ def test_set_env(
         mock_num_cpus.assert_called_once()
         mock_num_gpus.assert_called_once()
         mock_num_neurons.assert_called_once()
+        mock_read_distributed_json.assert_called_once()
+        mock_read_source_code_json.assert_called_once()
 
         with open(OUTPUT_FILE, "r") as f:
             env_file = f.read().strip()

tests/unit/sagemaker/modules/train/container_drivers/test_utils.py

@@ -12,11 +12,13 @@
 # language governing permissions and limitations under the License.
 """Container Utils Unit Tests."""
 from __future__ import absolute_import
+import os
 
 from sagemaker.modules.train.container_drivers.common.utils import (
     safe_deserialize,
     safe_serialize,
     hyperparameters_to_cli_args,
+    get_process_count,
 )
 
 SM_HPS = {
@@ -119,3 +121,18 @@ def test_safe_serialize_empty_data():
     assert safe_serialize("") == ""
     assert safe_serialize([]) == "[]"
     assert safe_serialize({}) == "{}"
+
+
+def test_get_process_count():
+    assert get_process_count() == 1
+    assert get_process_count(2) == 2
+    os.environ["SM_NUM_GPUS"] = "4"
+    assert get_process_count() == 4
+    os.environ["SM_NUM_GPUS"] = "0"
+    os.environ["SM_NUM_NEURONS"] = "8"
+    assert get_process_count() == 8
+    os.environ["SM_NUM_NEURONS"] = "0"
+    assert get_process_count() == 1
+    del os.environ["SM_NUM_GPUS"]
+    del os.environ["SM_NUM_NEURONS"]
+    assert get_process_count() == 1