chore: Do not specify PackageSpec.dependency_files_gcs_uri if extra_packages is empty

PiperOrigin-RevId: 673523227

Author: yeesian

tests/unit/vertex_langchain/test_reasoning_engines.py

@@ -352,25 +352,7 @@ def setup_method(self):
     def teardown_method(self):
         initializer.global_pool.shutdown(wait=True)
 
-    def test_prepare_create(
-        self,
-        cloud_storage_create_bucket_mock,
-        tarfile_open_mock,
-        cloudpickle_dump_mock,
-    ):
-        _reasoning_engines._prepare_create(
-            reasoning_engine=self.test_app,
-            requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
-            extra_packages=[],
-            project=_TEST_PROJECT,
-            location=_TEST_LOCATION,
-            staging_bucket=_TEST_STAGING_BUCKET,
-            gcs_dir_name=_TEST_GCS_DIR_NAME,
-        )
-        cloudpickle_dump_mock.assert_called()  # when preparing object.pkl
-        tarfile_open_mock.assert_called()  # when preparing extra_packages
-
-    def test_prepare_update_with_unspecified_extra_packages(
+    def test_prepare_with_unspecified_extra_packages(
         self,
         cloud_storage_create_bucket_mock,
         cloudpickle_dump_mock,
@@ -379,7 +361,7 @@ def test_prepare_update_with_unspecified_extra_packages(
             _reasoning_engines,
             "_upload_extra_packages",
         ) as upload_extra_packages_mock:
-            _reasoning_engines._prepare_update(
+            _reasoning_engines._prepare(
                 reasoning_engine=self.test_app,
                 requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
                 extra_packages=None,
@@ -390,7 +372,7 @@ def test_prepare_update_with_unspecified_extra_packages(
             )
             upload_extra_packages_mock.assert_not_called()
 
-    def test_prepare_update_with_empty_extra_packages(
+    def test_prepare_with_empty_extra_packages(
         self,
         cloud_storage_create_bucket_mock,
         cloudpickle_dump_mock,
@@ -399,7 +381,7 @@ def test_prepare_update_with_empty_extra_packages(
             _reasoning_engines,
             "_upload_extra_packages",
         ) as upload_extra_packages_mock:
-            _reasoning_engines._prepare_update(
+            _reasoning_engines._prepare(
                 reasoning_engine=self.test_app,
                 requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
                 extra_packages=[],
@@ -429,6 +411,7 @@ def test_create_reasoning_engine(
             self.test_app,
             display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
             requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
+            extra_packages=[_TEST_REASONING_ENGINE_EXTRA_PACKAGE_PATH],
         )
         # Manually set _gca_resource here to prevent the mocks from propagating.
         test_reasoning_engine._gca_resource = _TEST_REASONING_ENGINE_OBJ
@@ -494,6 +477,7 @@ def test_create_reasoning_engine_requirements_from_file(
                 self.test_app,
                 display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
                 requirements="requirements.txt",
+                extra_packages=[_TEST_REASONING_ENGINE_EXTRA_PACKAGE_PATH],
             )
         mock_file.assert_called_with("requirements.txt")
         # Manually set _gca_resource here to prevent the mocks from propagating.
@@ -668,6 +652,7 @@ def test_delete_after_create_reasoning_engine(
             self.test_app,
             display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
             requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
+            extra_packages=[_TEST_REASONING_ENGINE_EXTRA_PACKAGE_PATH],
         )
         # Manually set _gca_resource here to prevent the mocks from propagating.
         test_reasoning_engine._gca_resource = _TEST_REASONING_ENGINE_OBJ

vertexai/reasoning_engines/_reasoning_engines.py

@@ -209,7 +209,7 @@ def create(
         # This involves packaging and uploading the artifacts for
         # reasoning_engine, requirements and extra_packages to
         # `staging_bucket/gcs_dir_name`.
-        _prepare_create(
+        _prepare(
             reasoning_engine=reasoning_engine,
             requirements=requirements,
             project=sdk_resource.project,
@@ -226,12 +226,13 @@ def create(
                 gcs_dir_name,
                 _BLOB_FILENAME,
             ),
-            dependency_files_gcs_uri="{}/{}/{}".format(
+        )
+        if extra_packages:
+            package_spec.dependency_files_gcs_uri = "{}/{}/{}".format(
                 staging_bucket,
                 gcs_dir_name,
                 _EXTRA_PACKAGES_FILE,
-            ),
-        )
+            )
         if requirements:
             package_spec.requirements_gcs_uri = "{}/{}/{}".format(
                 staging_bucket,
@@ -377,7 +378,7 @@ def update(
         # This involves packaging and uploading the artifacts for
         # reasoning_engine, requirements and extra_packages to
         # `staging_bucket/gcs_dir_name`.
-        _prepare_update(
+        _prepare(
             reasoning_engine=reasoning_engine,
             requirements=requirements,
             project=self.project,
@@ -564,40 +565,7 @@ def _upload_extra_packages(
     _LOGGER.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}")
 
 
-def _prepare_create(
-    reasoning_engine: Queryable,
-    requirements: Sequence[str],
-    extra_packages: Sequence[str],
-    project: str,
-    location: str,
-    staging_bucket: str,
-    gcs_dir_name: str,
-) -> None:
-    """Prepares the reasoning engine for creation in Vertex AI.
-
-    This involves packaging and uploading artifacts to Cloud Storage. Note that
-    1. This does not actually create the Reasoning Engine in Vertex AI.
-    2. This will always generate and upload a pickled object.
-    3. This will always generate and upload the dependencies.tar.gz file.
-
-    Args:
-        reasoning_engine: The reasoning engine to be prepared.
-        requirements (Sequence[str]): The set of PyPI dependencies needed.
-        extra_packages (Sequence[str]): The set of extra user-provided packages.
-        project (str): The project for the staging bucket.
-        location (str): The location for the staging bucket.
-        staging_bucket (str): The staging bucket name in the form "gs://...".
-        gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to
-            use for staging the artifacts needed.
-    """
-    gcs_bucket = _get_gcs_bucket(project, location, staging_bucket)
-    _upload_reasoning_engine(reasoning_engine, gcs_bucket, gcs_dir_name)
-    if requirements:
-        _upload_requirements(requirements, gcs_bucket, gcs_dir_name)
-    _upload_extra_packages(extra_packages, gcs_bucket, gcs_dir_name)
-
-
-def _prepare_update(
+def _prepare(
     reasoning_engine: Optional[Queryable],
     requirements: Optional[Sequence[str]],
     extra_packages: Optional[Sequence[str]],
@@ -606,7 +574,7 @@ def _prepare_update(
     staging_bucket: str,
     gcs_dir_name: str,
 ) -> None:
-    """Prepares the reasoning engine for updates in Vertex AI.
+    """Prepares the reasoning engine for creation or updates in Vertex AI.
 
     This involves packaging and uploading artifacts to Cloud Storage. Note that
     1. This does not actually update the Reasoning Engine in Vertex AI.