def test__create_external_environment(self):
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
'localhost:50000',
})),
beam_runner_api_pb2.Environment(
urn=common_urns.environments.EXTERNAL.urn,
payload=beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(
url='localhost:50000')).SerializeToString()))
raw_config = ' {"url":"localhost:50000", "params":{"test":"test"}} '
for env_config in (raw_config, raw_config.lstrip(),
raw_config.strip()):
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
env_config,
})),
beam_runner_api_pb2.Environment(
urn=common_urns.environments.EXTERNAL.urn,
payload=beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(
url='localhost:50000'),
params={
"test": "test"
}).SerializeToString()))
with self.assertRaises(ValueError):
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
'{invalid}',
}))
with self.assertRaises(ValueError) as ctx:
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
'{"params":{"test":"test"}}',
}))
self.assertIn('External environment endpoint must be set.',
ctx.exception.args)
def to_runner_api_parameter(self, context):
# type: (PipelineContext) -> Tuple[str, beam_runner_api_pb2.ExternalPayload]
return (
common_urns.environments.EXTERNAL.urn,
beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(url=self.url),
params=self.params))
def _create_environment(options):
portable_options = options.view_as(PortableOptions)
environment_urn = common_urns.environments.DOCKER.urn
if portable_options.environment_type == 'DOCKER':
environment_urn = common_urns.environments.DOCKER.urn
elif portable_options.environment_type == 'PROCESS':
environment_urn = common_urns.environments.PROCESS.urn
elif portable_options.environment_type in ('EXTERNAL', 'LOOPBACK'):
environment_urn = common_urns.environments.EXTERNAL.urn
elif portable_options.environment_type:
if portable_options.environment_type.startswith('beam:env:'):
environment_urn = portable_options.environment_type
else:
raise ValueError(
'Unknown environment type: %s' % portable_options.environment_type)
if environment_urn == common_urns.environments.DOCKER.urn:
docker_image = (
portable_options.environment_config
or PortableRunner.default_docker_image())
return beam_runner_api_pb2.Environment(
url=docker_image,
urn=common_urns.environments.DOCKER.urn,
payload=beam_runner_api_pb2.DockerPayload(
container_image=docker_image
).SerializeToString())
elif environment_urn == common_urns.environments.PROCESS.urn:
config = json.loads(portable_options.environment_config)
return beam_runner_api_pb2.Environment(
urn=common_urns.environments.PROCESS.urn,
payload=beam_runner_api_pb2.ProcessPayload(
os=(config.get('os') or ''),
arch=(config.get('arch') or ''),
command=config.get('command'),
env=(config.get('env') or '')
).SerializeToString())
elif environment_urn == common_urns.environments.EXTERNAL.urn:
return beam_runner_api_pb2.Environment(
urn=common_urns.environments.EXTERNAL.urn,
payload=beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(
url=portable_options.environment_config)
).SerializeToString())
else:
return beam_runner_api_pb2.Environment(
urn=environment_urn,
payload=(portable_options.environment_config.encode('ascii')
if portable_options.environment_config else None))
def _create_environment(options):
portable_options = options.view_as(PortableOptions)
# Do not set a Runner. Otherwise this can cause problems in Java's
# PipelineOptions, i.e. ClassNotFoundException, if the corresponding Runner
# does not exist in the Java SDK. In portability, the entry point is clearly
# defined via the JobService.
portable_options.view_as(StandardOptions).runner = None
environment_urn = common_urns.environments.DOCKER.urn
if portable_options.environment_type == 'DOCKER':
environment_urn = common_urns.environments.DOCKER.urn
elif portable_options.environment_type == 'PROCESS':
environment_urn = common_urns.environments.PROCESS.urn
elif portable_options.environment_type in ('EXTERNAL', 'LOOPBACK'):
environment_urn = common_urns.environments.EXTERNAL.urn
elif portable_options.environment_type:
if portable_options.environment_type.startswith('beam:env:'):
environment_urn = portable_options.environment_type
else:
raise ValueError('Unknown environment type: %s' %
portable_options.environment_type)
if environment_urn == common_urns.environments.DOCKER.urn:
docker_image = (portable_options.environment_config
or PortableRunner.default_docker_image())
return beam_runner_api_pb2.Environment(
urn=common_urns.environments.DOCKER.urn,
payload=beam_runner_api_pb2.DockerPayload(
container_image=docker_image).SerializeToString())
elif environment_urn == common_urns.environments.PROCESS.urn:
config = json.loads(portable_options.environment_config)
return beam_runner_api_pb2.Environment(
urn=common_urns.environments.PROCESS.urn,
payload=beam_runner_api_pb2.ProcessPayload(
os=(config.get('os') or ''),
arch=(config.get('arch') or ''),
command=config.get('command'),
env=(config.get('env') or '')).SerializeToString())
elif environment_urn == common_urns.environments.EXTERNAL.urn:
def looks_like_json(environment_config):
import re
return re.match(r'\s*\{.*\}\s*$', environment_config)
if looks_like_json(portable_options.environment_config):
config = json.loads(portable_options.environment_config)
url = config.get('url')
if not url:
raise ValueError(
'External environment endpoint must be set.')
params = config.get('params')
else:
url = portable_options.environment_config
params = None
return beam_runner_api_pb2.Environment(
urn=common_urns.environments.EXTERNAL.urn,
payload=beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(url=url),
params=params).SerializeToString())
else:
return beam_runner_api_pb2.Environment(
urn=environment_urn,
payload=(portable_options.environment_config.encode('ascii')
if portable_options.environment_config else None))
