def testSuccessCreateVersion(self): with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \ as mock_hook: success_response = {'name': 'some-name', 'done': True} hook_instance = mock_hook.return_value hook_instance.create_version.return_value = success_response training_op = MLEngineVersionOperator(version=self.VERSION_INPUT, **self.VERSION_DEFAULT_ARGS) training_op.execute(None) mock_hook.assert_called_with(gcp_conn_id='google_cloud_default', delegate_to=None) # Make sure only 'create_version' is invoked on hook instance self.assertEquals(len(hook_instance.mock_calls), 1) hook_instance.create_version.assert_called_with( 'test-project', 'test-model', self.VERSION_INPUT)
def testSuccessCreateVersion(self): with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \ as mock_hook: success_response = {'name': 'some-name', 'done': True} hook_instance = mock_hook.return_value hook_instance.create_version.return_value = success_response training_op = MLEngineVersionOperator( version=self.VERSION_INPUT, **self.VERSION_DEFAULT_ARGS) training_op.execute(None) mock_hook.assert_called_with(gcp_conn_id='google_cloud_default', delegate_to=None) # Make sure only 'create_version' is invoked on hook instance self.assertEqual(len(hook_instance.mock_calls), 1) hook_instance.create_version.assert_called_with( 'test-project', 'test-model', self.VERSION_INPUT)
def deploy_tasks(model, parent_dag_name, child_dag_name, default_args, PROJECT_ID, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION): # Create inner dag dag = DAG("{0}.{1}".format(parent_dag_name, child_dag_name), default_args=default_args, schedule_interval=None) # Constants OTHER_VERSION_NAME = "v_{0}".format( datetime.datetime.now().strftime("%Y%m%d%H%M%S")[0:12]) # Create model on ML-Engine bash_ml_engine_models_list_op = BashOperator( task_id="bash_ml_engine_models_list_{}_task".format( model.replace(".", "_")), xcom_push=True, bash_command="gcloud ml-engine models list --filter='name:{0}'".format( MODEL_NAME + model.replace(".", "_")), dag=dag) def check_if_model_already_exists(templates_dict, **kwargs): cur_model = templates_dict["model"].replace(".", "_") ml_engine_models_list = kwargs["ti"].xcom_pull( task_ids="bash_ml_engine_models_list_{}_task".format(cur_model)) logging.info( "check_if_model_already_exists: {}: ml_engine_models_list = \n{}". format(cur_model, ml_engine_models_list)) create_model_task = "ml_engine_create_model_{}_task".format(cur_model) dont_create_model_task = "dont_create_model_dummy_branch_{}_task".format( cur_model) if len(ml_engine_models_list ) == 0 or ml_engine_models_list == "Listed 0 items.": return create_model_task return dont_create_model_task check_if_model_already_exists_op = BranchPythonOperator( task_id="check_if_model_already_exists_{}_task".format( model.replace(".", "_")), templates_dict={"model": model.replace(".", "_")}, python_callable=check_if_model_already_exists, provide_context=True, dag=dag) ml_engine_create_model_op = MLEngineModelOperator( task_id="ml_engine_create_model_{}_task".format(model.replace( ".", "_")), project_id=PROJECT_ID, model={"name": MODEL_NAME + model.replace(".", "_")}, operation="create", dag=dag) create_model_dummy_op = DummyOperator( task_id="create_model_dummy_{}_task".format(model.replace(".", "_")), trigger_rule="all_done", dag=dag) dont_create_model_dummy_branch_op = DummyOperator( task_id="dont_create_model_dummy_branch_{}_task".format( model.replace(".", "_")), dag=dag) dont_create_model_dummy_op = DummyOperator( task_id="dont_create_model_dummy_{}_task".format( model.replace(".", "_")), trigger_rule="all_done", dag=dag) # Create version of model on ML-Engine bash_ml_engine_versions_list_op = BashOperator( task_id="bash_ml_engine_versions_list_{}_task".format( model.replace(".", "_")), xcom_push=True, bash_command= "gcloud ml-engine versions list --model {0} --filter='name:{1}'". format(MODEL_NAME + model.replace(".", "_"), MODEL_VERSION), dag=dag) def check_if_model_version_already_exists(templates_dict, **kwargs): cur_model = templates_dict["model"].replace(".", "_") ml_engine_versions_list = kwargs["ti"].xcom_pull( task_ids="bash_ml_engine_versions_list_{}_task".format(cur_model)) logging.info( "check_if_model_version_already_exists: {}: ml_engine_versions_list = \n{}" .format(cur_model, ml_engine_versions_list)) create_version_task = "ml_engine_create_version_{}_task".format( cur_model) create_other_version_task = "ml_engine_create_other_version_{}_task".format( cur_model) if len(ml_engine_versions_list ) == 0 or ml_engine_versions_list == "Listed 0 items.": return create_version_task return create_other_version_task check_if_model_version_already_exists_op = BranchPythonOperator( task_id="check_if_model_version_already_exists_{}_task".format( model.replace(".", "_")), templates_dict={"model": model.replace(".", "_")}, python_callable=check_if_model_version_already_exists, provide_context=True, dag=dag) ml_engine_create_version_op = MLEngineVersionOperator( task_id="ml_engine_create_version_{}_task".format( model.replace(".", "_")), project_id=PROJECT_ID, model_name=MODEL_NAME + model.replace(".", "_"), version_name=MODEL_VERSION, version={ "name": MODEL_VERSION, "deploymentUri": MODEL_LOCATION + model.replace(".", "_"), "runtimeVersion": "1.13", "framework": "TENSORFLOW", "pythonVersion": "3.5", }, operation="create", dag=dag) ml_engine_create_other_version_op = MLEngineVersionOperator( task_id="ml_engine_create_other_version_{}_task".format( model.replace(".", "_")), project_id=PROJECT_ID, model_name=MODEL_NAME + model.replace(".", "_"), version_name=OTHER_VERSION_NAME, version={ "name": OTHER_VERSION_NAME, "deploymentUri": MODEL_LOCATION + model.replace(".", "_"), "runtimeVersion": "1.13", "framework": "TENSORFLOW", "pythonVersion": "3.5", }, operation="create", dag=dag) ml_engine_set_default_version_op = MLEngineVersionOperator( task_id="ml_engine_set_default_version_{}_task".format( model.replace(".", "_")), project_id=PROJECT_ID, model_name=MODEL_NAME + model.replace(".", "_"), version_name=MODEL_VERSION, version={"name": MODEL_VERSION}, operation="set_default", dag=dag) ml_engine_set_default_other_version_op = MLEngineVersionOperator( task_id="ml_engine_set_default_other_version_{}_task".format( model.replace(".", "_")), project_id=PROJECT_ID, model_name=MODEL_NAME + model.replace(".", "_"), version_name=OTHER_VERSION_NAME, version={"name": OTHER_VERSION_NAME}, operation="set_default", dag=dag) # Build dependency graph, set_upstream dependencies for all tasks check_if_model_already_exists_op.set_upstream( bash_ml_engine_models_list_op) ml_engine_create_model_op.set_upstream(check_if_model_already_exists_op) create_model_dummy_op.set_upstream(ml_engine_create_model_op) dont_create_model_dummy_branch_op.set_upstream( check_if_model_already_exists_op) dont_create_model_dummy_op.set_upstream(dont_create_model_dummy_branch_op) bash_ml_engine_versions_list_op.set_upstream( [dont_create_model_dummy_op, create_model_dummy_op]) check_if_model_version_already_exists_op.set_upstream( bash_ml_engine_versions_list_op) ml_engine_create_version_op.set_upstream( check_if_model_version_already_exists_op) ml_engine_create_other_version_op.set_upstream( check_if_model_version_already_exists_op) ml_engine_set_default_version_op.set_upstream(ml_engine_create_version_op) ml_engine_set_default_other_version_op.set_upstream( ml_engine_create_other_version_op) return dag
templates_dict={"model": model.replace(".", "_")}, python_callable=check_if_model_version_already_exists, provide_context=True, dag=dag) OTHER_VERSION_NAME = "v_{0}".format( datetime.datetime.now().strftime("%Y%m%d%H%M%S")[0:12]) ml_engine_create_version_op = MLEngineVersionOperator( task_id="ml_engine_create_version_{}_task".format( model.replace(".", "_")), project_id=PROJECT_ID, model_name=MODEL_NAME + model.replace(".", "_"), version_name=MODEL_VERSION, version={ "name": MODEL_VERSION, "deploymentUri": MODEL_LOCATION + model.replace(".", "_"), "runtimeVersion": "1.13", "framework": "TENSORFLOW", "pythonVersion": "3.5", }, operation="create", dag=dag) ml_engine_create_other_version_op = MLEngineVersionOperator( task_id="ml_engine_create_other_version_{}_task".format( model.replace(".", "_")), project_id=PROJECT_ID, model_name=MODEL_NAME + model.replace(".", "_"), version_name=OTHER_VERSION_NAME, version={
provide_context=True, trigger_rule="none_failed", dag=dag ) # MLEngineVersionOperator with operation set to "create" to create a new # version of our model ml_engine_create_version_op = MLEngineVersionOperator( task_id="ml_engine_create_version_{}_task" .format(model.replace(".", "_")), project_id=PROJECT_ID, model_name=MODEL_NAME, version_name=Variable.get("CURRENT_VERSION_NAME"), version={ "name": Variable.get("CURRENT_VERSION_NAME"), "deploymentUri": MODEL_LOCATION + model.replace(".", "_"), "runtimeVersion": "2.1", "framework": "TENSORFLOW", "pythonVersion": "3.7", }, operation="create", trigger_rule='none_failed', dag=dag ) # MLEngineVersionOperator with operation set to "set_default" to set our # newly deployed version to be the default version. ml_engine_set_default_version_op = MLEngineVersionOperator( task_id="ml_engine_set_default_version_{}_task" .format(model.replace(".", "_")), project_id=PROJECT_ID,