def testDeployModelForAIPPredictionWithCustomMachineType(self): self._setUpPredictionMocks() self._ai_platform_serving_args['machine_type'] = 'custom_machine_type' runner.deploy_model_for_aip_prediction(self._mock_api_client, self._serving_path, self._model_version, self._ai_platform_serving_args, self._job_labels) expected_versions_create_body = { 'name': self._model_version, 'deployment_uri': self._serving_path, 'machine_type': 'custom_machine_type', 'runtime_version': runner._get_tf_runtime_version(tf.__version__), 'python_version': runner._get_caip_python_version( runner._get_tf_runtime_version(tf.__version__)), 'labels': self._job_labels, } self._assertDeployModelMockCalls( expected_versions_create_body=expected_versions_create_body)
def _assertDeployModelMockCalls(self, expected_models_create_body=None, expected_versions_create_body=None, expect_set_default=True): if not expected_models_create_body: expected_models_create_body = { 'name': self._model_name, 'regions': [], } if not expected_versions_create_body: with telemetry_utils.scoped_labels({ telemetry_utils.LABEL_TFX_EXECUTOR: self._executor_class_path }): labels = telemetry_utils.get_labels_dict() expected_versions_create_body = { 'name': self._model_version, 'deployment_uri': self._serving_path, 'runtime_version': runner._get_tf_runtime_version(tf.__version__), 'python_version': runner._get_caip_python_version( runner._get_tf_runtime_version(tf.__version__)), 'labels': labels } self._mock_models_create.assert_called_with( body=mock.ANY, parent='projects/{}'.format(self._project_id), ) (_, models_create_kwargs) = self._mock_models_create.call_args self.assertDictEqual(expected_models_create_body, models_create_kwargs['body']) self._mock_versions_create.assert_called_with( body=mock.ANY, parent='projects/{}/models/{}'.format(self._project_id, self._model_name)) (_, versions_create_kwargs) = self._mock_versions_create.call_args self.assertDictEqual(expected_versions_create_body, versions_create_kwargs['body']) if not expect_set_default: return self._mock_set_default.assert_called_with( name='projects/{}/models/{}/versions/{}'.format( self._project_id, self._model_name, self._model_version)) self._mock_set_default_execute.assert_called_with()
def _assertDeployModelMockCalls(self, expected_models_create_body=None, expected_versions_create_body=None, expect_set_default=True): if not expected_models_create_body: expected_models_create_body = { 'name': self._model_name, 'regions': [], 'labels': self._job_labels } if not expected_versions_create_body: expected_versions_create_body = { 'name': self._model_version, 'deployment_uri': self._serving_path, 'runtime_version': runner._get_tf_runtime_version(tf.__version__), 'python_version': runner._get_caip_python_version( runner._get_tf_runtime_version(tf.__version__)), 'labels': self._job_labels } self._mock_models_create.assert_called_with( body=mock.ANY, parent='projects/{}'.format(self._project_id), ) models_create_kwargs = self._mock_models_create.call_args[1] self.assertDictEqual(expected_models_create_body, models_create_kwargs['body']) self._mock_versions_create.assert_called_with( body=mock.ANY, parent='projects/{}/models/{}'.format(self._project_id, self._model_name)) versions_create_kwargs = self._mock_versions_create.call_args[1] self.assertDictEqual(expected_versions_create_body, versions_create_kwargs['body']) if not expect_set_default: return self._mock_set_default.assert_called_with( name='projects/{}/models/{}/versions/{}'.format( self._project_id, self._model_name, self._model_version)) self._mock_set_default_execute.assert_called_with()
def testDeployModelForCMLEServing(self, mock_discovery): serving_path = os.path.join(self._output_data_dir, 'serving_path') model_version = 'model_version' mock_discovery.build.return_value = self._mock_api_client mock_create = mock.Mock() mock_create.return_value.execute.return_value = {'name': 'op_name'} self._mock_api_client.projects().models().versions( ).create = mock_create mock_get = mock.Mock() self._mock_api_client.projects().operations().get = mock_get mock_get.return_value.execute.return_value = { 'done': 'Done', } runner.deploy_model_for_cmle_serving(serving_path, model_version, self._cmle_serving_args) mock_create.assert_called_with(body=mock.ANY, parent='projects/{}/models/{}'.format( self._project_id, 'model_name')) (_, kwargs) = mock_create.call_args body = kwargs['body'] self.assertDictEqual( { 'name': 'v{}'.format(model_version), 'deployment_uri': serving_path, 'runtime_version': runner._get_tf_runtime_version(), 'python_version': runner._get_caip_python_version(), }, body) mock_get.assert_called_with(name='op_name')
def testGetTensorflowRuntime(self): self.assertEqual('1.14', runner._get_tf_runtime_version('1.14')) self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.0')) self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.1')) self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.0')) self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.1')) self.assertEqual('2.1', runner._get_tf_runtime_version('2.1.0')) self.assertEqual('2.2', runner._get_tf_runtime_version('2.2.0')) self.assertEqual('2.2', runner._get_tf_runtime_version('2.3.0'))
def testGetTensorflowRuntime(self): self.assertEqual('1.14', runner._get_tf_runtime_version('1.14')) self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.0')) self.assertEqual('1.15', runner._get_tf_runtime_version('1.15.1')) self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.0')) self.assertEqual('1.15', runner._get_tf_runtime_version('2.0.1')) self.assertEqual('2.1', runner._get_tf_runtime_version('2.1.0')) # TODO(b/157039850) Remove this once CAIP model support TF 2.2 runtime. self.assertEqual('2.1', runner._get_tf_runtime_version('2.2.0'))
def testDeployModelForAIPPredictionWithCustomRegion(self, mock_discovery): serving_path = os.path.join(self._output_data_dir, 'serving_path') model_version = 'model_version' mock_discovery.build.return_value = self._mock_api_client mock_create = mock.Mock() mock_create.return_value.execute.return_value = {'name': 'op_name'} self._mock_api_client.projects().models().versions( ).create = mock_create mock_get = mock.Mock() self._mock_api_client.projects().operations().get = mock_get mock_set_default = mock.Mock() self._mock_api_client.projects().models().versions( ).setDefault = mock_set_default mock_set_default_execute = mock.Mock() self._mock_api_client.projects().models().versions().setDefault( ).execute = mock_set_default_execute mock_get.return_value.execute.return_value = { 'done': 'Done', 'response': { 'name': model_version }, } self._ai_platform_serving_args['regions'] = ['custom-region'] runner.deploy_model_for_aip_prediction(serving_path, model_version, self._ai_platform_serving_args) mock_create.assert_called_with(body=mock.ANY, parent='projects/{}/models/{}'.format( self._project_id, 'model_name')) (_, kwargs) = mock_create.call_args body = kwargs['body'] self.assertDictEqual( { 'name': 'v{}'.format(model_version), 'regions': ['custom-region'], 'deployment_uri': serving_path, 'runtime_version': runner._get_tf_runtime_version( tf.__version__), 'python_version': runner._get_caip_python_version(), }, body) mock_get.assert_called_with(name='op_name') mock_set_default.assert_called_with( name='projects/{}/models/{}/versions/{}'.format( self._project_id, 'model_name', model_version)) mock_set_default_execute.assert_called_with()
def testDeployModelForAIPPredictionWithCustomRegion(self, mock_discovery): mock_discovery.build.return_value = self._mock_api_client self._setUpPredictionMocks() self._ai_platform_serving_args['regions'] = ['custom-region'] runner.deploy_model_for_aip_prediction(self._serving_path, self._model_version, self._ai_platform_serving_args, self._executor_class_path) self._mock_models_create.assert_called_with( body=mock.ANY, parent='projects/{}'.format(self._project_id), ) (_, models_create_kwargs) = self._mock_models_create.call_args models_create_body = models_create_kwargs['body'] self.assertDictEqual( { 'name': 'model_name', 'regions': ['custom-region'] }, models_create_body) self._mock_versions_create.assert_called_with( body=mock.ANY, parent='projects/{}/models/{}'.format(self._project_id, 'model_name')) (_, versions_create_kwargs) = self._mock_versions_create.call_args versions_create_body = versions_create_kwargs['body'] with telemetry_utils.scoped_labels( {telemetry_utils.TFX_EXECUTOR: self._executor_class_path}): labels = telemetry_utils.get_labels_dict() runtime_version = runner._get_tf_runtime_version(tf.__version__) self.assertDictEqual( { 'name': self._model_version, 'deployment_uri': self._serving_path, 'runtime_version': runtime_version, 'python_version': runner._get_caip_python_version(runtime_version), 'labels': labels, }, versions_create_body) self._mock_get.assert_called_with(name='op_name') self._mock_set_default.assert_called_with( name='projects/{}/models/{}/versions/{}'.format( self._project_id, 'model_name', self._model_version)) self._mock_set_default_execute.assert_called_with()