From 0b55762a8f762b1b46a8d0dcd93381d7870e5200 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Thu, 22 Feb 2024 19:27:35 -0800 Subject: [PATCH] fix: allow destination directory to exist before we copy. PiperOrigin-RevId: 609578936 --- google/cloud/aiplatform/utils/prediction_utils.py | 2 +- tests/unit/aiplatform/test_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/google/cloud/aiplatform/utils/prediction_utils.py b/google/cloud/aiplatform/utils/prediction_utils.py index ba601f43ee..c6a15d712c 100644 --- a/google/cloud/aiplatform/utils/prediction_utils.py +++ b/google/cloud/aiplatform/utils/prediction_utils.py @@ -150,4 +150,4 @@ def download_model_artifacts(artifact_uri: str) -> None: blob.download_to_filename(name_without_prefix) else: # Copy files to the current working directory. - shutil.copytree(artifact_uri, ".") + shutil.copytree(artifact_uri, ".", dirs_exist_ok=True) diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py index 9ea54554e2..e4f75d63be 100644 --- a/tests/unit/aiplatform/test_utils.py +++ b/tests/unit/aiplatform/test_utils.py @@ -965,7 +965,7 @@ def test_download_model_artifacts_not_gcs_uri( prediction_utils.download_model_artifacts(model_dir_name) assert not mock_storage_client.called - copy_tree_mock.assert_called_once_with(model_dir_name, ".") + copy_tree_mock.assert_called_once_with(model_dir_name, ".", dirs_exist_ok=True) @pytest.fixture(scope="function")