-
Notifications
You must be signed in to change notification settings - Fork 706
/
pipeline_config.py
71 lines (63 loc) · 3.04 KB
/
pipeline_config.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
# Lint as: python3
# Copyright 2019 Google LLC. All Rights Reserved.
#
# 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.
"""Settings for controlling how to run a pipeline."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from typing import Dict, List, Text, Type
from tfx.orchestration.config import base_component_config
from tfx.orchestration.launcher import base_component_launcher
from tfx.orchestration.launcher import in_process_component_launcher
class PipelineConfig(object):
"""Config class which controls how to run a pipeline.
Attributes
supported_launcher_classes: A list of component launcher classes that are
supported by the current pipeline. List sequence determines the order in
which launchers are chosen for each component being run.
default_component_configs: A list of default component configs which will
be used as default component config to run each component in the pipeline.
List sequence determines the order in which config are chosen for each
component being run.
component_config_overrides: component configs for customizing the launching
of each component. The key is the component ID.
"""
# TODO(hongyes): figure out the best practice to put the
# SUPPORTED_LAUNCHER_CLASSES.
def __init__(self,
supported_launcher_classes: List[Type[
base_component_launcher.BaseComponentLauncher]] = None,
default_component_configs: List[
base_component_config.BaseComponentConfig] = None,
component_config_overrides: Dict[
Text, base_component_config.BaseComponentConfig] = None):
self.supported_launcher_classes = supported_launcher_classes or [
in_process_component_launcher.InProcessComponentLauncher
]
self.default_component_configs = default_component_configs or []
self.component_config_overrides = component_config_overrides or {}
self._validate_configs()
def _validate_configs(self):
"""Validate the config settings."""
if len(self.supported_launcher_classes) > len(
set(self.supported_launcher_classes)):
raise ValueError(
'supported_launcher_classes must not have duplicate types')
default_component_config_classes = [
type(config) for config in self.default_component_configs
]
if len(default_component_config_classes) > len(
set(default_component_config_classes)):
raise ValueError(
'default_component_configs must not have configs with the same type')