def _injection(self):
"""Sample-specific image injection into yaml file."""
subs = {
'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-confusion-matrix:\w+':self._local_confusionmatrix_image,
'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-roc:\w+':self._local_roc_image
}
if self._test_name == 'xgboost_training_cm':
subs.update({
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-create-cluster:\w+':self._dataproc_create_cluster_image,
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-delete-cluster:\w+':self._dataproc_delete_cluster_image,
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-analyze:\w+':self._dataproc_analyze_image,
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-transform:\w+':self._dataproc_transform_image,
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-train:\w+':self._dataproc_train_image,
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-predict:\w+':self._dataproc_predict_image,
})
utils.file_injection('%s.yaml' % self._test_name,
'%s.yaml.tmp' % self._test_name,
subs)
else:
# Only the above two samples need injection for now.
pass
utils.file_injection('%s.yaml' % self._test_name,
'%s.yaml.tmp' % self._test_name,
subs)
def _injection(self):
"""Sample-specific image injection into yaml file."""
subs = { # Tag can look like 1.0.0-rc.3, so we need both "-" and "." in the regex.
'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-confusion-matrix:(\w+|[.-])+':
self._local_confusionmatrix_image,
'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-roc:(\w+|[.-])+':
self._local_roc_image
}
if self._test_name == 'xgboost_training_cm':
subs.update({
'gcr\.io/ml-pipeline/ml-pipeline-gcp:(\w|[.-])+':
self._dataproc_gcp_image
})
utils.file_injection('%s.py.yaml' % self._test_name,
'%s.py.yaml.tmp' % self._test_name, subs)
else:
# Only the above sample need injection for now.
pass
utils.file_injection('%s.py.yaml' % self._test_name,
'%s.py.yaml.tmp' % self._test_name, subs)
def _injection(self):
"""Sample-specific image injection into yaml file."""
subs = {
'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-confusion-matrix:\w+':
self._local_confusionmatrix_image,
'gcr\.io/ml-pipeline/ml-pipeline/ml-pipeline-local-roc:\w+':
self._local_roc_image
}
if self._test_name == 'xgboost_training_cm':
subs.update({
'gcr\.io/ml-pipeline/ml-pipeline-gcp:\w+':
self._dataproc_gcp_image
})
utils.file_injection('%s.yaml' % self._test_name,
'%s.yaml.tmp' % self._test_name, subs)
else:
# Only the above two samples need injection for now.
pass
utils.file_injection('%s.yaml' % self._test_name,
'%s.yaml.tmp' % self._test_name, subs)
def test_file_injection(self):
"""Test file_injection function."""
subs = {
'gcr\.io/ml-pipeline/ml-pipeline-local-confusion-matrix:\w+':
'gcr.io/ml-pipeline/LOCAL_CONFUSION_MATRIX_IMAGE',
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-analyze:\w+':
'gcr.io/ml-pipeline/DATAPROC_ANALYZE_IMAGE',
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-create-cluster:\w+':
'gcr.io/ml-pipeline/DATAPROC_CREATE_IMAGE',
'gcr\.io/ml-pipeline/ml-pipeline-dataproc-delete-cluster:\w+':
'gcr.io/ml-pipeline/DATAPROC_DELETE_IMAGE',
}
utils.file_injection(
os.path.join(_WORK_DIR, 'test_file_injection.yaml'),
os.path.join(_WORK_DIR, 'test_file_injection_tmp.yaml'), subs)
with open(os.path.join(_DATAPATH, 'test_file_injection_output.yaml'),
'r') as f:
golden = yaml.safe_load(f)
with open(os.path.join(_WORK_DIR, 'test_file_injection.yaml'),
'r') as f:
injected = yaml.safe_load(f)
self.assertEqual(golden, injected)