Skip to content

Commit b1079b7

Browse files
sararobcopybara-github
sararob
authored and
copybara-github
committed
chore: remove test_training_jobs dependency from test_custom_job
PiperOrigin-RevId: 505115749
1 parent bb27619 commit b1079b7

3 files changed

Lines changed: 100 additions & 42 deletions

File tree

‎tests/unit/aiplatform/conftest.py‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# -*- coding: utf-8 -*-
22

3-
# Copyright 2022 Google LLC
3+
# Copyright 2023 Google LLC
44
#
55
# Licensed under the Apache License, Version 2.0 (the "License");
66
# you may not use this file except in compliance with the License.
@@ -19,14 +19,28 @@
1919

2020
from google import auth
2121
from google.auth import credentials as auth_credentials
22-
from unittest.mock import patch
22+
from unittest import mock
23+
24+
from google.cloud.aiplatform.utils import source_utils
25+
import constants as test_constants
2326

2427

2528
@pytest.fixture(scope="module")
2629
def google_auth_mock():
27-
with patch.object(auth, "default") as google_auth_mock:
30+
with mock.patch.object(auth, "default") as google_auth_mock:
2831
google_auth_mock.return_value = (
2932
auth_credentials.AnonymousCredentials(),
3033
"test-project",
3134
)
3235
yield google_auth_mock
36+
37+
38+
@pytest.fixture
39+
def mock_python_package_to_gcs():
40+
with mock.patch.object(
41+
source_utils._TrainingScriptPythonPackager, "package_and_copy_to_gcs"
42+
) as mock_package_to_copy_gcs:
43+
mock_package_to_copy_gcs.return_value = (
44+
test_constants.TrainingJobConstants._TEST_OUTPUT_PYTHON_PACKAGE_PATH
45+
)
46+
yield mock_package_to_copy_gcs

‎tests/unit/aiplatform/constants.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2023 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
17+
# Use this file to store global variables that will be shared across multiple tests
18+
19+
import dataclasses
20+
21+
from google.cloud.aiplatform.utils import source_utils
22+
23+
24+
@dataclasses.dataclass(frozen=True)
25+
class TrainingJobConstants:
26+
"""Defines constants used by tests that create training jobs."""
27+
28+
_TEST_OUTPUT_PYTHON_PACKAGE_PATH = "gs://test-staging-bucket/trainer.tar.gz"
29+
_TEST_MODULE_NAME = (
30+
f"{source_utils._TrainingScriptPythonPackager._ROOT_MODULE}.task"
31+
)
32+
_TEST_LOCAL_SCRIPT_FILE_NAME = "____test____script.py"
33+
_TEST_REQUIREMENTS = ["pandas", "numpy", "tensorflow"]
34+
_TEST_ENVIRONMENT_VARIABLES = {
35+
"MY_PATH": "/path/to/my_path",
36+
}
37+
_TEST_REPLICA_COUNT = 1
38+
_TEST_MACHINE_TYPE = "n1-standard-4"
39+
_TEST_ACCELERATOR_TYPE = "NVIDIA_TESLA_K80"
40+
_TEST_ACCELERATOR_COUNT = 1
41+
_TEST_BOOT_DISK_TYPE = "pd-standard"
42+
_TEST_BOOT_DISK_SIZE_GB = 300
43+
_TEST_REDUCTION_SERVER_REPLICA_COUNT = 1
44+
_TEST_REDUCTION_SERVER_MACHINE_TYPE = "n1-highcpu-16"
45+
_TEST_REDUCTION_SERVER_CONTAINER_URI = (
46+
"us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest"
47+
)

‎tests/unit/aiplatform/test_custom_job.py‎

Lines changed: 36 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# -*- coding: utf-8 -*-
2-
# Copyright 2022 Google LLC
2+
# Copyright 2023 Google LLC
33
#
44
# Licensed under the Apache License, Version 2.0 (the "License");
55
# you may not use this file except in compliance with the License.
@@ -22,14 +22,11 @@
2222
from unittest import mock
2323
from unittest.mock import patch
2424

25+
import constants as test_constants
26+
2527
from google.protobuf import duration_pb2 # type: ignore
2628
from google.rpc import status_pb2
2729

