def testMainEmptyInputs(self):
"""Test executor class import under empty inputs/outputs."""
inputs = {
'x':
[types.Artifact(type_name='X'),
types.Artifact(type_name='X')]
}
outputs = {'y': [types.Artifact(type_name='Y')]}
exec_properties = {'a': 'b'}
args = [
'--executor_class_path=%s.%s' %
(FakeExecutor.__module__, FakeExecutor.__name__),
'--inputs=%s' % artifact_utils.jsonify_artifact_dict(inputs),
'--outputs=%s' % artifact_utils.jsonify_artifact_dict(outputs),
'--exec-properties=%s' % json.dumps(exec_properties),
]
with ArgsCapture() as args_capture:
run_executor.main(args)
# TODO(b/131417512): Add equal comparison to types.Artifact class so we
# can use asserters.
self.assertSetEqual(set(args_capture.input_dict.keys()),
set(inputs.keys()))
self.assertSetEqual(set(args_capture.output_dict.keys()),
set(outputs.keys()))
self.assertDictEqual(args_capture.exec_properties, exec_properties)
def generate_container_command(self, input_dict: Dict[str,
List[types.Artifact]],
output_dict: Dict[str, List[types.Artifact]],
exec_properties: Dict[str, Any],
executor_class_path: str) -> List[str]:
"""Generate container command to run executor."""
json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
logging.info('json_inputs=\'%s\'.', json_inputs)
json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
logging.info('json_outputs=\'%s\'.', json_outputs)
json_exec_properties = json.dumps(exec_properties, sort_keys=True)
logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
# We use custom containers to launch training on AI Platform, which invokes
# the specified image using the container's entrypoint. The default
# entrypoint for TFX containers is to call scripts/run_executor.py. The
# arguments below are passed to this run_executor entry to run the executor
# specified in `executor_class_path`.
return _CONTAINER_COMMAND + [
'--executor_class_path',
executor_class_path,
'--inputs',
json_inputs,
'--outputs',
json_outputs,
'--exec-properties',
json_exec_properties,
]
def _log_startup(self, inputs: Dict[Text, List[types.Artifact]],
outputs: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
"""Log inputs, outputs, and executor properties in a standard format."""
tf.logging.debug('Starting %s execution.', self.__class__.__name__)
tf.logging.debug('Inputs for %s are: %s', self.__class__.__name__,
artifact_utils.jsonify_artifact_dict(inputs))
tf.logging.debug('Outputs for %s are: %s', self.__class__.__name__,
artifact_utils.jsonify_artifact_dict(outputs))
tf.logging.debug('Execution properties for %s are: %s',
self.__class__.__name__, json.dumps(exec_properties))
def _log_startup(self, inputs: Dict[str, List[types.Artifact]],
outputs: Dict[str, List[types.Artifact]],
exec_properties: Dict[str, Any]) -> None:
"""Log inputs, outputs, and executor properties in a standard format."""
logging.debug('Starting %s execution.', self.__class__.__name__)
logging.debug('Inputs for %s are: %s', self.__class__.__name__,
artifact_utils.jsonify_artifact_dict(inputs))
logging.debug('Outputs for %s are: %s', self.__class__.__name__,
artifact_utils.jsonify_artifact_dict(outputs))
logging.debug(
'Execution properties for %s are: %s', self.__class__.__name__,
json.dumps(
data_types_utils.build_value_dict(
data_types_utils.build_metadata_value_dict(
exec_properties))))
def testCsvExampleGenWrapper(self):
input_base = types.Artifact(type_name='ExternalPath', split='')
input_base.uri = '/path/to/dataset'
with patch.object(executor, 'Executor', autospec=True) as _:
wrapper = executor_wrappers.CsvExampleGenWrapper(
argparse.Namespace(
exec_properties=json.dumps(self.exec_properties),
outputs=artifact_utils.jsonify_artifact_dict(
{'examples': self.examples}),
executor_class_path=
('tfx.components.example_gen.csv_example_gen.executor.Executor'
),
input_base=json.dumps([input_base.json_dict()])), )
wrapper.run(output_basedir=self.output_basedir)
# TODO(b/133011207): Validate arguments for executor and Do() method.
metadata_file = os.path.join(self.output_basedir,
'output/ml_metadata/examples')
expected_output_examples = types.Artifact(type_name='ExamplesPath',
split='dummy')
# Expect that span and path are resolved.
expected_output_examples.span = 1
expected_output_examples.uri = (
'/path/to/output/csv_example_gen/examples/mock_workflow_id/dummy/'
)
with tf.gfile.GFile(metadata_file) as f:
self.assertEqual([expected_output_examples.json_dict()],
json.loads(f.read()))
def testContainerOpArguments(self):
self.assertEqual(self.component.container_op.arguments[0],
'--exec_properties')
self.assertDictEqual(
{
'output_dir': 'output_dir',
'log_root': 'log_root',
'module_file': '/path/to/module.py'
}, json.loads(self.component.container_op.arguments[1]))
self.assertEqual(self.component.container_op.arguments[2:], [
'--outputs',
artifact_utils.jsonify_artifact_dict(self._output_dict),
'--executor_class_path',
'some.executor.Class',
'TFXComponent',
'--input_data',
'input-data-contents',
'--train_steps',
'300',
'--accuracy_threshold',
'0.3',
])
def start_aip_training(input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text,
Any], executor_class_path: Text,
training_inputs: Dict[Text,
Any], job_id: Optional[Text]):
"""Start a trainer job on AI Platform (AIP).
This is done by forwarding the inputs/outputs/exec_properties to the
tfx.scripts.run_executor module on a AI Platform training job interpreter.
Args:
input_dict: Passthrough input dict for tfx.components.Trainer.executor.
output_dict: Passthrough input dict for tfx.components.Trainer.executor.
exec_properties: Passthrough input dict for tfx.components.Trainer.executor.
executor_class_path: class path for TFX core default trainer.
training_inputs: Training input argment for AI Platform training job.
'pythonModule', 'pythonVersion' and 'runtimeVersion' will be inferred. For
the full set of parameters, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#TrainingInput
job_id: Job ID for AI Platform Training job. If not supplied,
system-determined unique ID is given. Refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#resource-job
Returns:
None
Raises:
RuntimeError: if the Google Cloud AI Platform training job failed.
"""
training_inputs = training_inputs.copy()
json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
absl.logging.info('json_inputs=\'%s\'.', json_inputs)
json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
absl.logging.info('json_outputs=\'%s\'.', json_outputs)
json_exec_properties = json.dumps(exec_properties, sort_keys=True)
absl.logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
# Configure AI Platform training job
api_client = discovery.build('ml', 'v1')
# We use custom containers to launch training on AI Platform, which invokes
# the specified image using the container's entrypoint. The default
# entrypoint for TFX containers is to call scripts/run_executor.py. The
# arguments below are passed to this run_executor entry to run the executor
# specified in `executor_class_path`.
job_args = [
'--executor_class_path', executor_class_path, '--inputs', json_inputs,
'--outputs', json_outputs, '--exec-properties', json_exec_properties
]
if not training_inputs.get('masterConfig'):
training_inputs['masterConfig'] = {
'imageUri': _TFX_IMAGE,
}
training_inputs['args'] = job_args
# Pop project_id so AIP doesn't complain about an unexpected parameter.
# It's been a stowaway in aip_args and has finally reached its destination.
project = training_inputs.pop('project')
project_id = 'projects/{}'.format(project)
# 'tfx_YYYYmmddHHMMSS' is the default job ID if not explicitly specified.
job_id = job_id or 'tfx_%s' % datetime.datetime.now().strftime('%Y%m%d%H%M%S')
job_spec = {'jobId': job_id, 'trainingInput': training_inputs}
# Submit job to AIP Training
absl.logging.info(
'Submitting job=\'{}\', project=\'{}\' to AI Platform.'.format(
job_id, project))
request = api_client.projects().jobs().create(
body=job_spec, parent=project_id)
request.execute()
# Wait for AIP Training job to finish
job_name = '{}/jobs/{}'.format(project_id, job_id)
request = api_client.projects().jobs().get(name=job_name)
response = request.execute()
while response['state'] not in ('SUCCEEDED', 'FAILED'):
time.sleep(_POLLING_INTERVAL_IN_SECONDS)
response = request.execute()
if response['state'] == 'FAILED':
err_msg = 'Job \'{}\' did not succeed. Detailed response {}.'.format(
job_name, response)
absl.logging.error(err_msg)
raise RuntimeError(err_msg)
# AIP training complete
absl.logging.info('Job \'{}\' successful.'.format(job_name))
def start_aip_training(input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text,
Any], executor_class_path: Text,
training_inputs: Dict[Text,
Any], job_id: Optional[Text]):
"""Start a trainer job on AI Platform (AIP).
This is done by forwarding the inputs/outputs/exec_properties to the
tfx.scripts.run_executor module on a AI Platform training job interpreter.
Args:
input_dict: Passthrough input dict for tfx.components.Trainer.executor.
output_dict: Passthrough input dict for tfx.components.Trainer.executor.
exec_properties: Passthrough input dict for tfx.components.Trainer.executor.
executor_class_path: class path for TFX core default trainer.
training_inputs: Training input argument for AI Platform training job.
'pythonModule', 'pythonVersion' and 'runtimeVersion' will be inferred. For
the full set of parameters, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#TrainingInput
job_id: Job ID for AI Platform Training job. If not supplied,
system-determined unique ID is given. Refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#resource-job
Returns:
None
Raises:
RuntimeError: if the Google Cloud AI Platform training job failed/cancelled.
"""
training_inputs = training_inputs.copy()
json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
logging.info('json_inputs=\'%s\'.', json_inputs)
json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
logging.info('json_outputs=\'%s\'.', json_outputs)
json_exec_properties = json.dumps(exec_properties, sort_keys=True)
logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
# Configure AI Platform training job
api_client = discovery.build('ml', 'v1')
# We use custom containers to launch training on AI Platform, which invokes
# the specified image using the container's entrypoint. The default
# entrypoint for TFX containers is to call scripts/run_executor.py. The
# arguments below are passed to this run_executor entry to run the executor
# specified in `executor_class_path`.
job_args = [
'--executor_class_path', executor_class_path, '--inputs', json_inputs,
'--outputs', json_outputs, '--exec-properties', json_exec_properties
]
if not training_inputs.get('masterConfig'):
training_inputs['masterConfig'] = {
'imageUri': _TFX_IMAGE,
}
training_inputs['args'] = job_args
# Pop project_id so AIP doesn't complain about an unexpected parameter.
# It's been a stowaway in aip_args and has finally reached its destination.
project = training_inputs.pop('project')
project_id = 'projects/{}'.format(project)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
job_labels = telemetry_utils.get_labels_dict()
# 'tfx_YYYYmmddHHMMSS' is the default job ID if not explicitly specified.
job_id = job_id or 'tfx_{}'.format(
datetime.datetime.now().strftime('%Y%m%d%H%M%S'))
job_spec = {
'jobId': job_id,
'trainingInput': training_inputs,
'labels': job_labels,
}
# Submit job to AIP Training
logging.info('Submitting job=\'%s\', project=\'%s\' to AI Platform.', job_id,
project)
request = api_client.projects().jobs().create(
body=job_spec, parent=project_id)
request.execute()
# Wait for AIP Training job to finish
job_name = '{}/jobs/{}'.format(project_id, job_id)
request = api_client.projects().jobs().get(name=job_name)
response = request.execute()
retry_count = 0
# Monitors the long-running operation by polling the job state periodically,
# and retries the polling when a transient connectivity issue is encountered.
#
# Long-running operation monitoring:
# The possible states of "get job" response can be found at
# https://cloud.google.com/ai-platform/training/docs/reference/rest/v1/projects.jobs#State
# where SUCCEEDED/FAILED/CANCELLED are considered to be final states.
# The following logic will keep polling the state of the job until the job
# enters a final state.
#
# During the polling, if a connection error was encountered, the GET request
# will be retried by recreating the Python API client to refresh the lifecycle
# of the connection being used. See
# https://github.com/googleapis/google-api-python-client/issues/218
# for a detailed description of the problem. If the error persists for
# _CONNECTION_ERROR_RETRY_LIMIT consecutive attempts, the function will exit
# with code 1.
while response['state'] not in ('SUCCEEDED', 'FAILED', 'CANCELLED'):
time.sleep(_POLLING_INTERVAL_IN_SECONDS)
try:
response = request.execute()
retry_count = 0
# Handle transient connection error.
except ConnectionError as err:
if retry_count < _CONNECTION_ERROR_RETRY_LIMIT:
retry_count += 1
logging.warning(
'ConnectionError (%s) encountered when polling job: %s. Trying to '
'recreate the API client.', err, job_id)
# Recreate the Python API client.
api_client = discovery.build('ml', 'v1')
request = api_client.projects().jobs().get(name=job_name)
else:
# TODO(b/158433873): Consider raising the error instead of exit with
# code 1 after CMLE supports configurable retry policy.
# Currently CMLE will automatically retry the job unless return code
# 1-128 is returned.
logging.error('Request failed after %s retries.',
_CONNECTION_ERROR_RETRY_LIMIT)
sys.exit(1)
if response['state'] in ('FAILED', 'CANCELLED'):
err_msg = 'Job \'{}\' did not succeed. Detailed response {}.'.format(
job_name, response)
logging.error(err_msg)
raise RuntimeError(err_msg)
# AIP training complete
logging.info('Job \'%s\' successful.', job_name)
def jsonify_tfx_type_dict(artifact_dict: Dict[Text, List[Artifact]]) -> Text:
return artifact_utils.jsonify_artifact_dict(artifact_dict)
def create_training_args(self, input_dict: Dict[Text,
List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any],
executor_class_path: Text,
training_inputs: Dict[Text, Any],
job_id: Optional[Text]) -> Dict[Text, Any]:
"""Get training args for runner._launch_aip_training.
The training args contain the inputs/outputs/exec_properties to the
tfx.scripts.run_executor module.
Args:
input_dict: Passthrough input dict for tfx.components.Trainer.executor.
output_dict: Passthrough input dict for tfx.components.Trainer.executor.
exec_properties: Passthrough input dict for
tfx.components.Trainer.executor.
executor_class_path: class path for TFX core default trainer.
training_inputs: Spec for CustomJob for AI Platform (Unified) custom
training job. See
https://cloud.google.com/ai-platform-unified/docs/reference/rest/v1/CustomJobSpec
for the detailed schema.
job_id: Display name for AI Platform (Unified) custom training job. If not
supplied, system-determined unique ID is given. Refer to
https://cloud.google.com/ai-platform-unified/docs/reference/rest/v1/projects.locations.customJobs
Returns:
A dict containing the training arguments
"""
training_inputs = training_inputs.copy()
json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
logging.info('json_inputs=\'%s\'.', json_inputs)
json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
logging.info('json_outputs=\'%s\'.', json_outputs)
json_exec_properties = json.dumps(exec_properties, sort_keys=True)
logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
# We use custom containers to launch training on AI Platform (unified),
# which invokes the specified image using the container's entrypoint. The
# default entrypoint for TFX containers is to call scripts/run_executor.py.
# The arguments below are passed to this run_executor entry to run the
# executor specified in `executor_class_path`.
container_command = _CONTAINER_COMMAND + [
'--executor_class_path',
executor_class_path,
'--inputs',
json_inputs,
'--outputs',
json_outputs,
'--exec-properties',
json_exec_properties,
]
if not training_inputs.get('worker_pool_specs'):
training_inputs['worker_pool_specs'] = [{}]
for worker_pool_spec in training_inputs['worker_pool_specs']:
if not worker_pool_spec.get('container_spec'):
worker_pool_spec['container_spec'] = {
'image_uri': _TFX_IMAGE,
}
# Always use our own entrypoint instead of relying on container default.
if 'command' in worker_pool_spec['container_spec']:
logging.warn(
'Overriding custom value of container_spec.command')
worker_pool_spec['container_spec']['command'] = container_command
# Pop project_id so AIP doesn't complain about an unexpected parameter.
# It's been a stowaway in aip_args and has finally reached its destination.
project = training_inputs.pop('project')
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
job_labels = telemetry_utils.get_labels_dict()
# 'tfx_YYYYmmddHHMMSS' is the default job display name if not explicitly
# specified.
job_id = job_id or 'tfx_{}'.format(
datetime.datetime.now().strftime('%Y%m%d%H%M%S'))
training_args = {
'job_id': job_id,
'project': project,
'training_input': training_inputs,
'job_labels': job_labels
}
return training_args
def create_training_args(self, input_dict: Dict[Text,
List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any],
executor_class_path: Text,
training_inputs: Dict[Text, Any],
job_id: Optional[Text]) -> Dict[Text, Any]:
"""Get training args for runner._launch_aip_training.
The training args contain the inputs/outputs/exec_properties to the
tfx.scripts.run_executor module.
Args:
input_dict: Passthrough input dict for tfx.components.Trainer.executor.
output_dict: Passthrough input dict for tfx.components.Trainer.executor.
exec_properties: Passthrough input dict for
tfx.components.Trainer.executor.
executor_class_path: class path for TFX core default trainer.
training_inputs: Training input argument for AI Platform training job.
'pythonModule', 'pythonVersion' and 'runtimeVersion' will be inferred.
For the full set of parameters, refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#TrainingInput
job_id: Job ID for AI Platform Training job. If not supplied,
system-determined unique ID is given. Refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#resource-job
Returns:
A dict containing the training arguments
"""
training_inputs = training_inputs.copy()
json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
logging.info('json_inputs=\'%s\'.', json_inputs)
json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
logging.info('json_outputs=\'%s\'.', json_outputs)
json_exec_properties = json.dumps(exec_properties, sort_keys=True)
logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
# We use custom containers to launch training on AI Platform, which invokes
# the specified image using the container's entrypoint. The default
# entrypoint for TFX containers is to call scripts/run_executor.py. The
# arguments below are passed to this run_executor entry to run the executor
# specified in `executor_class_path`.
container_command = _CONTAINER_COMMAND + [
'--executor_class_path',
executor_class_path,
'--inputs',
json_inputs,
'--outputs',
json_outputs,
'--exec-properties',
json_exec_properties,
]
if not training_inputs.get('masterConfig'):
training_inputs['masterConfig'] = {
'imageUri': _TFX_IMAGE,
}
# Always use our own entrypoint instead of relying on container default.
if 'containerCommand' in training_inputs['masterConfig']:
logging.warn('Overriding custom value of containerCommand')
training_inputs['masterConfig']['containerCommand'] = container_command
# Pop project_id so AIP doesn't complain about an unexpected parameter.
# It's been a stowaway in aip_args and has finally reached its destination.
project = training_inputs.pop('project')
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
job_labels = telemetry_utils.get_labels_dict()
# 'tfx_YYYYmmddHHMMSS' is the default job ID if not explicitly specified.
job_id = job_id or 'tfx_{}'.format(
datetime.datetime.now().strftime('%Y%m%d%H%M%S'))
training_args = {
'job_id': job_id,
'project': project,
'training_input': training_inputs,
'job_labels': job_labels
}
return training_args
def __new__(
cls,
component_name: Text,
input_dict: Dict[Text, Any],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any],
executor_class_path: Text,
pipeline_properties: PipelineProperties,
):
"""Creates a new component.
Args:
component_name: TFX component name.
input_dict: Dictionary of input names to TFX types, or
kfp.dsl.PipelineParam representing input parameters.
output_dict: Dictionary of output names to List of TFX types.
exec_properties: Execution properties.
executor_class_path: <module>.<class> for Python class of executor.
pipeline_properties: Pipeline level properties shared by all components.
Returns:
Newly constructed TFX Kubeflow component instance.
"""
outputs = output_dict.keys()
file_outputs = {
output: '/output/ml_metadata/{}'.format(output) for output in outputs
}
for k, v in pipeline_properties.exec_properties.items():
exec_properties[k] = v
arguments = [
'--exec_properties',
json.dumps(exec_properties),
'--outputs',
artifact_utils.jsonify_artifact_dict(output_dict),
'--executor_class_path',
executor_class_path,
component_name,
]
for k, v in input_dict.items():
if isinstance(v, float) or isinstance(v, int):
v = str(v)
arguments.append('--{}'.format(k))
arguments.append(v)
container_op = dsl.ContainerOp(
name=component_name,
command=_COMMAND,
image=pipeline_properties.tfx_image,
arguments=arguments,
file_outputs=file_outputs,
)
# Add the Argo workflow ID to the container's environment variable so it
# can be used to uniquely place pipeline outputs under the pipeline_root.
field_path = "metadata.labels['workflows.argoproj.io/workflow']"
container_op.add_env_variable(
k8s_client.V1EnvVar(
name='WORKFLOW_ID',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=field_path))))
named_outputs = {output: container_op.outputs[output] for output in outputs}
# This allows user code to refer to the ContainerOp 'op' output named 'x'
# as op.outputs.x
component_outputs = type('Output', (), named_outputs)
return type(component_name, (BaseComponent,), {
'container_op': container_op,
'outputs': component_outputs
})
def start_cmle_training(input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any],
executor_class_path: Text, training_inputs: Dict[Text,
Any]):
"""Start a trainer job on CMLE.
This is done by forwarding the inputs/outputs/exec_properties to the
tfx.scripts.run_executor module on a CMLE training job interpreter.
Args:
input_dict: Passthrough input dict for tfx.components.Trainer.executor.
output_dict: Passthrough input dict for tfx.components.Trainer.executor.
exec_properties: Passthrough input dict for tfx.components.Trainer.executor.
executor_class_path: class path for TFX core default trainer.
training_inputs: Training input for CMLE training job. 'pythonModule',
'pythonVersion' and 'runtimeVersion' will be inferred by the runner. For
the full set of parameters supported, refer to
https://cloud.google.com/ml-engine/docs/tensorflow/deploying-models#creating_a_model_version.
Returns:
None
Raises:
RuntimeError: if the Google Cloud AI Platform training job failed.
"""
training_inputs = training_inputs.copy()
# Remove cmle_args from exec_properties so CMLE trainer doesn't call itself
for gaip_training_key in ['cmle_training_args', 'gaip_training_args']:
if gaip_training_key in exec_properties.get('custom_config'):
exec_properties['custom_config'].pop(gaip_training_key)
json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
tf.logging.info('json_inputs=\'%s\'.', json_inputs)
json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
tf.logging.info('json_outputs=\'%s\'.', json_outputs)
json_exec_properties = json.dumps(exec_properties)
tf.logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
# Configure CMLE job
api_client = discovery.build('ml', 'v1')
job_args = [
'--executor_class_path', executor_class_path, '--inputs', json_inputs,
'--outputs', json_outputs, '--exec-properties', json_exec_properties
]
training_inputs['args'] = job_args
training_inputs['pythonModule'] = 'tfx.scripts.run_executor'
training_inputs['pythonVersion'] = _get_caip_python_version()
# runtimeVersion should be same as <major>.<minor> of currently
# installed tensorflow version.
training_inputs['runtimeVersion'] = _get_tf_runtime_version()
# Pop project_id so CMLE doesn't complain about an unexpected parameter.
# It's been a stowaway in cmle_args and has finally reached its destination.
project = training_inputs.pop('project')
project_id = 'projects/{}'.format(project)
package_uris = training_inputs.get('packageUris', [])
if package_uris:
tf.logging.info('Following packageUris \'%s\' are provided by user.',
package_uris)
else:
local_package = dependency_utils.build_ephemeral_package()
# TODO(b/125451545): Use a safe temp dir instead of jobDir.
cloud_package = os.path.join(training_inputs['jobDir'],
os.path.basename(local_package))
io_utils.copy_file(local_package, cloud_package, True)
training_inputs['packageUris'] = [cloud_package]
tf.logging.info('Package %s will be used',
training_inputs['packageUris'])
job_name = 'tfx_' + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
job_spec = {'jobId': job_name, 'trainingInput': training_inputs}
# Submit job to CMLE
tf.logging.info('Submitting job=\'{}\', project=\'{}\' to CMLE.'.format(
job_name, project))
request = api_client.projects().jobs().create(body=job_spec,
parent=project_id)
request.execute()
# Wait for CMLE job to finish
job_id = '{}/jobs/{}'.format(project_id, job_name)
request = api_client.projects().jobs().get(name=job_id)
response = request.execute()
while response['state'] not in ('SUCCEEDED', 'FAILED'):
time.sleep(_POLLING_INTERVAL_IN_SECONDS)
response = request.execute()
if response['state'] == 'FAILED':
err_msg = 'Job \'{}\' did not succeed. Detailed response {}.'.format(
job_name, response)
tf.logging.error(err_msg)
raise RuntimeError(err_msg)
# CMLE training complete
tf.logging.info('Job \'{}\' successful.'.format(job_name))
def start_cmle_training(input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any],
executor_class_path: Text, training_inputs: Dict[Text,
Any]):
"""Start a trainer job on CMLE.
This is done by forwarding the inputs/outputs/exec_properties to the
tfx.scripts.run_executor module on a CMLE training job interpreter.
Args:
input_dict: Passthrough input dict for tfx.components.Trainer.executor.
output_dict: Passthrough input dict for tfx.components.Trainer.executor.
exec_properties: Passthrough input dict for tfx.components.Trainer.executor.
executor_class_path: class path for TFX core default trainer.
training_inputs: Training input for CMLE training job. 'pythonModule',
'pythonVersion' and 'runtimeVersion' will be inferred by the runner. For
the full set of parameters supported, refer to
https://cloud.google.com/ml-engine/docs/tensorflow/deploying-models#creating_a_model_version.
Returns:
None
Raises:
RuntimeError: if the Google Cloud AI Platform training job failed.
"""
training_inputs = training_inputs.copy()
# Remove cmle_args from exec_properties so CMLE trainer doesn't call itself
for gaip_training_key in ['cmle_training_args', 'gaip_training_args']:
if gaip_training_key in exec_properties.get('custom_config'):
exec_properties['custom_config'].pop(gaip_training_key)
json_inputs = artifact_utils.jsonify_artifact_dict(input_dict)
absl.logging.info('json_inputs=\'%s\'.', json_inputs)
json_outputs = artifact_utils.jsonify_artifact_dict(output_dict)
absl.logging.info('json_outputs=\'%s\'.', json_outputs)
json_exec_properties = json.dumps(exec_properties)
absl.logging.info('json_exec_properties=\'%s\'.', json_exec_properties)
# Configure CMLE job
api_client = discovery.build('ml', 'v1')
# We use custom containers to launch training on AI Platform, which invokes
# the specified image using the container's entrypoint. The default
# entrypoint for TFX containers is to call scripts/run_executor.py. The
# arguments below are passed to this run_executor entry to run the executor
# specified in `executor_class_path`.
job_args = [
'--executor_class_path', executor_class_path, '--inputs', json_inputs,
'--outputs', json_outputs, '--exec-properties', json_exec_properties
]
if not training_inputs.get('masterConfig'):
training_inputs['masterConfig'] = {
'imageUri': _TFX_IMAGE,
}
training_inputs['args'] = job_args
# Pop project_id so CMLE doesn't complain about an unexpected parameter.
# It's been a stowaway in cmle_args and has finally reached its destination.
project = training_inputs.pop('project')
project_id = 'projects/{}'.format(project)
job_name = 'tfx_' + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
job_spec = {'jobId': job_name, 'trainingInput': training_inputs}
# Submit job to CMLE
absl.logging.info('Submitting job=\'{}\', project=\'{}\' to CMLE.'.format(
job_name, project))
request = api_client.projects().jobs().create(
body=job_spec, parent=project_id)
request.execute()
# Wait for CMLE job to finish
job_id = '{}/jobs/{}'.format(project_id, job_name)
request = api_client.projects().jobs().get(name=job_id)
response = request.execute()
while response['state'] not in ('SUCCEEDED', 'FAILED'):
time.sleep(_POLLING_INTERVAL_IN_SECONDS)
response = request.execute()
if response['state'] == 'FAILED':
err_msg = 'Job \'{}\' did not succeed. Detailed response {}.'.format(
job_name, response)
absl.logging.error(err_msg)
raise RuntimeError(err_msg)
# CMLE training complete
absl.logging.info('Job \'{}\' successful.'.format(job_name))
def _run_executor(args, pipeline_args) -> None:
r"""Select a particular executor and run it based on name.
# pylint: disable=line-too-long
_run_executor() is used to invoke a class subclassing
tfx.components.base.base_executor.BaseExecutor. This function can be used for
both invoking the executor on remote environments as well as for unit testing
of executors.
How to invoke an executor as standalone:
# TODO(b/132958430): Create utility script to generate arguments for run_executor.py
First, the input data needs to be prepared. An easy way to generate the test
data is to fully run the pipeline once. This will generate the data to be
used for testing as well as log the artifacts to be used as input parameters.
In each executed component, three log entries will be generated similar to the
below:
```
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,116] {base_executor.py:72} INFO - Starting Executor execution.
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,117] {base_executor.py:74} INFO - Inputs for Executor is: {"input_base": [{"artifact": {"id": "1", "typeId": "1", "uri": "/usr/local/google/home/khaas/taxi/data/simple", "properties": {"split": {"stringValue": ""}, "state": {"stringValue": "published"}, "span": {"intValue": "1"}, "type_name": {"stringValue": "ExternalPath"}}}, "artifact_type": {"id": "1", "name": "ExternalPath", "properties": {"span": "INT", "name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING"}}}]}
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,117] {base_executor.py:76} INFO - Outputs for Executor is: {"examples": [{"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/train/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "train"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}, {"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/eval/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "eval"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}]}
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,117] {base_executor.py:78} INFO - Execution properties for Executor is: {"output": "{ \"splitConfig\": {\"splits\": [{\"name\": \"train\", \"hashBuckets\": 2}, {\"name\": \"eval\",\"hashBuckets\": 1}]}}"}
```
Each of these map directly to the input parameters expected by run_executor():
```
python scripts/run_executor.py \
--executor_class_path=tfx.components.example_gen.big_query_example_gen.executor.Executor \
--inputs={"input_base": [{"artifact": {"id": "1", "typeId": "1", "uri": "/usr/local/google/home/khaas/taxi/data/simple", "properties": {"split": {"stringValue": ""}, "state": {"stringValue": "published"}, "span": {"intValue": "1"}, "type_name": {"stringValue": "ExternalPath"}}}, "artifact_type": {"id": "1", "name": "ExternalPath", "properties": {"span": "INT", "name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING"}}}]} \
--outputs={"examples": [{"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/train/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "train"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}, {"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/eval/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "eval"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}]} \
--exec-properties={"output": "{ \"splitConfig\": {\"splits\": [{\"name\": \"train\", \"hashBuckets\": 2}, {\"name\": \"eval\",\"hashBuckets\": 1}]}}"}
```
# pylint: disable=line-too-long
Args:
args:
- inputs: The input artifacts for this execution, serialized as JSON.
- outputs: The output artifacts to be generated by this execution,
serialized as JSON.
- exec_properties: The execution properties to be used by this execution,
serialized as JSON. Technically all the exec_properties values should be
a primitive, and nested exec_properties needs to be JSON-encoded as a
string. But as a convenience, the script allows you to feed in
non-serialized values of exec_properties, which is then automatically
serialized.
pipeline_args: Optional parameter that maps to the optional_pipeline_args
parameter in the pipeline, which provides additional configuration options
for apache-beam and tensorflow.logging.
Returns:
None
Raises:
None
"""
(inputs_str, outputs_str,
exec_properties_str) = (args.inputs
or base64.b64decode(args.inputs_base64),
args.outputs
or base64.b64decode(args.outputs_base64),
args.exec_properties
or base64.b64decode(args.exec_properties_base64))
inputs = artifact_utils.parse_artifact_dict(inputs_str)
outputs = artifact_utils.parse_artifact_dict(outputs_str)
exec_properties = json.loads(exec_properties_str)
# Technically exec_properties value can only be a primitive (e.g. string), and
# one of our convention is to use proto object by JSON-serializing it.
# Unfortunately, run_executor.py script accepts serialized exec_properties as
# an input, thus proto object value would be serialized twice. This is really
# inconvenient if you're manually constructing exec_properties, so we allow
# to feed in non-serialized values of exec_properties, and serialize them
# here.
for key, value in exec_properties.items():
if isinstance(value, (dict, list)):
exec_properties[key] = json.dumps(value)
logging.info(
'Executor %s do: inputs: %s, outputs: %s, exec_properties: %s',
args.executor_class_path, inputs, outputs, exec_properties)
executor_cls = import_utils.import_class_by_path(args.executor_class_path)
executor_context = base_executor.BaseExecutor.Context(
beam_pipeline_args=pipeline_args,
tmp_dir=args.temp_directory_path,
unique_id='')
executor = executor_cls(executor_context)
logging.info('Starting executor')
executor.Do(inputs, outputs, exec_properties)
# The last line of stdout will be pushed to xcom by Airflow.
if args.write_outputs_stdout:
print(artifact_utils.jsonify_artifact_dict(outputs))
def _run_executor(args, pipeline_args) -> None:
r"""Select a particular executor and run it based on name.
# pylint: disable=line-too-long
_run_executor() is used to invoke a class subclassing
tfx.components.base.base_executor.BaseExecutor. This function can be used for
both invoking the executor on remote environments as well as for unit testing
of executors.
How to invoke an executor as standalone:
# TODO(b/132958430): Create utility script to generate arguments for run_executor.py
First, the input data needs to be prepared. An easy way to generate the test
data is to fully run the pipeline once. This will generate the data to be
used for testing as well as log the artifacts to be used as input parameters.
In each executed component, three log entries will be generated similar to the
below:
```
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,116] {base_executor.py:72} INFO - Starting Executor execution.
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,117] {base_executor.py:74} INFO - Inputs for Executor is: {"input_base": [{"artifact": {"id": "1", "typeId": "1", "uri": "/usr/local/google/home/khaas/taxi/data/simple", "properties": {"split": {"stringValue": ""}, "state": {"stringValue": "published"}, "span": {"intValue": "1"}, "type_name": {"stringValue": "ExternalPath"}}}, "artifact_type": {"id": "1", "name": "ExternalPath", "properties": {"span": "INT", "name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING"}}}]}
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,117] {base_executor.py:76} INFO - Outputs for Executor is: {"examples": [{"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/train/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "train"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}, {"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/eval/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "eval"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}]}
[2019-05-16 08:59:27,117] {logging_mixin.py:95} INFO - [2019-05-16 08:59:27,117] {base_executor.py:78} INFO - Execution properties for Executor is: {"output": "{ \"splitConfig\": {\"splits\": [{\"name\": \"train\", \"hashBuckets\": 2}, {\"name\": \"eval\",\"hashBuckets\": 1}]}}"}
```
Each of these map directly to the input parameters expected by run_executor():
```
python scripts/run_executor.py \
--executor_class_path=tfx.components.example_gen.big_query_example_gen.executor.Executor \
--inputs={"input_base": [{"artifact": {"id": "1", "typeId": "1", "uri": "/usr/local/google/home/khaas/taxi/data/simple", "properties": {"split": {"stringValue": ""}, "state": {"stringValue": "published"}, "span": {"intValue": "1"}, "type_name": {"stringValue": "ExternalPath"}}}, "artifact_type": {"id": "1", "name": "ExternalPath", "properties": {"span": "INT", "name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING"}}}]} \
--outputs={"examples": [{"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/train/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "train"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}, {"artifact": {"uri": "/usr/local/google/home/khaas/tfx/pipelines/chicago_taxi_simple/CsvExampleGen/examples/1/eval/", "properties": {"type_name": {"stringValue": "ExamplesPath"}, "split": {"stringValue": "eval"}, "span": {"intValue": "1"}}}, "artifact_type": {"name": "ExamplesPath", "properties": {"name": "STRING", "type_name": "STRING", "split": "STRING", "state": "STRING", "span": "INT"}}}]} \
--exec-properties={"output": "{ \"splitConfig\": {\"splits\": [{\"name\": \"train\", \"hashBuckets\": 2}, {\"name\": \"eval\",\"hashBuckets\": 1}]}}"}
```
# pylint: disable=line-too-long
Args:
args:
- inputs: The input artifacts for this execution, serialized as JSON.
- outputs: The output artifacts to be generated by this execution,
serialized as JSON.
- exec_properties: The execution properties to be used by this execution,
serialized as JSON.
pipeline_args: Optional parameter that maps to the optional_pipeline_args
parameter in the pipeline, which provides additional configuration options
for apache-beam and tensorflow.logging.
Returns:
None
Raises:
None
"""
tf.logging.set_verbosity(tf.logging.INFO)
(inputs_str, outputs_str,
exec_properties_str) = (args.inputs
or base64.b64decode(args.inputs_base64),
args.outputs
or base64.b64decode(args.outputs_base64),
args.exec_properties
or base64.b64decode(args.exec_properties_base64))
inputs = artifact_utils.parse_artifact_dict(inputs_str)
outputs = artifact_utils.parse_artifact_dict(outputs_str)
exec_properties = json.loads(exec_properties_str)
tf.logging.info(
'Executor {} do: inputs: {}, outputs: {}, exec_properties: {}'.format(
args.executor_class_path, inputs, outputs, exec_properties))
executor_cls = import_utils.import_class_by_path(args.executor_class_path)
executor_context = base_executor.BaseExecutor.Context(
beam_pipeline_args=pipeline_args,
tmp_dir=args.temp_directory_path,
unique_id='')
executor = executor_cls(executor_context)
tf.logging.info('Starting executor')
executor.Do(inputs, outputs, exec_properties)
# The last line of stdout will be pushed to xcom by Airflow.
if args.write_outputs_stdout:
print(artifact_utils.jsonify_artifact_dict(outputs))
def __init__(
self,
component: tfx_base_component.BaseComponent,
component_launcher_class: Type[
base_component_launcher.BaseComponentLauncher],
depends_on: Set[dsl.ContainerOp],
pipeline: tfx_pipeline.Pipeline,
tfx_image: Text,
kubeflow_metadata_config: Optional[
kubeflow_pb2.KubeflowMetadataConfig],
):
"""Creates a new Kubeflow-based component.
This class essentially wraps a dsl.ContainerOp construct in Kubeflow
Pipelines.
Args:
component: The logical TFX component to wrap.
component_launcher_class: the class of the launcher to launch the
component.
depends_on: The set of upstream KFP ContainerOp components that this
component will depend on.
pipeline: The logical TFX pipeline to which this component belongs.
tfx_image: The container image to use for this component.
kubeflow_metadata_config: Configuration settings for
connecting to the MLMD store in a Kubeflow cluster.
"""
driver_class_path = '.'.join([
component.driver_class.__module__, component.driver_class.__name__
])
executor_spec = json_utils.dumps(component.executor_spec)
component_launcher_class_path = '.'.join([
component_launcher_class.__module__,
component_launcher_class.__name__
])
arguments = [
'--pipeline_name',
pipeline.pipeline_info.pipeline_name,
'--pipeline_root',
pipeline.pipeline_info.pipeline_root,
'--kubeflow_metadata_config',
json_format.MessageToJson(kubeflow_metadata_config),
'--additional_pipeline_args',
json.dumps(pipeline.additional_pipeline_args),
'--component_id',
component.component_id,
'--component_type',
component.component_type,
'--driver_class_path',
driver_class_path,
'--executor_spec',
executor_spec,
'--component_launcher_class_path',
component_launcher_class_path,
'--inputs',
artifact_utils.jsonify_artifact_dict(
_prepare_artifact_dict(component.inputs)),
'--outputs',
artifact_utils.jsonify_artifact_dict(
_prepare_artifact_dict(component.outputs)),
'--exec_properties',
json.dumps(component.exec_properties),
]
if pipeline.enable_cache:
arguments.append('--enable_cache')
self.container_op = dsl.ContainerOp(
name=component.component_id.replace('.', '_'),
command=_COMMAND,
image=tfx_image,
arguments=arguments,
)
tf.logging.info('Adding upstream dependencies for component {}'.format(
self.container_op.name))
for op in depends_on:
tf.logging.info(' -> Component: {}'.format(op.name))
self.container_op.after(op)
# TODO(b/140172100): Document the use of additional_pipeline_args.
if _WORKFLOW_ID_KEY in pipeline.additional_pipeline_args:
# Allow overriding pipeline's run_id externally, primarily for testing.
self.container_op.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value=pipeline.additional_pipeline_args[_WORKFLOW_ID_KEY]))
else:
# Add the Argo workflow ID to the container's environment variable so it
# can be used to uniquely place pipeline outputs under the pipeline_root.
field_path = "metadata.labels['workflows.argoproj.io/workflow']"
self.container_op.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=field_path))))
