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 _label_with_context_and_execution(self):
self.pod = podutils.get_pod(self.pod_name, self.pod_namespace)
labels = self.pod.metadata.labels
labels.setdefault(METADATA_EXECUTION_ID_LABEL_KEY,
str(self.execution.id))
labels.setdefault(METADATA_CONTEXT_ID_LABEL_KEY,
str(self.run_context.id))
podutils.patch_pod(self.pod_name, self.pod_namespace,
{"metadata": {
"labels": labels
}})
def find_poddefault_labels_on_server(request):
"""Find server's labels that correspond to poddefaults applied."""
request.log.info("Retrieving PodDefaults applied to server...")
applied_poddefaults = kfutils.find_applied_poddefaults(
podutils.get_pod(podutils.get_pod_name(), podutils.get_namespace()),
kfutils.list_poddefaults())
pd_names = [pd["metadata"]["name"] for pd in applied_poddefaults]
request.log.info("Retrieved applied PodDefaults: %s", pd_names)
labels = kfutils.get_poddefault_labels(applied_poddefaults)
request.log.info("PodDefault labels applied on server: %s",
", ".join(["%s: %s" % (k, v) for k, v in labels.items()]))
return labels
def _annotate_artifacts(self, annotation_key, ids):
self.pod = podutils.get_pod(self.pod_name, self.pod_namespace)
annotations = self.pod.metadata.annotations
all_ids_str = annotations.get(annotation_key, "[]")
all_ids = json.loads(all_ids_str)
all_ids.extend(ids)
all_ids.sort()
all_ids_str = json.dumps(all_ids)
annotations[annotation_key] = all_ids_str
podutils.patch_pod(self.pod_name, self.pod_namespace,
{"metadata": {
"annotations": annotations
}})
