def test_missing_version(self):
with self.assertRaises(AirflowException):
MLEngineCreateVersionOperator(
task_id="task-id",
project_id=TEST_PROJECT_ID,
model_name=TEST_MODEL_NAME,
version=None,
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
)
def test_success(self, mock_hook):
task = MLEngineCreateVersionOperator(
task_id="task-id",
project_id=TEST_PROJECT_ID,
model_name=TEST_MODEL_NAME,
version=TEST_VERSION,
gcp_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
task.execute(None)
mock_hook.assert_called_once_with(
delegate_to=TEST_DELEGATE_TO,
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.create_version.assert_called_once_with(
project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, version_spec=TEST_VERSION
)
# [START howto_operator_gcp_mlengine_print_model]
get_model_result = BashOperator(
bash_command=f"echo {get_model.output}",
task_id="get-model-result",
)
# [END howto_operator_gcp_mlengine_print_model]
# [START howto_operator_gcp_mlengine_create_version1]
create_version = MLEngineCreateVersionOperator(
task_id="create-version",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
version={
"name": "v1",
"description": "First-version",
"deployment_uri": f'{JOB_DIR}/keras_export/',
"runtime_version": "1.15",
"machineType": "mls1-c1-m2",
"framework": "TENSORFLOW",
"pythonVersion": "3.7",
},
)
# [END howto_operator_gcp_mlengine_create_version1]
# [START howto_operator_gcp_mlengine_create_version2]
create_version_2 = MLEngineCreateVersionOperator(
task_id="create-version-2",
project_id=PROJECT_ID,
model_name=MODEL_NAME,
version={
"name": "v2",
description='Deploys model to prediction service and function refresh',
schedule_interval=INTERVAL,
start_date=START_DATE,
catchup=False)
# Dummy tasks
begin = DummyOperator(task_id='begin', retries=1, dag=dag1)
end = DummyOperator(task_id='end', retries=1)
create_version = MLEngineCreateVersionOperator(
task_id="create-version",
project_id=PROJECT,
model_name=MODEL_NAME,
version={
"name": VERSION_NAME,
"deployment_uri": f'{MODEL_DIR}',
"runtime_version": "2.1",
"machineType": "mls1-c1-m2",
"framework": "TENSORFLOW",
"pythonVersion": "3.7",
},
)
set_defaults_version = MLEngineSetDefaultVersionOperator(
task_id="set-default-version",
project_id=PROJECT,
model_name=MODEL_NAME,
version_name=VERSION_NAME,
)
function_body = {
