def testPrepareExecution(self):
with kubeflow_metadata_adapter.KubeflowMetadataAdapter(
connection_config=self._connection_config) as m:
contexts = m.register_contexts_if_not_exists(
self._pipeline_info, self._component_info)
exec_properties = {'arg_one': 1}
os.environ['KFP_POD_NAME'] = 'fake_pod_name'
m.register_execution(exec_properties=exec_properties,
pipeline_info=self._pipeline_info,
component_info=self._component_info,
contexts=contexts)
[execution] = m.store.get_executions_by_context(contexts[0].id)
self.assertProtoEquals(
"""
id: 1
type_id: 4
properties {
key: "state"
value {
string_value: "new"
}
}
properties {
key: "pipeline_name"
value {
string_value: "fake_pipeline_name"
}
}
properties {
key: "pipeline_root"
value {
string_value: "/fake_pipeline_root"
}
}
properties {
key: "run_id"
value {
string_value: "fake_run_id"
}
}
properties {
key: "component_id"
value {
string_value: "fake_component_id"
}
}
properties {
key: "arg_one"
value {
string_value: "1"
}
}
properties {
key: "kfp_pod_name"
value {
string_value: "fake_pod_name"
}
}""", execution)
def main():
# Log to the container's stdout so Kubeflow Pipelines UI can display logs to
# the user.
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().setLevel(logging.INFO)
parser = argparse.ArgumentParser()
parser.add_argument('--pipeline_name', type=str, required=True)
parser.add_argument('--pipeline_root', type=str, required=True)
parser.add_argument('--kubeflow_metadata_config', type=str, required=True)
parser.add_argument('--beam_pipeline_args', type=str, required=True)
parser.add_argument('--additional_pipeline_args', type=str, required=True)
parser.add_argument('--component_launcher_class_path',
type=str,
required=True)
parser.add_argument('--enable_cache', action='store_true')
parser.add_argument('--serialized_component', type=str, required=True)
parser.add_argument('--component_config', type=str, required=True)
args = parser.parse_args()
component = json_utils.loads(args.serialized_component)
component_config = json_utils.loads(args.component_config)
component_launcher_class = import_utils.import_class_by_path(
args.component_launcher_class_path)
if not issubclass(component_launcher_class,
base_component_launcher.BaseComponentLauncher):
raise TypeError(
'component_launcher_class "%s" is not subclass of base_component_launcher.BaseComponentLauncher'
% component_launcher_class)
kubeflow_metadata_config = kubeflow_pb2.KubeflowMetadataConfig()
json_format.Parse(args.kubeflow_metadata_config, kubeflow_metadata_config)
metadata_connection = kubeflow_metadata_adapter.KubeflowMetadataAdapter(
_get_metadata_connection_config(kubeflow_metadata_config))
driver_args = data_types.DriverArgs(enable_cache=args.enable_cache)
beam_pipeline_args = _make_beam_pipeline_args(args.beam_pipeline_args)
additional_pipeline_args = json.loads(args.additional_pipeline_args)
launcher = component_launcher_class.create(
component=component,
pipeline_info=data_types.PipelineInfo(
pipeline_name=args.pipeline_name,
pipeline_root=args.pipeline_root,
run_id=os.environ['WORKFLOW_ID']),
driver_args=driver_args,
metadata_connection=metadata_connection,
beam_pipeline_args=beam_pipeline_args,
additional_pipeline_args=additional_pipeline_args,
component_config=component_config)
execution_info = launcher.launch()
# Dump the UI metadata.
_dump_ui_metadata(component, execution_info)
def testIsEligiblePreviousExecution(self):
with kubeflow_metadata_adapter.KubeflowMetadataAdapter(
connection_config=self._connection_config) as m:
contexts = m.register_contexts_if_not_exists(
self._pipeline_info, self._component_info)
exec_properties = {'arg_one': 1}
os.environ['KFP_POD_NAME'] = 'fake_pod_name1'
m.register_execution(exec_properties=exec_properties,
pipeline_info=self._pipeline_info,
component_info=self._component_info,
contexts=contexts)
os.environ['KFP_POD_NAME'] = 'fake_pod_name2'
m.register_execution(exec_properties=exec_properties,
pipeline_info=self._pipeline_info,
component_info=self._component_info,
contexts=contexts)
[execution1,
execution2] = m.store.get_executions_by_context(contexts[0].id)
self.assertTrue(
m._is_eligible_previous_execution(execution1, execution2))
def main():
# Log to the container's stdout so Kubeflow Pipelines UI can display logs to
# the user.
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().setLevel(logging.INFO)
parser = argparse.ArgumentParser()
parser.add_argument('--pipeline_name', type=str, required=True)
parser.add_argument('--pipeline_root', type=str, required=True)
parser.add_argument('--kubeflow_metadata_config', type=str, required=True)
parser.add_argument('--beam_pipeline_args', type=str, required=True)
parser.add_argument('--additional_pipeline_args', type=str, required=True)
parser.add_argument(
'--component_launcher_class_path', type=str, required=True)
parser.add_argument('--enable_cache', action='store_true')
parser.add_argument('--serialized_component', type=str, required=True)
parser.add_argument('--component_config', type=str, required=True)
args = parser.parse_args()
component = json_utils.loads(args.serialized_component)
component_config = json_utils.loads(args.component_config)
component_launcher_class = import_utils.import_class_by_path(
args.component_launcher_class_path)
if not issubclass(component_launcher_class,
base_component_launcher.BaseComponentLauncher):
raise TypeError(
'component_launcher_class "%s" is not subclass of base_component_launcher.BaseComponentLauncher'
% component_launcher_class)
kubeflow_metadata_config = kubeflow_pb2.KubeflowMetadataConfig()
json_format.Parse(args.kubeflow_metadata_config, kubeflow_metadata_config)
metadata_connection = kubeflow_metadata_adapter.KubeflowMetadataAdapter(
_get_metadata_connection_config(kubeflow_metadata_config))
driver_args = data_types.DriverArgs(enable_cache=args.enable_cache)
beam_pipeline_args = json.loads(args.beam_pipeline_args)
additional_pipeline_args = json.loads(args.additional_pipeline_args)
launcher = component_launcher_class.create(
component=component,
pipeline_info=data_types.PipelineInfo(
pipeline_name=args.pipeline_name,
pipeline_root=args.pipeline_root,
run_id=os.environ['WORKFLOW_ID']),
driver_args=driver_args,
metadata_connection=metadata_connection,
beam_pipeline_args=beam_pipeline_args,
additional_pipeline_args=additional_pipeline_args,
component_config=component_config)
# Attach necessary labels to distinguish different runner and DSL.
# TODO(zhitaoli): Pass this from KFP runner side when the same container
# entrypoint can be used by a different runner.
with telemetry_utils.scoped_labels({
telemetry_utils.LABEL_TFX_RUNNER: 'kfp',
}):
execution_info = launcher.launch()
# Dump the UI metadata.
_dump_ui_metadata(component, execution_info)
def testPrepareExecution(self):
with kubeflow_metadata_adapter.KubeflowMetadataAdapter(
connection_config=self._connection_config) as m:
contexts = m.register_pipeline_contexts_if_not_exists(self._pipeline_info)
exec_properties = {'arg_one': 1}
os.environ['KFP_POD_NAME'] = 'fake_pod_name'
m.register_execution(
exec_properties=exec_properties,
pipeline_info=self._pipeline_info,
component_info=self._component_info,
contexts=contexts)
[execution] = m.store.get_executions_by_context(contexts[0].id)
# Skip verifying time sensitive fields.
execution.ClearField('create_time_since_epoch')
execution.ClearField('last_update_time_since_epoch')
self.assertProtoEquals(
"""
id: 1
type_id: 3
last_known_state: RUNNING
properties {
key: "state"
value {
string_value: "new"
}
}
properties {
key: "pipeline_name"
value {
string_value: "fake_pipeline_name"
}
}
properties {
key: "pipeline_root"
value {
string_value: "/fake_pipeline_root"
}
}
properties {
key: "run_id"
value {
string_value: "fake_run_id"
}
}
properties {
key: "component_id"
value {
string_value: "fake_component_id"
}
}
properties {
key: "arg_one"
value {
string_value: "1"
}
}
properties {
key: "kfp_pod_name"
value {
string_value: "fake_pod_name"
}
}""", execution)
def main():
# Log to the container's stdout so Kubeflow Pipelines UI can display logs to
# the user.
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().setLevel(logging.INFO)
parser = argparse.ArgumentParser()
parser.add_argument('--pipeline_root', type=str, required=True)
parser.add_argument('--kubeflow_metadata_config', type=str, required=True)
parser.add_argument('--serialized_component', type=str, required=True)
parser.add_argument('--tfx_ir', type=str, required=True)
parser.add_argument('--node_id', type=str, required=True)
launcher._register_execution = _register_execution # pylint: disable=protected-access
args = parser.parse_args()
tfx_ir = pipeline_pb2.Pipeline()
json_format.Parse(args.tfx_ir, tfx_ir)
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
tfx_ir, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME:
os.environ['WORKFLOW_ID'],
})
deployment_config = runner_utils.extract_local_deployment_config(tfx_ir)
kubeflow_metadata_config = kubeflow_pb2.KubeflowMetadataConfig()
json_format.Parse(args.kubeflow_metadata_config, kubeflow_metadata_config)
metadata_connection = kubeflow_metadata_adapter.KubeflowMetadataAdapter(
_get_metadata_connection_config(kubeflow_metadata_config))
node_id = args.node_id
# Attach necessary labels to distinguish different runner and DSL.
# TODO(zhitaoli): Pass this from KFP runner side when the same container
# entrypoint can be used by a different runner.
with telemetry_utils.scoped_labels({
telemetry_utils.LABEL_TFX_RUNNER: 'kfp',
}):
custom_executor_operators = {
executable_spec_pb2.ContainerExecutableSpec:
kubernetes_executor_operator.KubernetesExecutorOperator
}
executor_spec = runner_utils.extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = runner_utils.extract_custom_driver_spec(
deployment_config, node_id)
pipeline_node = _get_pipeline_node(tfx_ir, node_id)
component_launcher = launcher.Launcher(
pipeline_node=pipeline_node,
mlmd_connection=metadata_connection,
pipeline_info=tfx_ir.pipeline_info,
pipeline_runtime_spec=tfx_ir.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec,
custom_executor_operators=custom_executor_operators)
logging.info('Component %s is running.', node_id)
execution_info = component_launcher.launch()
logging.info('Component %s is finished.', node_id)
# Dump the UI metadata.
_dump_ui_metadata(pipeline_node, execution_info)