fix(sdk): fix kfp component build only produces empty requirements.txt

sdk/python/kfp/cli/component.py

@@ -35,16 +35,18 @@
 from kfp.components import kfp_config
 from kfp.components import utils
 
-_REQUIREMENTS_TXT = 'requirements.txt'
+_REQUIREMENTS_TXT = 'runtime-requirements.txt'
 
 _DOCKERFILE = 'Dockerfile'
 
+# TODO: merge kfp_package_path into runtime-requirements.txt, once we have
+# kfp_runtime package that is dependency-free.
 _DOCKERFILE_TEMPLATE = '''
 FROM {base_image}
 
 WORKDIR {component_root_dir}
-COPY requirements.txt requirements.txt
-RUN pip install --no-cache-dir -r requirements.txt
+COPY {requirements_file} {requirements_file}
+RUN pip install --no-cache-dir -r {requirements_file}
 {maybe_copy_kfp_package}
 RUN pip install --no-cache-dir {kfp_package_path}
 COPY . .
@@ -235,8 +237,15 @@ def _maybe_write_file(self,
             f.write(f'# Generated by KFP.\n{contents}')
         logging.info(f'Generated file {filepath}.')
 
-    def maybe_generate_requirements_txt(self):
-        self._maybe_write_file(_REQUIREMENTS_TXT, '')
+    def generate_requirements_txt(self):
+        dependencies = set()
+        for component in self._components:
+            for package in component.packages_to_install or []:
+                dependencies.add(package)
+        self._maybe_write_file(
+            filename=_REQUIREMENTS_TXT,
+            contents='\n'.join(sorted(dependencies)),
+            overwrite=True)
 
     def maybe_generate_dockerignore(self):
         self._maybe_write_file(_DOCKERIGNORE, _DOCKERIGNORE_TEMPLATE)
@@ -265,12 +274,14 @@ def maybe_generate_dockerfile(self, overwrite_dockerfile: bool = False):
             base_image=self._base_image,
             maybe_copy_kfp_package=self._maybe_copy_kfp_package,
             component_root_dir=_COMPONENT_ROOT_DIR,
-            kfp_package_path=self._kfp_package_path)
+            kfp_package_path=self._kfp_package_path,
+            requirements_file=_REQUIREMENTS_TXT,
+        )
 
         self._maybe_write_file(_DOCKERFILE, dockerfile_contents,
                                overwrite_dockerfile)
 
-    def build_image(self, push_image: bool = True):
+    def build_image(self, push_image: bool):
         logging.info(f'Building image {self._target_image} using Docker...')
         client = docker.from_env()
 
@@ -394,7 +405,7 @@ def build(components_directory: str, component_filepattern: str, engine: str,
     builder.write_component_files()
     builder.generate_kfp_config()
 
-    builder.maybe_generate_requirements_txt()
+    builder.generate_requirements_txt()
     builder.maybe_generate_dockerignore()
     builder.maybe_generate_dockerfile(overwrite_dockerfile=overwrite_dockerfile)
     builder.build_image(push_image=push_image)

sdk/python/kfp/cli/component_test.py

@@ -31,19 +31,22 @@
 def _make_component(func_name: str,
                     base_image: Optional[str] = None,
                     target_image: Optional[str] = None,
+                    packages_to_install: Optional[List[str]] = None,
                     output_component_file: Optional[str] = None) -> str:
     return textwrap.dedent('''
     from kfp.dsl import *
 
     @component(
         base_image={base_image},
         target_image={target_image},
+        packages_to_install={packages_to_install},
         output_component_file={output_component_file})
     def {func_name}():
         pass
     ''').format(
         base_image=repr(base_image),
         target_image=repr(target_image),
+        packages_to_install=repr(packages_to_install),
         output_component_file=repr(output_component_file),
         func_name=func_name)
 
@@ -255,19 +258,39 @@ def test_component_filepattern_can_be_used_to_restrict_discovery(self):
 
             '''))
 
