Fix sagemaker config

src/accelerate/commands/config/config_args.py

@@ -241,3 +241,4 @@ class SageMakerConfig(BaseConfig):
     sagemaker_metrics_file: str = None
     additional_args: dict = None
     dynamo_config: dict = None
+    enable_cpu_affinity: bool = False

src/accelerate/commands/config/default.py

@@ -95,6 +95,7 @@ def write_basic_config(mixed_precision="no", save_location: str = default_json_c
         config["num_processes"] = 1
         config["distributed_type"] = "NO"
     config["debug"] = False
+    config["enable_cpu_affinity"] = False
     config = ClusterConfig(**config)
     config.to_json_file(path)
     return path

tests/test_cli.py

@@ -20,7 +20,7 @@
 import torch
 from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
 
-from accelerate.commands.config.config_args import BaseConfig, ClusterConfig, SageMakerConfig
+from accelerate.commands.config.config_args import BaseConfig, ClusterConfig, SageMakerConfig, load_config_from_file
 from accelerate.commands.estimate import estimate_command, estimate_command_parser, gather_data
 from accelerate.commands.launch import _validate_launch_command, launch_command_parser
 from accelerate.test_utils import execute_subprocess_async
@@ -73,8 +73,9 @@ def test_no_config(self):
         execute_subprocess_async(cmd, env=os.environ.copy())
 
     def test_config_compatibility(self):
+        invalid_configs = ["invalid", "mpi", "sagemaker"]
         for config in sorted(self.test_config_path.glob("**/*.yaml")):
-            if "invalid" in str(config) or "mpi" in str(config):
+            if any(invalid_config in str(config) for invalid_config in invalid_configs):
                 continue
             with self.subTest(config_file=config):
                 cmd = get_launch_command(config_file=config) + [self.test_file_path]
@@ -196,6 +197,8 @@ class ClusterConfigTester(unittest.TestCase):
     Test case for verifying the config dataclasses work
     """
 
+    test_config_path = Path("tests/test_configs")
+
     def test_base_config(self):
         # Tests that all the dataclasses can be initialized
         config = BaseConfig(
@@ -257,6 +260,8 @@ def test_sagemaker_config(self):
         assert config.ec2_instance_type == "MY_TYPE"
         assert config.iam_role_name == "MY_ROLE"
 
+        config = load_config_from_file(str(self.test_config_path / "0_30_0_sagemaker.yaml"))
+
 
 class TpuConfigTester(unittest.TestCase):
     """

tests/test_configs/0_30_0_sagemaker.yaml

@@ -0,0 +1,8 @@
+compute_environment: AMAZON_SAGEMAKER
+debug: false
+distributed_type: NO
+mixed_precision: fp16
+debug: false
+use_cpu: false
+ec2_instance_type: MY_TYPE
+iam_role_name: MY_ROLE