To create the custom executor, we start with a copy of the current Trainer executor and then make the modifications to our custom executor to initiate a training job on Google Cloud AI Platform. Much of the basic executor structure will remain the same as the inputs, outputs, and execution parameters will be the same. The changes will be in how the inputs are processed and the outputs are generated. This is achieved by creating a new Executor class that extends tfx.components.base.base_executor.BaseExecutor and implements Do(). class Executor(base_executor.BaseExecutor):
"""Start a trainer job on Google Cloud AI Platform."""
def Do(self, input_dict,
output_dict,
exec_properties):
"""Starts a trainer job on Google Cloud AI Platform.
Don’t forget to test it before moving on to the next step! We have created a convenience script for you to try out your executor before putting it into production. You should write similar code to exercise unit tests for your code. As with any production software deployment, when developing for TFX you should make sure to have good test coverage and a strong CI/CD framework.Override the executor used by the Trainer component
In order to do this, we will replace the default trainer executor used by TFX with the new custom executor which will create the training job on Google Cloud AI Platform. This is achieved with the optional executor_class component parameter. from tfx.extensions.google_cloud_ai_platform.trainer
import executor as ai_platform_trainer_executor
...
trainer = Trainer(
...,
custom_executor_spec = executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.Executor)
)
That’s it! Now when the Trainer component is called by the workflow engine, it will run the custom executor instead of the default executor, while creating and consuming the same ML Metadata artifacts as the default executor. Pass the component arguments to your trainer
TFX executors are self-contained binaries focused on running a single step of the ML pipeline. Custom executors require the same three parameters as all other TFX executors: input_dict, output_dict, exec_properties. More details on the semantics of these parameters can be found in the BaseExecutor class. When processing data flowing through your TFX pipeline you will often typically want to read input data from the artifact URIs in your input_dict, and often you might want to write your output to artifact URIs from your output_dict. This may include reading and writing to more than one split, as in the case of processing with train and eval splits.
from tfx.types import artifact_utils
train_input_examples_uri = artifact_utils.get_split_uri(
input_dict['my_input_data'], 'train')
eval_input_examples_uri = artifact_utils.get_split_uri(
input_dict['my_input_data'], 'eval')
train_output_examples_uri = artifact_utils.get_split_uri(
output_dict['my_output_data'], 'train')
eval_output_examples_uri = artifact_utils.get_split_uri(
output_dict[‘my_output_data'], 'eval')
In the example above the dictionary keys my_input_data and my_output_data are defined in the ComponentSpec for the component that you’re overriding the executor for.
class MyComponentSpec(tfx.types.ComponentSpec):
PARAMETERS = {
<...>
}
INPUTS = {
'my_input_data':
ChannelParameter(type=standard_artifacts.Examples),
}
OUTPUTS = {
'my_output_data':
ChannelParameter(type=standard_artifacts.Examples),
}
The splits are defined in the output Channel for the component that you’re overriding the executor for, typically in the constructor:
output_data = tfx.types.Channel(
type=standard_artifacts.Examples,
artifacts=[
standard_artifacts.Examples(split=split)
for split in artifact.DEFAULT_EXAMPLE_SPLITS
])
spec = LabelingComponentSpec(
input_data=input_data,
my_output_data=output_data)
Additional parameters are passed to your custom trainer executor using a custom_config dict. These can be retrieved by the custom executor using exec_properties.get(‘custom_config’).get(‘your_config_key’). In the example below, all of the additional arguments needed to submit a Google Cloud AI Platform training job can be found in _ai_platform_training_args. _ai_platform_training_args = {
'pythonModule': None, # Will be populated by TFX
'args': None, # Will be populated by TFX
'region': _gcp_region,
'jobDir': os.path.join(_output_bucket, 'tmp'),
'project': ‘your GCP project id’
}
...
trainer = Trainer(
...,
custom_config={'ai_platform_training_args':
_ai_platform_training_args})
}
Linking the custom trainer’s output to the expected output artifact
Hooking up the custom trainer to emit the expected outputs is essential to the success of downstream components. For the Google Cloud AI Platform custom trainer, we serialize the executor input parameters so they can be transmitted as part of the GCP training job. Because the Google Cloud AI Platform (CAIP) executor is redirecting the default TFX executor to run on Google Cloud AI Platform, both take the same {transformed examples, a transform_fn, and a schema} input parameters to create a TF model. The custom executor used in this example submits a CAIP training job that will invoke (via run_executor.py) the default TFX trainer as the CAIP python module, effectively opening a conduit from the local workstation to run the TFX trainer on CAIP. # Configure Google Cloud AI Platform job
training_inputs = exec_properties.get('custom_config',
{}).pop('ai_platform_training_args')
executor_class_path = '%s.%s' %
(tfx_trainer_executor.Executor.__module__,
tfx_trainer_executor.Executor.__name__)
# Start Google Cloud AI Platform job
return runner.start_cmle_training(input_dict, output_dict,
exec_properties, executor_class_path, training_inputs)
Running the pipeline remotely with your custom executor
So far we’ve been assuming that your pipeline is running locally, using code available in your $PYTHONPATH. An upcoming blog post will explain how to execute custom executors packaged in containers, or as PyPI packages.
Related topics
In addition to Trainer, the TFX ExampleGen component also supports executor-level customization. ExampleGen provides a generic component and a base executor which apply ML best practices, e.g. data shuffling and consistent/configurable partitions. If existing ExampleGen components don’t meet your needs, create a new Apache Beam PTransform for handling the conversion from an input split to TF examples and TFX will do the rest. The ExampleGen doc has more details. For more information
