def _create_file_stage_to_artifact(local_path, staged_name): return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=local_path).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=staged_name).SerializeToString())
def embedded_artifact(self, data, name=None): return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.EMBEDDED.urn, type_payload=beam_runner_api_pb2.EmbeddedFilePayload( data=data).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn if name else None, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=name).SerializeToString() if name else None)
def python_sdk_dependencies(options, tmp_dir=None): if tmp_dir is None: tmp_dir = tempfile.mkdtemp() return tuple( beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=local_path).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=staged_name).SerializeToString()) for local_path, staged_name in stager.Stager.create_job_resources(options, tmp_dir))
def extract_staging_tuple_iter( artifacts: List[beam_runner_api_pb2.ArtifactInformation]): for artifact in artifacts: if artifact.type_urn == common_urns.artifact_types.FILE.urn: file_payload = beam_runner_api_pb2.ArtifactFilePayload() file_payload.ParseFromString(artifact.type_payload) src = file_payload.path if artifact.role_urn == common_urns.artifact_roles.STAGING_TO.urn: role_payload = beam_runner_api_pb2.ArtifactStagingToRolePayload() role_payload.ParseFromString(artifact.role_payload) dst = role_payload.staged_name elif (artifact.role_urn == common_urns.artifact_roles.PIP_REQUIREMENTS_FILE.urn): dst = hashlib.sha256(artifact.SerializeToString()).hexdigest() else: raise RuntimeError("unknown role type: %s" % artifact.role_urn) yield (src, dst) else: raise RuntimeError("unknown artifact type: %s" % artifact.type_urn)