Skip to content

Commit d8fe643

Browse files
committed
fix: model deploy pipeline fails when PRE_PROD/PROD env vars are empty
This fixes two bugs in the model deployment seed code: 1. Pipeline fails with "Unable to parse environment specification aws:///" when PRE_PROD_ACCOUNT_ID, PRE_PROD_REGION, PROD_ACCOUNT_ID, or PROD_REGION are not provided. Now these values fall back to DEV environment values, allowing single-account deployments. 2. EventBridge rule name exceeds 64 character limit for long project names. Now truncates project name to fit within the limit. Also adds unit tests for the constants module to prevent regression.
1 parent c31182e commit d8fe643

4 files changed

Lines changed: 263 additions & 12 deletions

File tree

modules/sagemaker/sagemaker-templates/templates/model_deploy/seed_code/deploy_app/config/constants.py

Lines changed: 41 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
import json
1919
import os
20-
from typing import Union
20+
from typing import Any, List, Union
2121

2222

2323
def get_bool_env(name: str, default: Union[bool, str] = True) -> bool:
@@ -26,30 +26,60 @@ def get_bool_env(name: str, default: Union[bool, str] = True) -> bool:
2626
return os.getenv(name, default_str).lower() == "true"
2727

2828

29+
def get_env_with_fallback(key: str, fallback: Any, is_json_list: bool = False, warn: bool = False) -> Any:
30+
"""Get environment variable with fallback to provided value.
31+
32+
Args:
33+
key: Environment variable name
34+
fallback: Value to use if env var is not set or empty
35+
is_json_list: If True, parse as JSON list and fall back if empty list
36+
warn: If True, print a warning when fallback is used
37+
38+
Returns:
39+
The environment variable value, or fallback if not set/empty
40+
"""
41+
raw = os.environ.get(key)
42+
if is_json_list:
43+
value: List[str] = json.loads(raw) if raw else []
44+
if not value:
45+
if warn:
46+
print(f"INFO: {key} not provided, using fallback")
47+
return fallback
48+
return value
49+
if not raw:
50+
if warn:
51+
print(f"INFO: {key} not provided, using fallback")
52+
return fallback
53+
return raw
54+
55+
2956
MAX_NAME_LENGTH = 63
3057
REPOSITORY_TYPE = os.getenv("REPOSITORY_TYPE", "CodeCommit") # Default to CODECOMMIT if not set
3158
CODE_CONNECTION_ARN = os.getenv("CODE_CONNECTION_ARN", "")
3259
SOURCE_REPOSITORY = os.getenv("SOURCE_REPOSITORY", "")
3360
MODEL_BUCKET_ARN = os.environ["MODEL_BUCKET_ARN"]
3461
MODEL_PACKAGE_GROUP_NAME = os.getenv("MODEL_PACKAGE_GROUP_NAME", "")
3562

63+
# DEV environment configuration (required)
3664
DEV_ACCOUNT_ID = os.environ["DEV_ACCOUNT_ID"]
3765
DEV_REGION = os.environ["DEV_REGION"]
3866
DEV_VPC_ID = os.environ["DEV_VPC_ID"]
3967
DEV_SUBNET_IDS = json.loads(os.environ["DEV_SUBNET_IDS"])
4068
DEV_SECURITY_GROUP_IDS = json.loads(os.environ["DEV_SECURITY_GROUP_IDS"])
4169

42-
PRE_PROD_ACCOUNT_ID = os.environ["PRE_PROD_ACCOUNT_ID"]
43-
PRE_PROD_REGION = os.environ["PRE_PROD_REGION"]
44-
PRE_PROD_VPC_ID = os.environ["PRE_PROD_VPC_ID"]
45-
PRE_PROD_SUBNET_IDS = json.loads(os.environ["PRE_PROD_SUBNET_IDS"])
46-
PRE_PROD_SECURITY_GROUP_IDS = json.loads(os.environ["PRE_PROD_SECURITY_GROUP_IDS"])
70+
# PRE_PROD environment configuration (falls back to DEV if not provided)
71+
PRE_PROD_ACCOUNT_ID = get_env_with_fallback("PRE_PROD_ACCOUNT_ID", DEV_ACCOUNT_ID, warn=True)
72+
PRE_PROD_REGION = get_env_with_fallback("PRE_PROD_REGION", DEV_REGION, warn=True)
73+
PRE_PROD_VPC_ID = get_env_with_fallback("PRE_PROD_VPC_ID", DEV_VPC_ID)
74+
PRE_PROD_SUBNET_IDS = get_env_with_fallback("PRE_PROD_SUBNET_IDS", DEV_SUBNET_IDS, is_json_list=True)
75+
PRE_PROD_SECURITY_GROUP_IDS = get_env_with_fallback("PRE_PROD_SECURITY_GROUP_IDS", DEV_SECURITY_GROUP_IDS, is_json_list=True)
4776

