yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-25 00:48:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

build(ingest/boto3): Update boto3-stubs to fix CI (#8452)

Browse Source
This commit is contained in:
Andrew Sikowitz 2023-07-18 17:29:50 -04:00 committed by GitHub
parent 2a3f6c42d4
commit 48c1dc820e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 33 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
metadata-ingestion/setup.py
View File

@ -413,7 +413,7 @@ mypy_stubs = {
"types-cachetools",
# versions 0.1.13 and 0.1.14 seem to have issues
"types-click==0.1.12",
"boto3-stubs[s3,glue,sagemaker,sts]>=1.28.3",
"boto3-stubs[s3,glue,sagemaker,sts]>=1.28.4",
"types-tabulate",
# avrogen package requires this
"types-pytz",

12
metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/feature_groups.py
View File

@ -23,9 +23,9 @@ from datahub.metadata.schema_classes import (
if TYPE_CHECKING:
from mypy_boto3_sagemaker import SageMakerClient
from mypy_boto3_sagemaker.type_defs import (
DescribeFeatureGroupResponseOutputTypeDef,
DescribeFeatureGroupResponseTypeDef,
FeatureDefinitionOutputTypeDef,
FeatureGroupSummaryOutputTypeDef,
FeatureGroupSummaryTypeDef,
)
@ -35,7 +35,7 @@ class FeatureGroupProcessor:
env: str
report: SagemakerSourceReport
def get_all_feature_groups(self) -> List["FeatureGroupSummaryOutputTypeDef"]:
def get_all_feature_groups(self) -> List["FeatureGroupSummaryTypeDef"]:
"""
List all feature groups in SageMaker.
"""
@ -51,7 +51,7 @@ class FeatureGroupProcessor:
def get_feature_group_details(
self, feature_group_name: str
) -> "DescribeFeatureGroupResponseOutputTypeDef":
) -> "DescribeFeatureGroupResponseTypeDef":
"""
Get details of a feature group (including list of component features).
"""
@ -75,7 +75,7 @@ class FeatureGroupProcessor:
return feature_group
def get_feature_group_wu(
self, feature_group_details: "DescribeFeatureGroupResponseOutputTypeDef"
self, feature_group_details: "DescribeFeatureGroupResponseTypeDef"
) -> MetadataWorkUnit:
"""
Generate an MLFeatureTable workunit for a SageMaker feature group.
@ -146,7 +146,7 @@ class FeatureGroupProcessor:
def get_feature_wu(
self,
feature_group_details: "DescribeFeatureGroupResponseOutputTypeDef",
feature_group_details: "DescribeFeatureGroupResponseTypeDef",
feature: "FeatureDefinitionOutputTypeDef",
) -> MetadataWorkUnit:
"""

22
metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/lineage.py
View File

@ -9,10 +9,10 @@ from datahub.ingestion.source.aws.sagemaker_processors.common import (
if TYPE_CHECKING:
from mypy_boto3_sagemaker import SageMakerClient
from mypy_boto3_sagemaker.type_defs import (
ActionSummaryOutputTypeDef,
ArtifactSummaryOutputTypeDef,
AssociationSummaryOutputTypeDef,
ContextSummaryOutputTypeDef,
ActionSummaryTypeDef,
ArtifactSummaryTypeDef,
AssociationSummaryTypeDef,
ContextSummaryTypeDef,
)
@ -49,7 +49,7 @@ class LineageProcessor:
nodes: Dict[str, Dict[str, Any]] = field(default_factory=dict)
lineage_info: LineageInfo = field(default_factory=LineageInfo)
def get_all_actions(self) -> List["ActionSummaryOutputTypeDef"]:
def get_all_actions(self) -> List["ActionSummaryTypeDef"]:
"""
List all actions in SageMaker.
"""
@ -63,7 +63,7 @@ class LineageProcessor:
return actions
def get_all_artifacts(self) -> List["ArtifactSummaryOutputTypeDef"]:
def get_all_artifacts(self) -> List["ArtifactSummaryTypeDef"]:
"""
List all artifacts in SageMaker.
"""
@ -77,7 +77,7 @@ class LineageProcessor:
return artifacts
def get_all_contexts(self) -> List["ContextSummaryOutputTypeDef"]:
def get_all_contexts(self) -> List["ContextSummaryTypeDef"]:
"""
List all contexts in SageMaker.
"""
@ -91,9 +91,7 @@ class LineageProcessor:
return contexts
def get_incoming_edges(
self, node_arn: str
) -> List["AssociationSummaryOutputTypeDef"]:
def get_incoming_edges(self, node_arn: str) -> List["AssociationSummaryTypeDef"]:
"""
Get all incoming edges for a node in the lineage graph.
"""
@ -107,9 +105,7 @@ class LineageProcessor:
return edges
def get_outgoing_edges(
self, node_arn: str
) -> List["AssociationSummaryOutputTypeDef"]:
def get_outgoing_edges(self, node_arn: str) -> List["AssociationSummaryTypeDef"]:
"""
Get all outgoing edges for a node in the lineage graph.
"""

34
metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/models.py
View File

@ -45,12 +45,12 @@ from datahub.metadata.schema_classes import (
if TYPE_CHECKING:
from mypy_boto3_sagemaker import SageMakerClient
from mypy_boto3_sagemaker.type_defs import (
DescribeEndpointOutputOutputTypeDef,
DescribeModelOutputOutputTypeDef,
DescribeModelPackageGroupOutputOutputTypeDef,
EndpointSummaryOutputTypeDef,
ModelPackageGroupSummaryOutputTypeDef,
ModelSummaryOutputTypeDef,
DescribeEndpointOutputTypeDef,
DescribeModelOutputTypeDef,
DescribeModelPackageGroupOutputTypeDef,
EndpointSummaryTypeDef,
ModelPackageGroupSummaryTypeDef,
ModelSummaryTypeDef,
)
ENDPOINT_STATUS_MAP: Dict[str, str] = {
@ -91,7 +91,7 @@ class ModelProcessor:
group_arn_to_name: Dict[str, str] = field(default_factory=dict)
def get_all_models(self) -> List["ModelSummaryOutputTypeDef"]:
def get_all_models(self) -> List["ModelSummaryTypeDef"]:
"""
List all models in SageMaker.
"""
@ -105,7 +105,7 @@ class ModelProcessor:
return models
def get_model_details(self, model_name: str) -> "DescribeModelOutputOutputTypeDef":
def get_model_details(self, model_name: str) -> "DescribeModelOutputTypeDef":
"""
Get details of a model.
"""
@ -113,7 +113,7 @@ class ModelProcessor:
# see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_model
return self.sagemaker_client.describe_model(ModelName=model_name)
def get_all_groups(self) -> List["ModelPackageGroupSummaryOutputTypeDef"]:
def get_all_groups(self) -> List["ModelPackageGroupSummaryTypeDef"]:
"""
List all model groups in SageMaker.
"""
@ -128,7 +128,7 @@ class ModelProcessor:
def get_group_details(
self, group_name: str
) -> "DescribeModelPackageGroupOutputOutputTypeDef":
) -> "DescribeModelPackageGroupOutputTypeDef":
"""
Get details of a model group.
"""
@ -138,7 +138,7 @@ class ModelProcessor:
ModelPackageGroupName=group_name
)
def get_all_endpoints(self) -> List["EndpointSummaryOutputTypeDef"]:
def get_all_endpoints(self) -> List["EndpointSummaryTypeDef"]:
endpoints = []
# see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_endpoints
@ -151,7 +151,7 @@ class ModelProcessor:
def get_endpoint_details(
self, endpoint_name: str
) -> "DescribeEndpointOutputOutputTypeDef":
) -> "DescribeEndpointOutputTypeDef":
# see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_endpoint
return self.sagemaker_client.describe_endpoint(EndpointName=endpoint_name)
@ -171,7 +171,7 @@ class ModelProcessor:
return endpoint_status
def get_endpoint_wu(
self, endpoint_details: "DescribeEndpointOutputOutputTypeDef"
self, endpoint_details: "DescribeEndpointOutputTypeDef"
) -> MetadataWorkUnit:
"""a
Get a workunit for an endpoint.
@ -215,7 +215,7 @@ class ModelProcessor:
def get_model_endpoints(
self,
model_details: "DescribeModelOutputOutputTypeDef",
model_details: "DescribeModelOutputTypeDef",
endpoint_arn_to_name: Dict[str, str],
model_image: Optional[str],
model_uri: Optional[str],
@ -244,7 +244,7 @@ class ModelProcessor:
return model_endpoints_sorted
def get_group_wu(
self, group_details: "DescribeModelPackageGroupOutputOutputTypeDef"
self, group_details: "DescribeModelPackageGroupOutputTypeDef"
) -> MetadataWorkUnit:
"""
Get a workunit for a model group.
@ -294,7 +294,7 @@ class ModelProcessor:
return MetadataWorkUnit(id=group_name, mce=mce)
def match_model_jobs(
self, model_details: "DescribeModelOutputOutputTypeDef"
self, model_details: "DescribeModelOutputTypeDef"
) -> Tuple[Set[str], Set[str], List[MLHyperParamClass], List[MLMetricClass]]:
model_training_jobs: Set[str] = set()
model_downstream_jobs: Set[str] = set()
@ -387,7 +387,7 @@ class ModelProcessor:
def get_model_wu(
self,
model_details: "DescribeModelOutputOutputTypeDef",
model_details: "DescribeModelOutputTypeDef",
endpoint_arn_to_name: Dict[str, str],
) -> MetadataWorkUnit:
"""