28-
import test_training_jobs
29-
from test_training_jobs import ( # noqa: F401
30-
mock_python_package_to_gcs,
31-
)
32-
3330
from google.cloud import aiplatform
3431
from google.cloud.aiplatform import base
3532
from google.cloud.aiplatform.compat.types import (
@@ -79,8 +76,8 @@
7976

8077
_TEST_PYTHON_PACKAGE_SPEC = gca_custom_job_compat.PythonPackageSpec(
8178
executor_image_uri=_TEST_PREBUILT_CONTAINER_IMAGE,
82-
package_uris=[test_training_jobs._TEST_OUTPUT_PYTHON_PACKAGE_PATH],
83-
python_module=test_training_jobs._TEST_MODULE_NAME,
79+
package_uris=[test_constants.TrainingJobConstants._TEST_OUTPUT_PYTHON_PACKAGE_PATH],
80+
python_module=test_constants.TrainingJobConstants._TEST_MODULE_NAME,
8481
)
8582

8683
_TEST_CONTAINER_SPEC = gca_custom_job_compat.ContainerSpec(
@@ -89,10 +86,10 @@
8986
"sh",
9087
"-c",
9188
"pip install --upgrade pip && "
92-
+ f"pip3 install -q --user {test_training_jobs._TEST_OUTPUT_PYTHON_PACKAGE_PATH} && ".replace(
89+
+ f"pip3 install -q --user {test_constants.TrainingJobConstants._TEST_OUTPUT_PYTHON_PACKAGE_PATH} && ".replace(
9390
"gs://", "/gcs/"
9491
)
95-
+ f"python3 -m {test_training_jobs._TEST_MODULE_NAME}",
92+
+ f"python3 -m {test_constants.TrainingJobConstants._TEST_MODULE_NAME}",
9693
],
9794
)
9895

@@ -615,7 +612,7 @@ def test_create_from_local_script_prebuilt_container(
615612
# configuration on this is tested in test_training_jobs.py
616613
job = aiplatform.CustomJob.from_local_script(
617614
display_name=_TEST_DISPLAY_NAME,
618-
script_path=test_training_jobs._TEST_LOCAL_SCRIPT_FILE_NAME,
615+
script_path=test_constants.TrainingJobConstants._TEST_LOCAL_SCRIPT_FILE_NAME,
619616
container_uri=_TEST_PREBUILT_CONTAINER_IMAGE,
620617
base_output_dir=_TEST_BASE_OUTPUT_DIR,
621618
labels=_TEST_LABELS,
@@ -649,7 +646,7 @@ def test_create_from_local_script_custom_container(
649646
# configuration on this is tested in test_training_jobs.py
650647
job = aiplatform.CustomJob.from_local_script(
651648
display_name=_TEST_DISPLAY_NAME,
652-
script_path=test_training_jobs._TEST_LOCAL_SCRIPT_FILE_NAME,
649+
script_path=test_constants.TrainingJobConstants._TEST_LOCAL_SCRIPT_FILE_NAME,
653650
container_uri=_TEST_TRAINING_CONTAINER_IMAGE,
654651
base_output_dir=_TEST_BASE_OUTPUT_DIR,
655652
labels=_TEST_LABELS,
@@ -681,7 +678,7 @@ def test_create_from_local_script_raises_with_no_staging_bucket(
681678
# configuration on this is tested in test_training_jobs.py
682679
job = aiplatform.CustomJob.from_local_script( # noqa: F841
683680
display_name=_TEST_DISPLAY_NAME,
684-
script_path=test_training_jobs._TEST_LOCAL_SCRIPT_FILE_NAME,
681+
script_path=test_constants.TrainingJobConstants._TEST_LOCAL_SCRIPT_FILE_NAME,
685682
container_uri=_TEST_TRAINING_CONTAINER_IMAGE,
686683
)
687684

@@ -700,20 +697,20 @@ def test_create_from_local_script_prebuilt_container_with_all_args(
700697
# configuration on this is tested in test_training_jobs.py
701698
job = aiplatform.CustomJob.from_local_script(
702699
display_name=_TEST_DISPLAY_NAME,
703-
script_path=test_training_jobs._TEST_LOCAL_SCRIPT_FILE_NAME,
700+
script_path=test_constants.TrainingJobConstants._TEST_LOCAL_SCRIPT_FILE_NAME,
704701
container_uri=_TEST_PREBUILT_CONTAINER_IMAGE,
705702
args=_TEST_RUN_ARGS,
706-
requirements=test_training_jobs._TEST_REQUIREMENTS,
707-
environment_variables=test_training_jobs._TEST_ENVIRONMENT_VARIABLES,
708-
replica_count=test_training_jobs._TEST_REPLICA_COUNT,
709-
machine_type=test_training_jobs._TEST_MACHINE_TYPE,
710-
accelerator_type=test_training_jobs._TEST_ACCELERATOR_TYPE,
711-
accelerator_count=test_training_jobs._TEST_ACCELERATOR_COUNT,
712-
boot_disk_type=test_training_jobs._TEST_BOOT_DISK_TYPE,
713-
boot_disk_size_gb=test_training_jobs._TEST_BOOT_DISK_SIZE_GB,
714-
reduction_server_replica_count=test_training_jobs._TEST_REDUCTION_SERVER_REPLICA_COUNT,
715-
reduction_server_machine_type=test_training_jobs._TEST_REDUCTION_SERVER_MACHINE_TYPE,
716-
reduction_server_container_uri=test_training_jobs._TEST_REDUCTION_SERVER_CONTAINER_URI,
703+
requirements=test_constants.TrainingJobConstants._TEST_REQUIREMENTS,
704+
environment_variables=test_constants.TrainingJobConstants._TEST_ENVIRONMENT_VARIABLES,
705+
replica_count=test_constants.TrainingJobConstants._TEST_REPLICA_COUNT,
706+
machine_type=test_constants.TrainingJobConstants._TEST_MACHINE_TYPE,
707+
accelerator_type=test_constants.TrainingJobConstants._TEST_ACCELERATOR_TYPE,
708+
accelerator_count=test_constants.TrainingJobConstants._TEST_ACCELERATOR_COUNT,
709+
boot_disk_type=test_constants.TrainingJobConstants._TEST_BOOT_DISK_TYPE,
710+
boot_disk_size_gb=test_constants.TrainingJobConstants._TEST_BOOT_DISK_SIZE_GB,
711+
reduction_server_replica_count=test_constants.TrainingJobConstants._TEST_REDUCTION_SERVER_REPLICA_COUNT,
712+
reduction_server_machine_type=test_constants.TrainingJobConstants._TEST_REDUCTION_SERVER_MACHINE_TYPE,
713+
reduction_server_container_uri=test_constants.TrainingJobConstants._TEST_REDUCTION_SERVER_CONTAINER_URI,
717714
base_output_dir=_TEST_BASE_OUTPUT_DIR,
718715
labels=_TEST_LABELS,
719716
)
@@ -722,7 +719,7 @@ def test_create_from_local_script_prebuilt_container_with_all_args(
722719
expected_python_package_spec.args = _TEST_RUN_ARGS
723720
expected_python_package_spec.env = [
724721
{"name": key, "value": value}
725-
for key, value in test_training_jobs._TEST_ENVIRONMENT_VARIABLES.items()
722+
for key, value in test_constants.TrainingJobConstants._TEST_ENVIRONMENT_VARIABLES.items()
726723
]
727724

728725
assert (
@@ -752,20 +749,20 @@ def test_create_from_local_script_custom_container_with_all_args(
752749
# configuration on this is tested in test_training_jobs.py
753750
job = aiplatform.CustomJob.from_local_script(
754751
display_name=_TEST_DISPLAY_NAME,
755-
script_path=test_training_jobs._TEST_LOCAL_SCRIPT_FILE_NAME,
752+
script_path=test_constants.TrainingJobConstants._TEST_LOCAL_SCRIPT_FILE_NAME,
756753
container_uri=_TEST_TRAINING_CONTAINER_IMAGE,
757754
args=_TEST_RUN_ARGS,
758-
requirements=test_training_jobs._TEST_REQUIREMENTS,
759-
environment_variables=test_training_jobs._TEST_ENVIRONMENT_VARIABLES,
760-
replica_count=test_training_jobs._TEST_REPLICA_COUNT,
761-
machine_type=test_training_jobs._TEST_MACHINE_TYPE,
762-
accelerator_type=test_training_jobs._TEST_ACCELERATOR_TYPE,
763-
accelerator_count=test_training_jobs._TEST_ACCELERATOR_COUNT,
764-
boot_disk_type=test_training_jobs._TEST_BOOT_DISK_TYPE,
765-
boot_disk_size_gb=test_training_jobs._TEST_BOOT_DISK_SIZE_GB,
766-
reduction_server_replica_count=test_training_jobs._TEST_REDUCTION_SERVER_REPLICA_COUNT,
767-
reduction_server_machine_type=test_training_jobs._TEST_REDUCTION_SERVER_MACHINE_TYPE,
768-
reduction_server_container_uri=test_training_jobs._TEST_REDUCTION_SERVER_CONTAINER_URI,
755+
requirements=test_constants.TrainingJobConstants._TEST_REQUIREMENTS,
756+
environment_variables=test_constants.TrainingJobConstants._TEST_ENVIRONMENT_VARIABLES,
757+
replica_count=test_constants.TrainingJobConstants._TEST_REPLICA_COUNT,
758+
machine_type=test_constants.TrainingJobConstants._TEST_MACHINE_TYPE,
759+
accelerator_type=test_constants.TrainingJobConstants._TEST_ACCELERATOR_TYPE,
760+
accelerator_count=test_constants.TrainingJobConstants._TEST_ACCELERATOR_COUNT,
761+
boot_disk_type=test_constants.TrainingJobConstants._TEST_BOOT_DISK_TYPE,
762+
boot_disk_size_gb=test_constants.TrainingJobConstants._TEST_BOOT_DISK_SIZE_GB,
763+
reduction_server_replica_count=test_constants.TrainingJobConstants._TEST_REDUCTION_SERVER_REPLICA_COUNT,
764+
reduction_server_machine_type=test_constants.TrainingJobConstants._TEST_REDUCTION_SERVER_MACHINE_TYPE,
765+
reduction_server_container_uri=test_constants.TrainingJobConstants._TEST_REDUCTION_SERVER_CONTAINER_URI,
769766
base_output_dir=_TEST_BASE_OUTPUT_DIR,
770767
labels=_TEST_LABELS,
771768
)
@@ -774,7 +771,7 @@ def test_create_from_local_script_custom_container_with_all_args(
774771
expected_container_spec.args = _TEST_RUN_ARGS
775772
expected_container_spec.env = [
776773
{"name": key, "value": value}
777-
for key, value in test_training_jobs._TEST_ENVIRONMENT_VARIABLES.items()
774+
for key, value in test_constants.TrainingJobConstants._TEST_ENVIRONMENT_VARIABLES.items()
778775
]
779776

780777
assert (

0 commit comments

Comments
 (0)