From 255dabc77c647ef3ac33a10b06b3a36db122118a Mon Sep 17 00:00:00 2001
From: Yeesian Ng <ysian@google.com>
Date: Wed, 10 Apr 2024 21:48:58 -0700
Subject: [PATCH] fix: Add validation check for extra_packages when creating a
 reasoning engine.

PiperOrigin-RevId: 623704048
---
 tests/unit/vertexai/test_reasoning_engines.py | 22 +++++++++++++++----
 .../reasoning_engines/_reasoning_engines.py   | 15 +++++++++----
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/tests/unit/vertexai/test_reasoning_engines.py b/tests/unit/vertexai/test_reasoning_engines.py
index ffa8c65613..7158838771 100644
--- a/tests/unit/vertexai/test_reasoning_engines.py
+++ b/tests/unit/vertexai/test_reasoning_engines.py
@@ -82,7 +82,7 @@ def query(self, unused_arbitrary_string_name: str) -> str:
     "google-cloud-aiplatform==1.29.0",
     "langchain",
 ]
-_TEST_REASONING_ENGINE_EXTRA_PACKAGES = [
+_TEST_REASONING_ENGINE_INVALID_EXTRA_PACKAGES = [
     "lib",
     "main.py",
 ]
@@ -282,7 +282,6 @@ def test_create_reasoning_engine(
             reasoning_engine_name=_TEST_REASONING_ENGINE_RESOURCE_NAME,
             display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
             requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
-            extra_packages=_TEST_REASONING_ENGINE_EXTRA_PACKAGES,
         )
         # Manually set _gca_resource here to prevent the mocks from propagating.
         test_reasoning_engine._gca_resource = _TEST_REASONING_ENGINE_OBJ
@@ -405,10 +404,26 @@ def test_create_reasoning_engine_unsupported_sys_version(
                 reasoning_engine_name=_TEST_REASONING_ENGINE_RESOURCE_NAME,
                 display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
                 requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
-                extra_packages=_TEST_REASONING_ENGINE_EXTRA_PACKAGES,
                 sys_version="2.6",
             )
 
+    def test_create_reasoning_engine_nonexistent_extra_packages(
+        self,
+        create_reasoning_engine_mock,
+        cloud_storage_create_bucket_mock,
+        tarfile_open_mock,
+        cloudpickle_dump_mock,
+        get_reasoning_engine_mock,
+    ):
+        with pytest.raises(FileNotFoundError, match="not found"):
+            reasoning_engines.ReasoningEngine.create(
+                self.test_app,
+                reasoning_engine_name=_TEST_REASONING_ENGINE_RESOURCE_NAME,
+                display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
+                requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
+                extra_packages=_TEST_REASONING_ENGINE_INVALID_EXTRA_PACKAGES,
+            )
+
     def test_create_reasoning_engine_with_invalid_query_method(
         self,
         create_reasoning_engine_mock,
@@ -423,7 +438,6 @@ def test_create_reasoning_engine_with_invalid_query_method(
                 reasoning_engine_name=_TEST_REASONING_ENGINE_RESOURCE_NAME,
                 display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
                 requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
-                extra_packages=_TEST_REASONING_ENGINE_EXTRA_PACKAGES,
             )
 
 
diff --git a/vertexai/reasoning_engines/_reasoning_engines.py b/vertexai/reasoning_engines/_reasoning_engines.py
index 751fac22fc..16e1e99174 100644
--- a/vertexai/reasoning_engines/_reasoning_engines.py
+++ b/vertexai/reasoning_engines/_reasoning_engines.py
@@ -162,10 +162,12 @@ def create(
 
         Raises:
             ValueError: If `sys.version` is not supported by ReasoningEngine.
-            ValueError: If the project was not set using vertexai.init.
-            ValueError: If the location was not set using vertexai.init.
-            ValueError: If the staging_bucket was not set using vertexai.init.
-            ValueError: If the staging_bucket does not start with "gs://".
+            ValueError: If the `project` was not set using `vertexai.init`.
+            ValueError: If the `location` was not set using `vertexai.init`.
+            ValueError: If the `staging_bucket` was not set using vertexai.init.
+            ValueError: If the `staging_bucket` does not start with "gs://".
+            FileNotFoundError: If `extra_packages` includes a file or directory
+            that does not exist.
         """
         if not sys_version:
             sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
@@ -202,6 +204,11 @@ def create(
             ) from err
         requirements = requirements or []
         extra_packages = extra_packages or []
+        for extra_package in extra_packages:
+            if not os.path.exists(extra_package):
+                raise FileNotFoundError(
+                    f"Extra package specified but not found: {extra_package=}"
+                )
         # Prepares the Reasoning Engine for creation in Vertex AI.
         # This involves packaging and uploading the artifacts for
         # reasoning_engine, requirements and extra_packages to