/**
 * Copyright 2021 The Kubeflow Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */ import { Struct } from 'google-protobuf/google/protobuf/struct_pb' ; import { logger } from 'src/lib/Utils' ; import { isV2Pipeline } from 'src/lib/v2/WorkflowUtils' ; import { Api , ArtifactCustomProperties , ArtifactProperties , ExecutionCustomProperties , ExecutionProperties , getResourceProperty , getResourcePropertyViaFallBack , } from 'src/mlmd/library' ; import { Artifact , ArtifactType , Context , Event , Execution , GetArtifactsByIDRequest , GetArtifactsByIDResponse , GetArtifactTypesRequest , GetArtifactTypesResponse , GetContextByTypeAndNameRequest , GetEventsByExecutionIDsRequest , GetEventsByExecutionIDsResponse , GetExecutionsByContextRequest , } from 'src/third_party/mlmd' ; import { Workflow } from 'src/third_party/mlmd/argo_template' ; async function getContext ( { type , name } ) { if ( type === '' ) { throw new Error ( 'Failed to getContext: type is empty.' ) ; } if ( name === '' ) { throw new Error ( 'Failed to getContext: name is empty.' ) ; } const request = new GetContextByTypeAndNameRequest ( ) ; request . setTypeName ( type ) ; request . setContextName ( name ) ; try { const res = await Api . getInstance ( ) . metadataStoreService . getContextByTypeAndName ( request ) ; const context = res . getContext ( ) ; if ( context == null ) { throw new Error ( 'Cannot find specified context' ) ; } return context ; } catch ( err ) { err . message = `Cannot find context with ${ JSON . stringify ( request . toObject ( ) ) }: ` + err . message ; throw err ; } } throws async function getTfxRunContext ( argoWorkflowName ) { return await getContext ( { name : argoWorkflowName , type : 'pipeline_run' } ) ; } throws async function getKfpRunContext ( argoWorkflowName ) { return await getContext ( { name : argoWorkflowName , type : 'KfpRun' } ) ; } async function getKfpV2RunContext ( runID ) { return await getContext ( { name : runID , type : 'system.PipelineRun' } ) ; } export async function getRunContext ( workflow , runID ) { console . log ( workflow , runID ) ; const workflowName = workflow ?. metadata ?. name || '' ; if ( isV2Pipeline ( workflow ) ) { return await getKfpV2RunContext ( runID ) ; } try { return await getTfxRunContext ( workflowName ) ; } catch ( err ) { logger . warn ( `Cannot find tfx run context (this is expected for non tfx runs)` , err ) ; return await getKfpRunContext ( workflowName ) ; } } throws export async function getExecutionsFromContext ( context ) { const request = new GetExecutionsByContextRequest ( ) ; request . setContextId ( context . getId ( ) ) ; try { const res = await Api . getInstance ( ) . metadataStoreService . getExecutionsByContext ( request ) ; const list = res . getExecutionsList ( ) ; if ( list == null ) { throw new Error ( 'response.getExecutionsList() is empty' ) ; } return list ; } catch ( err ) { err . message = `Cannot find executions by context ${ context . getId ( ) } with name ${ context . getName ( ) }: ` + err . message ; throw err ; } } export enum KfpExecutionProperties { KFP_POD_NAME = 'kfp_pod_name' , POD_NAME = 'pod_name' , DISPLAY_NAME = 'display_name' , TASK_NAME = 'task_name' } const EXECUTION_PROPERTY_REPOS = [ ExecutionProperties , ExecutionCustomProperties ] ; export const ExecutionHelpers = { getWorkspace ( execution ) { return ( getResourcePropertyViaFallBack ( execution , EXECUTION_PROPERTY_REPOS , [ 'RUN_ID' ] ) || getStringProperty ( execution , ExecutionCustomProperties . WORKSPACE , true ) || getStringProperty ( execution , ExecutionProperties . PIPELINE_NAME ) || undefined ) ; } , getName ( execution ) { return `${ getStringProperty ( execution , KfpExecutionProperties . DISPLAY_NAME , true ) || getStringProperty ( execution , KfpExecutionProperties . TASK_NAME , true ) || getStringProperty ( execution , ExecutionProperties . NAME ) || getStringProperty ( execution , ExecutionProperties . COMPONENT_ID ) || getStringProperty ( execution , ExecutionCustomProperties . TASK_ID , true ) || getStringProperty ( execution , KfpExecutionProperties . KFP_POD_NAME , true ) || '(No name)' }` ; } , getState ( execution ) { return getStringProperty ( execution , ExecutionProperties . STATE ) || undefined ; } , getKfpPod ( execution ) { return ( getStringProperty ( execution , KfpExecutionProperties . POD_NAME , true ) || getStringProperty ( execution , KfpExecutionProperties . KFP_POD_NAME ) || getStringProperty ( execution , KfpExecutionProperties . KFP_POD_NAME , true ) || undefined ) ; } , } ; export enum KfpArtifactProperties { DISPLAY_NAME = 'display_name' } export const ArtifactHelpers = { getName ( a ) { const name = getResourceProperty ( a , KfpArtifactProperties . DISPLAY_NAME , true ) || getResourceProperty ( a , ArtifactProperties . NAME ) || getResourceProperty ( a , ArtifactCustomProperties . NAME , true ) || '(No name)' ; return `${ name }` ; } , } ; function getStringProperty ( resource , propertyName , fromCustomProperties = false ) { const value = getResourceProperty ( resource , propertyName , fromCustomProperties ) ; return getStringValue ( value ) ; } function getStringValue ( value ? ) { if ( typeof value != 'string' ) { return undefined ; } return value ; } async function getEventByExecution ( execution ) { const executionId = execution . getId ( ) ; if ( ! executionId ) { throw new Error ( 'Execution must have an ID' ) ; } const request = new GetEventsByExecutionIDsRequest ( ) . addExecutionIds ( executionId ) ; let response ; try { response = await Api . getInstance ( ) . metadataStoreService . getEventsByExecutionIDs ( request ) ; } catch ( err ) { err . message = 'Failed to getEventsByExecutionIDs: ' + err . message ; throw err ; } return response . getEventsList ( ) ; } export interface LinkedArtifact { event : Event ; artifact : Artifact ; } export async function getLinkedArtifactsByEvents ( events ) { const artifactIds = events . filter ( event => event . getArtifactId ( ) ) . map ( event => event . getArtifactId ( ) ) ; const artifactsRequest = new GetArtifactsByIDRequest ( ) . setArtifactIdsList ( artifactIds ) ; let artifactsRes ; try { artifactsRes = await Api . getInstance ( ) . metadataStoreService . getArtifactsByID ( artifactsRequest ) ; } catch ( artifactsErr ) { artifactsErr . message = 'Failed to getArtifactsByID: ' + artifactsErr . message ; throw artifactsErr ; } const artifactMap = new Map ( ) ; for ( const [ , artifactEntry ] of Object . entries ( artifactsRes . getArtifactsList ( ) ) ) { artifactMap . set ( artifactEntry . getId ( ) , artifactEntry ) ; } return events . map ( event => { const artifact = artifactMap . get ( event . getArtifactId ( ) ) ; return { event : event , artifact : artifact } ; } ) ; } export async function getLinkedArtifactsByExecution ( execution ) { const event = await getEventByExecution ( execution ) ; return getLinkedArtifactsByEvents ( event ) ; } export function filterEventWithInputArtifact ( linkedArtifact ) { return linkedArtifact . filter ( obj => obj . event . getType ( ) === Event . Type . INPUT ) ; } export function filterEventWithOutputArtifact ( linkedArtifact ) { return linkedArtifact . filter ( obj => obj . event . getType ( ) === Event . Type . OUTPUT ) ; } throws export async function getOutputArtifactsInExecution ( execution ) { const linkedArtifacts = await getLinkedArtifactsByExecution ( execution ) ; return filterEventWithOutputArtifact ( linkedArtifacts ) . map ( linkedArtifact => linkedArtifact . artifact ) ; } export async function getOutputLinkedArtifactsInExecution ( execution ) { const linkedArtifacts = await getLinkedArtifactsByExecution ( execution ) ; return filterEventWithOutputArtifact ( linkedArtifacts ) ; } export async function getArtifactTypes ( ) { const request = new GetArtifactTypesRequest ( ) ; let res ; try { res = await Api . getInstance ( ) . metadataStoreService . getArtifactTypes ( request ) ; } catch ( err ) { err . message = 'Failed to getArtifactTypes: ' + err . message ; throw err ; } return res . getArtifactTypesList ( ) ; } export function filterArtifactsByType ( artifactTypeName , artifactTypes , artifacts ) { const artifactTypeIds = artifactTypes . filter ( artifactType => artifactType . getName ( ) === artifactTypeName ) . map ( artifactType => artifactType . getId ( ) ) ; return artifacts . filter ( artifact => artifactTypeIds . includes ( artifact . getTypeId ( ) ) ) ; } export function filterLinkedArtifactsByType ( artifactTypeName , artifactTypes , artifacts ) { const artifactTypeIds = artifactTypes . filter ( artifactType => artifactType . getName ( ) === artifactTypeName ) . map ( artifactType => artifactType . getId ( ) ) ; return artifacts . filter ( x => artifactTypeIds . includes ( x . artifact . getTypeId ( ) ) ) ; } export function getArtifactName ( linkedArtifact ) { return linkedArtifact . event . getPath ( ) ?. getStepsList ( ) [ 0 ] . getKey ( ) ; }