def do_list_versions(**kwargs):
""" Check if a version with the name exists using Hooks instead of operators.
Uses xcom_push to pass it to the next step. Could use return too if no key.
"""
mle = MLEngineHook()
model_name = kwargs['dag_run'].conf.get('model_name')
model_versions = mle.list_versions(PROJECT, model_name)
kwargs['ti'].xcom_push(key='model_versions', value=model_versions)
def execute(self, context):
hook = MLEngineHook(
gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
if self._operation == 'create':
return hook.create_model(self._project_id, self._model)
elif self._operation == 'get':
return hook.get_model(self._project_id, self._model['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
def execute(self, context):
hook = MLEngineHook(gcp_conn_id=self._gcp_conn_id,
delegate_to=self._delegate_to)
if self._operation == 'create':
return hook.create_model(self._project_id, self._model)
elif self._operation == 'get':
return hook.get_model(self._project_id, self._model['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
def execute(self, context):
job_id = _normalize_mlengine_job_id(self._job_id)
prediction_request = {
'jobId': job_id,
'predictionInput': {
'dataFormat': self._data_format,
'inputPaths': self._input_paths,
'outputPath': self._output_path,
'region': self._region
}
}
if self._uri:
prediction_request['predictionInput']['uri'] = self._uri
elif self._model_name:
origin_name = 'projects/{}/models/{}'.format(
self._project_id, self._model_name)
if not self._version_name:
prediction_request['predictionInput'][
'modelName'] = origin_name
else:
prediction_request['predictionInput']['versionName'] = \
origin_name + '/versions/{}'.format(self._version_name)
if self._max_worker_count:
prediction_request['predictionInput'][
'maxWorkerCount'] = self._max_worker_count
if self._runtime_version:
prediction_request['predictionInput'][
'runtimeVersion'] = self._runtime_version
if self._signature_name:
prediction_request['predictionInput'][
'signatureName'] = self._signature_name
hook = MLEngineHook(self._gcp_conn_id, self._delegate_to)
# Helper method to check if the existing job's prediction input is the
# same as the request we get here.
def check_existing_job(existing_job):
return existing_job.get('predictionInput', None) == \
prediction_request['predictionInput']
try:
finished_prediction_job = hook.create_job(self._project_id,
prediction_request,
check_existing_job)
except HttpError:
raise
if finished_prediction_job['state'] != 'SUCCEEDED':
self.log.error('MLEngine batch prediction job failed: %s',
str(finished_prediction_job))
raise RuntimeError(finished_prediction_job['errorMessage'])
return finished_prediction_job['predictionOutput']
def do_check_model(**kwargs):
""" Check if a model with the name exists using Hooks instead of operators.
Uses xcom_push to pass it to the next step. Could use return too if no key.
"""
# pushes an XCom without a specific target, just by returning it
mle = MLEngineHook()
model_name = kwargs['dag_run'].conf.get('model_name')
# return bool(mle.get_model(PROJECT, MODEL_DNN_NAME))
project = mle.get_model(PROJECT, model_name)
kwargs['ti'].xcom_push(key='is_project', value=bool(project))
def execute(self, context):
job_id = _normalize_mlengine_job_id(self._job_id)
training_request = {
'jobId': job_id,
'trainingInput': {
'scaleTier': self._scale_tier,
'packageUris': self._package_uris,
'pythonModule': self._training_python_module,
'region': self._region,
'args': self._training_args,
}
}
if self._runtime_version:
training_request['trainingInput'][
'runtimeVersion'] = self._runtime_version
if self._python_version:
training_request['trainingInput'][
'pythonVersion'] = self._python_version
if self._job_dir:
training_request['trainingInput']['jobDir'] = self._job_dir
if self._scale_tier is not None and self._scale_tier.upper(
) == "CUSTOM":
training_request['trainingInput']['masterType'] = self._master_type
if self._mode == 'DRY_RUN':
self.log.info('In dry_run mode.')
self.log.info('MLEngine Training job request is: {}'.format(
training_request))
return
hook = MLEngineHook(gcp_conn_id=self._gcp_conn_id,
delegate_to=self._delegate_to)
# Helper method to check if the existing job's training input is the
# same as the request we get here.
def check_existing_job(existing_job):
return existing_job.get('trainingInput', None) == \
training_request['trainingInput']
try:
finished_training_job = hook.create_job(self._project_id,
training_request,
check_existing_job)
except errors.HttpError:
raise
if finished_training_job['state'] != 'SUCCEEDED':
self.log.error('MLEngine training job failed: {}'.format(
str(finished_training_job)))
raise RuntimeError(finished_training_job['errorMessage'])
def execute(self, context):
job_id = _normalize_mlengine_job_id(self._job_id)
prediction_request = {
'jobId': job_id,
'predictionInput': {
'dataFormat': self._data_format,
'inputPaths': self._input_paths,
'outputPath': self._output_path,
'region': self._region
}
}
if self._uri:
prediction_request['predictionInput']['uri'] = self._uri
elif self._model_name:
origin_name = 'projects/{}/models/{}'.format(
self._project_id, self._model_name)
if not self._version_name:
prediction_request['predictionInput'][
'modelName'] = origin_name
else:
prediction_request['predictionInput']['versionName'] = \
origin_name + '/versions/{}'.format(self._version_name)
if self._max_worker_count:
prediction_request['predictionInput'][
'maxWorkerCount'] = self._max_worker_count
if self._runtime_version:
prediction_request['predictionInput'][
'runtimeVersion'] = self._runtime_version
hook = MLEngineHook(self._gcp_conn_id, self._delegate_to)
# Helper method to check if the existing job's prediction input is the
# same as the request we get here.
def check_existing_job(existing_job):
return existing_job.get('predictionInput', None) == \
prediction_request['predictionInput']
try:
finished_prediction_job = hook.create_job(
self._project_id, prediction_request, check_existing_job)
except HttpError:
raise
if finished_prediction_job['state'] != 'SUCCEEDED':
self.log.error(
'MLEngine batch prediction job failed: %s', str(finished_prediction_job)
)
raise RuntimeError(finished_prediction_job['errorMessage'])
return finished_prediction_job['predictionOutput']
def execute(self, context):
job_id = _normalize_mlengine_job_id(self._job_id)
training_request = {
'jobId': job_id,
'trainingInput': {
'scaleTier': self._scale_tier,
'packageUris': self._package_uris,
'pythonModule': self._training_python_module,
'region': self._region,
'args': self._training_args,
}
}
if self._runtime_version:
training_request['trainingInput']['runtimeVersion'] = self._runtime_version
if self._python_version:
training_request['trainingInput']['pythonVersion'] = self._python_version
if self._job_dir:
training_request['trainingInput']['jobDir'] = self._job_dir
if self._scale_tier is not None and self._scale_tier.upper() == "CUSTOM":
training_request['trainingInput']['masterType'] = self._master_type
if self._mode == 'DRY_RUN':
self.log.info('In dry_run mode.')
self.log.info('MLEngine Training job request is: {}'.format(
training_request))
return
hook = MLEngineHook(
gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
# Helper method to check if the existing job's training input is the
# same as the request we get here.
def check_existing_job(existing_job):
return existing_job.get('trainingInput', None) == \
training_request['trainingInput']
try:
finished_training_job = hook.create_job(
self._project_id, training_request, check_existing_job)
except HttpError:
raise
if finished_training_job['state'] != 'SUCCEEDED':
self.log.error('MLEngine training job failed: {}'.format(
str(finished_training_job)))
raise RuntimeError(finished_training_job['errorMessage'])
def do_create_model(**kwargs):
""" Creates a model only if one with the same name did not exist.
It leverages the check from the previous task pushed using xcom.
"""
model_params = {
'name': kwargs['dag_run'].conf.get('model_name'),
'description': 'A custom DNN regressor model',
'regions': [REGION]
}
ti = kwargs['ti']
is_model = ti.xcom_pull(key='is_project', task_ids='check_model')
if not is_model:
mle = MLEngineHook()
mle.create_model(PROJECT, model_params)
def execute(self, context):
conf = context['dag_run'].conf
job_id = conf['jobName']
print('jobID: {}'.format(job_id))
training_request = {
'jobId': job_id,
'trainingInput': {
'scaleTier':
'BASIC_GPU',
'packageUris': [
'gs://elvos/cloud-ml/blueno-0.1.0.tar.gz',
'gs://elvos/cloud-ml/c3d/cloudml-c3d-0.0.2.tar.gz'
],
'pythonModule':
'trainer.task',
'region':
'us-east1',
'args':
'',
'runtimeVersion':
'1.4',
'pythonVersion':
'3.5',
}
}
hook = MLEngineHook(gcp_conn_id=self._gcp_conn_id,
delegate_to=self._delegate_to)
# Helper method to check if the existing job's training input is the
# same as the request we get here.
def check_existing_job(existing_job):
return existing_job.get('trainingInput', None) == \
training_request['trainingInput']
try:
finished_training_job = hook.create_job(self._project_id,
training_request,
check_existing_job)
except errors.HttpError:
raise
if finished_training_job['state'] != 'SUCCEEDED':
self.log.error('MLEngine training job failed: {}'.format(
str(finished_training_job)))
raise RuntimeError(finished_training_job['errorMessage'])
def execute(self, context):
hook = MLEngineHook(self.gcp_conn_id, self.delegate_to)
def check_existing_job(existing_job):
return existing_job.get('predictionInput', None) == \
self.prediction_job_request['predictionInput']
try:
finished_prediction_job = hook.create_job(
self.project_id, self.prediction_job_request,
check_existing_job)
except errors.HttpError:
raise
if finished_prediction_job['state'] != 'SUCCEEDED':
self.log.error('Batch prediction job failed: %s',
str(finished_prediction_job))
raise RuntimeError(finished_prediction_job['errorMessage'])
return finished_prediction_job['predictionOutput']
def execute(self, context):
if 'name' not in self._version:
self._version['name'] = self._version_name
hook = MLEngineHook(
gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
if self._operation == 'create':
assert self._version is not None
return hook.create_version(self._project_id, self._model_name,
self._version)
elif self._operation == 'set_default':
return hook.set_default_version(self._project_id, self._model_name,
self._version['name'])
elif self._operation == 'list':
return hook.list_versions(self._project_id, self._model_name)
elif self._operation == 'delete':
return hook.delete_version(self._project_id, self._model_name,
self._version['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
def execute(self, context):
hook = MLEngineHook(self.gcp_conn_id, self.delegate_to)
def check_existing_job(existing_job):
return existing_job.get('predictionInput', None) == \
self.prediction_job_request['predictionInput']
try:
finished_prediction_job = hook.create_job(
self.project_id,
self.prediction_job_request,
check_existing_job)
except errors.HttpError:
raise
if finished_prediction_job['state'] != 'SUCCEEDED':
self.log.error(
'Batch prediction job failed: %s',
str(finished_prediction_job))
raise RuntimeError(finished_prediction_job['errorMessage'])
return finished_prediction_job['predictionOutput']
def do_create_version(**kwargs):
""" Creates a new version or overwrite if existing one. It leverages the
check from the previous task pushed using xcom.
"""
version_params = {
"name": kwargs['dag_run'].conf.get('model_version'),
"description": 'Version 1',
"runtimeVersion": kwargs['dag_run'].conf.get('tf_version'),
"deploymentUri": 'gs://{}/{}'.format(COMPOSER_BUCKET_NAME, PREFIX_FINAL_MODEL)
}
ti = kwargs['ti']
mle = MLEngineHook()
model_name = kwargs['dag_run'].conf.get('model_name')
model_versions = ti.xcom_pull(key='model_versions', task_ids='list_versions')
version_path = 'projects/{}/models/{}/versions/{}'.format(PROJECT,
model_name,
version_params['name'])
if version_path in [v['name'] for v in model_versions]:
logging.info("Delete previously version of the model to overwrite.")
mle.delete_version(PROJECT, model_name, version_params['name'])
mle.create_version(PROJECT, model_name, version_params)
def execute(self, context):
if 'name' not in self._version:
self._version['name'] = self._version_name
hook = MLEngineHook(
gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
if self._operation == 'create':
if not self._version:
raise ValueError("version attribute of {} could not "
"be empty".format(self.__class__.__name__))
return hook.create_version(self._project_id, self._model_name,
self._version)
elif self._operation == 'set_default':
return hook.set_default_version(self._project_id, self._model_name,
self._version['name'])
elif self._operation == 'list':
return hook.list_versions(self._project_id, self._model_name)
elif self._operation == 'delete':
return hook.delete_version(self._project_id, self._model_name,
self._version['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
def execute(self, context):
if 'name' not in self._version:
self._version['name'] = self._version_name
hook = MLEngineHook(gcp_conn_id=self._gcp_conn_id,
delegate_to=self._delegate_to)
if self._operation == 'create':
assert self._version is not None
return hook.create_version(self._project_id, self._model_name,
self._version)
elif self._operation == 'set_default':
return hook.set_default_version(self._project_id, self._model_name,
self._version['name'])
elif self._operation == 'list':
return hook.list_versions(self._project_id, self._model_name)
elif self._operation == 'delete':
return hook.delete_version(self._project_id, self._model_name,
self._version['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
def execute(self, context):
if 'name' not in self._version:
self._version['name'] = self._version_name
hook = MLEngineHook(gcp_conn_id=self._gcp_conn_id,
delegate_to=self._delegate_to)
if self._operation == 'create':
if not self._version:
raise ValueError("version attribute of {} could not "
"be empty".format(self.__class__.__name__))
return hook.create_version(self._project_id, self._model_name,
self._version)
elif self._operation == 'set_default':
return hook.set_default_version(self._project_id, self._model_name,
self._version['name'])
elif self._operation == 'list':
return hook.list_versions(self._project_id, self._model_name)
elif self._operation == 'delete':
return hook.delete_version(self._project_id, self._model_name,
self._version['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
