reformat files

Author: zhaoqizqwang

src/sagemaker/__init__.py

@@ -63,7 +63,7 @@
 from sagemaker.automl.automl import AutoML, AutoMLJob, AutoMLInput  # noqa: F401
 from sagemaker.automl.candidate_estimator import CandidateEstimator, CandidateStep  # noqa: F401
 from sagemaker.automl.automlv2 import (  # noqa: F401
-    AutoMLV2, 
+    AutoMLV2,
     AutoMLJobV2,
     LocalAutoMLDataChannel,
     AutoMLDataChannel,

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

@@ -110,24 +110,24 @@ def test_load_base_recipe_types(
 
     if recipe_type == "sagemaker":
         # Mock the file check to simulate recipe exists
-        with patch("os.path.isfile", return_value=True), \
-             patch("shutil.copy") as mock_copy:
+        with patch("os.path.isfile", return_value=True), patch("shutil.copy") as mock_copy:
             # Create a temporary recipe file for the copy operation
             import tempfile
             import os
-            with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+
+            with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
                 yaml.dump({"trainer": {"num_nodes": 1}}, f)
                 temp_path = f.name
-            
+
             def copy_side_effect(src, dst):
                 # Read from temp file and write to destination
-                with open(temp_path, 'r') as src_file:
+                with open(temp_path, "r") as src_file:
                     content = src_file.read()
-                with open(dst, 'w') as dst_file:
+                with open(dst, "w") as dst_file:
                     dst_file.write(content)
-            
+
             mock_copy.side_effect = copy_side_effect
-            
+
             load_recipe = _load_base_recipe(
                 training_recipe="training/llama/p4_hf_llama3_70b_seq8k_gpu",
                 recipe_overrides=None,
@@ -136,7 +136,7 @@ def copy_side_effect(src, dst):
             assert load_recipe is not None
             assert "trainer" in load_recipe
             assert mock_clone.call_args.args[0] == training_recipes_cfg.get("launcher_repo")
-            
+
             # Clean up
             os.unlink(temp_path)
 

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

@@ -934,13 +934,15 @@ def mock_upload_data(path, bucket, key_prefix):
 @patch("sagemaker.modules.train.model_trainer._load_base_recipe")
 def test_model_trainer_gpu_recipe_full_init(mock_load_recipe, modules_session):
     from omegaconf import OmegaConf
-    
+
     # Mock the recipe loading to return a valid GPU recipe structure
-    mock_load_recipe.return_value = OmegaConf.create({
-        "trainer": {"num_nodes": 2},
-        "model": {"model_type": "llama_v3"},
-    })
-    
+    mock_load_recipe.return_value = OmegaConf.create(
+        {
+            "trainer": {"num_nodes": 2},
+            "model": {"model_type": "llama_v3"},
+        }
+    )
+
     training_recipe = "training/llama/p4_hf_llama3_70b_seq8k_gpu"
     recipe_overrides = {"run": {"results_dir": "/opt/ml/model"}}
     compute = Compute(instance_type="ml.p4d.24xlarge", instance_count="2")

tests/unit/test_pytorch.py

@@ -900,23 +900,27 @@ def test_training_recipe_for_cpu(sagemaker_session):
 )
 @patch("sagemaker.pytorch.estimator.PyTorch._recipe_load")
 @patch("sagemaker.pytorch.estimator._get_training_recipe_gpu_script")
-def test_training_recipe_for_gpu(mock_gpu_script, mock_recipe_load, sagemaker_session, recipe, model):
+def test_training_recipe_for_gpu(
+    mock_gpu_script, mock_recipe_load, sagemaker_session, recipe, model
+):
     from omegaconf import OmegaConf
-    
+
     # Mock the GPU script function to return the expected entry point
     mock_gpu_script.return_value = f"{model}_pretrain.py"
-    
+
     # Mock the recipe structure that would be loaded
-    mock_recipe = OmegaConf.create({
-        "trainer": {
-            "num_nodes": 1,
-        },
-        "model": {
-            "model_type": model,
-        },
-    })
+    mock_recipe = OmegaConf.create(
+        {
+            "trainer": {
+                "num_nodes": 1,
+            },
+            "model": {
+                "model_type": model,
+            },
+        }
+    )
     mock_recipe_load.return_value = (recipe, mock_recipe)
-    
+
     container_log_level = '"logging.INFO"'
 
     recipe_overrides = {
@@ -968,21 +972,23 @@ def test_training_recipe_for_gpu(mock_gpu_script, mock_recipe_load, sagemaker_se
 @patch("sagemaker.pytorch.estimator._get_training_recipe_gpu_script")
 def test_training_recipe_with_override(mock_gpu_script, mock_recipe_load, sagemaker_session):
     from omegaconf import OmegaConf
-    
+
     # Mock the GPU script function to return the expected entry point
     mock_gpu_script.return_value = "mistral_pretrain.py"
-    
+
     # Mock the recipe structure that would be loaded
-    mock_recipe = OmegaConf.create({
-        "trainer": {
-            "num_nodes": 1,
-        },
-        "model": {
-            "model_type": "mistral",
-        },
-    })
+    mock_recipe = OmegaConf.create(
+        {
+            "trainer": {
+                "num_nodes": 1,
+            },
+            "model": {
+                "model_type": "mistral",
+            },
+        }
+    )
     mock_recipe_load.return_value = ("hf_llama3_8b_seq8k_gpu_p5x16_pretrain", mock_recipe)
-    
+
     container_log_level = '"logging.INFO"'
 
     recipe_overrides = {
@@ -1021,23 +1027,27 @@ def test_training_recipe_with_override(mock_gpu_script, mock_recipe_load, sagema
 
 @patch("sagemaker.pytorch.estimator.PyTorch._recipe_load")
 @patch("sagemaker.pytorch.estimator._get_training_recipe_gpu_script")
-def test_training_recipe_gpu_custom_source_dir(mock_gpu_script, mock_recipe_load, sagemaker_session):
+def test_training_recipe_gpu_custom_source_dir(
+    mock_gpu_script, mock_recipe_load, sagemaker_session
+):
     from omegaconf import OmegaConf
-    
+
     # Mock the GPU script function to return the expected entry point
     mock_gpu_script.return_value = "mistral_pretrain.py"
-    
+
     # Mock the recipe structure that would be loaded
-    mock_recipe = OmegaConf.create({
-        "trainer": {
-            "num_nodes": 1,
-        },
-        "model": {
-            "model_type": "mistral",
-        },
-    })
+    mock_recipe = OmegaConf.create(
+        {
+            "trainer": {
+                "num_nodes": 1,
+            },
+            "model": {
+                "model_type": "mistral",
+            },
+        }
+    )
     mock_recipe_load.return_value = ("hf_llama3_8b_seq8k_gpu_p5x16_pretrain", mock_recipe)
-    
+
     container_log_level = '"logging.INFO"'
 
     recipe_overrides = {