Source code for kfp.dsl._pipeline

# Copyright 2018-2019 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

import enum
from typing import Callable, Optional, Union

from kubernetes.client.models import V1PodDNSConfig
from kfp.dsl import _container_op
from kfp.dsl import _resource_op
from kfp.dsl import _ops_group
from kfp.dsl import _component_bridge
from kfp.components import _components
from kfp.components import _naming
import sys

# This handler is called whenever the @pipeline decorator is applied.
# It can be used by command-line DSL compiler to inject code that runs for every
# pipeline definition.
_pipeline_decorator_handler = None

[docs]class PipelineExecutionMode(enum.Enum): # Compile to Argo YAML without support for metadata-enabled components. V1_LEGACY = 1 # Compiles to Argo YAML with support for metadata-enabled components. # Pipelines compiled using this mode aim to be compatible with v2 semantics. V2_COMPATIBLE = 2 # Compiles to KFP v2 IR for execution using the v2 engine. # This option is unsupported right now. V2_ENGINE = 3
[docs]def pipeline(name: Optional[str] = None, description: Optional[str] = None, pipeline_root: Optional[str] = None): """Decorator of pipeline functions. Example :: @pipeline( name='my-pipeline', description='My ML Pipeline.' pipeline_root='gs://my-bucket/my-output-path' ) def my_pipeline(a: PipelineParam, b: PipelineParam): ... Args: name: The pipeline name. Default to a sanitized version of the function name. description: Optionally, a human-readable description of the pipeline. pipeline_root: The root directory to generate input/output URI under this pipeline. This is required if input/output URI placeholder is used in this pipeline. """ def _pipeline(func: Callable): if name: func._component_human_name = name if description: func._component_description = description if pipeline_root: func.pipeline_root = pipeline_root if _pipeline_decorator_handler: return _pipeline_decorator_handler(func) or func else: return func return _pipeline
[docs]class PipelineConf(): """PipelineConf contains pipeline level settings.""" def __init__(self): self.image_pull_secrets = [] self.timeout = 0 self.ttl_seconds_after_finished = -1 self._pod_disruption_budget_min_available = None self.op_transformers = [] self.default_pod_node_selector = {} self.image_pull_policy = None self.parallelism = None self._data_passing_method = None self.dns_config = None
[docs] def set_image_pull_secrets(self, image_pull_secrets): """Configures the pipeline level imagepullsecret. Args: image_pull_secrets: a list of Kubernetes V1LocalObjectReference For detailed description, check Kubernetes V1LocalObjectReference definition """ self.image_pull_secrets = image_pull_secrets return self
[docs] def set_timeout(self, seconds: int): """Configures the pipeline level timeout. Args: seconds: number of seconds for timeout """ self.timeout = seconds return self
[docs] def set_parallelism(self, max_num_pods: int): """Configures the max number of total parallel pods that can execute at the same time in a workflow. Args: max_num_pods: max number of total parallel pods. """ if max_num_pods < 1: raise ValueError( 'Pipeline max_num_pods set to < 1, allowed values are > 0') self.parallelism = max_num_pods return self
[docs] def set_ttl_seconds_after_finished(self, seconds: int): """Configures the ttl after the pipeline has finished. Args: seconds: number of seconds for the workflow to be garbage collected after it is finished. """ self.ttl_seconds_after_finished = seconds return self
[docs] def set_pod_disruption_budget(self, min_available: Union[int, str]): """PodDisruptionBudget holds the number of concurrent disruptions that you allow for pipeline Pods. Args: min_available (Union[int, str]): An eviction is allowed if at least "minAvailable" pods selected by "selector" will still be available after the eviction, i.e. even in the absence of the evicted pod. So for example you can prevent all voluntary evictions by specifying "100%". "minAvailable" can be either an absolute number or a percentage. """ self._pod_disruption_budget_min_available = min_available return self
[docs] def set_default_pod_node_selector(self, label_name: str, value: str): """Add a constraint for nodeSelector for a pipeline. Each constraint is a key-value pair label. For the container to be eligible to run on a node, the node must have each of the constraints appeared as labels. Args: label_name: The name of the constraint label. value: The value of the constraint label. """ self.default_pod_node_selector[label_name] = value return self
[docs] def set_image_pull_policy(self, policy: str): """Configures the default image pull policy. Args: policy: the pull policy, has to be one of: Always, Never, IfNotPresent. For more info: """ self.image_pull_policy = policy return self
[docs] def add_op_transformer(self, transformer): """Configures the op_transformers which will be applied to all ops in the pipeline. The ops can be ResourceOp, VolumeOp, or ContainerOp. Args: transformer: A function that takes a kfp Op as input and returns a kfp Op """ self.op_transformers.append(transformer)
[docs] def set_dns_config(self, dns_config: V1PodDNSConfig): """Set the dnsConfig to be given to each pod. Args: dns_config: Kubernetes V1PodDNSConfig For detailed description, check Kubernetes V1PodDNSConfig definition Example: :: import kfp from kubernetes.client.models import V1PodDNSConfig, V1PodDNSConfigOption pipeline_conf = kfp.dsl.PipelineConf() pipeline_conf.set_dns_config(dns_config=V1PodDNSConfig( nameservers=[""], options=[V1PodDNSConfigOption(name="ndots", value="2")], )) """ self.dns_config = dns_config
@property def data_passing_method(self): return self._data_passing_method @data_passing_method.setter def data_passing_method(self, value): """Sets the object representing the method used for intermediate data passing. Example: :: from kfp.dsl import PipelineConf, data_passing_methods from kubernetes.client.models import V1Volume, V1PersistentVolumeClaimVolumeSource pipeline_conf = PipelineConf() pipeline_conf.data_passing_method = data_passing_methods.KubernetesVolume( volume=V1Volume( name='data', persistent_volume_claim=V1PersistentVolumeClaimVolumeSource('data-volume'), ), path_prefix='artifact_data/', ) """ self._data_passing_method = value
[docs]def get_pipeline_conf(): """Configure the pipeline level setting to the current pipeline Note: call the function inside the user defined pipeline function. """ return Pipeline.get_default_pipeline().conf
# TODO: Pipeline is in fact an opsgroup, refactor the code. class Pipeline(): """A pipeline contains a list of operators. This class is not supposed to be used by pipeline authors since pipeline authors can use pipeline functions (decorated with @pipeline) to reference their pipelines. This class is useful for implementing a compiler. For example, the compiler can use the following to get the pipeline object and its ops: Example: :: with Pipeline() as p: pipeline_func(*args_list) traverse(p.ops) """ # _default_pipeline is set when it (usually a compiler) runs "with Pipeline()" _default_pipeline = None @staticmethod def get_default_pipeline(): """Get default pipeline.""" return Pipeline._default_pipeline @staticmethod def add_pipeline(name, description, func): """Add a pipeline function with the specified name and description.""" # Applying the @pipeline decorator to the pipeline function func = pipeline(name=name, description=description)(func) def __init__(self, name: str): """Create a new instance of Pipeline. Args: name: the name of the pipeline. Once deployed, the name will show up in Pipeline System UI. """ = name self.ops = {} # Add the root group. self.groups = [_ops_group.OpsGroup('pipeline', name=name)] self.group_id = 0 self.conf = PipelineConf() self._metadata = None def __enter__(self): if Pipeline._default_pipeline: raise Exception('Nested pipelines are not allowed.') Pipeline._default_pipeline = self self._old_container_task_constructor = ( _components._container_task_constructor) _components._container_task_constructor = ( _component_bridge._create_container_op_from_component_and_arguments) def register_op_and_generate_id(op): return self.add_op(op, op.is_exit_handler) self._old__register_op_handler = _container_op._register_op_handler _container_op._register_op_handler = register_op_and_generate_id return self def __exit__(self, *args): Pipeline._default_pipeline = None _container_op._register_op_handler = self._old__register_op_handler _components._container_task_constructor = ( self._old_container_task_constructor) def add_op(self, op: _container_op.BaseOp, define_only: bool): """Add a new operator. Args: op: An operator of ContainerOp, ResourceOp or their inherited types. Returns op_name: a unique op name. """ # Sanitizing the op name. # Technically this could be delayed to the compilation stage, but string # serialization of PipelineParams make unsanitized names problematic. op_name = _naming._sanitize_python_function_name(op.human_name).replace( '_', '-') #If there is an existing op with this name then generate a new name. op_name = _naming._make_name_unique_by_adding_index( op_name, list(self.ops.keys()), '-') if op_name == '': op_name = _naming._make_name_unique_by_adding_index( 'task', list(self.ops.keys()), '-') self.ops[op_name] = op if not define_only: self.groups[-1].ops.append(op) return op_name def push_ops_group(self, group: _ops_group.OpsGroup): """Push an OpsGroup into the stack. Args: group: An OpsGroup. Typically it is one of ExitHandler, Branch, and Loop. """ self.groups[-1].groups.append(group) self.groups.append(group) def pop_ops_group(self): """Remove the current OpsGroup from the stack.""" del self.groups[-1] def remove_op_from_groups(self, op): for group in self.groups: group.remove_op_recursive(op) def get_next_group_id(self): """Get next id for a new group.""" self.group_id += 1 return self.group_id def _set_metadata(self, metadata): """_set_metadata passes the containerop the metadata information. Args: metadata (ComponentMeta): component metadata """ self._metadata = metadata