Skip to content

Commit d959a9d

Browse files
nargokulnargokul
committed
fixes
1 parent ebcab5c commit d959a9d

File tree

5 files changed

+255
-58
lines changed

5 files changed

+255
-58
lines changed

‎sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_telemetry_integration.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# Licensed under the Apache License, Version 2.0
3+
4+
"""Tests for telemetry integration in feature store interfaces."""
5+
6+
import pytest
7+
from unittest.mock import Mock, patch
8+
import pandas as pd
9+
10+
from sagemaker.core.telemetry import Feature
11+
from sagemaker.mlops.feature_store.dataset_builder import DatasetBuilder
12+
from sagemaker.mlops.feature_store.ingestion_manager_pandas import IngestionManagerPandas
13+
from sagemaker.mlops.feature_store.athena_query import AthenaQuery
14+
15+
16+
class TestFeatureStoreTelemetry:
17+
"""Test telemetry integration for feature store interfaces."""
18+
19+
@patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request")
20+
def test_dataset_builder_telemetry(self, mock_send_telemetry):
21+
"""Test that DatasetBuilder methods emit telemetry."""
22+
mock_session = Mock()
23+
mock_session.sagemaker_config = None
24+
mock_session.local_mode = False
25+
26+
# Create a simple DataFrame
27+
df = pd.DataFrame({"feature1": [1, 2, 3], "feature2": ["a", "b", "c"]})
28+
29+
builder = DatasetBuilder.create(
30+
base=df,
31+
output_path="s3://test-bucket/output",
32+
session=mock_session,
33+
record_identifier_feature_name="feature1",
34+
event_time_identifier_feature_name="feature2"
35+
)
36+
37+
# Mock the internal methods to avoid actual S3/Athena calls
38+
with patch.object(builder, '_to_csv_from_dataframe', return_value=("s3://test/file.csv", "SELECT * FROM test")):
39+
result = builder.to_csv_file()
40+
41+
# Verify telemetry was called
42+
assert mock_send_telemetry.called
43+
# Verify the feature code for FEATURE_STORE was used
44+
call_args = mock_send_telemetry.call_args[0]
45+
assert 17 in call_args[1] # FEATURE_STORE code
46+
47+
@patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request")
48+
def test_ingestion_manager_telemetry(self, mock_send_telemetry):
49+
"""Test that IngestionManagerPandas.run emits telemetry."""
50+
mock_session = Mock()
51+
mock_session.sagemaker_config = None
52+
mock_session.local_mode = False
53+
54+
feature_definitions = {
55+
"feature1": {"FeatureType": "Integral", "CollectionType": None},
56+
"feature2": {"FeatureType": "String", "CollectionType": None}
57+
}
58+
59+
manager = IngestionManagerPandas(
60+
feature_group_name="test-fg",
61+
feature_definitions=feature_definitions,
62+
max_workers=1,
63+
max_processes=1
64+
)
65+
66+
# Add session to manager for telemetry
67+
manager.sagemaker_session = mock_session
68+
69+
df = pd.DataFrame({"feature1": [1, 2, 3], "feature2": ["a", "b", "c"]})
70+
71+
# Mock the internal methods to avoid actual ingestion
72+
with patch.object(manager, '_run_single_process_single_thread'):
73+
manager.run(df, wait=True)
74+
75+
# Verify telemetry was called
76+
assert mock_send_telemetry.called
77+
# Verify the feature code for FEATURE_STORE was used
78+
call_args = mock_send_telemetry.call_args[0]
79+
assert 17 in call_args[1] # FEATURE_STORE code
80+
81+
@patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request")
82+
def test_athena_query_telemetry(self, mock_send_telemetry):
83+
"""Test that AthenaQuery methods emit telemetry."""
84+
mock_session = Mock()
85+
mock_session.sagemaker_config = None
86+
mock_session.local_mode = False
87+
88+
query = AthenaQuery(
89+
catalog="test_catalog",
90+
database="test_db",
91+
table_name="test_table",
92+
sagemaker_session=mock_session
93+
)
94+
95+
# Mock the internal methods to avoid actual Athena calls
96+
with patch("sagemaker.mlops.feature_store.feature_utils.start_query_execution",
97+
return_value={"QueryExecutionId": "test-id"}):
98+
query.run("SELECT * FROM test", "s3://test-bucket/output")
99+
100+
# Verify telemetry was called
101+
assert mock_send_telemetry.called
102+
# Verify the feature code for FEATURE_STORE was used
103+
call_args = mock_send_telemetry.call_args[0]
104+
assert 17 in call_args[1] # FEATURE_STORE code

