def is_kfp_step() -> bool:
"""Detect if running inside a KFP step.
The detection involves two steps:
1. Auto-detect if the current Pod is part of an Argo workflow
2. Read one of the annotations that the KFP API Server sets in the
workflow object (one-off runs and recurring ones have different
annotations).
"""
log.info("Checking if running inside a KFP step...")
try:
namespace = podutils.get_namespace()
workflow = workflowutils.get_workflow(
workflowutils.get_workflow_name(podutils.get_pod_name(),
namespace), namespace)
annotations = workflow["metadata"]["annotations"]
try:
_ = annotations[KFP_RUN_NAME_ANNOTATION_KEY]
except KeyError:
_ = annotations[KFP_SWF_NAME_ANNOTATION_KEY]
except Exception:
log.info("Not in a KFP step.")
return False
log.info("Running in a KFP step.")
return True
def __init__(self):
log.info("%s Initializing MLMD context... %s", "-" * 10, "-" * 10)
log.info("Connecting to MLMD...")
self.store = self._connect()
log.info("Successfully connected to MLMD")
log.info("Getting step details...")
log.info("Getting pod name...")
self.pod_name = podutils.get_pod_name()
log.info("Successfully retrieved pod name: %s", self.pod_name)
log.info("Getting pod namespace...")
self.pod_namespace = podutils.get_namespace()
log.info("Successfully retrieved pod namespace: %s",
self.pod_namespace)
log.info("Getting pod...")
self.pod = podutils.get_pod(self.pod_name, self.pod_namespace)
log.info("Successfully retrieved pod")
log.info("Getting workflow name from pod...")
self.workflow_name = self.pod.metadata.labels.get(
workflowutils.ARGO_WORKFLOW_LABEL_KEY)
log.info("Successfully retrieved workflow name: %s",
self.workflow_name)
log.info("Getting workflow...")
self.workflow = workflowutils.get_workflow(self.workflow_name,
self.pod_namespace)
log.info("Successfully retrieved workflow")
workflow_labels = self.workflow["metadata"].get("labels", {})
self.run_uuid = workflow_labels.get(podutils.KFP_RUN_ID_LABEL_KEY,
self.workflow_name)
log.info("Successfully retrieved KFP run ID: %s", self.run_uuid)
workflow_annotations = self.workflow["metadata"].get("annotations", {})
pipeline_spec = json.loads(
workflow_annotations.get("pipelines.kubeflow.org/pipeline_spec",
"{}"))
self.pipeline_name = pipeline_spec.get("name", self.workflow_name)
if self.pipeline_name:
log.info("Successfully retrieved KFP pipeline_name: %s",
self.pipeline_name)
else:
log.info("Could not retrieve KFP pipeline name")
self.component_id = podutils.compute_component_id(self.pod)
self.execution_hash = self.pod.metadata.annotations.get(
MLMD_EXECUTION_HASH_PROPERTY_KEY)
if self.execution_hash:
log.info("Successfully retrieved execution hash: %s",
self.execution_hash)
else:
self.execution_hash = utils.random_string(10)
log.info(
"Failed to retrieve execution hash."
" Generating random string...: %s", self.execution_hash)
self.run_context = self._get_or_create_run_context()
self.execution = self._create_execution_in_run_context()
self._label_with_context_and_execution()
log.info("%s Successfully initialized MLMD context %s", "-" * 10,
"-" * 10)
def detect_run_uuid() -> str:
"""Get the workflow's UUID form inside a pipeline step."""
namespace = podutils.get_namespace()
workflow = workflowutils.get_workflow(
workflowutils.get_workflow_name(podutils.get_pod_name(), namespace),
namespace)
run_uuid = (workflow["metadata"].get("labels",
{}).get(KFP_RUN_ID_LABEL_KEY, None))
# KFP api-server adds run UUID as label to workflows for KFP>=0.1.26.
# Return run UUID if available. Else return workflow UUID to maintain
# backwards compatibility.
return run_uuid or workflow["metadata"]["uid"]
