Skip to content

Commit 78eddfd

Browse files
authored
infra: add TF batch transform integ test with KMS and network isolation (#1413)
1 parent 9f81006 commit 78eddfd

File tree

4 files changed

+861
-3
lines changed

4 files changed

+861
-3
lines changed

src/sagemaker/tensorflow/estimator.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -590,6 +590,9 @@ def create_model(
590590
if "name" not in kwargs:
591591
kwargs["name"] = self._current_job_name
592592

593+
if "enable_network_isolation" not in kwargs:
594+
kwargs["enable_network_isolation"] = self.enable_network_isolation()
595+
593596
if endpoint_type == "tensorflow-serving" or self._script_mode_enabled():
594597
return self._create_tfs_model(
595598
role=role,
@@ -630,7 +633,6 @@ def _create_tfs_model(
630633
entry_point=entry_point,
631634
source_dir=source_dir,
632635
dependencies=dependencies,
633-
enable_network_isolation=self.enable_network_isolation(),
634636
**kwargs
635637
)
636638

@@ -660,7 +662,6 @@ def _create_default_model(
660662
sagemaker_session=self.sagemaker_session,
661663
vpc_config=self.get_vpc_config(vpc_config_override),
662664
dependencies=dependencies or self.dependencies,
663-
enable_network_isolation=self.enable_network_isolation(),
664665
**kwargs
665666
)
666667

0 commit comments

Comments
 (0)