Fixed sagemaker either (#15156)

This commit is contained in:
Onkar Ravgan 2024-02-13 10:17:30 +05:30 committed by GitHub
parent 75eb54a0c8
commit 3cc074ba5a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -24,6 +24,9 @@ from metadata.generated.schema.entity.data.mlmodel import (
from metadata.generated.schema.entity.services.connections.mlmodel.sageMakerConnection import ( from metadata.generated.schema.entity.services.connections.mlmodel.sageMakerConnection import (
SageMakerConnection, SageMakerConnection,
) )
from metadata.generated.schema.entity.services.ingestionPipelines.status import (
StackTraceError,
)
from metadata.generated.schema.metadataIngestion.workflow import ( from metadata.generated.schema.metadataIngestion.workflow import (
Source as WorkflowSource, Source as WorkflowSource,
) )
@ -128,14 +131,23 @@ class SagemakerSource(MlModelServiceSource):
""" """
Prepare the Request model Prepare the Request model
""" """
mlmodel_request = CreateMlModelRequest( try:
name=model.name, mlmodel_request = CreateMlModelRequest(
algorithm=self._get_algorithm(), # Setting this to a constant name=model.name,
mlStore=self._get_ml_store(model.name), algorithm=self._get_algorithm(), # Setting this to a constant
service=self.context.mlmodel_service, mlStore=self._get_ml_store(model.name),
) service=self.context.mlmodel_service,
yield mlmodel_request )
self.register_record(mlmodel_request=mlmodel_request) yield Either(right=mlmodel_request)
self.register_record(mlmodel_request=mlmodel_request)
except Exception as exc: # pylint: disable=broad-except
yield Either(
left=StackTraceError(
name=model.name,
error=f"Error creating mlmodel: {exc}",
stackTrace=traceback.format_exc(),
)
)
def _get_ml_store( # pylint: disable=arguments-differ def _get_ml_store( # pylint: disable=arguments-differ
self, self,