Skip to content

fix(core): read container image from the Model resource in Transformer #6418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AkshayShah03 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix(core): read container image from the Model resource in Transformer #6418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions sagemaker-core/src/sagemaker/core/transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -442,13 +442,11 @@ def _retrieve_image_uri(self):
)
if not model:
return None
model_desc = model.__dict__

primary_container = getattr(model_desc, "primary_container", None)
primary_container = getattr(model, "primary_container", None)
if primary_container:
return getattr(primary_container, "image", None)

containers = getattr(model_desc, "containers", None)
containers = getattr(model, "containers", None)
if containers:
return getattr(containers[0], "image", None)

Expand Down
42 changes: 10 additions & 32 deletions sagemaker-core/tests/unit/test_transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import pytest
from unittest.mock import Mock, patch
from sagemaker.core.transformer import Transformer
from sagemaker.core.resources import Model
from sagemaker.core.shapes import ContainerDefinition


@pytest.fixture
Expand Down Expand Up @@ -271,22 +273,10 @@ def test_prepare_data_processing_partial(self, mock_session):
@patch("sagemaker.core.transformer.Model")
def test_retrieve_image_uri_success(self, mock_model_class, mock_session):
"""Test _retrieve_image_uri with successful model retrieval"""
mock_primary_container = Mock()
mock_primary_container.image = "test-image:latest"

class DictWithAttrs(dict):
"""A dict that also supports attribute access"""

def __getattr__(self, name):
return self.get(name)

class MockModel:
def __init__(self):
self.__dict__ = DictWithAttrs()
self.__dict__["primary_container"] = mock_primary_container
self.__dict__["containers"] = None

mock_model = MockModel()
mock_model = Model(
model_name="test-model",
primary_container=ContainerDefinition(image="test-image:latest"),
)
mock_model_class.get.return_value = mock_model

transformer = Transformer(
Expand All @@ -302,22 +292,10 @@ def __init__(self):
@patch("sagemaker.core.transformer.Model")
def test_retrieve_image_uri_with_containers(self, mock_model_class, mock_session):
"""Test _retrieve_image_uri with containers instead of primary_container"""
mock_container = Mock()
mock_container.image = "container-image:latest"

class DictWithAttrs(dict):
"""A dict that also supports attribute access"""

def __getattr__(self, name):
return self.get(name)

class MockModel:
def __init__(self):
self.__dict__ = DictWithAttrs()
self.__dict__["primary_container"] = None
self.__dict__["containers"] = [mock_container]

mock_model = MockModel()
mock_model = Model(
model_name="test-model",
containers=[ContainerDefinition(image="container-image:latest")],
)
mock_model_class.get.return_value = mock_model

transformer = Transformer(
Expand Down
Loading