48-
PROD_ACCOUNT_ID = os.environ["PROD_ACCOUNT_ID"]
49-
PROD_REGION = os.environ["PROD_REGION"]
50-
PROD_VPC_ID = os.environ["PROD_VPC_ID"]
51-
PROD_SUBNET_IDS = json.loads(os.environ["PROD_SUBNET_IDS"])
52-
PROD_SECURITY_GROUP_IDS = json.loads(os.environ["PROD_SECURITY_GROUP_IDS"])
77+
# PROD environment configuration (falls back to DEV if not provided)
78+
PROD_ACCOUNT_ID = get_env_with_fallback("PROD_ACCOUNT_ID", DEV_ACCOUNT_ID, warn=True)
79+
PROD_REGION = get_env_with_fallback("PROD_REGION", DEV_REGION, warn=True)
80+
PROD_VPC_ID = get_env_with_fallback("PROD_VPC_ID", DEV_VPC_ID)
81+
PROD_SUBNET_IDS = get_env_with_fallback("PROD_SUBNET_IDS", DEV_SUBNET_IDS, is_json_list=True)
82+
PROD_SECURITY_GROUP_IDS = get_env_with_fallback("PROD_SECURITY_GROUP_IDS", DEV_SECURITY_GROUP_IDS, is_json_list=True)
5383

5484
PROJECT_NAME = os.getenv("PROJECT_NAME", "")
5585
PROJECT_ID = os.getenv("PROJECT_ID", "")

modules/sagemaker/sagemaker-templates/templates/model_deploy/seed_code/deploy_app/deploy_app/pipeline_stack.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,10 +223,15 @@ def __init__(self, scope: Construct, construct_id: str, **kwargs: Any) -> None:
223223

