def Prepare(self, request, context=None): # For now, just use the job name as the job id. logging.debug('Got Prepare request.') preparation_id = '%s-%s' % (request.job_name, uuid.uuid4()) provision_info = fn_api_runner.ExtendedProvisionInfo( beam_provision_api_pb2.ProvisionInfo( job_id=preparation_id, job_name=request.job_name, pipeline_options=request.pipeline_options, retrieval_token=self._artifact_service.retrieval_token( preparation_id)), self._staging_dir) self._jobs[preparation_id] = BeamJob(preparation_id, request.pipeline_options, request.pipeline, provision_info) logging.debug("Prepared job '%s' as '%s'", request.job_name, preparation_id) # TODO(angoenka): Pass an appropriate staging_session_token. The token can # be obtained in PutArtifactResponse from JobService if not self._artifact_staging_endpoint: # The front-end didn't try to stage anything, but the worker may # request what's here so we should at least store an empty manifest. self._artifact_service.CommitManifest( beam_artifact_api_pb2.CommitManifestRequest( staging_session_token=preparation_id, manifest=beam_artifact_api_pb2.Manifest())) return beam_job_api_pb2.PrepareJobResponse( preparation_id=preparation_id, artifact_staging_endpoint=self._artifact_staging_endpoint, staging_session_token=preparation_id)
def create_beam_job( self, preparation_id, # stype: str job_name, # type: str pipeline, # type: beam_runner_api_pb2.Pipeline options # type: struct_pb2.Struct ): # type: (...) -> BeamJob # TODO(angoenka): Pass an appropriate staging_session_token. The token can # be obtained in PutArtifactResponse from JobService if not self._artifact_staging_endpoint: # The front-end didn't try to stage anything, but the worker may # request what's here so we should at least store an empty manifest. self._artifact_service.CommitManifest( beam_artifact_api_pb2.CommitManifestRequest( staging_session_token=preparation_id, manifest=beam_artifact_api_pb2.Manifest())) provision_info = fn_api_runner.ExtendedProvisionInfo( beam_provision_api_pb2.ProvisionInfo( job_id=preparation_id, job_name=job_name, pipeline_options=options, retrieval_token=self._artifact_service.retrieval_token( preparation_id)), self._staging_dir) return BeamJob(preparation_id, pipeline, options, provision_info, self._artifact_staging_endpoint)