‎sagemaker-serve/src/sagemaker/serve/model_builder.py‎

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -804,6 +804,14 @@ def _resolve_compute_requirements_from_config(
804804
)
805805
default_cpus = actual_cpus
806806

807+
# Adjust memory if it exceeds instance capacity
808+
if actual_memory_mb and default_memory_mb > actual_memory_mb:
809+
logger.warning(
810+
f"Default requirements request {default_memory_mb} MB memory but {instance_type} has {actual_memory_mb} MB. "
811+
f"Adjusting to {actual_memory_mb} MB."
812+
)
813+
default_memory_mb = actual_memory_mb
814+
807815
# Initialize with defaults
808816
final_cpus = default_cpus
809817
final_min_memory = default_memory_mb
@@ -852,21 +860,23 @@ def _resolve_compute_requirements_from_config(
852860
)
853861

854862
# Validate requirements are compatible with instance type
855-
if actual_cpus and final_cpus > actual_cpus:
856-
raise ValueError(
857-
f"Resource requirements incompatible with instance type '{instance_type}'.\n"
858-
f"Requested: {final_cpus} CPUs\n"
859-
f"Available: {actual_cpus} CPUs\n"
860-
f"Please reduce CPU requirements or select a larger instance type."
861-
)
863+
# Only validate user-provided requirements (defaults are already adjusted above)
864+
if user_resource_requirements:
865+
if actual_cpus and user_resource_requirements.num_cpus is not None and user_resource_requirements.num_cpus > actual_cpus:
866+
raise ValueError(
867+
f"Resource requirements incompatible with instance type '{instance_type}'.\n"
868+
f"Requested: {user_resource_requirements.num_cpus} CPUs\n"
869+
f"Available: {actual_cpus} CPUs\n"
870+
f"Please reduce CPU requirements or select a larger instance type."
871+
)
862872

863-
if actual_memory_mb and final_min_memory > actual_memory_mb:
864-
raise ValueError(
865-
f"Resource requirements incompatible with instance type '{instance_type}'.\n"
866-
f"Requested: {final_min_memory} MB memory\n"
867-
f"Available: {actual_memory_mb} MB memory\n"
868-
f"Please reduce memory requirements or select a larger instance type."
869-
)
873+
if actual_memory_mb and user_resource_requirements.min_memory is not None and user_resource_requirements.min_memory > actual_memory_mb:
874+
raise ValueError(
875+
f"Resource requirements incompatible with instance type '{instance_type}'.\n"
876+
f"Requested: {user_resource_requirements.min_memory} MB memory\n"
877+
f"Available: {actual_memory_mb} MB memory\n"
878+
f"Please reduce memory requirements or select a larger instance type."
879+
)
870880

871881
# Create and return InferenceComponentComputeResourceRequirements
872882
requirements = InferenceComponentComputeResourceRequirements(

‎sagemaker-serve/tests/unit/test_artifact_path_resolution.py‎

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,8 @@ def test_base_model_artifact_uri_retrieval(
7575
},
7676
mode=Mode.SAGEMAKER_ENDPOINT,
7777
role_arn="arn:aws:iam::123456789012:role/TestRole",
78-
sagemaker_session=self.mock_session
78+
sagemaker_session=self.mock_session,
79+
instance_type="ml.m5.xlarge"
7980
)
8081

8182
# Execute
@@ -119,7 +120,8 @@ def test_fine_tuned_adapter_artifact_location(
119120
},
120121
mode=Mode.SAGEMAKER_ENDPOINT,
121122
role_arn="arn:aws:iam::123456789012:role/TestRole",
122-
sagemaker_session=self.mock_session
123+
sagemaker_session=self.mock_session,
124+
instance_type="ml.m5.xlarge"
123125
)
124126

