def generate_pod_spec_for_task():
pod_spec = generated_pb2.PodSpec()
secondary_container = generated_pb2.Container(
name="secondary",
image="alpine",
)
secondary_container.command.extend(["/bin/sh"])
secondary_container.args.extend(
["-c", "echo hi sidecar world > /data/message.txt"])
shared_volume_mount = generated_pb2.VolumeMount(
name="shared-data",
mountPath="/data",
)
secondary_container.volumeMounts.extend([shared_volume_mount])
primary_container = generated_pb2.Container(name="primary")
primary_container.volumeMounts.extend([shared_volume_mount])
pod_spec.volumes.extend([
generated_pb2.Volume(name="shared-data",
volumeSource=generated_pb2.VolumeSource(
emptyDir=generated_pb2.EmptyDirVolumeSource(
medium="Memory", )))
])
pod_spec.containers.extend([primary_container, secondary_container])
return pod_spec
def get_pod_spec():
a_container = generated_pb2.Container(name="a container",)
a_container.command.extend(["fee", "fi", "fo", "fum"])
a_container.volumeMounts.extend([generated_pb2.VolumeMount(name="volume mount", mountPath="some/where",)])
pod_spec = generated_pb2.PodSpec(restartPolicy="OnFailure",)
pod_spec.containers.extend([a_container, generated_pb2.Container(name="another container")])
return pod_spec
def get_pod_spec():
a_container = generated_pb2.Container(name="main")
a_container.command.extend(["foo", "bar"])
a_container.volumeMounts.extend(
[generated_pb2.VolumeMount(
name="scratch",
mountPath="/scratch",
)])
pod_spec = generated_pb2.PodSpec(restartPolicy="Never", )
pod_spec.containers.extend(
[a_container, generated_pb2.Container(name="sidecar")])
return pod_spec
def test_sidecar_task():
pod_spec = generated_pb2.PodSpec()
container = generated_pb2.Container(name="containery")
pod_spec.containers.extend([container])
obj = task.SidecarJob(pod_spec=pod_spec, primary_container_name="primary")
assert obj.primary_container_name == "primary"
assert len(obj.pod_spec.containers) == 1
assert obj.pod_spec.containers[0].name == "containery"
obj2 = task.SidecarJob.from_flyte_idl(obj.to_flyte_idl())
assert obj2 == obj
def generate_pod_spec_for_task():
pod_spec = generated_pb2.PodSpec()
# Primary containers do not require us to specify an image, the default image built for flyte tasks will get used.
primary_container = generated_pb2.Container(name="primary")
# Note: for non-primary containers we must specify an image.
secondary_container = generated_pb2.Container(
name="secondary",
image="alpine",
)
secondary_container.command.extend(["/bin/sh"])
secondary_container.args.extend(
["-c", "echo hi pod world > {}".format(_SHARED_DATA_PATH)])
resources = generated_pb2.ResourceRequirements()
resources.limits["cpu"].CopyFrom(Quantity(string="1"))
resources.requests["cpu"].CopyFrom(Quantity(string="1"))
resources.limits["memory"].CopyFrom(Quantity(string="100Mi"))
resources.requests["memory"].CopyFrom(Quantity(string="100Mi"))
primary_container.resources.CopyFrom(resources)
secondary_container.resources.CopyFrom(resources)
shared_volume_mount = generated_pb2.VolumeMount(
name="shared-data",
mountPath="/data",
)
secondary_container.volumeMounts.extend([shared_volume_mount])
primary_container.volumeMounts.extend([shared_volume_mount])
pod_spec.volumes.extend([
generated_pb2.Volume(
name="shared-data",
volumeSource=generated_pb2.VolumeSource(
emptyDir=generated_pb2.EmptyDirVolumeSource(medium="Memory",
)),
)
])
pod_spec.containers.extend([primary_container, secondary_container])
return pod_spec
def generate_pod_spec_for_task():
pod_spec = generated_pb2.PodSpec()
primary_container = generated_pb2.Container(name="primary")
secondary_container = generated_pb2.Container(
name="secondary",
image="alpine",
)
secondary_container.command.extend(["/bin/sh"])
secondary_container.args.extend(
["-c", "echo hi sidecar world > {}".format(_SHARED_DATA_PATH)])
resources = generated_pb2.ResourceRequirements()
resources.limits["cpu"].CopyFrom(Quantity(string="1"))
resources.requests["cpu"].CopyFrom(Quantity(string="1"))
resources.limits["memory"].CopyFrom(Quantity(string="100Mi"))
resources.requests["memory"].CopyFrom(Quantity(string="100Mi"))
primary_container.resources.CopyFrom(resources)
secondary_container.resources.CopyFrom(resources)
shared_volume_mount = generated_pb2.VolumeMount(
name="shared-data",
mountPath="/data",
)
secondary_container.volumeMounts.extend([shared_volume_mount])
primary_container.volumeMounts.extend([shared_volume_mount])
pod_spec.volumes.extend([
generated_pb2.Volume(name="shared-data",
volumeSource=generated_pb2.VolumeSource(
emptyDir=generated_pb2.EmptyDirVolumeSource(
medium="Memory", )))
])
pod_spec.containers.extend([primary_container, secondary_container])
return pod_spec
def test_sidecar_task():
pod_spec = generated_pb2.PodSpec()
container = generated_pb2.Container(name="containery")
pod_spec.containers.extend([container])
obj = task.SidecarJob(
pod_spec=pod_spec,
primary_container_name="primary",
annotations={"a1": "a1"},
labels={"b1": "b1"},
)
assert obj.primary_container_name == "primary"
assert len(obj.pod_spec.containers) == 1
assert obj.pod_spec.containers[0].name == "containery"
assert obj.annotations["a1"] == "a1"
assert obj.labels["b1"] == "b1"
obj2 = task.SidecarJob.from_flyte_idl(obj.to_flyte_idl())
assert obj2 == obj
def generate_simple_pod_spec_for_task():
pod_spec = generated_pb2.PodSpec()
primary_container = generated_pb2.Container(name="primary")
pod_spec.containers.extend([primary_container])
return pod_spec
def reconcile_partial_pod_spec_and_task(self, pod_spec,
primary_container_name):
"""
Assigns the custom field as a the reconciled primary container and pod spec defintion.
:param _sdk_runnable.SdkRunnableTask sdk_runnable_task:
:param generated_pb2.PodSpec pod_spec:
:param Text primary_container_name:
:rtype: SdkSidecarTask
"""
# First, insert a placeholder primary container if it is not defined in the pod spec.
containers = pod_spec.containers
primary_exists = False
for container in containers:
if container.name == primary_container_name:
primary_exists = True
break
if not primary_exists:
containers.extend(
[_k8s_pb2.Container(name=primary_container_name)])
final_containers = []
for container in containers:
# In the case of the primary container, we overwrite specific container attributes with the default values
# used in an SDK runnable task.
if container.name == primary_container_name:
container.image = self._container.image
# clear existing commands
del container.command[:]
container.command.extend(self._container.command)
# also clear existing args
del container.args[:]
container.args.extend(self._container.args)
resource_requirements = _k8s_pb2.ResourceRequirements()
for resource in self._container.resources.limits:
resource_requirements.limits[
_core_task.Resources.ResourceName.Name(
resource.name).lower()].CopyFrom(
_resource_pb2.Quantity(string=resource.value))
for resource in self._container.resources.requests:
resource_requirements.requests[
_core_task.Resources.ResourceName.Name(
resource.name).lower()].CopyFrom(
_resource_pb2.Quantity(string=resource.value))
if resource_requirements.ByteSize():
# Important! Only copy over resource requirements if they are non-empty.
container.resources.CopyFrom(resource_requirements)
del container.env[:]
container.env.extend([
_k8s_pb2.EnvVar(name=key, value=val)
for key, val in _six.iteritems(self._container.env)
])
final_containers.append(container)
del pod_spec.containers[:]
pod_spec.containers.extend(final_containers)
sidecar_job_plugin = _task_models.SidecarJob(
pod_spec=pod_spec,
primary_container_name=primary_container_name,
).to_flyte_idl()
self.assign_custom_and_return(_MessageToDict(sidecar_job_plugin))