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 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): 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 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): 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))