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)