Add an option for user to remove inputs and container artifacts when using local model trainer

Author: pintaoz-aws

src/sagemaker/modules/local_core/local_container.py

@@ -108,6 +108,8 @@ class _LocalContainer(BaseModel):
     container_entrypoint: Optional[List[str]]
     container_arguments: Optional[List[str]]
 
+    _temperary_folders: List[str] = []
+
     def model_post_init(self, __context: Any):
         """Post init method to perform custom validation and set default values."""
         self.hosts = [f"algo-{i}" for i in range(1, self.instance_count + 1)]
@@ -146,12 +148,15 @@ def model_post_init(self, __context: Any):
     def train(
         self,
         wait: bool,
+        remove_inputs_and_container_artifacts: Optional[bool] = True,
     ) -> str:
         """Run a training job locally using docker-compose.
 
         Args:
             wait (bool):
                 Whether to wait the training output before exiting.
+            remove_inputs_and_container_artifacts (Optional[bool]):
+                Whether to remove inputs and container artifacts after training.
         """
         # create output/data folder since sagemaker-containers 2.0 expects it
         os.makedirs(os.path.join(self.container_root, "output", "data"), exist_ok=True)
@@ -201,6 +206,13 @@ def train(
 
         # Print our Job Complete line
         logger.info("Local training job completed, output artifacts saved to %s", artifacts)
+
+        if remove_inputs_and_container_artifacts:
+            shutil.rmtree(os.path.join(self.container_root, "input"))
+            for host in self.hosts:
+                shutil.rmtree(os.path.join(self.container_root, host))
+            for folder in self._temperary_folders:
+                shutil.rmtree(os.path.join(self.container_root, folder))
         return artifacts
 
     def retrieve_artifacts(
@@ -540,6 +552,7 @@ def _get_data_source_local_path(self, data_source: DataSource):
             uri = data_source.s3_data_source.s3_uri
             parsed_uri = urlparse(uri)
             local_dir = TemporaryDirectory(prefix=os.path.join(self.container_root + "/")).name
+            self._temperary_folders.append(local_dir)
             download_folder(parsed_uri.netloc, parsed_uri.path, local_dir, self.sagemaker_session)
             return local_dir
         else:

src/sagemaker/modules/train/model_trainer.py

@@ -203,6 +203,8 @@ class ModelTrainer(BaseModel):
         local_container_root (Optional[str]):
             The local root directory to store artifacts from a training job launched in
             "LOCAL_CONTAINER" mode.
+        remove_inputs_and_container_artifacts (Optional[bool]):
+                Whether to remove inputs and container artifacts after training.
     """
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
@@ -227,6 +229,7 @@ class ModelTrainer(BaseModel):
     hyperparameters: Optional[Dict[str, Any]] = {}
     tags: Optional[List[Tag]] = None
     local_container_root: Optional[str] = os.getcwd()
+    remove_inputs_and_container_artifacts: Optional[bool] = True
 
     # Created Artifacts
     _latest_training_job: Optional[resources.TrainingJob] = PrivateAttr(default=None)
@@ -646,7 +649,7 @@ def train(
                 hyper_parameters=string_hyper_parameters,
                 environment=self.environment,
             )
-            local_container.train(wait)
+            local_container.train(wait, self.remove_inputs_and_container_artifacts)
 
     def create_input_data_channel(
         self, channel_name: str, data_source: DataSourceType, key_prefix: Optional[str] = None

tests/integ/sagemaker/modules/train/test_local_model_trainer.py

@@ -92,18 +92,15 @@ def test_single_container_local_mode_local_data(modules_sagemaker_session):
                 "compressed_artifacts",
                 "artifacts",
                 "model",
-                "shared",
-                "input",
                 "output",
-                "algo-1",
             ]
 
             for directory in directories:
                 path = os.path.join(CWD, directory)
                 delete_local_path(path)
 
 
-def test_single_container_local_mode_s3_data(modules_sagemaker_session):
+def test_single_container_local_mode_s3_data_remove_input(modules_sagemaker_session):
     with lock.lock(LOCK_PATH):
         try:
             # upload local data to s3
@@ -145,6 +142,70 @@ def test_single_container_local_mode_s3_data(modules_sagemaker_session):
                 training_mode=Mode.LOCAL_CONTAINER,
             )
 
+            model_trainer.train()
+            assert os.path.exists(os.path.join(CWD, "compressed_artifacts/model.tar.gz"))
+        finally:
+            subprocess.run(["docker", "compose", "down", "-v"])
+
+            assert not os.path.exists(os.path.join(CWD, "shared"))
+            assert not os.path.exists(os.path.join(CWD, "input"))
+            assert not os.path.exists(os.path.join(CWD, "algo-1"))
+
+            directories = [
+                "compressed_artifacts",
+                "artifacts",
+                "model",
+                "output",
+            ]
+
+            for directory in directories:
+                path = os.path.join(CWD, directory)
+                delete_local_path(path)
+
+
+def test_single_container_local_mode_s3_data_not_remove_input(modules_sagemaker_session):
+    with lock.lock(LOCK_PATH):
+        try:
+            # upload local data to s3
+            session = modules_sagemaker_session
+            bucket = session.default_bucket()
+            session.upload_data(
+                path=os.path.join(SOURCE_DIR, "data/train/"),
+                bucket=bucket,
+                key_prefix="data/train",
+            )
+            session.upload_data(
+                path=os.path.join(SOURCE_DIR, "data/test/"),
+                bucket=bucket,
+                key_prefix="data/test",
+            )
+
+            source_code = SourceCode(
+                source_dir=SOURCE_DIR,
+                entry_script="local_training_script.py",
+            )
+
+            compute = Compute(
+                instance_type="local_cpu",
+                instance_count=1,
+            )
+
+            # read input data from s3
+            train_data = InputData(channel_name="train", data_source=f"s3://{bucket}/data/train/")
+
+            test_data = InputData(channel_name="test", data_source=f"s3://{bucket}/data/test/")
+
+            model_trainer = ModelTrainer(
+                training_image=DEFAULT_CPU_IMAGE,
+                sagemaker_session=modules_sagemaker_session,
+                source_code=source_code,
+                compute=compute,
+                input_data_config=[train_data, test_data],
+                base_job_name="local_mode_single_container_s3_data",
+                training_mode=Mode.LOCAL_CONTAINER,
+                remove_inputs_and_container_artifacts=False,
+            )
+
             model_trainer.train()
             assert os.path.exists(os.path.join(CWD, "compressed_artifacts/model.tar.gz"))
         finally:
@@ -213,11 +274,7 @@ def test_multi_container_local_mode(modules_sagemaker_session):
                 "compressed_artifacts",
                 "artifacts",
                 "model",
-                "shared",
-                "input",
                 "output",
-                "algo-1",
-                "algo-2",
             ]
 
             for directory in directories: