def endpoint_context(self):
"""Retrieves the lineage context object representing the endpoint.
Examples:
.. code-block:: python
predictor = Predictor()
...
context = predictor.endpoint_context()
models = context.models()
Returns:
ContextEndpoint: The context for the endpoint.
"""
if self._context:
return self._context
# retrieve endpoint by name to get arn
response = self.sagemaker_session.sagemaker_client.describe_endpoint(
EndpointName=self.endpoint_name
)
endpoint_arn = response["EndpointArn"]
# list context by source uri using arn
contexts = list(
EndpointContext.list(sagemaker_session=self.sagemaker_session, source_uri=endpoint_arn)
)
if len(contexts) != 0:
# create endpoint context object
self._context = EndpointContext.load(
sagemaker_session=self.sagemaker_session, context_name=contexts[0].context_name
)
return self._context
def to_lineage_object(self):
"""Convert the ``Vertex`` object to its corresponding lineage object.
Returns:
A ``Vertex`` object to its corresponding ``Artifact``,``Action``, ``Context``
or ``TrialComponent`` object.
"""
from sagemaker.lineage.context import Context, EndpointContext
from sagemaker.lineage.action import Action
from sagemaker.lineage.lineage_trial_component import LineageTrialComponent
if self.lineage_entity == LineageEntityEnum.CONTEXT.value:
resource_name = get_resource_name_from_arn(self.arn)
if self.lineage_source == LineageSourceEnum.ENDPOINT.value:
return EndpointContext.load(context_name=resource_name,
sagemaker_session=self._session)
return Context.load(context_name=resource_name,
sagemaker_session=self._session)
if self.lineage_entity == LineageEntityEnum.ARTIFACT.value:
return self._artifact_to_lineage_object()
if self.lineage_entity == LineageEntityEnum.ACTION.value:
return Action.load(action_name=self.arn.split("/")[1],
sagemaker_session=self._session)
if self.lineage_entity == LineageEntityEnum.TRIAL_COMPONENT.value:
trial_component_name = get_resource_name_from_arn(self.arn)
return LineageTrialComponent.load(
trial_component_name=trial_component_name,
sagemaker_session=self._session)
raise ValueError("Vertex cannot be converted to a lineage object.")
def to_lineage_object(self):
"""Convert the ``Vertex`` object to its corresponding ``Artifact`` or ``Context`` object."""
from sagemaker.lineage.artifact import Artifact, ModelArtifact
from sagemaker.lineage.context import Context, EndpointContext
from sagemaker.lineage.artifact import DatasetArtifact
if self.lineage_entity == LineageEntityEnum.CONTEXT.value:
resource_name = get_resource_name_from_arn(self.arn)
if self.lineage_source == LineageSourceEnum.ENDPOINT.value:
return EndpointContext.load(
context_name=resource_name, sagemaker_session=self._session
)
return Context.load(context_name=resource_name, sagemaker_session=self._session)
if self.lineage_entity == LineageEntityEnum.ARTIFACT.value:
if self.lineage_source == LineageSourceEnum.MODEL.value:
return ModelArtifact.load(artifact_arn=self.arn, sagemaker_session=self._session)
if self.lineage_source == LineageSourceEnum.DATASET.value:
return DatasetArtifact.load(artifact_arn=self.arn, sagemaker_session=self._session)
return Artifact.load(artifact_arn=self.arn, sagemaker_session=self._session)
raise ValueError("Vertex cannot be converted to a lineage object.")