224224
# Add EventBridge rule to trigger pipeline when model is approved in Model Registry
225225
if constants.ENABLE_EVENTBRIDGE_TRIGGER:
226+
# Truncate project name to fit within 64 char limit for EventBridge rule names
227+
rule_suffix = "-model-approval-trigger"
228+
max_prefix_len = constants.MAX_NAME_LENGTH - len(rule_suffix)
229+
truncated_project_name = constants.PROJECT_NAME[:max_prefix_len]
230+
226231
events.Rule(
227232
self,
228233
"ModelApprovalEventRule",
229-
rule_name=f"{constants.PROJECT_NAME}-model-approval-trigger",
234+
rule_name=f"{truncated_project_name}{rule_suffix}",
230235
event_pattern=events.EventPattern(
231236
source=["aws.sagemaker"],
232237
detail_type=["SageMaker Model Package State Change"],
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import os
5+
import sys
6+
from unittest import mock
7+
8+
import pytest
9+
10+
11+
@pytest.fixture(scope="function")
12+
def clean_constants_module():
13+
"""Remove cached constants module to ensure fresh imports with mocked env vars."""
14+
modules_to_remove = [k for k in sys.modules if k.startswith("config")]
15+
for mod in modules_to_remove:
16+
del sys.modules[mod]
17+
18+
yield
19+
20+
modules_to_remove = [k for k in sys.modules if k.startswith("config")]
21+
for mod in modules_to_remove:
22+
del sys.modules[mod]
23+
24+
25+
@pytest.fixture(scope="function")
26+
def base_env_vars():
27+
"""Base environment variables required for constants module to load."""
28+
return {
29+
"DEV_ACCOUNT_ID": "111111111111",
30+
"DEV_REGION": "us-east-1",
31+
"DEV_VPC_ID": "vpc-dev",
32+
"DEV_SUBNET_IDS": '["subnet-dev"]',
33+
"DEV_SECURITY_GROUP_IDS": '["sg-dev"]',
34+
"MODEL_BUCKET_ARN": "arn:aws:s3:::test-bucket",
35+
}
36+
37+
38+
# Tests for get_env_with_fallback helper function
39+
40+
41+
def test_get_env_with_fallback_returns_value_when_set(clean_constants_module, base_env_vars):
42+
"""Should return the environment variable value when it's set."""
43+
with mock.patch.dict(os.environ, {**base_env_vars, "TEST_VAR": "test_value"}, clear=True):
44+
from config.constants import get_env_with_fallback
45+
46+
result = get_env_with_fallback("TEST_VAR", "fallback")
47+
assert result == "test_value"
48+
49+
50+
def test_get_env_with_fallback_returns_fallback_when_not_set(clean_constants_module, base_env_vars):
51+
"""Should return fallback when environment variable is not set."""
52+
with mock.patch.dict(os.environ, base_env_vars, clear=True):
53+
from config.constants import get_env_with_fallback
54+
55+
result = get_env_with_fallback("TEST_VAR", "fallback")
56+
assert result == "fallback"
57+
58+
59+
def test_get_env_with_fallback_returns_fallback_when_empty_string(clean_constants_module, base_env_vars):
60+
"""Should return fallback when environment variable is empty string."""
61+
with mock.patch.dict(os.environ, {**base_env_vars, "TEST_VAR": ""}, clear=True):
62+
from config.constants import get_env_with_fallback
63+
64+
result = get_env_with_fallback("TEST_VAR", "fallback")
65+
assert result == "fallback"
66+
67+
68+
def test_get_env_with_fallback_parses_json_list(clean_constants_module, base_env_vars):
69+
"""Should parse JSON list when is_json_list=True."""
70+
with mock.patch.dict(os.environ, {**base_env_vars, "TEST_LIST": '["a", "b", "c"]'}, clear=True):
71+
from config.constants import get_env_with_fallback
72+
73+
result = get_env_with_fallback("TEST_LIST", ["default"], is_json_list=True)
74+
assert result == ["a", "b", "c"]
75+
76+
77+
def test_get_env_with_fallback_json_list_fallback_when_empty(clean_constants_module, base_env_vars):
78+
"""Should return fallback when JSON list is empty."""
79+
with mock.patch.dict(os.environ, {**base_env_vars, "TEST_LIST": "[]"}, clear=True):
80+
from config.constants import get_env_with_fallback
81+
82+
result = get_env_with_fallback("TEST_LIST", ["default"], is_json_list=True)
83+
assert result == ["default"]
84+
85+
86+
def test_get_env_with_fallback_warns_when_enabled(clean_constants_module, base_env_vars, capsys):
87+
"""Should print warning when warn=True and fallback is used."""
88+
with mock.patch.dict(os.environ, base_env_vars, clear=True):
89+
from config.constants import get_env_with_fallback
90+
91+
get_env_with_fallback("TEST_VAR", "fallback", warn=True)
92+
captured = capsys.readouterr()
93+
assert "INFO: TEST_VAR not provided, using fallback" in captured.out
94+
95+
96+
def test_get_env_with_fallback_no_warning_by_default(clean_constants_module, base_env_vars, capsys):
97+
"""Should not print warning when warn=False (default)."""
98+
with mock.patch.dict(os.environ, base_env_vars, clear=True):
99+
from config.constants import get_env_with_fallback
100+
101+
# Clear any output from module load (PRE_PROD/PROD fallback warnings)
102+
capsys.readouterr()
103+
104+
get_env_with_fallback("TEST_VAR", "fallback")
105+
captured = capsys.readouterr()
106+
assert "TEST_VAR" not in captured.out
107+
108+
109+
# Tests for PRE_PROD/PROD fallback behavior
110+
111+
112+
def test_preprod_falls_back_to_dev_when_not_set(clean_constants_module, base_env_vars):
113+
"""PRE_PROD should use DEV values when not provided."""
114+
with mock.patch.dict(os.environ, base_env_vars, clear=True):
115+
import config.constants as constants
116+
117+
assert constants.PRE_PROD_ACCOUNT_ID == "111111111111"
118+
assert constants.PRE_PROD_REGION == "us-east-1"
119+
assert constants.PRE_PROD_VPC_ID == "vpc-dev"
120+
assert constants.PRE_PROD_SUBNET_IDS == ["subnet-dev"]
121+
assert constants.PRE_PROD_SECURITY_GROUP_IDS == ["sg-dev"]
122+
123+
124+
def test_prod_falls_back_to_dev_when_not_set(clean_constants_module, base_env_vars):
125+
"""PROD should use DEV values when not provided."""
126+
with mock.patch.dict(os.environ, base_env_vars, clear=True):
127+
import config.constants as constants
128+
129+
assert constants.PROD_ACCOUNT_ID == "111111111111"
130+
assert constants.PROD_REGION == "us-east-1"
131+
assert constants.PROD_VPC_ID == "vpc-dev"
132+
assert constants.PROD_SUBNET_IDS == ["subnet-dev"]
133+
assert constants.PROD_SECURITY_GROUP_IDS == ["sg-dev"]
134+
135+
136+
def test_preprod_uses_own_values_when_set(clean_constants_module, base_env_vars):
137+
"""PRE_PROD should use its own values when provided."""
138+
env = {
139+
**base_env_vars,
140+
"PRE_PROD_ACCOUNT_ID": "222222222222",
141+
"PRE_PROD_REGION": "us-west-2",
142+
"PRE_PROD_VPC_ID": "vpc-preprod",
143+
"PRE_PROD_SUBNET_IDS": '["subnet-preprod"]',
144+
"PRE_PROD_SECURITY_GROUP_IDS": '["sg-preprod"]',
145+
}
146+
with mock.patch.dict(os.environ, env, clear=True):
147+
import config.constants as constants
148+
149+
assert constants.PRE_PROD_ACCOUNT_ID == "222222222222"
150+
assert constants.PRE_PROD_REGION == "us-west-2"
151+
assert constants.PRE_PROD_VPC_ID == "vpc-preprod"
152+
assert constants.PRE_PROD_SUBNET_IDS == ["subnet-preprod"]
153+
assert constants.PRE_PROD_SECURITY_GROUP_IDS == ["sg-preprod"]
154+
155+
156+
def test_prod_uses_own_values_when_set(clean_constants_module, base_env_vars):
157+
"""PROD should use its own values when provided."""
158+
env = {
159+
**base_env_vars,
160+
"PROD_ACCOUNT_ID": "333333333333",
161+
"PROD_REGION": "eu-west-1",
162+
"PROD_VPC_ID": "vpc-prod",
163+
"PROD_SUBNET_IDS": '["subnet-prod"]',
164+
"PROD_SECURITY_GROUP_IDS": '["sg-prod"]',
165+
}
166+
with mock.patch.dict(os.environ, env, clear=True):
167+
import config.constants as constants
168+
169+
assert constants.PROD_ACCOUNT_ID == "333333333333"
170+
assert constants.PROD_REGION == "eu-west-1"
171+
assert constants.PROD_VPC_ID == "vpc-prod"
172+
assert constants.PROD_SUBNET_IDS == ["subnet-prod"]
173+
assert constants.PROD_SECURITY_GROUP_IDS == ["sg-prod"]
174+
175+
176+
def test_empty_string_treated_as_not_set(clean_constants_module, base_env_vars):
177+
"""Empty string values should fall back to DEV."""
178+
env = {
179+
**base_env_vars,
180+
"PRE_PROD_ACCOUNT_ID": "",
181+
"PRE_PROD_REGION": "",
182+
"PROD_ACCOUNT_ID": "",
183+
"PROD_REGION": "",
184+
}
185+
with mock.patch.dict(os.environ, env, clear=True):
186+
import config.constants as constants
187+
188+
assert constants.PRE_PROD_ACCOUNT_ID == "111111111111"
189+
assert constants.PRE_PROD_REGION == "us-east-1"
190+
assert constants.PROD_ACCOUNT_ID == "111111111111"
191+
assert constants.PROD_REGION == "us-east-1"
192+
193+
194+
def test_warns_for_account_and_region_fallback(clean_constants_module, base_env_vars, capsys):
195+
"""Should print warnings when ACCOUNT_ID and REGION fall back."""
196+
with mock.patch.dict(os.environ, base_env_vars, clear=True):
197+
import config.constants # noqa: F401
198+
199+
captured = capsys.readouterr()
200+
assert "PRE_PROD_ACCOUNT_ID not provided" in captured.out
201+
assert "PRE_PROD_REGION not provided" in captured.out
202+
assert "PROD_ACCOUNT_ID not provided" in captured.out
203+
assert "PROD_REGION not provided" in captured.out
204+
205+
206+
def test_no_warning_for_vpc_subnet_sg_fallback(clean_constants_module, base_env_vars, capsys):
207+
"""Should NOT print warnings when VPC/subnet/SG fall back."""
208+
with mock.patch.dict(os.environ, base_env_vars, clear=True):
209+
import config.constants # noqa: F401
210+
211+
captured = capsys.readouterr()
212+
assert "VPC_ID not provided" not in captured.out
213+
assert "SUBNET_IDS not provided" not in captured.out
214+
assert "SECURITY_GROUP_IDS not provided" not in captured.out

0 commit comments

Comments
 (0)