File tree Expand file tree Collapse file tree 1 file changed +23
-0
lines changed Expand file tree Collapse file tree 1 file changed +23
-0
lines changed Original file line number Diff line number Diff line change @@ -283,6 +283,29 @@ def models(self) -> List[association.Association]:
283
283
]
284
284
return model_list
285
285
286
+ def models_v2 (
287
+ self , direction : LineageQueryDirectionEnum = LineageQueryDirectionEnum .BOTH
288
+ ) -> List [Artifact ]:
289
+ """Get artifacts representing models from the context lineage.
290
+
291
+ Returns:
292
+ list of Artifacts: Artifacts representing a model.
293
+
294
+ """
295
+ query_filter = LineageFilter (
296
+ entities = [LineageEntityEnum .ARTIFACT ], sources = [LineageSourceEnum .MODEL ]
297
+ )
298
+ query_result = LineageQuery (self .sagemaker_session ).query (
299
+ start_arns = [self .context_arn ],
300
+ query_filter = query_filter ,
301
+ direction = direction ,
302
+ include_edges = False ,
303
+ )
304
+ model_artifacts = []
305
+ for vertex in query_result .vertices :
306
+ model_artifacts .append (vertex .to_lineage_object ())
307
+ return model_artifacts
308
+
286
309
def dataset_artifacts (
287
310
self , direction : LineageQueryDirectionEnum = LineageQueryDirectionEnum .ASCENDANTS
288
311
) -> List [Artifact ]:
You can’t perform that action at this time.
0 commit comments