Skip to content

Commit 34ef5a3

Browse files
yeesiancopybara-github
authored andcommitted
fix: Always upload the pickled object and dependencies tarball when creating ReasoningEngine
PiperOrigin-RevId: 665166642
1 parent 45e4251 commit 34ef5a3

File tree

2 files changed

+122
-23
lines changed

2 files changed

+122
-23
lines changed

tests/unit/vertex_langchain/test_reasoning_engines.py

+58
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,64 @@ def setup_method(self):
352352
def teardown_method(self):
353353
initializer.global_pool.shutdown(wait=True)
354354

355+
def test_prepare_create(
356+
self,
357+
cloud_storage_create_bucket_mock,
358+
tarfile_open_mock,
359+
cloudpickle_dump_mock,
360+
):
361+
_reasoning_engines._prepare_create(
362+
reasoning_engine=self.test_app,
363+
requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
364+
extra_packages=[],
365+
project=_TEST_PROJECT,
366+
location=_TEST_LOCATION,
367+
staging_bucket=_TEST_STAGING_BUCKET,
368+
gcs_dir_name=_TEST_GCS_DIR_NAME,
369+
)
370+
cloudpickle_dump_mock.assert_called() # when preparing object.pkl
371+
tarfile_open_mock.assert_called() # when preparing extra_packages
372+
373+
def test_prepare_update_with_unspecified_extra_packages(
374+
self,
375+
cloud_storage_create_bucket_mock,
376+
cloudpickle_dump_mock,
377+
):
378+
with mock.patch.object(
379+
_reasoning_engines,
380+
"_upload_extra_packages",
381+
) as upload_extra_packages_mock:
382+
_reasoning_engines._prepare_update(
383+
reasoning_engine=self.test_app,
384+
requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
385+
extra_packages=None,
386+
project=_TEST_PROJECT,
387+
location=_TEST_LOCATION,
388+
staging_bucket=_TEST_STAGING_BUCKET,
389+
gcs_dir_name=_TEST_GCS_DIR_NAME,
390+
)
391+
upload_extra_packages_mock.assert_not_called()
392+
393+
def test_prepare_update_with_empty_extra_packages(
394+
self,
395+
cloud_storage_create_bucket_mock,
396+
cloudpickle_dump_mock,
397+
):
398+
with mock.patch.object(
399+
_reasoning_engines,
400+
"_upload_extra_packages",
401+
) as upload_extra_packages_mock:
402+
_reasoning_engines._prepare_update(
403+
reasoning_engine=self.test_app,
404+
requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
405+
extra_packages=[],
406+
project=_TEST_PROJECT,
407+
location=_TEST_LOCATION,
408+
staging_bucket=_TEST_STAGING_BUCKET,
409+
gcs_dir_name=_TEST_GCS_DIR_NAME,
410+
)
411+
upload_extra_packages_mock.assert_called() # user wants to override
412+
355413
def test_get_reasoning_engine(self, get_reasoning_engine_mock):
356414
reasoning_engines.ReasoningEngine(_TEST_RESOURCE_ID)
357415
get_reasoning_engine_mock.assert_called_with(

vertexai/reasoning_engines/_reasoning_engines.py

+64-23
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ def create(
209209
# This involves packaging and uploading the artifacts for
210210
# reasoning_engine, requirements and extra_packages to
211211
# `staging_bucket/gcs_dir_name`.
212-
_prepare(
212+
_prepare_create(
213213
reasoning_engine=reasoning_engine,
214214
requirements=requirements,
215215
project=sdk_resource.project,
@@ -304,11 +304,15 @@ def update(
304304
305305
Args:
306306
reasoning_engine (ReasoningEngineInterface):
307-
Optional. The Reasoning Engine to be replaced.
307+
Optional. The Reasoning Engine to be replaced. If it is not
308+
specified, the existing Reasoning Engine will be used.
308309
requirements (Union[str, Sequence[str]]):
309310
Optional. The set of PyPI dependencies needed. It can either be
310311
the path to a single file (requirements.txt), or an ordered list
311312
of strings corresponding to each line of the requirements file.
313+
If it is not specified, the existing requirements will be used.
314+
If it is set to an empty string or list, the existing
315+
requirements will be removed.
312316
display_name (str):
313317
Optional. The user-defined name of the Reasoning Engine.
314318
The name can be up to 128 characters long and can comprise any
@@ -322,7 +326,10 @@ def update(
322326
Optional. The Python system version used. Currently updating
323327
sys version is not supported.
324328
extra_packages (Sequence[str]):
325-
Optional. The set of extra user-provided packages (if any).
329+
Optional. The set of extra user-provided packages (if any). If
330+
it is not specified, the existing extra packages will be used.
331+
If it is set to an empty list, the existing extra packages will
332+
be removed.
326333
327334
Returns:
328335
ReasoningEngine: The Reasoning Engine that was updated.
@@ -359,18 +366,18 @@ def update(
359366
)
360367
if sys_version:
361368
_LOGGER.warning("Updated sys_version is not supported.")
362-
if requirements:
369+
if requirements is not None:
363370
requirements = _validate_requirements_or_raise(requirements)
364-
if extra_packages:
371+
if extra_packages is not None:
365372
extra_packages = _validate_extra_packages_or_raise(extra_packages)
366-
if reasoning_engine:
373+
if reasoning_engine is not None:
367374
reasoning_engine = _validate_reasoning_engine_or_raise(reasoning_engine)
368375

369-
# Prepares the Reasoning Engine for creation in Vertex AI.
376+
# Prepares the Reasoning Engine for update in Vertex AI.
370377
# This involves packaging and uploading the artifacts for
371378
# reasoning_engine, requirements and extra_packages to
372379
# `staging_bucket/gcs_dir_name`.
373-
_prepare(
380+
_prepare_update(
374381
reasoning_engine=reasoning_engine,
375382
requirements=requirements,
376383
project=self.project,
@@ -392,14 +399,11 @@ def update(
392399
operation_future = self.api_client.update_reasoning_engine(
393400
request=update_request
394401
)
395-
_LOGGER.log_create_with_lro(ReasoningEngine, operation_future)
396-
created_resource = operation_future.result()
397-
_LOGGER.log_create_complete(
398-
ReasoningEngine,
399-
created_resource,
400-
self._resource_noun,
401-
module_name="vertexai.preview.reasoning_engines",
402+
_LOGGER.info(
403+
f"Update ReasoningEngine backing LRO: {operation_future.operation.name}"
402404
)
405+
created_resource = operation_future.result()
406+
_LOGGER.info(f"ReasoningEngine updated. Resource name: {created_resource.name}")
403407
self._operation_schemas = None
404408
return self
405409

@@ -560,7 +564,7 @@ def _upload_extra_packages(
560564
_LOGGER.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}")
561565

562566

563-
def _prepare(
567+
def _prepare_create(
564568
reasoning_engine: Queryable,
565569
requirements: Sequence[str],
566570
extra_packages: Sequence[str],
@@ -571,7 +575,10 @@ def _prepare(
571575
) -> None:
572576
"""Prepares the reasoning engine for creation in Vertex AI.
573577
574-
This involves packaging and uploading the artifacts to Cloud Storage.
578+
This involves packaging and uploading artifacts to Cloud Storage. Note that
579+
1. This does not actually create the Reasoning Engine in Vertex AI.
580+
2. This will always generate and upload a pickled object.
581+
3. This will always generate and upload the dependencies.tar.gz file.
575582
576583
Args:
577584
reasoning_engine: The reasoning engine to be prepared.
@@ -584,11 +591,45 @@ def _prepare(
584591
use for staging the artifacts needed.
585592
"""
586593
gcs_bucket = _get_gcs_bucket(project, location, staging_bucket)
587-
if reasoning_engine:
588-
_upload_reasoning_engine(reasoning_engine, gcs_bucket, gcs_dir_name)
594+
_upload_reasoning_engine(reasoning_engine, gcs_bucket, gcs_dir_name)
589595
if requirements:
590596
_upload_requirements(requirements, gcs_bucket, gcs_dir_name)
591-
if extra_packages:
597+
_upload_extra_packages(extra_packages, gcs_bucket, gcs_dir_name)
598+
599+
600+
def _prepare_update(
601+
reasoning_engine: Optional[Queryable],
602+
requirements: Optional[Sequence[str]],
603+
extra_packages: Optional[Sequence[str]],
604+
project: str,
605+
location: str,
606+
staging_bucket: str,
607+
gcs_dir_name: str,
608+
) -> None:
609+
"""Prepares the reasoning engine for updates in Vertex AI.
610+
611+
This involves packaging and uploading artifacts to Cloud Storage. Note that
612+
1. This does not actually update the Reasoning Engine in Vertex AI.
613+
2. This will only generate and upload a pickled object if specified.
614+
3. This will only generate and upload the dependencies.tar.gz file if
615+
extra_packages is non-empty.
616+
617+
Args:
618+
reasoning_engine: The reasoning engine to be prepared.
619+
requirements (Sequence[str]): The set of PyPI dependencies needed.
620+
extra_packages (Sequence[str]): The set of extra user-provided packages.
621+
project (str): The project for the staging bucket.
622+
location (str): The location for the staging bucket.
623+
staging_bucket (str): The staging bucket name in the form "gs://...".
624+
gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to
625+
use for staging the artifacts needed.
626+
"""
627+
gcs_bucket = _get_gcs_bucket(project, location, staging_bucket)
628+
if reasoning_engine is not None:
629+
_upload_reasoning_engine(reasoning_engine, gcs_bucket, gcs_dir_name)
630+
if requirements is not None:
631+
_upload_requirements(requirements, gcs_bucket, gcs_dir_name)
632+
if extra_packages is not None:
592633
_upload_extra_packages(extra_packages, gcs_bucket, gcs_dir_name)
593634

594635

@@ -607,23 +648,23 @@ def _generate_update_request_or_raise(
607648
update_masks: List[str] = []
608649
reasoning_engine_spec = types.ReasoningEngineSpec()
609650
package_spec = types.ReasoningEngineSpec.PackageSpec()
610-
if requirements:
651+
if requirements is not None:
611652
is_spec_update = True
612653
update_masks.append("spec.package_spec.requirements_gcs_uri")
613654
package_spec.requirements_gcs_uri = "{}/{}/{}".format(
614655
staging_bucket,
615656
gcs_dir_name,
616657
_REQUIREMENTS_FILE,
617658
)
618-
if extra_packages:
659+
if extra_packages is not None:
619660
is_spec_update = True
620661
update_masks.append("spec.package_spec.dependency_files_gcs_uri")
621662
package_spec.dependency_files_gcs_uri = "{}/{}/{}".format(
622663
staging_bucket,
623664
gcs_dir_name,
624665
_EXTRA_PACKAGES_FILE,
625666
)
626-
if reasoning_engine:
667+
if reasoning_engine is not None:
627668
is_spec_update = True
628669
update_masks.append("spec.package_spec.pickle_object_gcs_uri")
629670
package_spec.pickle_object_gcs_uri = "{}/{}/{}".format(

0 commit comments

Comments
 (0)