125127
# Execute
@@ -148,7 +150,8 @@ def test_lora_adapter_returns_none(
148150
},
149151
mode=Mode.SAGEMAKER_ENDPOINT,
150152
role_arn="arn:aws:iam::123456789012:role/TestRole",
151-
sagemaker_session=self.mock_session
153+
sagemaker_session=self.mock_session,
154+
instance_type="ml.m5.xlarge"
152155
)
153156

154157
# Execute
@@ -194,7 +197,8 @@ def test_missing_hosting_artifact_uri_returns_none(
194197
},
195198
mode=Mode.SAGEMAKER_ENDPOINT,
196199
role_arn="arn:aws:iam::123456789012:role/TestRole",
197-
sagemaker_session=self.mock_session
200+
sagemaker_session=self.mock_session,
201+
instance_type="ml.m5.xlarge"
198202
)
199203

200204
# Execute
@@ -240,7 +244,8 @@ def test_hub_document_fetch_exception_handling(
240244
},
241245
mode=Mode.SAGEMAKER_ENDPOINT,
242246
role_arn="arn:aws:iam::123456789012:role/TestRole",
243-
sagemaker_session=self.mock_session
247+
sagemaker_session=self.mock_session,
248+
instance_type="ml.m5.xlarge"
244249
)
245250

246251
# Execute - should not raise exception
@@ -262,7 +267,8 @@ def test_non_model_customization_returns_none(
262267
model="my-model",
263268
mode=Mode.SAGEMAKER_ENDPOINT,
264269
role_arn="arn:aws:iam::123456789012:role/TestRole",
265-
sagemaker_session=self.mock_session
270+
sagemaker_session=self.mock_session,
271+
instance_type="ml.m5.xlarge" # Provide instance_type to avoid auto-detection
266272
)
267273

268274
# Execute
@@ -294,7 +300,8 @@ def test_missing_model_package_returns_none(
294300
},
295301
mode=Mode.SAGEMAKER_ENDPOINT,
296302
role_arn="arn:aws:iam::123456789012:role/TestRole",
297-
sagemaker_session=self.mock_session
303+
sagemaker_session=self.mock_session,
304+
instance_type="ml.m5.xlarge"
298305
)
299306

300307
# Execute
@@ -329,7 +336,8 @@ def test_missing_inference_specification_returns_none(
329336
},
330337
mode=Mode.SAGEMAKER_ENDPOINT,
331338
role_arn="arn:aws:iam::123456789012:role/TestRole",
332-
sagemaker_session=self.mock_session
339+
sagemaker_session=self.mock_session,
340+
instance_type="ml.m5.xlarge"
333341
)
334342

335343
# Execute
@@ -365,7 +373,8 @@ def test_empty_containers_list_returns_none(
365373
},
366374
mode=Mode.SAGEMAKER_ENDPOINT,
367375
role_arn="arn:aws:iam::123456789012:role/TestRole",
368-
sagemaker_session=self.mock_session
376+
sagemaker_session=self.mock_session,
377+
instance_type="ml.m5.xlarge"
369378
)
370379

371380
# Execute
@@ -406,7 +415,8 @@ def test_base_model_without_base_model_attribute_returns_none(
406415
},
407416
mode=Mode.SAGEMAKER_ENDPOINT,
408417
role_arn="arn:aws:iam::123456789012:role/TestRole",
409-
sagemaker_session=self.mock_session
418+
sagemaker_session=self.mock_session,
419+
instance_type="ml.m5.xlarge"
410420
)
411421

412422
# Execute
@@ -453,7 +463,8 @@ def test_fine_tuned_model_with_nested_s3_data_source(
453463
},
454464
mode=Mode.SAGEMAKER_ENDPOINT,
455465
role_arn="arn:aws:iam::123456789012:role/TestRole",
456-
sagemaker_session=self.mock_session
466+
sagemaker_session=self.mock_session,
467+
instance_type="ml.m5.xlarge"
457468
)
458469

459470
# Execute
@@ -502,7 +513,8 @@ def test_base_model_with_multiple_hosting_artifact_uris(
502513
},
503514
mode=Mode.SAGEMAKER_ENDPOINT,
504515
role_arn="arn:aws:iam::123456789012:role/TestRole",
505-
sagemaker_session=self.mock_session
516+
sagemaker_session=self.mock_session,
517+
instance_type="ml.m5.xlarge"
506518
)
507519

508520
# Execute

0 commit comments

Comments
 (0)