-    def test_empty_requirements_txt_file_is_generated(self):
+    def test_emptry_requirements_txt_file_is_generated(self):
         component = _make_component(
             func_name='train', target_image='custom-image')
         _write_components('components.py', component)
 
         result = self.runner.invoke(self.cli, ['build', str(self._working_dir)])
         self.assertEqual(result.exit_code, 0)
-        self.assert_file_exists_and_contains('requirements.txt',
+        self.assert_file_exists_and_contains('runtime-requirements.txt',
                                              '# Generated by KFP.\n')
 
-    def test_existing_requirements_txt_file_is_unchanged(self):
+    def test_existing_requirements_txt_file_is_updated(self):
         component = _make_component(
-            func_name='train', target_image='custom-image')
+            func_name='train',
+            target_image='custom-image',
+            packages_to_install=[
+                'tensorflow==2.10.0', 'google-cloud-pipeline-components'
+            ])
+        _write_components('components.py', component)
+
+        _write_file('runtime-requirements.txt', 'Some pre-existing content')
+
+        result = self.runner.invoke(self.cli, ['build', str(self._working_dir)])
+        self.assertEqual(result.exit_code, 0)
+        self.assert_file_exists_and_contains(
+            'runtime-requirements.txt',
+            '# Generated by KFP.\ngoogle-cloud-pipeline-components\ntensorflow==2.10.0'
+        )
+
+    def test_user_requirements_txt_file_is_unchanged(self):
+        component = _make_component(
+            func_name='train',
+            target_image='custom-image',
+            packages_to_install=['tensorflow==2.10.0'])
         _write_components('components.py', component)
 
         _write_file('requirements.txt', 'Some pre-existing content')
@@ -397,8 +420,8 @@ def test_docker_file_is_created_correctly(self):
                 FROM python:3.7
 
                 WORKDIR /usr/local/src/kfp/components
-                COPY requirements.txt requirements.txt
-                RUN pip install --no-cache-dir -r requirements.txt
+                COPY runtime-requirements.txt runtime-requirements.txt
+                RUN pip install --no-cache-dir -r runtime-requirements.txt
 
                 RUN pip install --no-cache-dir kfp==1.2.3
                 COPY . .
@@ -440,8 +463,8 @@ def test_existing_dockerfile_can_be_overwritten(self):
                 FROM python:3.7
 
                 WORKDIR /usr/local/src/kfp/components
-                COPY requirements.txt requirements.txt
-                RUN pip install --no-cache-dir -r requirements.txt
+                COPY runtime-requirements.txt runtime-requirements.txt
+                RUN pip install --no-cache-dir -r runtime-requirements.txt
 
                 RUN pip install --no-cache-dir kfp==1.2.3
                 COPY . .
@@ -472,8 +495,8 @@ def test_dockerfile_can_contain_custom_kfp_package(self):
                 FROM python:3.7
 
                 WORKDIR /usr/local/src/kfp/components
-                COPY requirements.txt requirements.txt
-                RUN pip install --no-cache-dir -r requirements.txt
+                COPY runtime-requirements.txt runtime-requirements.txt
+                RUN pip install --no-cache-dir -r runtime-requirements.txt
                 ''')
         self.assertTrue(contents.startswith(file_start))
         self.assertRegex(contents, 'RUN pip install --no-cache-dir kfp-*')

sdk/python/kfp/components/component_factory.py

@@ -49,6 +49,7 @@ class ComponentInfo():
     component_spec: structures.ComponentSpec
     output_component_file: Optional[str] = None
     base_image: str = _DEFAULT_BASE_IMAGE
+    packages_to_install: Optional[List[str]] = None
 
 
 # A map from function_name to components.  This is always populated when a
@@ -428,7 +429,8 @@ def create_component_from_func(
         module_path=module_path,
         component_spec=component_spec,
         output_component_file=output_component_file,
-        base_image=base_image)
+        base_image=base_image,
+        packages_to_install=packages_to_install)
 
     if REGISTERED_MODULES is not None:
         REGISTERED_MODULES[component_name] = component_info