178 lines
7.9 KiB
Python
178 lines
7.9 KiB
Python
# Copyright 2018 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from kfp.containers._component_builder import _generate_dockerfile, _dependency_to_requirements, VersionedDependency, DependencyHelper
|
|
|
|
import os
|
|
import unittest
|
|
import yaml
|
|
import tarfile
|
|
from pathlib import Path
|
|
import inspect
|
|
from collections import OrderedDict
|
|
from typing import NamedTuple
|
|
|
|
class TestVersionedDependency(unittest.TestCase):
|
|
|
|
def test_version(self):
|
|
""" test version overrides min_version and max_version """
|
|
version = VersionedDependency(name='tensorflow', version='0.3.0', min_version='0.1.0', max_version='0.4.0')
|
|
self.assertTrue(version.min_version == '0.3.0')
|
|
self.assertTrue(version.max_version == '0.3.0')
|
|
self.assertTrue(version.has_versions())
|
|
self.assertTrue(version.name == 'tensorflow')
|
|
|
|
def test_minmax_version(self):
|
|
""" test if min_version and max_version are configured when version is not given """
|
|
version = VersionedDependency(name='tensorflow', min_version='0.1.0', max_version='0.4.0')
|
|
self.assertTrue(version.min_version == '0.1.0')
|
|
self.assertTrue(version.max_version == '0.4.0')
|
|
self.assertTrue(version.has_versions())
|
|
|
|
def test_min_or_max_version(self):
|
|
""" test if min_version and max_version are configured when version is not given """
|
|
version = VersionedDependency(name='tensorflow', min_version='0.1.0')
|
|
self.assertTrue(version.min_version == '0.1.0')
|
|
self.assertTrue(version.has_versions())
|
|
version = VersionedDependency(name='tensorflow', max_version='0.3.0')
|
|
self.assertTrue(version.max_version == '0.3.0')
|
|
self.assertTrue(version.has_versions())
|
|
|
|
def test_no_version(self):
|
|
""" test the no version scenario """
|
|
version = VersionedDependency(name='tensorflow')
|
|
self.assertFalse(version.has_min_version())
|
|
self.assertFalse(version.has_max_version())
|
|
self.assertFalse(version.has_versions())
|
|
|
|
class TestDependencyHelper(unittest.TestCase):
|
|
|
|
def test_generate_requirement(self):
|
|
""" Test generating requirement file """
|
|
|
|
# prepare
|
|
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
|
|
temp_file = os.path.join(test_data_dir, 'test_requirements.tmp')
|
|
|
|
dependency_helper = DependencyHelper()
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='tensorflow', min_version='0.10.0', max_version='0.11.0'))
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='kubernetes', min_version='0.6.0'))
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='pytorch', max_version='0.3.0'))
|
|
dependency_helper.generate_pip_requirements(temp_file)
|
|
|
|
golden_requirement_payload = '''\
|
|
tensorflow >= 0.10.0, <= 0.11.0
|
|
kubernetes >= 0.6.0
|
|
pytorch <= 0.3.0
|
|
'''
|
|
with open(temp_file, 'r') as f:
|
|
target_requirement_payload = f.read()
|
|
self.assertEqual(target_requirement_payload, golden_requirement_payload)
|
|
os.remove(temp_file)
|
|
|
|
def test_add_python_package(self):
|
|
""" Test add_python_package """
|
|
|
|
# prepare
|
|
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
|
|
temp_file = os.path.join(test_data_dir, 'test_requirements.tmp')
|
|
|
|
dependency_helper = DependencyHelper()
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='tensorflow', min_version='0.10.0', max_version='0.11.0'))
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='kubernetes', min_version='0.6.0'))
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='tensorflow', min_version='0.12.0'), override=True)
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='kubernetes', min_version='0.8.0'), override=False)
|
|
dependency_helper.add_python_package(dependency=VersionedDependency(name='pytorch', version='0.3.0'))
|
|
dependency_helper.generate_pip_requirements(temp_file)
|
|
golden_requirement_payload = '''\
|
|
tensorflow >= 0.12.0
|
|
kubernetes >= 0.6.0
|
|
pytorch >= 0.3.0, <= 0.3.0
|
|
'''
|
|
with open(temp_file, 'r') as f:
|
|
target_requirement_payload = f.read()
|
|
self.assertEqual(target_requirement_payload, golden_requirement_payload)
|
|
os.remove(temp_file)
|
|
|
|
|
|
|
|
class TestGenerator(unittest.TestCase):
|
|
def test_generate_dockerfile(self):
|
|
""" Test generate dockerfile """
|
|
# prepare
|
|
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
|
|
target_dockerfile = os.path.join(test_data_dir, 'component.temp.dockerfile')
|
|
golden_dockerfile_payload_one = '''\
|
|
FROM gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0
|
|
RUN apt-get update -y && apt-get install --no-install-recommends -y -q python3 python3-pip python3-setuptools
|
|
ADD main.py /ml/main.py
|
|
'''
|
|
golden_dockerfile_payload_two = '''\
|
|
FROM gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0
|
|
RUN apt-get update -y && apt-get install --no-install-recommends -y -q python3 python3-pip python3-setuptools
|
|
ADD requirements.txt /ml/requirements.txt
|
|
RUN python3 -m pip install -r /ml/requirements.txt
|
|
ADD main.py /ml/main.py
|
|
'''
|
|
|
|
golden_dockerfile_payload_three = '''\
|
|
FROM gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0
|
|
RUN apt-get update -y && apt-get install --no-install-recommends -y -q python python-pip python-setuptools
|
|
ADD requirements.txt /ml/requirements.txt
|
|
RUN python -m pip install -r /ml/requirements.txt
|
|
ADD main.py /ml/main.py
|
|
'''
|
|
# check
|
|
_generate_dockerfile(filename=target_dockerfile, base_image='gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0',
|
|
python_version='python3', add_files={'main.py': '/ml/main.py'})
|
|
with open(target_dockerfile, 'r') as f:
|
|
target_dockerfile_payload = f.read()
|
|
self.assertEqual(target_dockerfile_payload, golden_dockerfile_payload_one)
|
|
_generate_dockerfile(filename=target_dockerfile, base_image='gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0',
|
|
python_version='python3', requirement_filename='requirements.txt', add_files={'main.py': '/ml/main.py'})
|
|
with open(target_dockerfile, 'r') as f:
|
|
target_dockerfile_payload = f.read()
|
|
self.assertEqual(target_dockerfile_payload, golden_dockerfile_payload_two)
|
|
_generate_dockerfile(filename=target_dockerfile, base_image='gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0',
|
|
python_version='python2', requirement_filename='requirements.txt', add_files={'main.py': '/ml/main.py'})
|
|
with open(target_dockerfile, 'r') as f:
|
|
target_dockerfile_payload = f.read()
|
|
self.assertEqual(target_dockerfile_payload, golden_dockerfile_payload_three)
|
|
|
|
self.assertRaises(ValueError, _generate_dockerfile, filename=target_dockerfile,
|
|
base_image='gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0',
|
|
python_version='python4', requirement_filename='requirements.txt', add_files={'main.py': '/ml/main.py'})
|
|
|
|
# clean up
|
|
os.remove(target_dockerfile)
|
|
|
|
def test_generate_requirement(self):
|
|
# prepare
|
|
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
|
|
temp_file = os.path.join(test_data_dir, 'test_requirements.tmp')
|
|
|
|
dependencies = [
|
|
VersionedDependency(name='tensorflow', min_version='0.10.0', max_version='0.11.0'),
|
|
VersionedDependency(name='kubernetes', min_version='0.6.0'),
|
|
]
|
|
_dependency_to_requirements(dependencies, filename=temp_file)
|
|
golden_payload = '''\
|
|
tensorflow >= 0.10.0, <= 0.11.0
|
|
kubernetes >= 0.6.0
|
|
'''
|
|
with open(temp_file, 'r') as f:
|
|
target_payload = f.read()
|
|
self.assertEqual(target_payload, golden_payload)
|
|
os.remove(temp_file)
|