def test_run_raises_if_pipeline_fails(
self, mock_pipeline_service_create_and_get_with_fail,
mock_dataset_tabular, sync):
aiplatform.init(project=_TEST_PROJECT,
staging_bucket=_TEST_BUCKET_NAME)
job = training_jobs.AutoMLTabularTrainingJob(
display_name=_TEST_DISPLAY_NAME,
optimization_prediction_type=
_TEST_TRAINING_OPTIMIZATION_PREDICTION_TYPE,
optimization_objective=_TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME,
column_transformations=_TEST_TRAINING_COLUMN_TRANSFORMATIONS,
optimization_objective_recall_value=None,
optimization_objective_precision_value=None,
)
with pytest.raises(RuntimeError):
job.run(
model_display_name=_TEST_MODEL_DISPLAY_NAME,
dataset=mock_dataset_tabular,
target_column=_TEST_TRAINING_TARGET_COLUMN,
training_fraction_split=_TEST_TRAINING_FRACTION_SPLIT,
validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT,
test_fraction_split=_TEST_TEST_FRACTION_SPLIT,
sync=sync,
)
if not sync:
job.wait()
with pytest.raises(RuntimeError):
job.get_model()
def test_raises_before_run_is_called(self, mock_pipeline_service_create):
aiplatform.init(project=_TEST_PROJECT,
staging_bucket=_TEST_BUCKET_NAME)
job = training_jobs.AutoMLTabularTrainingJob(
display_name=_TEST_DISPLAY_NAME,
optimization_prediction_type=
_TEST_TRAINING_OPTIMIZATION_PREDICTION_TYPE,
optimization_objective=_TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME,
column_transformations=_TEST_TRAINING_COLUMN_TRANSFORMATIONS,
optimization_objective_recall_value=None,
optimization_objective_precision_value=None,
)
with pytest.raises(RuntimeError):
job.get_model()
with pytest.raises(RuntimeError):
job.has_failed
with pytest.raises(RuntimeError):
job.state
def test_run_call_pipeline_if_no_model_display_name(
self,
mock_pipeline_service_create,
mock_dataset_tabular,
mock_model_service_get,
sync,
):
aiplatform.init(project=_TEST_PROJECT,
staging_bucket=_TEST_BUCKET_NAME)
job = training_jobs.AutoMLTabularTrainingJob(
display_name=_TEST_DISPLAY_NAME,
optimization_objective=_TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME,
optimization_prediction_type=
_TEST_TRAINING_OPTIMIZATION_PREDICTION_TYPE,
column_transformations=_TEST_TRAINING_COLUMN_TRANSFORMATIONS,
optimization_objective_recall_value=None,
optimization_objective_precision_value=None,
training_encryption_spec_key_name=
_TEST_PIPELINE_ENCRYPTION_KEY_NAME,
model_encryption_spec_key_name=_TEST_MODEL_ENCRYPTION_KEY_NAME,
)
model_from_job = job.run(
dataset=mock_dataset_tabular,
target_column=_TEST_TRAINING_TARGET_COLUMN,
training_fraction_split=_TEST_TRAINING_FRACTION_SPLIT,
validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT,
test_fraction_split=_TEST_TEST_FRACTION_SPLIT,
weight_column=_TEST_TRAINING_WEIGHT_COLUMN,
budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS,
disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING,
)
if not sync:
model_from_job.wait()
true_fraction_split = gca_training_pipeline.FractionSplit(
training_fraction=_TEST_TRAINING_FRACTION_SPLIT,
validation_fraction=_TEST_VALIDATION_FRACTION_SPLIT,
test_fraction=_TEST_TEST_FRACTION_SPLIT,
)
# Test that if defaults to the job display name
true_managed_model = gca_model.Model(
display_name=_TEST_DISPLAY_NAME,
encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC)
true_input_data_config = gca_training_pipeline.InputDataConfig(
fraction_split=true_fraction_split,
dataset_id=mock_dataset_tabular.name,
)
true_training_pipeline = gca_training_pipeline.TrainingPipeline(
display_name=_TEST_DISPLAY_NAME,
training_task_definition=schema.training_job.definition.
automl_tabular,
training_task_inputs=_TEST_TRAINING_TASK_INPUTS,
model_to_upload=true_managed_model,
input_data_config=true_input_data_config,
encryption_spec=_TEST_PIPELINE_ENCRYPTION_SPEC,
)
mock_pipeline_service_create.assert_called_once_with(
parent=initializer.global_config.common_location_path(),
training_pipeline=true_training_pipeline,
)