diff --git a/components/ibm-components/ffdl/train/src/train.py b/components/ibm-components/ffdl/train/src/train.py index ee86f080555..f43ad5dfeb8 100644 --- a/components/ibm-components/ffdl/train/src/train.py +++ b/components/ibm-components/ffdl/train/src/train.py @@ -78,7 +78,7 @@ ''' Update FfDL manifest with the corresponding object storage credentials ''' f = open('manifest.yml', 'r') - manifest = yaml.load(f.read()) + manifest = yaml.safe_load(f.read()) f.close() manifest['data_stores'][0]['connection']['auth_url'] = s3_url diff --git a/components/kubeflow/katib-launcher/src/launch_study_job.py b/components/kubeflow/katib-launcher/src/launch_study_job.py index 8f901a00d82..54e614abd4d 100644 --- a/components/kubeflow/katib-launcher/src/launch_study_job.py +++ b/components/kubeflow/katib-launcher/src/launch_study_job.py @@ -35,7 +35,7 @@ def yamlOrJsonStr(str): try: return json.loads(str) except: - return yaml.load(str) + return yaml.safe_load(str) def strToList(str): return str.split(",") @@ -50,7 +50,7 @@ def _generate_studyjob_yaml(src_filename, name, namespace, optimizationtype, obj metricsnames, parameterconfigs, nasConfig, workertemplatepath, mcollectortemplatepath, suggestionspec): """_generate_studyjob_yaml generates studyjob yaml file based on hp.template.yaml""" with open(src_filename, 'r') as f: - content = yaml.load(f) + content = yaml.safe_load(f) content['metadata']['name'] = name content['metadata']['namespace'] = namespace diff --git a/components/kubeflow/launcher/src/launch_tf_job.py b/components/kubeflow/launcher/src/launch_tf_job.py index a728ae76df9..414c15fd7c5 100644 --- a/components/kubeflow/launcher/src/launch_tf_job.py +++ b/components/kubeflow/launcher/src/launch_tf_job.py @@ -49,7 +49,7 @@ def _generate_train_yaml(src_filename, tfjob_ns, workers, pss, trainer_image, command): """_generate_train_yaml generates train yaml files based on train.template.yaml""" with open(src_filename, 'r') as f: - content = yaml.load(f) + content = yaml.safe_load(f) content['metadata']['generateName'] = 'trainer-' content['metadata']['namespace'] = tfjob_ns diff --git a/components/kubeflow/launcher/test/test_launcher.py b/components/kubeflow/launcher/test/test_launcher.py index c9447f980a8..4b1f0ab5029 100644 --- a/components/kubeflow/launcher/test/test_launcher.py +++ b/components/kubeflow/launcher/test/test_launcher.py @@ -38,7 +38,7 @@ def test_yaml_generation_basic(self): args_list.append('--learning-rate=0.1') generated_yaml = train._generate_train_yaml(train_template_file, tfjob_ns, worker, pss, args_list) with open(os.path.join(test_data_dir, 'train_basic.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) self.assertEqual(golden, generated_yaml) def test_yaml_generation_advanced(self): @@ -53,7 +53,7 @@ def test_yaml_generation_advanced(self): args_list.append('--learning-rate=0.1') generated_yaml = train._generate_train_yaml(train_template_file, tfjob_ns, worker, pss, args_list) with open(os.path.join(test_data_dir, 'train_zero_worker.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) self.assertEqual(golden, generated_yaml) if __name__ == '__main__': diff --git a/sdk/python/kfp/_client.py b/sdk/python/kfp/_client.py index 784ca427f21..c70d946ca4e 100644 --- a/sdk/python/kfp/_client.py +++ b/sdk/python/kfp/_client.py @@ -181,7 +181,7 @@ def _extract_pipeline_yaml(self, package_file): raise ValueError('Invalid package. Multiple yaml files in the package.') with tar.extractfile(all_yaml_files[0]) as f: - return yaml.load(f) + return yaml.safe_load(f) elif package_file.endswith('.zip'): with zipfile.ZipFile(package_file, 'r') as zip: all_yaml_files = [m for m in zip.namelist() if @@ -193,10 +193,10 @@ def _extract_pipeline_yaml(self, package_file): raise ValueError('Invalid package. Multiple yaml files in the package.') with zip.open(all_yaml_files[0]) as f: - return yaml.load(f) + return yaml.safe_load(f) elif package_file.endswith('.yaml') or package_file.endswith('.yml'): with open(package_file, 'r') as f: - return yaml.load(f) + return yaml.safe_load(f) else: raise ValueError('The package_file '+ package_file + ' should ends with one of the following formats: [.tar.gz, .tgz, .zip, .yaml, .yml]') diff --git a/sdk/python/kfp/components/_yaml_utils.py b/sdk/python/kfp/components/_yaml_utils.py index 71edad6b454..22be0899998 100644 --- a/sdk/python/kfp/components/_yaml_utils.py +++ b/sdk/python/kfp/components/_yaml_utils.py @@ -18,7 +18,7 @@ def load_yaml(stream): #!!! Yaml should only be loaded using this function. Otherwise the dict ordering may be broken in Python versions prior to 3.6 #See https://stackoverflow.com/questions/5121931/in-python-how-can-you-load-yaml-mappings-as-ordereddicts/21912744#21912744 - def ordered_load(stream, Loader=yaml.Loader, object_pairs_hook=OrderedDict): + def ordered_load(stream, Loader=yaml.SafeLoader, object_pairs_hook=OrderedDict): class OrderedLoader(Loader): pass def construct_mapping(loader, node): diff --git a/sdk/python/tests/compiler/compiler_tests.py b/sdk/python/tests/compiler/compiler_tests.py index 7f2f588f684..93f6839635c 100644 --- a/sdk/python/tests/compiler/compiler_tests.py +++ b/sdk/python/tests/compiler/compiler_tests.py @@ -122,11 +122,11 @@ def test_operator_to_template(self): def _get_yaml_from_zip(self, zip_file): with zipfile.ZipFile(zip_file, 'r') as zip: with open(zip.extract(zip.namelist()[0]), 'r') as yaml_file: - return yaml.load(yaml_file) + return yaml.safe_load(yaml_file) def _get_yaml_from_tar(self, tar_file): with tarfile.open(tar_file, 'r:gz') as tar: - return yaml.load(tar.extractfile(tar.getmembers()[0])) + return yaml.safe_load(tar.extractfile(tar.getmembers()[0])) def test_basic_workflow(self): """Test compiling a basic workflow.""" @@ -139,7 +139,7 @@ def test_basic_workflow(self): try: compiler.Compiler().compile(basic.save_most_frequent_word, package_path) with open(os.path.join(test_data_dir, 'basic.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) compiled = self._get_yaml_from_zip(package_path) self.maxDiff = None @@ -166,7 +166,7 @@ def test_composing_workflow(self): compose_package_path = os.path.join(tmpdir, 'compose.zip') compiler.Compiler().compile(compose.download_save_most_frequent_word, compose_package_path) with open(os.path.join(test_data_dir, 'compose.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) compiled = self._get_yaml_from_zip(compose_package_path) self.maxDiff = None @@ -193,7 +193,7 @@ def test_package_compile(self): 'dsl-compile', '--package', package_path, '--namespace', 'mypipeline', '--output', target_zip, '--function', 'download_save_most_frequent_word']) with open(os.path.join(test_data_dir, 'compose.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) compiled = self._get_yaml_from_zip(target_zip) self.maxDiff = None @@ -211,7 +211,7 @@ def _test_py_compile_zip(self, file_base_name): subprocess.check_call([ 'dsl-compile', '--py', py_file, '--output', target_zip]) with open(os.path.join(test_data_dir, file_base_name + '.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) compiled = self._get_yaml_from_zip(target_zip) self.maxDiff = None @@ -228,7 +228,7 @@ def _test_py_compile_targz(self, file_base_name): subprocess.check_call([ 'dsl-compile', '--py', py_file, '--output', target_tar]) with open(os.path.join(test_data_dir, file_base_name + '.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) compiled = self._get_yaml_from_tar(target_tar) self.maxDiff = None self.assertEqual(golden, compiled) @@ -244,10 +244,10 @@ def _test_py_compile_yaml(self, file_base_name): subprocess.check_call([ 'dsl-compile', '--py', py_file, '--output', target_yaml]) with open(os.path.join(test_data_dir, file_base_name + '.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) with open(os.path.join(test_data_dir, target_yaml), 'r') as f: - compiled = yaml.load(f) + compiled = yaml.safe_load(f) self.maxDiff = None self.assertEqual(golden, compiled) diff --git a/sdk/python/tests/compiler/component_builder_test.py b/sdk/python/tests/compiler/component_builder_test.py index 54d159a424f..abb6dd9e5dd 100644 --- a/sdk/python/tests/compiler/component_builder_test.py +++ b/sdk/python/tests/compiler/component_builder_test.py @@ -313,7 +313,7 @@ def test_generate_kaniko_yaml(self): generated_yaml = builder._generate_kaniko_spec(namespace='default', arc_dockerfile_name='dockerfile', gcs_path='gs://mlpipeline/kaniko_build.tar.gz', target_image='gcr.io/mlpipeline/kaniko_image:latest') with open(os.path.join(test_data_dir, 'kaniko.basic.yaml'), 'r') as f: - golden = yaml.load(f) + golden = yaml.safe_load(f) self.assertEqual(golden, generated_yaml)