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)
