Skip to content

Commit 06b87fb

Browse files
bradmirobradmiro
committed
updating tests and add region tags to setup
1 parent 5b72337 commit 06b87fb

File tree

6 files changed

+28
-22
lines changed

6 files changed

+28
-22
lines changed

‎data-science-onramp/ai-platform/modules/setup.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
# [START ai_platform_setup]
1516
from setuptools import find_packages
1617
from setuptools import setup
1718

@@ -20,8 +21,9 @@
2021
version="0.1",
2122
packages=find_packages(),
2223
include_package_data=True,
23-
description="Data Science Onramp trainer application",
24+
description="Tutorial Package",
2425
install_requires=[
2526
"gcsfs==0.7.2"
2627
]
2728
)
29+
# [END ai_platform_setup]

‎data-science-onramp/ai-platform/modules/trainer/sklearn_model/task.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,14 +46,20 @@ def get_args() -> argparse.Namespace:
4646
type=float,
4747
help="Regularization strength, default=0 (Standard Regression)",
4848
)
49+
parser.add_argument(
50+
"--model_dir",
51+
type=str,
52+
help="Output directory for the model.",
53+
default=os.environ["AIP_MODEL_DIR"],
54+
)
4955
return parser.parse_args()
5056
# [END ai_platform_sklearn_task_args]
5157

5258

5359
# [START ai_platform_sklearn_task_fit]
5460
def fit_model(
5561
input_path: str,
56-
job_dir: str,
62+
model_dir: str,
5763
degree: int = 1,
5864
alpha: int = 0
5965
) -> None:
@@ -82,7 +88,7 @@ def fit_model(
8288
# [START ai_platform_sklearn_task_export]
8389
# Save model to GCS
8490
print("Saving model")
85-
matches = re.match("gs://(.*?)/(.*)", job_dir)
91+
matches = re.match("gs://(.*?)/(.*)", model_dir)
8692
bucket = matches.group(1)
8793
blob = matches.group(2)
8894

@@ -99,14 +105,11 @@ def fit_model(
99105
if __name__ == "__main__":
100106
args = get_args()
101107

102-
input_path = args.input_path
103-
job_dir = os.environ["AIP_MODEL_DIR"]
104-
105108
kwargs = {}
106109
if args.degree:
107110
kwargs["degree"] = args.degree
108111
if args.alpha:
109112
kwargs["alpha"] = args.alpha
110113

111-
fit_model(input_path, job_dir, **kwargs)
114+
fit_model(args.input_path, args.model_dir, **kwargs)
112115
# [END ai_platform_sklearn_task]

‎data-science-onramp/ai-platform/modules/trainer/tfkeras_model/task.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,12 @@ def get_args() -> argparse.Namespace:
5656
choices=["DEBUG", "ERROR", "FATAL", "INFO", "WARN"],
5757
default="INFO",
5858
)
59+
parser.add_argument(
60+
"--model_dir",
61+
type=str,
62+
help="Output directory for the model.",
63+
default=os.environ["AIP_MODEL_DIR"],
64+
)
5965
return parser.parse_args()
6066
# [END ai_platform_tfkeras_task_args]
6167

@@ -64,7 +70,7 @@ def get_args() -> argparse.Namespace:
6470
# [START ai_platform_tfkeras_task_train_and_evaluate_load]
6571
def train_and_evaluate(
6672
input_path: str,
67-
job_dir: str,
73+
model_dir: str,
6874
num_epochs: int = 5,
6975
batch_size: int = 128,
7076
learning_rate: float = 0.01
@@ -124,7 +130,7 @@ def train_and_evaluate(
124130

125131
# Setup TensorBoard callback.
126132
tensorboard_cb = tf.keras.callbacks.TensorBoard(
127-
os.path.join(job_dir, "keras_tensorboard"), histogram_freq=1
133+
os.path.join(model_dir, "keras_tensorboard"), histogram_freq=1
128134
)
129135
# [END ai_platform_tfkeras_task_train_and_evaluate_tensorboard]
130136

@@ -141,17 +147,14 @@ def train_and_evaluate(
141147
)
142148

143149
# Export model
144-
export_path = os.path.join(job_dir, "tfkeras_model/")
145-
tf.keras.models.save_model(keras_model, export_path)
146-
print(f"Model exported to: {export_path}")
150+
keras_model.save(model_dir)
151+
print(f"Model exported to: {model_dir}")
147152
# [END ai_platform_tfkeras_task_train_and_evaluate_fit_export]
148153
# [END ai_platform_tfkeras_task_train_and_evaluate]
149154

150155

151156
if __name__ == "__main__":
152157
args = get_args()
153-
input_path = args.input_path
154-
job_dir = os.environ["AIP_MODEL_DIR"]
155158

156159
kwargs = {}
157160
if args.num_epochs:
@@ -163,5 +166,5 @@ def train_and_evaluate(
163166

164167
tf.compat.v1.logging.set_verbosity(args.verbosity)
165168

166-
train_and_evaluate(input_path, job_dir, **kwargs)
169+
train_and_evaluate(args.input_path, args.model_dir, **kwargs)
167170
# [END ai_platform_tfkeras_task]

‎data-science-onramp/ai-platform/requirements.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11

2-
google-cloud-aiplatform==0.5.1
3-
google-cloud-storage==1.36.1
2+
google-cloud-aiplatform==0.5.0
3+
google-cloud-storage==1.36.0

‎data-science-onramp/ai-platform/sklearn_test.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,8 +103,7 @@ def test_sklearn(
103103
"package_uris": [f"gs://{STAGING_BUCKET}/{TRAINER_TAR}"],
104104
"python_module": "trainer.sklearn_model.task",
105105
"args": [
106-
"--input-path",
107-
TRAIN_DATA_PATH
106+
f"--input-path={TRAIN_DATA_PATH}"
108107
]
109108
}
110109
}]

‎data-science-onramp/ai-platform/tfkeras_test.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,7 @@ def test_tfkeras(
104104
"package_uris": [f"gs://{STAGING_BUCKET}/{TRAINER_TAR}"],
105105
"python_module": "trainer.tfkeras_model.task",
106106
"args": [
107-
"--input-path",
108-
TRAIN_DATA_PATH
107+
f"--input-path={TRAIN_DATA_PATH}"
109108
]
110109
}
111110
}]
@@ -124,4 +123,4 @@ def test_tfkeras(
124123
time.sleep(60)
125124
response = aip_job_client.get_custom_job(name=resource_name)
126125

127-
assert bucket.blob(f"{MODEL_DIR}/tfkeras_model/saved_model.pb").exists()
126+
assert bucket.blob(f"{MODEL_DIR}/saved_model.pb").exists()

0 commit comments

Comments
 (0)