def _search(
cls,
search_resource,
search_item_factory,
boto_next_token_name="NextToken",
sagemaker_session=None,
**kwargs
):
"""Search for objects with the SageMaker API."""
sagemaker_session = sagemaker_session or _utils.default_session()
sagemaker_client = sagemaker_session.sagemaker_client
next_token = None
try:
while True:
search_request_kwargs = _boto_functions.to_boto(
kwargs, cls._custom_boto_names, cls._custom_boto_types
)
search_request_kwargs["Resource"] = search_resource
if next_token:
search_request_kwargs[boto_next_token_name] = next_token
search_method = getattr(sagemaker_client, "search")
search_method_response = search_method(**search_request_kwargs)
search_items = search_method_response.get("Results", [])
next_token = search_method_response.get(boto_next_token_name)
for item in search_items:
if cls.__name__ in item:
yield search_item_factory(item[cls.__name__])
if not next_token:
break
except StopIteration:
return
def _list(
cls,
boto_list_method,
list_item_factory,
boto_list_items_name,
boto_next_token_name="NextToken",
sagemaker_session=None,
**kwargs
):
"""List objects from the SageMaker API."""
sagemaker_session = sagemaker_session or _utils.default_session()
sagemaker_client = sagemaker_session.sagemaker_client
next_token = None
try:
while True:
list_request_kwargs = _boto_functions.to_boto(
kwargs, cls._custom_boto_names, cls._custom_boto_types
)
if next_token:
list_request_kwargs[boto_next_token_name] = next_token
list_method = getattr(sagemaker_client, boto_list_method)
list_method_response = list_method(**list_request_kwargs)
list_items = list_method_response.get(boto_list_items_name, [])
next_token = list_method_response.get(boto_next_token_name)
for item in list_items:
yield list_item_factory(item)
if not next_token:
break
except StopIteration:
return
def to_boto(cls, obj):
"""Convert an object to a boto representation.
Args:
obj (dict): The object to convert to boto.
"""
if not isinstance(obj, dict):
var_dict = vars(obj)
else:
var_dict = obj
return _boto_functions.to_boto(var_dict, cls._custom_boto_names, cls._custom_boto_types)
