def test_success(self): mock_client = _create_mock_client() new_name = "a new name" response = write_service_pb2.UpdateExperimentResponse() mock_client.UpdateExperiment.return_value = response uploader_lib.update_experiment_metadata(mock_client, "123", name=new_name) expected_request = write_service_pb2.UpdateExperimentRequest( experiment=experiment_pb2.Experiment(experiment_id="123", name=new_name), experiment_mask=experiment_pb2.ExperimentMask(name=True), ) mock_client.UpdateExperiment.assert_called_once() (args, _) = mock_client.UpdateExperiment.call_args self.assertEqual(args[0], expected_request)
def update_experiment_metadata(writer_client, experiment_id, name=None, description=None): """Modifies user data associated with an experiment. Args: writer_client: a TensorBoardWriterService stub instance experiment_id: string ID of the experiment to modify name: If provided, modifies name of experiment to this value. description: If provided, modifies the description of the experiment to this value Raises: ExperimentNotFoundError: If no such experiment exists. PermissionDeniedError: If the user is not authorized to modify this experiment. InvalidArgumentError: If the server rejected the name or description, if, for instance, the size limits have changed on the server. """ logger.info("Modifying experiment %r", experiment_id) request = write_service_pb2.UpdateExperimentRequest() request.experiment.experiment_id = experiment_id if name is not None: logger.info("Setting exp %r name to %r", experiment_id, name) request.experiment.name = name request.experiment_mask.name = True if description is not None: logger.info("Setting exp %r description to %r", experiment_id, description) request.experiment.description = description request.experiment_mask.description = True try: grpc_util.call_with_retries(writer_client.UpdateExperiment, request) except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: raise ExperimentNotFoundError() if e.code() == grpc.StatusCode.PERMISSION_DENIED: raise PermissionDeniedError() if e.code() == grpc.StatusCode.INVALID_ARGUMENT: raise InvalidArgumentError(e.details()) raise