def check_if_job_exists(gcp_resources) -> Optional[str]: """Check if the BigQuery job already created. Return the job url if created. Return None otherwise """ if path.exists(gcp_resources) and os.stat(gcp_resources).st_size != 0: with open(gcp_resources) as f: serialized_gcp_resources = f.read() job_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) # Resources should only contain one item. if len(job_resources.resources) != 1: raise ValueError( f'gcp_resources should contain one resource, found {len(job_resources.resources)}' ) # Validate the format of the resource uri. job_name_pattern = re.compile(_BQ_JOB_NAME_TEMPLATE) match = job_name_pattern.match( job_resources.resources[0].resource_uri) try: project = match.group('project') job = match.group('job') except AttributeError as err: raise ValueError( 'Invalid bigquery job uri: {}. Expect: {}.'.format( job_resources.resources[0].resource_uri, 'https://www.googleapis.com/bigquery/v2/projects/[projectId]/jobs/[jobId]?location=[location]' )) return job_resources.resources[0].resource_uri else: return None
def check_if_job_exists(self) -> Optional[str]: """Check if the job already exists.""" if path.exists(self.gcp_resources) and os.stat( self.gcp_resources).st_size != 0: with open(self.gcp_resources) as f: serialized_gcp_resources = f.read() job_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) # Resources should only contain one item. if len(job_resources.resources) != 1: raise ValueError( f'gcp_resources should contain one resource, found {len(job_resources.resources)}' ) job_name_group = re.findall( job_resources.resources[0].resource_uri, f'{self.job_uri_prefix}(.*)') if not job_name_group or not job_name_group[0]: raise ValueError( 'Job Name in gcp_resource is not formatted correctly or is empty.' ) job_name = job_name_group[0] logging.info( '%s name already exists: %s. Continue polling the status', self.job_type, job_name) return job_name else: return None
def test_hptuning_job_remote_runner_raises_exception_with_more_than_one_resource_in_gcp_resources( self, _, mock_job_service_client): job_client = mock.Mock() mock_job_service_client.return_value = job_client create_hptuning_job_response = mock.Mock() job_client.create_hyperparameter_tuning_job.return_value = create_hptuning_job_response create_hptuning_job_response.name = self._hptuning_job_name get_hptuning_job_response_success = mock.Mock() get_hptuning_job_response_success.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED job_client.get_hyperparameter_tuning_job.side_effect = [ get_hptuning_job_response_success ] # Write the job proto to output hptuning_job_resources = GcpResources() hptuning_job_resource_1 = hptuning_job_resources.resources.add() hptuning_job_resource_1.resource_type = "HyperparameterTuningJob" hptuning_job_resource_1.resource_uri = f"{self._hptuning_job_uri_prefix}{self._hptuning_job_name}" hptuning_job_resource_2 = hptuning_job_resources.resources.add() hptuning_job_resource_2.resource_type = "HyperparameterTuningJob" hptuning_job_resource_2.resource_uri = f"{self._hptuning_job_uri_prefix}{self._hptuning_job_name}" with open(self._gcp_resources, "w") as f: f.write(json_format.MessageToJson(hptuning_job_resources)) with self.assertRaisesRegex( ValueError, "gcp_resources should contain one resource, found 2"): hyperparameter_tuning_job_remote_runner.create_hyperparameter_tuning_job( self._type, self._project, self._location, self._payload, self._gcp_resources)
def test_model_deploy_remote_runner_succeeded(self, mock_post_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, "project"] deploy_model_lro = mock.Mock() deploy_model_lro.json.return_value = { 'name': self._lro_name, 'done': True, } mock_post_requests.return_value = deploy_model_lro deploy_model_remote_runner.deploy_model(self._type, '', '', self._payload, self._gcp_resouces_path) mock_post_requests.assert_called_once_with( url= f'{self._uri_prefix}projects/test_project/locations/test_region/endpoints/e12:deployModel', data=self._payload, headers={ 'Content-type': 'application/json', 'Authorization': 'Bearer fake_token', 'User-Agent': 'google-cloud-pipeline-components' }) with open(self._gcp_resouces_path) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto lro_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) self.assertEqual(len(lro_resources.resources), 1) self.assertEqual(lro_resources.resources[0].resource_uri, self._uri_prefix + self._lro_name)
def create_lro(self, create_url: str, request_body: str, gcp_resources: str, http_request: str = 'post') -> Any: """call the create API and get a LRO""" # Currently we don't check if operation already exists and continue from there # If this is desirable to the user and improves the reliability, we could do the following # ``` # from google.api_core import operations_v1, grpc_helpers # channel = grpc_helpers.create_channel(location + '-aiplatform.googleapis.com') # api = operations_v1.OperationsClient(channel) # current_status = api.get_operation(lro.operation.name) # ``` lro = self.request(request_url=create_url, request_body=request_body, http_request=http_request) lro_name = lro['name'] get_operation_uri = f"{self.vertex_uri_prefix}{lro_name}" # Write the lro to the gcp_resources output parameter long_running_operations = GcpResources() long_running_operation = long_running_operations.resources.add() long_running_operation.resource_type = "VertexLro" long_running_operation.resource_uri = get_operation_uri with open(gcp_resources, 'w') as f: f.write(json_format.MessageToJson(long_running_operations)) return lro
def GetTrialsOp(gcp_resources: str) -> list: """Retrieves the best trial from the trials. Args: gcp_resources (str): Proto tracking the hyperparameter tuning job. Returns: List of strings representing the intermediate JSON representation of the trials from the hyperparameter tuning job. """ from google.cloud import aiplatform from google_cloud_pipeline_components.proto.gcp_resources_pb2 import GcpResources from google.protobuf.json_format import Parse from google.cloud.aiplatform_v1.types import study api_endpoint_suffix = '-aiplatform.googleapis.com' gcp_resources_proto = Parse(gcp_resources, GcpResources()) gcp_resources_split = gcp_resources_proto.resources[ 0].resource_uri.partition('projects') resource_name = gcp_resources_split[1] + gcp_resources_split[2] prefix_str = gcp_resources_split[0] prefix_str = prefix_str[:prefix_str.find(api_endpoint_suffix)] api_endpoint = prefix_str[(prefix_str.rfind('//') + 2):] + api_endpoint_suffix client_options = {'api_endpoint': api_endpoint} job_client = aiplatform.gapic.JobServiceClient( client_options=client_options) response = job_client.get_hyperparameter_tuning_job(name=resource_name) return [study.Trial.to_json(trial) for trial in response.trials]
def test_custom_job_remote_runner_raises_exception_empty_URI_in_gcp_resources( self, mock_time_sleep, mock_job_service_client): job_client = mock.Mock() mock_job_service_client.return_value = job_client create_custom_job_response = mock.Mock() job_client.create_custom_job.return_value = create_custom_job_response create_custom_job_response.name = self._custom_job_name get_custom_job_response_success = mock.Mock() get_custom_job_response_success.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED job_client.get_custom_job.side_effect = [ get_custom_job_response_success ] # Write the job proto to output custom_job_resources = GcpResources() custom_job_resource_1 = custom_job_resources.resources.add() custom_job_resource_1.resource_type = "CustomJob" custom_job_resource_1.resource_uri = "" with open(self._gcp_resources, "w") as f: f.write(json_format.MessageToJson(custom_job_resources)) with self.assertRaisesRegex( ValueError, "Job Name in gcp_resource is not formatted correctly or is empty." ): custom_job_remote_runner.create_custom_job(self._type, self._project, self._location, self._payload, self._gcp_resources)
def test_custom_job_remote_runner_returns_gcp_resources( self, mock_time_sleep, mock_path_exists, mock_job_service_client): job_client = mock.Mock() mock_job_service_client.return_value = job_client create_custom_job_response = mock.Mock() job_client.create_custom_job.return_value = create_custom_job_response create_custom_job_response.name = self._custom_job_name get_custom_job_response_success = mock.Mock() get_custom_job_response_success.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED job_client.get_custom_job.side_effect = [ get_custom_job_response_success ] mock_path_exists.return_value = False custom_job_remote_runner.create_custom_job(self._type, self._project, self._location, self._payload, self._gcp_resources) with open(self._gcp_resources) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto custom_job_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) self.assertEqual(len(custom_job_resources.resources), 1) custom_job_name = custom_job_resources.resources[0].resource_uri[ len(self._custom_job_uri_prefix):] self.assertEqual(custom_job_name, self._custom_job_name)
def test_batch_prediction_job_remote_runner_succeeded_output_bq_table( self, mock_path_exists, mock_job_service_client): job_client = mock.Mock() mock_job_service_client.return_value = job_client create_batch_prediction_job_response = mock.Mock() job_client.create_batch_prediction_job.return_value = create_batch_prediction_job_response create_batch_prediction_job_response.name = self._batch_prediction_job_name get_batch_prediction_job_response = mock.Mock() job_client.get_batch_prediction_job.return_value = get_batch_prediction_job_response get_batch_prediction_job_response.state = gca_job_state.JobState.JOB_STATE_SUCCEEDED get_batch_prediction_job_response.name = 'job1' get_batch_prediction_job_response.output_info.bigquery_output_table = 'bigquery_output_table' get_batch_prediction_job_response.output_info.bigquery_output_dataset = 'bq://bq_project.bigquery_output_dataset' get_batch_prediction_job_response.output_info.gcs_output_directory = '' mock_path_exists.return_value = False batch_prediction_job_remote_runner.create_batch_prediction_job( self._job_type, self._project, self._location, self._payload, self._gcp_resources, self._executor_input) mock_job_service_client.assert_called_once_with(client_options={ 'api_endpoint': 'test_region-aiplatform.googleapis.com' }, client_info=mock.ANY) expected_parent = f'projects/{self._project}/locations/{self._location}' expected_job_spec = json.loads(self._payload, strict=False) job_client.create_batch_prediction_job.assert_called_once_with( parent=expected_parent, batch_prediction_job=expected_job_spec) with open(self._gcp_resources) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto batch_prediction_job_resources = json_format.Parse( serialized_gcp_resources, GcpResources()) self.assertEqual(len(batch_prediction_job_resources.resources), 1) batch_prediction_job_name = batch_prediction_job_resources.resources[ 0].resource_uri[len(self._batch_prediction_job_uri_prefix):] self.assertEqual(batch_prediction_job_name, self._batch_prediction_job_name) with open(self._output_file_path) as f: executor_output = json.load(f, strict=False) self.assertEqual( executor_output, json.loads('{"artifacts": {\ "batchpredictionjob": {"artifacts": [{"metadata": {"resourceName": "job1", "bigqueryOutputDataset": "bq://bq_project.bigquery_output_dataset","bigqueryOutputTable": "bigquery_output_table","gcsOutputDirectory": ""}, "name": "foobar", "type": {"schemaTitle": "google.VertexBatchPredictionJob"}, "uri": "https://test_region-aiplatform.googleapis.com/v1/job1"}]},\ "bigquery_output_table": {"artifacts": [{"metadata": {"projectId": "bq_project", "datasetId": "bigquery_output_dataset", "tableId": "bigquery_output_table"}, "name": "bq_table", "type": {"schemaTitle": "google.BQTable"}, "uri": "https://www.googleapis.com/bigquery/v2/projects/bq_project/datasets/bigquery_output_dataset/tables/bigquery_output_table"}]}}}' ))
def test_model_upload_remote_runner_append_unmanaged_model_succeeded( self, mock_post_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, 'project'] upload_model_lro = mock.Mock() upload_model_lro.json.return_value = { 'name': self._lro_name, 'done': True, 'response': { 'model': self._model_name } } mock_post_requests.return_value = upload_model_lro self._executor_input = ( '{"inputs":{"artifacts":{"unmanaged_container_model":{"artifacts":[{"metadata":{"predictSchemata":{"instanceSchemaUri":"instance_a"},' ' ' '"containerSpec":{"imageUri":"image_foo"}},"name":"unmanaged_container_model","type":{"schemaTitle":"google.UnmanagedContainerModel"},"uri":"gs://abc"}]}}},"outputs":{"artifacts":{"model":{"artifacts":[{"metadata":{},"name":"foobar","type":{"schemaTitle":"google.VertexModel"},"uri":"gs://abc"}]}},"outputFile":"' ) + self._output_file_path + '"}}' upload_model_remote_runner.upload_model(self._type, self._project, self._location, self._payload, self._gcp_resources_path, self._executor_input) mock_post_requests.assert_called_once_with( url= f'{self._uri_prefix}projects/test_project/locations/test_region/models:upload', data= '{"model": {"display_name": "model1", "predict_schemata": {"instance_schema_uri": "instance_a"}, "container_spec": {"image_uri": "image_foo"}, "artifact_uri": "gs://abc"}}', headers={ 'Content-type': 'application/json', 'Authorization': 'Bearer fake_token', 'User-Agent': 'google-cloud-pipeline-components' }) with open(self._output_file_path) as f: executor_output = json.load(f, strict=False) self.assertEqual( executor_output, json.loads( '{"artifacts": {"model": {"artifacts": [{"metadata": {"resourceName": "projects/test_project/locations/test_region/models/123"}, "name": "foobar", "type": {"schemaTitle": "google.VertexModel"}, "uri": "https://test_region-aiplatform.googleapis.com/v1/projects/test_project/locations/test_region/models/123"}]}}}' )) with open(self._gcp_resources_path) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto lro_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) self.assertEqual(len(lro_resources.resources), 1) self.assertEqual(lro_resources.resources[0].resource_uri, self._uri_prefix + self._lro_name)
def _create_job(job_type, project, location, job_request_json, creds, gcp_resources) -> str: """Create a new BigQuery job Args: job_type: BigQuery job type. project: Project to launch the job. location: location to launch the job. For more details, see https://cloud.google.com/bigquery/docs/locations#specifying_your_location job_request_json: A json object of Job proto. For more details, see https://cloud.google.com/bigquery/docs/reference/rest/v2/Job creds: Google auth credential. gcp_resources: File path for storing `gcp_resources` output parameter. Returns: The URI of the BigQuery Job. """ # Overrides the location if location: if 'jobReference' not in job_request_json: job_request_json['jobReference'] = {} job_request_json['jobReference']['location'] = location creds.refresh(google.auth.transport.requests.Request()) headers = { 'Content-type': 'application/json', 'Authorization': 'Bearer ' + creds.token, 'User-Agent': 'google-cloud-pipeline-components' } insert_job_url = f'https://www.googleapis.com/bigquery/v2/projects/{project}/jobs' job = requests.post( url=insert_job_url, data=json.dumps(job_request_json), headers=headers).json() if 'selfLink' not in job: raise RuntimeError( 'BigQquery Job failed. Cannot retrieve the job name. Response: {}.' .format(job)) # Write the bigquey job uri to gcp resource. job_uri = job['selfLink'] job_resources = GcpResources() job_resource = job_resources.resources.add() job_resource.resource_type = job_type job_resource.resource_uri = job_uri with open(gcp_resources, 'w') as f: f.write(json_format.MessageToJson(job_resources)) return job_uri
def GetTrialsOp(gcp_resources: str, region: str) -> list: from google.cloud import aiplatform from google_cloud_pipeline_components.proto.gcp_resources_pb2 import GcpResources from google.protobuf.json_format import Parse from google.cloud.aiplatform_v1.types import study client_options = {'api_endpoint': region + '-aiplatform.googleapis.com'} job_client = aiplatform.gapic.JobServiceClient( client_options=client_options) gcp_resources_proto = Parse(gcp_resources, GcpResources()) gcp_resources_split = gcp_resources_proto.resources[ 0].resource_uri.partition('projects') resource_name = gcp_resources_split[1] + gcp_resources_split[2] response = job_client.get_hyperparameter_tuning_job(name=resource_name) return [study.Trial.to_json(trial) for trial in response.trials]
def test_create_endpoint_remote_runner_succeeded(self, mock_post_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, "project"] create_endpoint_lro = mock.Mock() create_endpoint_lro.json.return_value = { 'name': self._lro_name, 'done': True, 'response': { 'name': self._endpoint_name } } mock_post_requests.return_value = create_endpoint_lro create_endpoint_remote_runner.create_endpoint( self._type, self._project, self._location, self._payload, self._gcp_resouces_path, self._executor_input) mock_post_requests.assert_called_once_with( url= f'{self._uri_prefix}projects/{self._project}/locations/{self._location}/endpoints', data=json.dumps(self._create_endpoint_request), headers={ 'Content-type': 'application/json', 'Authorization': 'Bearer fake_token', 'User-Agent': 'google-cloud-pipeline-components' }) with open(self._output_file_path) as f: executor_output = json.load(f, strict=False) self.assertEqual( executor_output, json.loads( '{"artifacts": {"endpoint": {"artifacts": [{"metadata": {"resourceName": "projects/test_project/locations/test_region/endpoints/123"}, "name": "foobar", "type": {"schemaTitle": "system.Endpoint"}, "uri": "https://test_region-aiplatform.googleapis.com/v1/projects/test_project/locations/test_region/endpoints/123"}]}}}' )) with open(self._gcp_resouces_path) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto lro_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) self.assertEqual(len(lro_resources.resources), 1) self.assertEqual(lro_resources.resources[0].resource_uri, self._uri_prefix + self._lro_name)
def create_job(self, create_job_fn, payload) -> str: """Create a job.""" parent = f'projects/{self.project}/locations/{self.location}' # TODO(kevinbnaughton) remove empty fields from the spec temporarily. job_spec = json_util.recursive_remove_empty( json.loads(payload, strict=False)) create_job_response = create_job_fn(self.job_client, parent, job_spec) job_name = create_job_response.name # Write the job proto to output. job_resources = GcpResources() job_resource = job_resources.resources.add() job_resource.resource_type = self.job_type job_resource.resource_uri = f'{self.job_uri_prefix}{job_name}' with open(self.gcp_resources, 'w') as f: f.write(json_format.MessageToJson(job_resources)) return job_name
def test_model_export_remote_runner_succeeded(self, mock_post_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, 'project'] export_model_lro = mock.Mock() export_model_lro.json.return_value = { 'name': self._lro_name, 'done': True, 'metadata': { 'outputInfo': self._output_info_content } } mock_post_requests.return_value = export_model_lro export_model_remote_runner.export_model(self._type, '', '', self._payload, self._gcp_resources_path, self._output_info) mock_post_requests.assert_called_once_with( url= f'{self._uri_prefix}projects/test_project/locations/test_region/models/m12:export', data=self._payload, headers={ 'Content-type': 'application/json', 'Authorization': 'Bearer fake_token', 'User-Agent': 'google-cloud-pipeline-components' }) with open(self._output_info) as f: self.assertEqual(f.read(), json.dumps(self._output_info_content)) with open(self._gcp_resources_path) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto lro_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) self.assertEqual(len(lro_resources.resources), 1) self.assertEqual(lro_resources.resources[0].resource_uri, self._uri_prefix + self._lro_name)
def test_import_model_evaluation_gcp_resources(self, mock_api): import_model_evaluation_response = mock.Mock() mock_api.return_value = import_model_evaluation_response import_model_evaluation_response.name = self._model_name main([ '--metrics', self.metrics_path, '--problem_type', 'classification', '--model_name', self._model_name, '--gcp_resources', self._gcp_resources ]) with open(self._gcp_resources) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto model_evaluation_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) self.assertLen(model_evaluation_resources.resources, 1) model_evaluation_name = model_evaluation_resources.resources[ 0].resource_uri[len(self._model_evaluation_uri_prefix):] self.assertEqual(model_evaluation_name, self._model_name)
def create_job(job_type, project, location, payload, creds, gcp_resources) -> str: """Create a new BigQuery job""" job_configuration = json.loads(payload, strict=False) # Always use standard SQL instead of legacy SQL. job_configuration['query']['useLegacySql'] = False job_request = { # TODO(IronPan) temporarily remove the empty fields from the spec 'configuration': json_util.recursive_remove_empty(job_configuration), } if location is not None: if 'jobReference' not in job_request: job_request['jobReference'] = {} job_request['jobReference']['location'] = location creds.refresh(google.auth.transport.requests.Request()) headers = { 'Content-type': 'application/json', 'Authorization': 'Bearer ' + creds.token, 'User-Agent': 'google-cloud-pipeline-components' } insert_job_url = f'https://www.googleapis.com/bigquery/v2/projects/{project}/jobs' job = requests.post(url=insert_job_url, data=json.dumps(job_request), headers=headers).json() if 'selfLink' not in job: raise RuntimeError( 'BigQquery Job failed. Cannot retrieve the job name. Response: {}.' .format(job)) # Write the bigquey job uri to gcp resource. job_uri = job['selfLink'] job_resources = GcpResources() job_resource = job_resources.resources.add() job_resource.resource_type = job_type job_resource.resource_uri = job_uri with open(gcp_resources, 'w') as f: f.write(json_format.MessageToJson(job_resources)) return job_uri
def create_lro(self, create_url: str, request_body: str, gcp_resources: str) -> Any: """call the create API and get a LRO""" # Currently we don't check if operation already exists and continue from there # If this is desirable to the user and improves the reliability, we could do the following # ``` # from google.api_core import operations_v1, grpc_helpers # channel = grpc_helpers.create_channel(location + '-aiplatform.googleapis.com') # api = operations_v1.OperationsClient(channel) # current_status = api.get_operation(lro.operation.name) # ``` self.creds.refresh(google.auth.transport.requests.Request()) headers = { 'Content-type': 'application/json', 'Authorization': 'Bearer ' + self.creds.token, 'User-Agent': 'google-cloud-pipeline-components' } lro = requests.post(url=create_url, data=request_body, headers=headers).json() if "error" in lro and lro["error"]["code"]: raise RuntimeError( "Failed to create the resource. Error: {}".format( lro["error"])) lro_name = lro['name'] get_operation_uri = f"{self.vertex_uri_prefix}{lro_name}" # Write the lro to the gcp_resources output parameter long_running_operations = GcpResources() long_running_operation = long_running_operations.resources.add() long_running_operation.resource_type = "VertexLro" long_running_operation.resource_uri = get_operation_uri with open(gcp_resources, 'w') as f: f.write(json_format.MessageToJson(long_running_operations)) return lro
def wait_gcp_resources( type, project, location, payload, gcp_resources, ): """ Poll the gcp resources till it reaches a final state. """ input_gcp_resources = Parse(payload, GcpResources()) if len(input_gcp_resources.resources) != 1: raise ValueError( "Invalid payload: %s. Wait component support waiting on only one resource at this moment." % payload) if input_gcp_resources.resources[0].resource_type != 'DataflowJob': raise ValueError( "Invalid payload: %s. Wait component only support waiting on Dataflow job at this moment." % payload) dataflow_job_uri = input_gcp_resources.resources[0].resource_uri uri_pattern = re.compile(_DATAFLOW_URI_TEMPLATE) match = uri_pattern.match(dataflow_job_uri) # Get the project and location from the job URI instead from the parameter. try: project = match.group('project') location = match.group('location') job_id = match.group('jobid') except AttributeError as err: # TODO(ruifang) propagate the error. raise ValueError('Invalid dataflow resource URI: {}. Expect: {}.'.format( dataflow_job_uri, 'https://dataflow.googleapis.com/v1b3/projects/[project_id]/locations/[location]/jobs/[job_id]' )) # Propagate the GCP resources as the output of the wait component with open(gcp_resources, 'w') as f: f.write(payload) with execution_context.ExecutionContext(on_cancel=partial( _send_cancel_request, project, job_id, location, )): # Poll the job status retry_count = 0 while True: try: df_client = discovery.build('dataflow', 'v1b3', cache_discovery=False) job = df_client.projects().locations().jobs().get( projectId=project, jobId=job_id, location=location, view=None).execute() retry_count = 0 except ConnectionError as err: retry_count += 1 if retry_count <= _CONNECTION_ERROR_RETRY_LIMIT: logging.warning( 'ConnectionError (%s) encountered when polling job: %s. Retrying.', err, job_id) else: logging.error('Request failed after %s retries.', _CONNECTION_ERROR_RETRY_LIMIT) # TODO(ruifang) propagate the error. raise job_state = job.get('currentState', None) # Write the job details as gcp_resources if job_state in _JOB_SUCCESSFUL_STATES: logging.info( 'GetDataflowJob response state =%s. Job completed', job_state) return elif job_state in _JOB_TERMINATED_STATES: # TODO(ruifang) propagate the error. raise RuntimeError( 'Job {} failed with error state: {}.'.format( job_id, job_state)) else: logging.info( 'Job %s is in a non-final state %s. Waiting for %s seconds for next poll.', job_id, job_state, _POLLING_INTERVAL_IN_SECONDS) time.sleep(_POLLING_INTERVAL_IN_SECONDS)
def test_bigquery_query_job_remote_runner_succeeded( self, mock_time_sleep, mock_get_requests, mock_post_requests, _, mock_auth): creds = mock.Mock() creds.token = 'fake_token' mock_auth.return_value = [creds, 'project'] mock_created_bq_job = mock.Mock() mock_created_bq_job.json.return_value = {'selfLink': self._job_uri} mock_post_requests.return_value = mock_created_bq_job mock_polled_bq_job = mock.Mock() mock_polled_bq_job.json.return_value = { 'selfLink': self._job_uri, 'status': { 'state': 'DONE' }, 'configuration': { 'query': { 'destinationTable': { 'projectId': 'test_project', 'datasetId': 'test_dataset', 'tableId': 'test_table' } } } } mock_get_requests.return_value = mock_polled_bq_job bigquery_query_job_remote_runner.create_bigquery_job( self._job_type, self._project, self._location, self._payload, self._gcp_resources, self._executor_input) mock_post_requests.assert_called_once_with( url= f'https://www.googleapis.com/bigquery/v2/projects/{self._project}/jobs', data= ('{"configuration": {"query": {"query": "CREATE OR REPLACE MODEL bqml_tutorial.penguins_model OPTIONS (model_type=\'linear_reg\', input_label_cols=[\'body_mass_g\']) AS SELECT * FROM `bigquery-public-data.ml_datasets.penguins` WHERE body_mass_g IS NOT NULL"}}, "jobReference": {"location": "US"}}' ), headers={ 'Content-type': 'application/json', 'Authorization': 'Bearer fake_token', 'User-Agent': 'google-cloud-pipeline-components' }) with open(self._output_file_path) as f: self.assertEqual( f.read(), '{"artifacts": {"destinationTable": {"artifacts": [{"metadata": {}, "name": "foobar", "type": {"schemaTitle": "google.BQTable"}, "uri": "https://www.googleapis.com/bigquery/v2/projects/test_project/datasets/test_dataset/tables/test_table"}]}}}' ) with open(self._gcp_resources) as f: serialized_gcp_resources = f.read() # Instantiate GCPResources Proto bq_job_resources = json_format.Parse(serialized_gcp_resources, GcpResources()) self.assertEqual(len(bq_job_resources.resources), 1) self.assertEqual( bq_job_resources.resources[0].resource_uri, 'https://www.googleapis.com/bigquery/v2/projects/test_project/jobs/fake_job?location=US' ) self.assertEqual(mock_post_requests.call_count, 1) self.assertEqual(mock_time_sleep.call_count, 1) self.assertEqual(mock_get_requests.call_count, 1)
def main(argv): """Calls ModelService.ImportModelEvaluation.""" parser = argparse.ArgumentParser( prog='Vertex Model Service evaluation importer', description='') parser.add_argument('--metrics', dest='metrics', type=str, required=True, default=argparse.SUPPRESS) parser.add_argument('--metrics_explanation', dest='metrics_explanation', type=str, default=None) parser.add_argument('--explanation', dest='explanation', type=str, default=None) parser.add_argument('--problem_type', dest='problem_type', type=str, required=True, default=argparse.SUPPRESS) parser.add_argument('--model_name', dest='model_name', type=str, required=True, default=argparse.SUPPRESS) parser.add_argument('--gcp_resources', dest='gcp_resources', type=_make_parent_dirs_and_return_path, required=True, default=argparse.SUPPRESS) parsed_args, _ = parser.parse_known_args(argv) _, project_id, _, location, _, model_id = parsed_args.model_name.split('/') api_endpoint = location + '-aiplatform.googleapis.com' resource_uri_prefix = f'https://{api_endpoint}/v1/' with open(parsed_args.metrics) as metrics_file: model_evaluation = { 'metrics': to_value( next( iter( json.loads(metrics_file.read())['slicedMetrics'][0] ['metrics'].values()))), 'metrics_schema_uri': PROBLEM_TYPE_TO_SCHEMA_URI.get(parsed_args.problem_type), } if parsed_args.explanation and parsed_args.explanation == "{{$.inputs.artifacts['explanation'].metadata['explanation_gcs_path']}}": # metrics_explanation must contain explanation_gcs_path when provided. logging.error('"explanation" must contain explanations when provided.') sys.exit(13) elif parsed_args.explanation: explanation_file_name = parsed_args.explanation if not parsed_args.explanation.startswith( 'gs://') else '/gcs' + parsed_args.explanation[4:] elif parsed_args.metrics_explanation and parsed_args.metrics_explanation != "{{$.inputs.artifacts['metrics'].metadata['explanation_gcs_path']}}": explanation_file_name = parsed_args.metrics_explanation if not parsed_args.metrics_explanation.startswith( 'gs://') else '/gcs' + parsed_args.metrics_explanation[4:] else: explanation_file_name = None if explanation_file_name: with open(explanation_file_name) as explanation_file: model_evaluation['model_explanation'] = { 'mean_attributions': [{ 'feature_attributions': to_value( json.loads(explanation_file.read())['explanation'] ['attributions'][0]['featureAttributions']) }] } import_model_evaluation_response = aiplatform.gapic.ModelServiceClient( client_info=gapic_v1.client_info.ClientInfo( user_agent='google-cloud-pipeline-components', ), client_options={ 'api_endpoint': api_endpoint, }).import_model_evaluation( parent=parsed_args.model_name, model_evaluation=model_evaluation, ) model_evaluation_name = import_model_evaluation_response.name # Write the model evaluation resource to GcpResources output. model_eval_resources = GcpResources() model_eval_resource = model_eval_resources.resources.add() model_eval_resource.resource_type = RESOURCE_TYPE model_eval_resource.resource_uri = f'{resource_uri_prefix}{model_evaluation_name}' with open(parsed_args.gcp_resources, 'w') as f: f.write(json_format.MessageToJson(model_eval_resources))