Override region for draft model data source

Jonathan Makunga · Jonathan Makunga · commit 4dca186fe473 · 2024-07-09T12:49:23.000-07:00
diff --git a/src/sagemaker/jumpstart/factory/model.py b/src/sagemaker/jumpstart/factory/model.py
@@ -697,7 +697,7 @@ def _add_additional_model_data_sources_to_kwargs(
     # Append speculative decoding data source from metadata
     speculative_decoding_data_sources = specs.get_speculative_decoding_s3_data_sources()
     for data_source in speculative_decoding_data_sources:
-        data_source.s3_data_source.set_bucket(get_neo_content_bucket())
+        data_source.s3_data_source.set_bucket(get_neo_content_bucket(region=kwargs.region))
     api_shape_additional_model_data_sources = (
         [
             camel_case_to_pascal_case(data_source.to_json())

Original file line number	Diff line number	Diff line change
`@@ -697,7 +697,7 @@ def _add_additional_model_data_sources_to_kwargs(`
`697`	`697`	`# Append speculative decoding data source from metadata`
`698`	`698`	`speculative_decoding_data_sources = specs.get_speculative_decoding_s3_data_sources()`
`699`	`699`	`for data_source in speculative_decoding_data_sources:`
`700`		`- data_source.s3_data_source.set_bucket(get_neo_content_bucket())`
	`700`	`+ data_source.s3_data_source.set_bucket(get_neo_content_bucket(region=kwargs.region))`
`701`	`701`	`api_shape_additional_model_data_sources = (`
`702`	`702`	`[`
`703`	`703`	`camel_case_to_pascal_case(data_source.to_json())`