Skip to content

Commit

Permalink
fix: Add validation check for extra_packages when creating a reasonin…
Browse files Browse the repository at this point in the history
…g engine.

PiperOrigin-RevId: 623704048
  • Loading branch information
Yeesian Ng authored and Copybara-Service committed Apr 11, 2024
1 parent c9b6b8b commit 255dabc
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 8 deletions.
22 changes: 18 additions & 4 deletions tests/unit/vertexai/test_reasoning_engines.py
Expand Up @@ -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",
]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
)


Expand Down
15 changes: 11 additions & 4 deletions vertexai/reasoning_engines/_reasoning_engines.py
Expand Up @@ -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}"
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 255dabc

Please sign in to comment.