def add_kfp_pod_env(op: BaseOp) -> BaseOp:
"""Adds KFP pod environment info to the specified ContainerOp.
"""
if not isinstance(op, ContainerOp):
warnings.warn(
'Trying to add default KFP environment variables to an Op that is '
'not a ContainerOp. Ignoring request.')
return op
op.container.add_env_variable(
k8s_client.V1EnvVar(name='KFP_POD_NAME',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path='metadata.name')))
).add_env_variable(
k8s_client.V1EnvVar(name='KFP_NAMESPACE',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path='metadata.namespace')))
).add_env_variable(
k8s_client.V1EnvVar(
name='WORKFLOW_ID',
value_from=k8s_client.
V1EnvVarSource(field_ref=k8s_client.V1ObjectFieldSelector(
field_path="metadata.labels['workflows.argoproj.io/workflow']")
)))
return op
def add_pod_env(op: BaseOp) -> BaseOp:
"""Adds pod environment info to ContainerOp.
"""
if isinstance(op, ContainerOp) and op.pod_labels and op.pod_labels['add-pod-env'] == 'true':
from kubernetes import client as k8s_client
op.container.add_env_variable(
k8s_client.V1EnvVar(
name='KFP_POD_NAME',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path='metadata.name'
)
)
)
).add_env_variable(
k8s_client.V1EnvVar(
name='KFP_NAMESPACE',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path='metadata.namespace'
)
)
)
)
return op
def downstream_api():
print_workflow_info().add_env_variable(
k8s_client.V1EnvVar(
name='KFP_RUN_NAME',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=
"metadata.annotations['pipelines.kubeflow.org/run_name']"))
)).add_env_variable(
k8s_client.V1EnvVar(
name='KFP_RUN_ID',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path="metadata.labels['pipeline/runid']"))))
def add_default_env(k8s_client, cop):
cop.container.add_env_variable(
k8s_client.V1EnvVar(
"MLRUN_NAMESPACE",
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path="metadata.namespace")),
))
if config.httpdb.api_url:
cop.container.add_env_variable(
k8s_client.V1EnvVar(name="MLRUN_DBPATH",
value=config.httpdb.api_url))
if config.mpijob_crd_version:
cop.container.add_env_variable(
k8s_client.V1EnvVar(name="MLRUN_MPIJOB_CRD_VERSION",
value=config.mpijob_crd_version))
if "MLRUN_AUTH_SESSION" in os.environ or "V3IO_ACCESS_KEY" in os.environ:
cop.container.add_env_variable(
k8s_client.V1EnvVar(
name="MLRUN_AUTH_SESSION",
value=os.environ.get("MLRUN_AUTH_SESSION")
or os.environ.get("V3IO_ACCESS_KEY"),
))
def _mount_v3iod(task):
from kubernetes import client as k8s_client
def add_vol(name, mount_path, host_path):
vol = k8s_client.V1Volume(
name=name,
host_path=k8s_client.V1HostPathVolumeSource(path=host_path, type=''),
)
task.add_volume(vol).add_volume_mount(
k8s_client.V1VolumeMount(mount_path=mount_path, name=name)
)
add_vol(name='shm', mount_path='/dev/shm', host_path='/dev/shm/' + namespace)
add_vol(
name='v3iod-comm',
mount_path='/var/run/iguazio/dayman',
host_path='/var/run/iguazio/dayman/' + namespace,
)
vol = k8s_client.V1Volume(
name='daemon-health', empty_dir=k8s_client.V1EmptyDirVolumeSource()
)
task.add_volume(vol).add_volume_mount(
k8s_client.V1VolumeMount(
mount_path='/var/run/iguazio/daemon_health', name='daemon-health'
)
)
vol = k8s_client.V1Volume(
name='v3io-config',
config_map=k8s_client.V1ConfigMapVolumeSource(
name=v3io_config_configmap, default_mode=420
),
)
task.add_volume(vol).add_volume_mount(
k8s_client.V1VolumeMount(mount_path='/etc/config/v3io', name='v3io-config')
)
# vol = k8s_client.V1Volume(name='v3io-auth',
# secret=k8s_client.V1SecretVolumeSource(secret_name=v3io_auth_secret,
# default_mode=420))
# task.add_volume(vol).add_volume_mount(k8s_client.V1VolumeMount(mount_path='/igz/.igz', name='v3io-auth'))
task.add_env_variable(
k8s_client.V1EnvVar(
name='CURRENT_NODE_IP',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
api_version='v1', field_path='status.hostIP'
)
),
)
)
task.add_env_variable(
k8s_client.V1EnvVar(
name='IGZ_DATA_CONFIG_FILE', value='/igz/java/conf/v3io.conf'
)
)
return task
def modify_pod_hook(spawner, pod):
pod.spec.containers[0].env.append(
client.V1EnvVar(
"MY_POD_IP", None,
client.V1EnvVarSource(
None, client.V1ObjectFieldSelector(None, "status.podIP"))))
return pod
def get_run_pod_env_vars(run_context):
config = SigOptConfig()
config.set_context_entry(GlobalRunContext(run_context))
env = [
k8s_client.V1EnvVar(
name="SIGOPT_API_TOKEN",
value=os.environ["SIGOPT_API_TOKEN"],
),
k8s_client.V1EnvVar(
name="SIGOPT_API_URL",
value=os.environ["SIGOPT_API_URL"],
),
k8s_client.V1EnvVar(
name="SIGOPT_PROJECT",
value=os.environ["SIGOPT_PROJECT"],
),
k8s_client.V1EnvVar(
name="SIGOPT_RUN_ID",
value=run_context.run.id,
),
k8s_client.V1EnvVar(
name="SIGOPT_RUN_NAME",
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path="metadata.name", ), ),
),
*(k8s_client.V1EnvVar(
name=key,
value=value.decode("ascii"),
) for key, value in config.get_environment_context().items()),
]
return env
def kube_env(self):
kube_env = []
for key, value in self.environment.items():
kube_env.append(client.V1EnvVar(name=key, value=value))
pod_name_value = client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(field_path='metadata.name'))
kube_env.append(
client.V1EnvVar(name='POD_NAME', value_from=pod_name_value))
return kube_env
def run_agent_deployment(agent_type,
replicas,
deploy_name='pymada-agents-deployment',
template_label={'app': 'pymada-agent'},
agent_port=5001,
container_name='pymada-single-agent',
auth_token=None,
no_agents_on_master_node=True,
pod_limits=None,
config_path=None):
env_vars = [
client.V1EnvVar("MASTER_URL", "http://pymadamaster:8000"),
client.V1EnvVar("AGENT_PORT", str(agent_port)),
client.V1EnvVar("AGENT_ADDR",
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path="status.podIP")))
]
if auth_token is not None:
env_vars.append(client.V1EnvVar("PYMADA_TOKEN_AUTH", auth_token))
agent_container_ports = [client.V1ContainerPort(container_port=agent_port)]
pod_node_selector = None
if no_agents_on_master_node:
pod_node_selector = {'pymada-role': 'agent'}
if agent_type == 'node_puppeteer':
agent_image_name = 'pymada/node-puppeteer'
pod_spec = create_general_pod_spec(agent_image_name, container_name,
agent_container_ports, env_vars,
pod_node_selector, pod_limits)
elif agent_type == 'python_selenium_firefox':
pod_spec = create_selenium_pod_spec('firefox', container_name,
agent_container_ports, env_vars,
pod_node_selector, pod_limits)
elif agent_type == 'python_selenium_chrome':
pod_spec = create_selenium_pod_spec('chrome', container_name,
agent_container_ports, env_vars,
pod_node_selector, pod_limits)
elif agent_type == 'python_agent':
agent_image_name = 'pymada/python-agent'
pod_spec = create_general_pod_spec(agent_image_name, container_name,
agent_container_ports, env_vars,
pod_node_selector, pod_limits)
run_deployment(pod_spec,
replicas,
deploy_name,
template_label,
config_path=config_path)
def _mount_v3iod(task):
from kubernetes import client as k8s_client
def add_vol(name, mount_path, host_path):
vol = k8s_client.V1Volume(
name=name,
host_path=k8s_client.V1HostPathVolumeSource(path=host_path,
type=""),
)
task.add_volume(vol).add_volume_mount(
k8s_client.V1VolumeMount(mount_path=mount_path, name=name))
add_vol(name="shm",
mount_path="/dev/shm",
host_path="/dev/shm/" + namespace)
add_vol(
name="v3iod-comm",
mount_path="/var/run/iguazio/dayman",
host_path="/var/run/iguazio/dayman/" + namespace,
)
vol = k8s_client.V1Volume(
name="daemon-health",
empty_dir=k8s_client.V1EmptyDirVolumeSource())
task.add_volume(vol).add_volume_mount(
k8s_client.V1VolumeMount(
mount_path="/var/run/iguazio/daemon_health",
name="daemon-health"))
vol = k8s_client.V1Volume(
name="v3io-config",
config_map=k8s_client.V1ConfigMapVolumeSource(
name=v3io_config_configmap, default_mode=420),
)
task.add_volume(vol).add_volume_mount(
k8s_client.V1VolumeMount(mount_path="/etc/config/v3io",
name="v3io-config"))
task.add_env_variable(
k8s_client.V1EnvVar(
name="CURRENT_NODE_IP",
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
api_version="v1", field_path="status.hostIP")),
))
task.add_env_variable(
k8s_client.V1EnvVar(name="IGZ_DATA_CONFIG_FILE",
value="/igz/java/conf/v3io.conf"))
return task
def __init__(self,
component: base_component.BaseComponent,
input_dict: Optional[Dict] = None):
output_dict = dict(
(k, v.get()) for k, v in component.outputs.get_all().items())
outputs = output_dict.keys()
file_outputs = {
output: '/output/ml_metadata/{}'.format(output)
for output in outputs
}
arguments = [
'--output_dir',
self._output_dir,
'--project_id',
self._project_id,
'--gcp_region',
self._gcp_region,
'--beam_runner',
self._beam_runner,
component.component_name,
]
if input_dict:
for k, v in input_dict.items():
if isinstance(v, float) or isinstance(v, int):
v = str(v)
arguments.append('--{}'.format(k))
arguments.append(v)
super().__init__(
name=component.component_name,
image=_IMAGE,
arguments=arguments,
file_outputs=file_outputs,
)
self.apply(gcp.use_gcp_secret('user-gcp-sa'))
field_path = "metadata.labels['workflows.argoproj.io/workflow']"
self.add_env_variable(
k8s_client.V1EnvVar(name='WORKFLOW_ID',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=field_path))))
class TestCastValue:
"""Tests for kubetest.manifest.cast_value"""
@pytest.mark.parametrize(
'value,t,expected',
[
# builtin types
(11, 'int', int(11)),
('11', 'int', int(11)),
(11.0, 'int', int(11)),
(11, 'float', float(11)),
(11, 'str', '11'),
# casting to object should result in no change
(11, 'object', 11),
('11', 'object', '11'),
# kubernetes types
({
'apiVersion': 'apps/v1',
'kind': 'Namespace'
}, 'V1Namespace',
client.V1Namespace(kind='Namespace', api_version='apps/v1')),
({
'fieldRef': {
'apiVersion': 'apps/v1beta1',
'fieldPath': 'foobar'
}
}, 'V1EnvVarSource',
client.V1EnvVarSource(field_ref=client.V1ObjectFieldSelector(
api_version='apps/v1beta1', field_path='foobar'))),
({
'finalizers': ['a', 'b', 'c']
}, 'V1ObjectMeta',
client.V1ObjectMeta(finalizers=['a', 'b', 'c'])),
])
def test_ok(self, value, t, expected):
"""Test casting values to the specified type successfully."""
actual = manifest.cast_value(value, t)
assert type(actual) == type(expected)
assert actual == expected
def _use_v3io_cred(task):
from kubernetes import client as k8s_client
from os import environ
web_api = api or environ.get('V3IO_API')
_user = user or environ.get('V3IO_USERNAME')
_access_key = access_key or environ.get('V3IO_ACCESS_KEY')
return (task.add_env_variable(
k8s_client.V1EnvVar(
name='V3IO_API', value=web_api)).add_env_variable(
k8s_client.V1EnvVar(
name='V3IO_USERNAME', value=_user)).add_env_variable(
k8s_client.V1EnvVar(name='V3IO_ACCESS_KEY',
value=_access_key)).
add_env_variable(
k8s_client.V1EnvVar(
name='CURRENT_NODE_IP',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
api_version='v1', field_path='status.hostIP')))
).add_env_variable(
k8s_client.V1EnvVar(name='IGZ_DATA_CONFIG_FILE',
value='/igz/java/conf/v3io.conf')))
def _create_container_object(name, image, always_pull, **kwargs):
# Set up environment variables
# Copy any passed in environment variables
env = kwargs.get('env') or {}
env_vars = [client.V1EnvVar(name=k, value=env[k]) for k in env]
# Add POD_IP with the IP address of the pod running the container
pod_ip = client.V1EnvVarSource(field_ref=client.V1ObjectFieldSelector(
field_path="status.podIP"))
env_vars.append(client.V1EnvVar(name="POD_IP", value_from=pod_ip))
# If a health check is specified, create a readiness/liveness probe
# (For an HTTP-based check, we assume it's at the first container port)
readiness = kwargs.get('readiness')
liveness = kwargs.get('liveness')
resources = kwargs.get('resources')
container_ports = kwargs.get('container_ports') or []
hc_port = container_ports[0][0] if container_ports else None
probe = _create_probe(readiness, hc_port) if readiness else None
live_probe = _create_probe(liveness, hc_port) if liveness else None
resources_obj = _create_resources(resources) if resources else None
port_objs = [
client.V1ContainerPort(container_port=port, protocol=proto)
for port, proto in container_ports
]
# Define container for pod
return client.V1Container(
name=name,
image=image,
image_pull_policy='Always' if always_pull else 'IfNotPresent',
env=env_vars,
ports=port_objs,
volume_mounts=kwargs.get('volume_mounts') or [],
resources=resources_obj,
readiness_probe=probe,
liveness_probe=live_probe)
def run(provider, provider_kwargs, cluster=None, job=None, storage=None):
# TODO, temp fix
s3 = storage["s3"]
_validate_fields(
provider=provider_kwargs, cluster=cluster, job=job, storage=storage, s3=s3
)
_required_run_arguments(provider_kwargs, cluster, job, storage, s3)
response = {"job": {}}
if "name" not in job["meta"] or not job["meta"]["name"]:
since_epoch = int(time.time())
job["meta"]["name"] = "{}-{}".format(JOB_DEFAULT_NAME, since_epoch)
if "bucket_name" not in s3 or not s3["bucket_name"]:
s3["bucket_name"] = job["meta"]["name"]
container_engine_client = new_client(
ContainerEngineClient,
composite_class=ContainerEngineClientCompositeOperations,
name=provider_kwargs["profile"]["name"],
)
compute_cluster = get_cluster_by_name(
container_engine_client,
provider_kwargs["profile"]["compartment_id"],
name=cluster["name"],
)
if not compute_cluster:
response["msg"] = "Failed to find a cluster with name: {}".format(
cluster["name"]
)
return False, response
refreshed = refresh_kube_config(
compute_cluster.id, name=provider_kwargs["profile"]["name"]
)
if not refreshed:
response["msg"] = "Failed to refresh the kubernetes config"
return False, response
node_manager = NodeManager()
if not node_manager.discover():
response["msg"] = "Failed to discover any nodes to schedule jobs on"
return False, response
node = node_manager.select()
if not node:
response["msg"] = "Failed to select a node to schedule on"
return False, response
# Ensure we have the newest config
scheduler = KubenetesScheduler()
jobio_args = [
"jobio",
"run",
]
jobio_args.extend(job["commands"])
jobio_args.extend(["--job-meta-name", job["meta"]["name"]])
if "output_path" in job:
jobio_args.extend(
["--job-output-path", job["output_path"],]
)
if "capture" in job and job["capture"]:
jobio_args.append("--job-capture")
if "debug" in job["meta"]:
jobio_args.append("--job-meta-debug")
if "env_override" in job["meta"]:
jobio_args.append("--job-meta-env-override")
# Maintained by the pod
volumes = []
# Maintained by the container
volume_mounts = []
# Environment to pass to the container
envs = []
# Prepare config for the scheduler
scheduler_config = {}
if storage and storage["enable"]:
validate_dict_values(storage, required_storage_fields, throw=True)
jobio_args.append("--storage-enable")
# Means that results should be exported to the specified storage
# Create kubernetes secrets
core_api = client.CoreV1Api()
# storage_api = client.StorageV1Api()
# Storage endpoint credentials secret (Tied to a profile and job)
secret_profile_name = "{}-{}-{}".format(
STORAGE_CREDENTIALS_NAME, s3["name"], job["meta"]["name"]
)
try:
storage_credentials_secret = core_api.read_namespaced_secret(
secret_profile_name, KUBERNETES_NAMESPACE
)
except ApiException:
storage_credentials_secret = None
# volumes
secret_volume_source = V1SecretVolumeSource(secret_name=secret_profile_name)
secret_volume = V1Volume(name=secret_profile_name, secret=secret_volume_source)
volumes.append(secret_volume)
# Where the storage credentials should be mounted
# in the compute unit
secret_mount = V1VolumeMount(
name=secret_profile_name,
mount_path=storage["credentials_path"],
read_only=True,
)
volume_mounts.append(secret_mount)
if s3:
validate_dict_values(s3, required_staging_values, verbose=True, throw=True)
jobio_args.append("--storage-s3")
# S3 storage
# Look for s3 credentials and config files
s3_config = load_aws_config(
s3["config_file"], s3["credentials_file"], profile_name=s3["name"],
)
s3_config["endpoint_url"] = storage["endpoint"]
if not storage_credentials_secret:
secret_data = dict(
aws_access_key_id=s3_config["aws_access_key_id"],
aws_secret_access_key=s3_config["aws_secret_access_key"],
)
secret_metadata = V1ObjectMeta(name=secret_profile_name)
secrets_config = dict(metadata=secret_metadata, string_data=secret_data)
scheduler_config.update(dict(secret_kwargs=secrets_config))
# If `access_key`
# TODO, unify argument endpoint, with s3 config endpoint'
s3_resource = boto3.resource("s3", **s3_config)
bucket = bucket_exists(s3_resource.meta.client, s3["bucket_name"])
if not bucket:
bucket = s3_resource.create_bucket(
Bucket=s3["bucket_name"],
CreateBucketConfiguration={
"LocationConstraint": s3_config["region_name"]
},
)
if "upload_path" in storage and storage["upload_path"]:
# Upload local path to the bucket as designated input for the job
uploaded = None
if os.path.exists(storage["upload_path"]):
if os.path.isdir(storage["upload_path"]):
uploaded = upload_directory_to_s3(
s3_resource.meta.client,
storage["upload_path"],
s3["bucket_name"],
s3_prefix=s3["bucket_input_prefix"],
)
elif os.path.isfile(storage["upload_path"]):
s3_path = os.path.basename(storage["upload_path"])
if s3["bucket_input_prefix"]:
s3_path = os.path.join(s3["bucket_input_prefix"], s3_path)
# Upload
uploaded = upload_to_s3(
s3_resource.meta.client,
storage["upload_path"],
s3_path,
s3["bucket_name"],
)
if not uploaded:
response[
"msg"
] = "Failed to local path: {} in the upload folder to s3".format(
storage["upload_path"]
)
return False, response
jobio_args.extend(
[
"--s3-region-name",
s3_config["region_name"],
"--storage-secrets-dir",
storage["credentials_path"],
"--storage-endpoint",
storage["endpoint"],
"--storage-input-path",
storage["input_path"],
"--storage-output-path",
storage["output_path"],
"--bucket-name",
s3["bucket_name"],
"--bucket-input-prefix",
s3["bucket_input_prefix"],
"--bucket-output-prefix",
s3["bucket_output_prefix"],
]
)
# Provide a way to allow pod specific output prefixes
field_ref = client.V1ObjectFieldSelector(field_path="metadata.name")
env_var_source = client.V1EnvVarSource(field_ref=field_ref)
# HACK, Set the output prefix in the bucket to the name of the pod
env_output_prefix = client.V1EnvVar(
name="JOBIO_BUCKET_OUTPUT_PREFIX", value_from=env_var_source
)
envs.append(env_output_prefix)
if scheduler_config:
prepared = scheduler.prepare(**scheduler_config)
if not prepared:
response["msg"] = "Failed to prepare the scheduler"
return False, response
container_spec = dict(
name=job["meta"]["name"],
image=cluster["image"],
env=envs,
args=jobio_args,
volume_mounts=volume_mounts,
)
# If the working directory does not exist inside the container
# It will set permissions where it will be unable to expand the
# s3 bucket if the user doesn't have root permissions
if "working_dir" in job:
container_spec.update({"working_dir": job["working_dir"]})
# If the container requires a specific set of resources
resources = {}
if "min_cores" in job:
resources["requests"] = {"cpu": job["min_cores"]}
if "max_cores" in job:
resources["limits"] = {"cpu": job["max_cores"]}
if "min_memory" in job:
resources["requests"].update({"memory": job["min_memory"]})
if "max_memory" in job:
resources["limits"].update({"memory": job["max_memory"]})
if resources:
resource_req = client.V1ResourceRequirements(**resources)
container_spec.update({"resources": resource_req})
# args=jobio_args,
pod_spec = dict(node_name=node.metadata.name, volumes=volumes, dns_policy="Default")
job_spec = dict(
backoff_limit=2,
parallelism=job["meta"]["num_parallel"],
completions=job["meta"]["num_jobs"],
)
task = dict(
container_kwargs=container_spec,
pod_spec_kwargs=pod_spec,
job_spec_kwargs=job_spec,
)
job = scheduler.submit(**task)
if not job:
response["msg"] = "Failed to submit the job"
return False, response
response["job"] = job
response["msg"] = "Job submitted"
return True, response
def build_op(
name,
function=None,
func_url=None,
image=None,
base_image=None,
commands: list = None,
secret_name="",
with_mlrun=True,
skip_deployed=False,
):
"""build Docker image."""
from kfp import dsl
from os import environ
from kubernetes import client as k8s_client
cmd = ["python", "-m", "mlrun", "build", "--kfp"]
if function:
if not hasattr(function, "to_dict"):
raise ValueError("function must specify a function runtime object")
cmd += ["-r", str(function.to_dict())]
elif not func_url:
raise ValueError("function object or func_url must be specified")
commands = commands or []
if image:
cmd += ["-i", image]
if base_image:
cmd += ["-b", base_image]
if secret_name:
cmd += ["--secret-name", secret_name]
if with_mlrun:
cmd += ["--with_mlrun"]
if skip_deployed:
cmd += ["--skip"]
for c in commands:
cmd += ["-c", c]
if func_url and not function:
cmd += [func_url]
cop = dsl.ContainerOp(
name=name,
image=config.kfp_image,
command=cmd,
file_outputs={"state": "/tmp/state", "image": "/tmp/image"},
)
if config.httpdb.builder.docker_registry:
cop.container.add_env_variable(
k8s_client.V1EnvVar(
name="MLRUN_HTTPDB__BUILDER__DOCKER_REGISTRY",
value=config.httpdb.builder.docker_registry,
)
)
if "IGZ_NAMESPACE_DOMAIN" in environ:
cop.container.add_env_variable(
k8s_client.V1EnvVar(
name="IGZ_NAMESPACE_DOMAIN",
value=os.environ.get("IGZ_NAMESPACE_DOMAIN"),
)
)
is_v3io = function.spec.build.source and function.spec.build.source.startswith(
"v3io"
)
if "V3IO_ACCESS_KEY" in environ and is_v3io:
cop.container.add_env_variable(
k8s_client.V1EnvVar(
name="V3IO_ACCESS_KEY", value=os.environ.get("V3IO_ACCESS_KEY")
)
)
cop.container.add_env_variable(
k8s_client.V1EnvVar(
"MLRUN_NAMESPACE",
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path="metadata.namespace"
)
),
)
)
return cop
def createStatefulSet(cls, cluster_object: V1MongoClusterConfiguration) -> client.V1beta1StatefulSet:
"""
Creates a the stateful set configuration for the given cluster.
:param cluster_object: The cluster object from the YAML file.
:return: The stateful set object.
"""
# Parse cluster data object.
name = cluster_object.metadata.name
namespace = cluster_object.metadata.namespace
replicas = cluster_object.spec.mongodb.replicas
storage_mount_path = cluster_object.spec.mongodb.host_path or cls.DEFAULT_STORAGE_MOUNT_PATH
host_path = cluster_object.spec.mongodb.host_path
cpu_limit = cluster_object.spec.mongodb.cpu_limit or cls.DEFAULT_CPU_LIMIT
memory_limit = cluster_object.spec.mongodb.memory_limit or cls.DEFAULT_MEMORY_LIMIT
run_as_user = cluster_object.spec.mongodb.run_as_user or cls.DEFAULT_RUN_AS_USER
service_account = cluster_object.spec.mongodb.service_account or cls.DEFAULT_SERVICE_ACCOUNT
wired_tiger_cache_size = cluster_object.spec.mongodb.wired_tiger_cache_size or cls.DEFAULT_CACHE_SIZE
secret_name = cls.ADMIN_SECRET_NAME_FORMAT.format(name)
# create container
mongo_container = client.V1Container(
name=name,
env=[client.V1EnvVar(
name="POD_IP",
value_from=client.V1EnvVarSource(
field_ref = client.V1ObjectFieldSelector(
api_version = "v1",
field_path = "status.podIP"
)
)
),
client.V1EnvVar(
name="MONGODB_PASSWORD",
value_from=client.V1EnvVarSource(
secret_key_ref=client.V1SecretKeySelector(
key="database-password",
name=secret_name
)
)
),
client.V1EnvVar(
name="MONGODB_USER",
value_from=client.V1EnvVarSource(
secret_key_ref=client.V1SecretKeySelector(
key="database-user",
name=secret_name
)
)
),
client.V1EnvVar(
name="MONGODB_DATABASE",
value_from=client.V1EnvVarSource(
secret_key_ref=client.V1SecretKeySelector(
key="database-name",
name=secret_name
)
)
),
client.V1EnvVar(
name="MONGODB_ADMIN_PASSWORD",
value_from=client.V1EnvVarSource(
secret_key_ref=client.V1SecretKeySelector(
key="database-admin-password",
name=secret_name
)
)
),
client.V1EnvVar(
name="WIREDTIGER_CACHE_SIZE",
value=wired_tiger_cache_size
),
client.V1EnvVar(
name="MONGODB_REPLICA_NAME",
value=name
),
client.V1EnvVar(
name="MONGODB_SERVICE_NAME",
value="svc-" + name + "-internal"
),
client.V1EnvVar(
name="MONGODB_KEYFILE_VALUE",
value="supersecretkeyfile123"
)],
liveness_probe=client.V1Probe(failure_threshold=3,
initial_delay_seconds=30,
period_seconds=30,
success_threshold=1,
tcp_socket=client.V1TCPSocketAction(port=cls.MONGO_PORT),
timeout_seconds=1
),
command=cls.MONGO_COMMAND.split(),
image=cls.MONGO_IMAGE,
image_pull_policy="Always",
ports=[client.V1ContainerPort(
name="mongodb",
container_port=cls.MONGO_PORT,
protocol="TCP"
)],
readiness_probe=client.V1Probe(_exec=client.V1ExecAction(command=["/bin/sh", "-i", "-c", "mongo 127.0.0.1:27017/$MONGODB_DATABASE -u $MONGODB_USER -p $MONGODB_PASSWORD --eval=\"quit()\""]),
failure_threshold=3,
initial_delay_seconds=10,
period_seconds=10,
success_threshold=1,
timeout_seconds=1
),
security_context=client.V1SecurityContext(
run_as_user=int(run_as_user),
se_linux_options=client.V1SELinuxOptions(
level="s0",
type="spc_t"
)
),
termination_message_path="/dev/termination-log",
volume_mounts=[client.V1VolumeMount(
name="mongo-data",
read_only=False,
mount_path=storage_mount_path
)],
resources=client.V1ResourceRequirements(
limits={"cpu": cpu_limit, "memory": memory_limit},
requests={"cpu": cpu_limit, "memory": memory_limit}
)
)
#create affinity rules
affinity = client.V1Affinity(
pod_anti_affinity=client.V1PodAntiAffinity(
required_during_scheduling_ignored_during_execution=[
client.V1PodAffinityTerm(label_selector=client.V1LabelSelector(
match_expressions=[client.V1LabelSelectorRequirement(
key="app",
operator="In",
values=[name]
)]
),
topology_key="kubernetes.io/hostname")
]
)
)
volumes = [client.V1Volume(
name="mongo-data",
host_path=client.V1HostPathVolumeSource(path=host_path)
)]
# Create stateful set.
return client.V1beta1StatefulSet(
metadata = client.V1ObjectMeta(annotations={"service.alpha.kubernetes.io/tolerate-unready-endpoints": "true"},
name=name,
namespace=namespace,
labels=cls.createDefaultLabels(name)),
spec = client.V1beta1StatefulSetSpec(
replicas = replicas,
service_name = "svc-" + name + "-internal",
template = client.V1PodTemplateSpec(
metadata = client.V1ObjectMeta(labels=cls.createDefaultLabels(name)),
spec = client.V1PodSpec(affinity = affinity,
containers=[mongo_container],
node_selector={"compute":"mongodb"},
service_account=service_account,
#restart_policy="Never",
volumes=volumes
)
),
),
)
def get_from_field_ref(name, field_path):
field_ref = client.V1ObjectFieldSelector(field_path=field_path)
value_from = client.V1EnvVarSource(field_ref=field_ref)
return client.V1EnvVar(name=name, value_from=value_from)
def mlrun_op(
name: str = "",
project: str = "",
function=None,
func_url=None,
image: str = "",
runobj=None,
command: str = "",
secrets: list = None,
params: dict = None,
job_image=None,
hyperparams: dict = None,
param_file: str = "",
labels: dict = None,
selector: str = "",
inputs: dict = None,
outputs: list = None,
in_path: str = "",
out_path: str = "",
rundb: str = "",
mode: str = "",
handler: str = "",
more_args: list = None,
tuning_strategy=None,
verbose=None,
scrape_metrics=False,
):
"""mlrun KubeFlow pipelines operator, use to form pipeline steps
when using kubeflow pipelines, each step is wrapped in an mlrun_op
one step can pass state and data to the next step, see example below.
:param name: name used for the step
:param project: optional, project name
:param image: optional, run container image (will be executing the step)
the container should host all requiered packages + code
for the run, alternatively user can mount packages/code via
shared file volumes like v3io (see example below)
:param function: optional, function object
:param func_url: optional, function object url
:param command: exec command (or URL for functions)
:param secrets: extra secrets specs, will be injected into the runtime
e.g. ['file=<filename>', 'env=ENV_KEY1,ENV_KEY2']
:param params: dictionary of run parameters and values
:param hyperparams: dictionary of hyper parameters and list values, each
hyperparam holds a list of values, the run will be
executed for every parameter combination (GridSearch)
:param param_file: a csv file with parameter combinations, first row hold
the parameter names, following rows hold param values
:param selector: selection criteria for hyperparams e.g. "max.accuracy"
:param tuning_strategy: selection strategy for hyperparams e.g. list, grid, random
:param labels: labels to tag the job/run with ({key:val, ..})
:param inputs: dictionary of input objects + optional paths (if path is
omitted the path will be the in_path/key.
:param outputs: dictionary of input objects + optional paths (if path is
omitted the path will be the out_path/key.
:param in_path: default input path/url (prefix) for inputs
:param out_path: default output path/url (prefix) for artifacts
:param rundb: path for rundb (or use 'MLRUN_DBPATH' env instead)
:param mode: run mode, e.g. 'noctx' for pushing params as args
:param handler code entry-point/hanfler name
:param job_image name of the image user for the job
:param verbose: add verbose prints/logs
:param scrape_metrics: whether to add the `mlrun/scrape-metrics` label to this run's resources
:returns: KFP step operation
Example:
from kfp import dsl
from mlrun import mlrun_op
from mlrun.platforms import mount_v3io
def mlrun_train(p1, p2):
return mlrun_op('training',
command = '/User/kubeflow/training.py',
params = {'p1':p1, 'p2':p2},
outputs = {'model.txt':'', 'dataset.csv':''},
out_path ='v3io:///projects/my-proj/mlrun/{{workflow.uid}}/',
rundb = '/User/kubeflow')
# use data from the first step
def mlrun_validate(modelfile):
return mlrun_op('validation',
command = '/User/kubeflow/validation.py',
inputs = {'model.txt':modelfile},
out_path ='v3io:///projects/my-proj/{{workflow.uid}}/',
rundb = '/User/kubeflow')
@dsl.pipeline(
name='My MLRUN pipeline', description='Shows how to use mlrun.'
)
def mlrun_pipeline(
p1 = 5 , p2 = '"text"'
):
# run training, mount_v3io will mount "/User" into the pipeline step
train = mlrun_train(p1, p2).apply(mount_v3io())
# feed 1st step results into the secound step
validate = mlrun_validate(
train.outputs['model-txt']).apply(mount_v3io())
"""
from kfp import dsl
from os import environ
from kubernetes import client as k8s_client
secrets = [] if secrets is None else secrets
params = {} if params is None else params
hyperparams = {} if hyperparams is None else hyperparams
inputs = {} if inputs is None else inputs
outputs = [] if outputs is None else outputs
labels = {} if labels is None else labels
rundb = rundb or get_or_set_dburl()
cmd = [
"python",
"-m",
"mlrun",
"run",
"--kfp",
"--from-env",
"--workflow",
"{{workflow.uid}}",
]
file_outputs = {}
runtime = None
code_env = None
function_name = ""
if function:
if not func_url:
if function.kind in ["", "local"]:
image = image or function.spec.image
command = command or function.spec.command
more_args = more_args or function.spec.args
mode = mode or function.spec.mode
rundb = rundb or function.spec.rundb
code_env = str(function.spec.build.functionSourceCode)
else:
runtime = str(function.to_dict())
function_name = function.metadata.name
if function.kind == "dask":
image = image or function.spec.kfp_image or config.dask_kfp_image
image = image or config.kfp_image
if runobj:
handler = handler or runobj.spec.handler_name
params = params or runobj.spec.parameters
hyperparams = hyperparams or runobj.spec.hyperparams
param_file = param_file or runobj.spec.param_file
tuning_strategy = tuning_strategy or runobj.spec.tuning_strategy
selector = selector or runobj.spec.selector
inputs = inputs or runobj.spec.inputs
outputs = outputs or runobj.spec.outputs
in_path = in_path or runobj.spec.input_path
out_path = out_path or runobj.spec.output_path
secrets = secrets or runobj.spec.secret_sources
project = project or runobj.metadata.project
labels = runobj.metadata.labels or labels
verbose = verbose or runobj.spec.verbose
scrape_metrics = scrape_metrics or runobj.spec.scrape_metrics
if not name:
if not function_name:
raise ValueError("name or function object must be specified")
name = function_name
if handler:
name += "-" + handler
if hyperparams or param_file:
outputs.append("iteration_results")
if "run_id" not in outputs:
outputs.append("run_id")
params = params or {}
hyperparams = hyperparams or {}
inputs = inputs or {}
secrets = secrets or []
if "V3IO_USERNAME" in environ and "v3io_user" not in labels:
labels["v3io_user"] = os.environ.get("V3IO_USERNAME")
if "owner" not in labels:
labels["owner"] = os.environ.get("V3IO_USERNAME") or getpass.getuser()
if name:
cmd += ["--name", name]
if func_url:
cmd += ["-f", func_url]
for secret in secrets:
cmd += ["-s", f"{secret['kind']}={secret['source']}"]
for param, val in params.items():
cmd += ["-p", f"{param}={val}"]
for xpram, val in hyperparams.items():
cmd += ["-x", f"{xpram}={val}"]
for input_param, val in inputs.items():
cmd += ["-i", f"{input_param}={val}"]
for label, val in labels.items():
cmd += ["--label", f"{label}={val}"]
for output in outputs:
cmd += ["-o", str(output)]
file_outputs[output.replace(".", "_")] = os.path.join("/tmp", output)
if project:
cmd += ["--project", project]
if handler:
cmd += ["--handler", handler]
if runtime:
cmd += ["--runtime", runtime]
if in_path:
cmd += ["--in-path", in_path]
if out_path:
cmd += ["--out-path", out_path]
if param_file:
cmd += ["--param-file", param_file]
if tuning_strategy:
cmd += ["--tuning-strategy", tuning_strategy]
if selector:
cmd += ["--selector", selector]
if job_image:
cmd += ["--image", job_image]
if mode:
cmd += ["--mode", mode]
if verbose:
cmd += ["--verbose"]
if scrape_metrics:
cmd += ["--scrape-metrics"]
if more_args:
cmd += more_args
registry = get_default_reg()
if image and image.startswith("."):
if registry:
image = f"{registry}/{image[1:]}"
else:
raise ValueError("local image registry env not found")
cop = dsl.ContainerOp(
name=name,
image=image,
command=cmd + [command],
file_outputs=file_outputs,
output_artifact_paths={
"mlpipeline-ui-metadata": "/mlpipeline-ui-metadata.json",
"mlpipeline-metrics": "/mlpipeline-metrics.json",
},
)
# if rundb:
# cop.container.add_env_variable(k8s_client.V1EnvVar(
# name='MLRUN_DBPATH', value=rundb))
if code_env:
cop.container.add_env_variable(
k8s_client.V1EnvVar(name="MLRUN_EXEC_CODE", value=code_env)
)
if registry:
cop.container.add_env_variable(
k8s_client.V1EnvVar(
name="MLRUN_HTTPDB__BUILDER__DOCKER_REGISTRY", value=registry
)
)
cop.container.add_env_variable(
k8s_client.V1EnvVar(
"MLRUN_NAMESPACE",
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path="metadata.namespace"
)
),
)
)
if config.mpijob_crd_version:
cop.container.add_env_variable(
k8s_client.V1EnvVar(
name="MLRUN_MPIJOB_CRD_VERSION", value=config.mpijob_crd_version
)
)
return cop
def create_job_object(name: str,
container_image: str,
env_list: dict,
command: List[str],
command_args: List[str],
volumes: List[Dict],
init_containers: List[Dict],
output: Output,
namespace: str = "stackl",
container_name: str = "jobcontainer",
api_version: str = "batch/v1",
image_pull_policy: str = "Always",
ttl_seconds_after_finished: int = 3600,
restart_policy: str = "Never",
backoff_limit: int = 0,
active_deadline_seconds: int = 3600,
service_account: str = "stackl-agent-stackl-agent",
image_pull_secrets: List[str] = [],
labels=None) -> client.V1Job:
# pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements
"""Creates a Job object using the Kubernetes client
:param name: Job name affix
:type name: str
:param container_image: automation container image
:type container_image: str
:param env_list: Dict with key/values for the environment inside the automation container
:type env_list: dict
:param command: entrypoint command
:type command: List[str]
:param command_args: command arguments
:type command_args: List[str]
:param volumes: volumes and volumemounts
:type volumes: List[Dict]
:param image_pull_secrets: secrets to pull images
:type image_pull_secrets: List[str]
:param init_containers: list with init_containers
:type init_containers: List[Dict]
:param output: output Object
:type output: Output
:param namespace: Kubernetes namespace, defaults to "stackl"
:type namespace: str, optional
:param container_name: name of automation container, defaults to "jobcontainer"
:type container_name: str, optional
:param api_version: Job api version, defaults to "batch/v1"
:type api_version: str, optional
:param image_pull_policy: always pull latest images, defaults to "Always"
:type image_pull_policy: str, optional
:param ttl_seconds_after_finished: Remove jobs after execution with ttl, defaults to 600
:type ttl_seconds_after_finished: int, optional
:param restart_policy: Restart the pod on the same node after failure, defaults to "Never"
:type restart_policy: str, optional
:param backoff_limit: Retries after failure, defaults to 0
:type backoff_limit: int, optional
:param active_deadline_seconds: Timeout on a job, defaults to 3600 seconds
:type active_deadline_seconds: int, optional
:param service_account: Kubernetes service account, defaults to "stackl-agent-stackl-agent"
:type service_account: str, optional
:param labels: metadata labels, defaults to {}
:type labels: dict, optional
:return: automation Job object
:rtype: client.V1Job
"""
id_job = id_generator()
name = name + "-" + id_job
body = client.V1Job(api_version=api_version, kind="Job")
body.metadata = client.V1ObjectMeta(namespace=namespace, name=name)
body.status = client.V1JobStatus()
template = client.V1PodTemplate()
template.template = client.V1PodTemplateSpec()
k8s_volumes = []
cms = []
logging.debug(f"volumes: {volumes}")
# create a k8s volume for each element in volumes
for vol in volumes:
vol_name = name + "-" + vol["name"]
k8s_volume = client.V1Volume(name=vol_name)
if vol["type"] == "config_map":
config_map = client.V1ConfigMapVolumeSource()
config_map.name = vol_name
k8s_volume.config_map = config_map
cms.append(create_cm(vol_name, namespace, vol['data']))
vol['name'] = vol_name
if vol["type"] == "empty_dir":
k8s_volume.empty_dir = client.V1EmptyDirVolumeSource(
medium="Memory")
vol['name'] = vol_name
k8s_volumes.append(k8s_volume)
logging.debug(f"Volumes created for job {name}: {k8s_volumes}")
# create a volume mount for each element in volumes
k8s_volume_mounts = []
for vol in volumes:
if vol["mount_path"]:
volume_mount = client.V1VolumeMount(name=vol["name"],
mount_path=vol["mount_path"])
if "sub_path" in vol:
volume_mount.sub_path = vol["sub_path"]
k8s_volume_mounts.append(volume_mount)
logging.debug(f"Volume mounts created for job {name}: {k8s_volume_mounts}")
# create an environment list
k8s_env_list = []
if env_list:
for key, value in env_list.items():
if isinstance(value, dict):
if 'config_map_key_ref' in value:
k8s_env_from = client.V1EnvVar(
name=key,
value_from=client.V1EnvVarSource(
config_map_key_ref=client.V1ConfigMapKeySelector(
name=value['config_map_key_ref']["name"],
key=value['config_map_key_ref']["key"])))
k8s_env_list.append(k8s_env_from)
elif 'field_ref' in value:
k8s_env_from = client.V1EnvVar(
name=key,
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path=value['field_ref'])))
k8s_env_list.append(k8s_env_from)
else:
k8s_env = client.V1EnvVar(name=key, value=value)
k8s_env_list.append(k8s_env)
k8s_env_from_list = []
# if env_from:
# for env in env_from:
# if 'config_map_ref' in env:
# k8s_env_from = client.V1EnvFromSource(
# config_map_ref=env['config_map_ref'])
# k8s_env_from_list.append(k8s_env_from)
# elif 'secret_ref' in env:
# k8s_env_from = client.V1EnvFromSource(
# secret_ref=env['secret_ref'])
# k8s_env_from_list.append(k8s_env_from)
logging.debug(f"Environment list created for job {name}: {k8s_env_list}")
print(f"Environment list created for job {name}: {k8s_env_list}")
container = client.V1Container(name=container_name,
image=container_image,
env=k8s_env_list,
volume_mounts=k8s_volume_mounts,
image_pull_policy=image_pull_policy,
command=command,
args=command_args,
env_from=k8s_env_from_list)
k8s_init_containers = []
logging.debug(f"Init containers for job {name}: {init_containers}")
for c in init_containers:
k8s_c = client.V1Container(name=c['name'],
image=c['image'],
volume_mounts=k8s_volume_mounts,
env=k8s_env_list)
if 'args' in c:
k8s_c.args = c['args']
k8s_init_containers.append(k8s_c)
k8s_secrets = []
for secret in image_pull_secrets:
k8s_secrets.append(client.V1LocalObjectReference(name=secret))
logging.debug(f"Secret list created for job {name}: {k8s_secrets}")
containers = [container]
if output:
output.volume_mounts = k8s_volume_mounts
output.env = k8s_env_list
output_containers = output.containers
containers = containers + output_containers
template.template.metadata = client.V1ObjectMeta(labels=labels)
template.template.spec = client.V1PodSpec(
containers=containers,
restart_policy=restart_policy,
image_pull_secrets=k8s_secrets,
volumes=k8s_volumes,
init_containers=k8s_init_containers,
service_account_name=service_account)
template.template = client.V1PodTemplateSpec(
metadata=template.template.metadata, spec=template.template.spec)
body.spec = client.V1JobSpec(
ttl_seconds_after_finished=ttl_seconds_after_finished,
template=template.template,
backoff_limit=backoff_limit,
active_deadline_seconds=active_deadline_seconds)
return body, cms
def export_deployment(self):
# Configureate Pod template container
volume_mounts = []
containers = []
volumes = []
volume_mounts.append(
client.V1VolumeMount(mount_path='/docker/logs', name='logs'))
volumes.append(
client.V1Volume(name='logs',
host_path=client.V1HostPathVolumeSource(
path='/opt/logs', type='DirectoryOrCreate')))
if self.mounts:
for path in self.mounts:
volume_mounts.append(
client.V1VolumeMount(mount_path=path,
name=self.mounts[path]))
volumes.append(
client.V1Volume(name=self.mounts[path],
host_path=client.V1HostPathVolumeSource(
path=path, type='DirectoryOrCreate')))
liveness_probe = client.V1Probe(initial_delay_seconds=15,
tcp_socket=client.V1TCPSocketAction(
port=int(self.container_port[0])))
readiness_probe = client.V1Probe(initial_delay_seconds=15,
tcp_socket=client.V1TCPSocketAction(
port=int(self.container_port[0])))
if self.healthcheck:
liveness_probe = client.V1Probe(initial_delay_seconds=15,
http_get=client.V1HTTPGetAction(
path=self.healthcheck,
port=int(
self.container_port[0])))
readiness_probe = client.V1Probe(initial_delay_seconds=15,
http_get=client.V1HTTPGetAction(
path=self.healthcheck,
port=int(
self.container_port[0])))
Env = [
client.V1EnvVar(name='LANG', value='en_US.UTF-8'),
client.V1EnvVar(name='LC_ALL', value='en_US.UTF-8'),
client.V1EnvVar(name='POD_NAME',
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path='metadata.name'))),
client.V1EnvVar(name='POD_IP',
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path='status.podIP'))),
]
container = client.V1Container(
name=self.dm_name,
image=self.image,
ports=[
client.V1ContainerPort(container_port=int(port))
for port in self.container_port
],
image_pull_policy='Always',
env=Env,
resources=client.V1ResourceRequirements(limits=self.re_limits,
requests=self.re_requests),
volume_mounts=volume_mounts,
liveness_probe=liveness_probe,
readiness_probe=readiness_probe)
containers.append(container)
if self.sidecar:
sidecar_container = client.V1Container(
name='sidecar-%s' % self.dm_name,
image=self.sidecar,
image_pull_policy='Always',
env=Env,
resources=client.V1ResourceRequirements(
limits=self.re_limits, requests=self.re_requests),
volume_mounts=volume_mounts)
containers.append(sidecar_container)
# Create and configurate a spec section
secrets = client.V1LocalObjectReference('registrysecret')
template = client.V1PodTemplateSpec(
metadata=client.V1ObjectMeta(labels={"project": self.dm_name}),
spec=client.V1PodSpec(
containers=containers,
image_pull_secrets=[secrets],
volumes=volumes,
affinity=client.V1Affinity(node_affinity=client.V1NodeAffinity(
preferred_during_scheduling_ignored_during_execution=[
client.V1PreferredSchedulingTerm(
preference=client.V1NodeSelectorTerm(
match_expressions=[
client.V1NodeSelectorRequirement(
key='project',
operator='In',
values=['moji'])
]),
weight=30),
client.V1PreferredSchedulingTerm(
preference=client.V1NodeSelectorTerm(
match_expressions=[
client.V1NodeSelectorRequirement(
key='deploy',
operator='In',
values=[self.dm_name])
]),
weight=70)
]))))
selector = client.V1LabelSelector(
match_labels={"project": self.dm_name})
# Create the specification of deployment
spec = client.ExtensionsV1beta1DeploymentSpec(replicas=int(
self.replicas),
template=template,
selector=selector,
min_ready_seconds=3)
# Instantiate the deployment object
deployment = client.ExtensionsV1beta1Deployment(
api_version="extensions/v1beta1",
kind="Deployment",
metadata=client.V1ObjectMeta(name=self.dm_name),
spec=spec)
return deployment
def __init__(self, **kwargs):
self.log = logging.getLogger("distributed.deploy.adaptive")
# Configure ourselves using environment variables
if 'KUBECONFIG' in os.environ:
config.load_kube_config()
else:
config.load_incluster_config()
# Switch off SSL host name verification for now (the JAP Python is old... :-()
client.configuration.assert_hostname = False
self.api = client.CoreV1Api()
# Read the environment variables for configuration
self.namespace = os.environ.get('NAMESPACE', 'dask')
self.worker_labels = os.environ.get('WORKER_LABELS',
'app=dask,component=worker')
dask_scheduler_service = os.environ.get('DASK_SCHEDULER_SERVICE',
'dask-scheduler')
worker_name_prefix = os.environ.get('WORKER_NAME_PREFIX',
'dask-worker-')
worker_image = os.environ.get('WORKER_IMAGE', 'daskdev/dask:latest')
worker_image_pull_policy = os.environ.get('WORKER_IMAGE_PULL_POLICY',
'')
# Worker resources should be given as a JSON string, e.g.
# "{'requests': {'cpu':'100m','memory':'1Gi'}}"
# We need them as a dict
worker_resources = json.loads(
os.environ.get('WORKER_RESOURCES', '\{\}'))
# Build the pod template once for use later
# Note that because we use generate_name rather than name, this is reusable
self.pod_template = client.V1Pod(
metadata=client.V1ObjectMeta(
generate_name=worker_name_prefix,
# Convert comma-separated 'k=v' pairs to a dict
labels=dict(l.strip().split('=')
for l in self.worker_labels.split(',')
if l.strip())),
spec=client.V1PodSpec(
# Don't attempt to restart failed workers as this causes funny things
# to happen when dask kills workers legitimately
restart_policy='Never',
containers=[
client.V1Container(
name='dask-worker',
image=worker_image,
image_pull_policy=worker_image_pull_policy,
env=[
client.V1EnvVar(
name='POD_IP',
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path='status.podIP'))),
client.V1EnvVar(
name='POD_NAME',
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path='metadata.name'))),
],
args=[
'dask-worker',
dask_scheduler_service,
'--nprocs',
'1',
'--nthreads',
'1',
'--host',
'$(POD_IP)',
'--name',
'$(POD_NAME)',
])
]))
# If resource requests were given, add them to the pod template
if worker_resources:
resources = client.V1ResourceRequirements(**worker_resources)
self.pod_template.spec.containers[0].resources = resources
def create_deployment_old(config_file):
"""
Create IBM Spectrum Scale CSI Operator deployment object in operator namespace using
deployment_operator_image_for_crd and deployment_driver_image_for_crd parameters from
config.json file
Args:
param1: config_file - configuration json file
Returns:
None
Raises:
Raises an exception on kubernetes client api failure and asserts
"""
deployment_apps_api_instance = client.AppsV1Api()
deployment_labels = {
"app.kubernetes.io/instance": "ibm-spectrum-scale-csi-operator",
"app.kubernetes.io/managed-by": "ibm-spectrum-scale-csi-operator",
"app.kubernetes.io/name": "ibm-spectrum-scale-csi-operator",
"product": "ibm-spectrum-scale-csi",
"release": "ibm-spectrum-scale-csi-operator"
}
deployment_annotations = {
"productID": "ibm-spectrum-scale-csi-operator",
"productName": "IBM Spectrum Scale CSI Operator",
"productVersion": "2.0.0"
}
deployment_metadata = client.V1ObjectMeta(
name="ibm-spectrum-scale-csi-operator",
labels=deployment_labels,
namespace=namespace_value)
deployment_selector = client.V1LabelSelector(
match_labels={
"app.kubernetes.io/name": "ibm-spectrum-scale-csi-operator"
})
podtemplate_metadata = client.V1ObjectMeta(
labels=deployment_labels, annotations=deployment_annotations)
pod_affinity = client.V1Affinity(node_affinity=client.V1NodeAffinity(
required_during_scheduling_ignored_during_execution=client.
V1NodeSelector(node_selector_terms=[
client.V1NodeSelectorTerm(match_expressions=[
client.V1NodeSelectorRequirement(key="beta.kubernetes.io/arch",
operator="Exists")
])
])))
ansible_pod_container = client.V1Container(
image=config_file["deployment_operator_image_for_crd"],
command=[
"/usr/local/bin/ao-logs", "/tmp/ansible-operator/runner", "stdout"
],
liveness_probe=client.V1Probe(
_exec=client.V1ExecAction(command=["/health_check.sh"]),
initial_delay_seconds=10,
period_seconds=30),
readiness_probe=client.V1Probe(
_exec=client.V1ExecAction(command=["/health_check.sh"]),
initial_delay_seconds=3,
period_seconds=1),
name="ansible",
image_pull_policy="IfNotPresent",
security_context=client.V1SecurityContext(
capabilities=client.V1Capabilities(drop=["ALL"])),
volume_mounts=[
client.V1VolumeMount(mount_path="/tmp/ansible-operator/runner",
name="runner",
read_only=True)
],
env=[
client.V1EnvVar(
name="CSI_DRIVER_IMAGE",
value=config_file["deployment_driver_image_for_crd"])
])
operator_pod_container = client.V1Container(
image=config_file["deployment_operator_image_for_crd"],
name="operator",
image_pull_policy="IfNotPresent",
liveness_probe=client.V1Probe(
_exec=client.V1ExecAction(command=["/health_check.sh"]),
initial_delay_seconds=10,
period_seconds=30),
readiness_probe=client.V1Probe(
_exec=client.V1ExecAction(command=["/health_check.sh"]),
initial_delay_seconds=3,
period_seconds=1),
security_context=client.V1SecurityContext(
capabilities=client.V1Capabilities(drop=["ALL"])),
env=[
client.V1EnvVar(name="WATCH_NAMESPACE",
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path="metadata.namespace"))),
client.V1EnvVar(name="POD_NAME",
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path="metadata.name"))),
client.V1EnvVar(name="OPERATOR_NAME",
value="ibm-spectrum-scale-csi-operator"),
client.V1EnvVar(
name="CSI_DRIVER_IMAGE",
value=config_file["deployment_driver_image_for_crd"])
],
volume_mounts=[
client.V1VolumeMount(mount_path="/tmp/ansible-operator/runner",
name="runner")
])
pod_spec = client.V1PodSpec(
affinity=pod_affinity,
containers=[ansible_pod_container, operator_pod_container],
service_account_name="ibm-spectrum-scale-csi-operator",
volumes=[
client.V1Volume(
empty_dir=client.V1EmptyDirVolumeSource(medium="Memory"),
name="runner")
])
podtemplate_spec = client.V1PodTemplateSpec(metadata=podtemplate_metadata,
spec=pod_spec)
deployment_spec = client.V1DeploymentSpec(replicas=1,
selector=deployment_selector,
template=podtemplate_spec)
body_dep = client.V1Deployment(kind='Deployment',
api_version='apps/v1',
metadata=deployment_metadata,
spec=deployment_spec)
try:
LOGGER.info("creating deployment for operator")
deployment_apps_api_response = deployment_apps_api_instance.create_namespaced_deployment(
namespace=namespace_value, body=body_dep)
LOGGER.debug(str(deployment_apps_api_response))
except ApiException as e:
LOGGER.error(
f"Exception when calling RbacAuthorizationV1Api->create_namespaced_deployment: {e}"
)
assert False
def __new__(
cls,
component_name: Text,
input_dict: Dict[Text, Any],
output_dict: Dict[Text, List[types.TfxArtifact]],
exec_properties: Dict[Text, Any],
executor_class_path: Text,
pipeline_properties: PipelineProperties,
):
"""Creates a new component.
Args:
component_name: TFX component name.
input_dict: Dictionary of input names to TFX types, or
kfp.dsl.PipelineParam representing input parameters.
output_dict: Dictionary of output names to List of TFX types.
exec_properties: Execution properties.
executor_class_path: <module>.<class> for Python class of executor.
pipeline_properties: Pipeline level properties shared by all components.
Returns:
Newly constructed TFX Kubeflow component instance.
"""
outputs = output_dict.keys()
file_outputs = {
output: '/output/ml_metadata/{}'.format(output)
for output in outputs
}
for k, v in pipeline_properties.exec_properties.items():
exec_properties[k] = v
arguments = [
'--exec_properties',
json.dumps(exec_properties),
'--outputs',
types.jsonify_tfx_type_dict(output_dict),
'--executor_class_path',
executor_class_path,
component_name,
]
for k, v in input_dict.items():
if isinstance(v, float) or isinstance(v, int):
v = str(v)
arguments.append('--{}'.format(k))
arguments.append(v)
container_op = dsl.ContainerOp(
name=component_name,
command=_COMMAND,
image=pipeline_properties.tfx_image,
arguments=arguments,
file_outputs=file_outputs,
)
# Add the Argo workflow ID to the container's environment variable so it
# can be used to uniquely place pipeline outputs under the pipeline_root.
field_path = "metadata.labels['workflows.argoproj.io/workflow']"
container_op.add_env_variable(
k8s_client.V1EnvVar(name='WORKFLOW_ID',
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=field_path))))
named_outputs = {
output: container_op.outputs[output]
for output in outputs
}
# This allows user code to refer to the ContainerOp 'op' output named 'x'
# as op.outputs.x
component_outputs = type('Output', (), named_outputs)
return type(component_name, (BaseComponent, ), {
'container_op': container_op,
'outputs': component_outputs
})
def export_deployment(self):
# Configureate Pod template container
volume_mounts = []
containers = []
volumes = []
ports = []
liveness_probe = None
readiness_probe = None
volume_mounts.append(
client.V1VolumeMount(mount_path='/docker/logs', name='logs'))
volumes.append(
client.V1Volume(name='logs',
host_path=client.V1HostPathVolumeSource(
path='/opt/logs', type='DirectoryOrCreate')))
if self.mounts:
for path in self.mounts:
volume_mounts.append(
client.V1VolumeMount(mount_path=path,
name=self.mounts[path]))
volumes.append(
client.V1Volume(name=self.mounts[path],
host_path=client.V1HostPathVolumeSource(
path=path, type='DirectoryOrCreate')))
if self.container_port:
ports = [
client.V1ContainerPort(container_port=int(port))
for port in self.container_port
]
liveness_probe = client.V1Probe(
initial_delay_seconds=15,
tcp_socket=client.V1TCPSocketAction(
port=int(self.container_port[0])))
readiness_probe = client.V1Probe(
initial_delay_seconds=15,
tcp_socket=client.V1TCPSocketAction(
port=int(self.container_port[0])))
if self.healthcheck:
liveness_probe = client.V1Probe(
initial_delay_seconds=15,
http_get=client.V1HTTPGetAction(
path=self.healthcheck,
port=int(self.container_port[0])))
readiness_probe = client.V1Probe(
initial_delay_seconds=15,
http_get=client.V1HTTPGetAction(
path=self.healthcheck,
port=int(self.container_port[0])))
Env = [
client.V1EnvVar(name='LANG', value='en_US.UTF-8'),
client.V1EnvVar(name='LC_ALL', value='en_US.UTF-8'),
client.V1EnvVar(name='POD_NAME',
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path='metadata.name'))),
client.V1EnvVar(name='POD_IP',
value_from=client.V1EnvVarSource(
field_ref=client.V1ObjectFieldSelector(
field_path='status.podIP'))),
]
container = client.V1Container(name=self.dm_name,
image=self.image,
ports=ports,
image_pull_policy='Always',
env=Env,
resources=client.V1ResourceRequirements(
limits=self.re_limits,
requests=self.re_requests),
volume_mounts=volume_mounts)
if liveness_probe and readiness_probe:
container = client.V1Container(
name=self.dm_name,
image=self.image,
ports=ports,
image_pull_policy='Always',
env=Env,
resources=client.V1ResourceRequirements(
limits=self.re_limits, requests=self.re_requests),
volume_mounts=volume_mounts,
liveness_probe=liveness_probe,
readiness_probe=readiness_probe)
containers.append(container)
if self.sidecar:
sidecar_container = client.V1Container(
name='sidecar-%s' % self.dm_name,
image=self.sidecar,
image_pull_policy='Always',
env=Env,
resources=client.V1ResourceRequirements(
limits=self.re_limits, requests=self.re_requests),
volume_mounts=volume_mounts)
containers.append(sidecar_container)
# Create and configurate a spec section
secrets = client.V1LocalObjectReference('registrysecret')
preference_key = self.dm_name
project_values = ['xxxx']
host_aliases = []
db_docker_hosts = db_op.docker_hosts
values = db_docker_hosts.query.with_entities(
db_docker_hosts.ip, db_docker_hosts.hostname).filter(
and_(db_docker_hosts.deployment == self.dm_name,
db_docker_hosts.context == self.context)).all()
db_op.DB.session.remove()
if values:
ips = []
for value in values:
try:
ip, hostname = value
key = "op_docker_hosts_%s" % ip
Redis.lpush(key, hostname)
ips.append(ip)
except Exception as e:
logging.error(e)
for ip in set(ips):
try:
key = "op_docker_hosts_%s" % ip
if Redis.exists(key):
hostnames = Redis.lrange(key, 0, -1)
if hostnames:
host_aliases.append(
client.V1HostAlias(hostnames=hostnames, ip=ip))
Redis.delete(key)
except Exception as e:
logging.error(e)
if self.labels:
if 'deploy' in self.labels:
preference_key = self.labels['deploy']
if 'project' in self.labels:
project_values = [self.labels['project']]
template = client.V1PodTemplateSpec(
metadata=client.V1ObjectMeta(labels={"project": self.dm_name}),
spec=client.V1PodSpec(
containers=containers,
image_pull_secrets=[secrets],
volumes=volumes,
host_aliases=host_aliases,
affinity=client.V1Affinity(node_affinity=client.V1NodeAffinity(
preferred_during_scheduling_ignored_during_execution=[
client.V1PreferredSchedulingTerm(
preference=client.V1NodeSelectorTerm(
match_expressions=[
client.V1NodeSelectorRequirement(
key=preference_key,
operator='In',
values=['mark'])
]),
weight=100)
],
required_during_scheduling_ignored_during_execution=client.
V1NodeSelector(node_selector_terms=[
client.V1NodeSelectorTerm(match_expressions=[
client.V1NodeSelectorRequirement(
key='project',
operator='In',
values=project_values)
])
])))))
selector = client.V1LabelSelector(
match_labels={"project": self.dm_name})
# Create the specification of deployment
spec = client.ExtensionsV1beta1DeploymentSpec(replicas=int(
self.replicas),
template=template,
selector=selector,
min_ready_seconds=3)
# Instantiate the deployment object
deployment = client.ExtensionsV1beta1Deployment(
api_version="extensions/v1beta1",
kind="Deployment",
metadata=client.V1ObjectMeta(name=self.dm_name),
spec=spec)
return deployment
def __init__(self,
component: tfx_base_node.BaseNode,
depends_on: Set[dsl.ContainerOp],
pipeline: tfx_pipeline.Pipeline,
pipeline_root: dsl.PipelineParam,
tfx_image: str,
kubeflow_metadata_config: kubeflow_pb2.KubeflowMetadataConfig,
tfx_ir: pipeline_pb2.Pipeline,
pod_labels_to_attach: Dict[str, str],
runtime_parameters: List[data_types.RuntimeParameter],
metadata_ui_path: str = '/mlpipeline-ui-metadata.json'):
"""Creates a new Kubeflow-based component.
This class essentially wraps a dsl.ContainerOp construct in Kubeflow
Pipelines.
Args:
component: The logical TFX component to wrap.
depends_on: The set of upstream KFP ContainerOp components that this
component will depend on.
pipeline: The logical TFX pipeline to which this component belongs.
pipeline_root: The pipeline root specified, as a dsl.PipelineParam
tfx_image: The container image to use for this component.
kubeflow_metadata_config: Configuration settings for connecting to the
MLMD store in a Kubeflow cluster.
tfx_ir: The TFX intermedia representation of the pipeline.
pod_labels_to_attach: Dict of pod labels to attach to the GKE pod.
runtime_parameters: Runtime parameters of the pipeline.
metadata_ui_path: File location for metadata-ui-metadata.json file.
"""
utils.replace_placeholder(component)
arguments = [
'--pipeline_root',
pipeline_root,
'--kubeflow_metadata_config',
json_format.MessageToJson(message=kubeflow_metadata_config,
preserving_proto_field_name=True),
'--node_id',
component.id,
# TODO(b/182220464): write IR to pipeline_root and let
# container_entrypoint.py read it back to avoid future issue that IR
# exeeds the flag size limit.
'--tfx_ir',
json_format.MessageToJson(tfx_ir),
'--metadata_ui_path',
metadata_ui_path,
]
for param in runtime_parameters:
arguments.append('--runtime_parameter')
arguments.append(_encode_runtime_parameter(param))
self.container_op = dsl.ContainerOp(
name=component.id,
command=_COMMAND,
image=tfx_image,
arguments=arguments,
output_artifact_paths={
'mlpipeline-ui-metadata': metadata_ui_path,
},
)
logging.info('Adding upstream dependencies for component %s',
self.container_op.name)
for op in depends_on:
logging.info(' -> Component: %s', op.name)
self.container_op.after(op)
# TODO(b/140172100): Document the use of additional_pipeline_args.
if _WORKFLOW_ID_KEY in pipeline.additional_pipeline_args:
# Allow overriding pipeline's run_id externally, primarily for testing.
self.container_op.container.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value=pipeline.additional_pipeline_args[_WORKFLOW_ID_KEY]))
else:
# Add the Argo workflow ID to the container's environment variable so it
# can be used to uniquely place pipeline outputs under the pipeline_root.
field_path = "metadata.labels['workflows.argoproj.io/workflow']"
self.container_op.container.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=field_path))))
if pod_labels_to_attach:
for k, v in pod_labels_to_attach.items():
self.container_op.add_pod_label(k, v)
def __init__(self,
component: tfx_base_node.BaseNode,
component_launcher_class: Type[
base_component_launcher.BaseComponentLauncher],
depends_on: Set[dsl.ContainerOp],
pipeline: tfx_pipeline.Pipeline,
pipeline_name: Text,
pipeline_root: dsl.PipelineParam,
tfx_image: Text,
kubeflow_metadata_config: Optional[
kubeflow_pb2.KubeflowMetadataConfig],
component_config: base_component_config.BaseComponentConfig,
pod_labels_to_attach: Optional[Dict[Text, Text]] = None):
"""Creates a new Kubeflow-based component.
This class essentially wraps a dsl.ContainerOp construct in Kubeflow
Pipelines.
Args:
component: The logical TFX component to wrap.
component_launcher_class: the class of the launcher to launch the
component.
depends_on: The set of upstream KFP ContainerOp components that this
component will depend on.
pipeline: The logical TFX pipeline to which this component belongs.
pipeline_name: The name of the TFX pipeline.
pipeline_root: The pipeline root specified, as a dsl.PipelineParam
tfx_image: The container image to use for this component.
kubeflow_metadata_config: Configuration settings for connecting to the
MLMD store in a Kubeflow cluster.
component_config: Component config to launch the component.
pod_labels_to_attach: Optional dict of pod labels to attach to the
GKE pod.
"""
component_launcher_class_path = '.'.join([
component_launcher_class.__module__,
component_launcher_class.__name__
])
serialized_component = utils.replace_placeholder(
json_utils.dumps(node_wrapper.NodeWrapper(component)))
arguments = [
'--pipeline_name',
pipeline_name,
'--pipeline_root',
pipeline_root,
'--kubeflow_metadata_config',
json_format.MessageToJson(message=kubeflow_metadata_config,
preserving_proto_field_name=True),
'--beam_pipeline_args',
json.dumps(pipeline.beam_pipeline_args),
'--additional_pipeline_args',
json.dumps(pipeline.additional_pipeline_args),
'--component_launcher_class_path',
component_launcher_class_path,
'--serialized_component',
serialized_component,
'--component_config',
json_utils.dumps(component_config),
]
if component.enable_cache or (component.enable_cache is None
and pipeline.enable_cache):
arguments.append('--enable_cache')
self.container_op = dsl.ContainerOp(
name=component.id.replace('.', '_'),
command=_COMMAND,
image=tfx_image,
arguments=arguments,
output_artifact_paths={
'mlpipeline-ui-metadata': '/mlpipeline-ui-metadata.json',
},
)
absl.logging.info(
'Adding upstream dependencies for component {}'.format(
self.container_op.name))
for op in depends_on:
absl.logging.info(' -> Component: {}'.format(op.name))
self.container_op.after(op)
# TODO(b/140172100): Document the use of additional_pipeline_args.
if _WORKFLOW_ID_KEY in pipeline.additional_pipeline_args:
# Allow overriding pipeline's run_id externally, primarily for testing.
self.container_op.container.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value=pipeline.additional_pipeline_args[_WORKFLOW_ID_KEY]))
else:
# Add the Argo workflow ID to the container's environment variable so it
# can be used to uniquely place pipeline outputs under the pipeline_root.
field_path = "metadata.labels['workflows.argoproj.io/workflow']"
self.container_op.container.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=field_path))))
if pod_labels_to_attach:
for k, v in pod_labels_to_attach.items():
self.container_op.add_pod_label(k, v)
class TestCastValue:
"""Tests for kubetest.manifest.cast_value"""
@pytest.mark.parametrize(
'value,t,expected',
[
# builtin types
(11, 'int', int(11)),
('11', 'int', int(11)),
(11.0, 'int', int(11)),
(11, 'float', float(11)),
(11, 'str', '11'),
# casting to object should result in no change
(11, 'object', 11),
('11', 'object', '11'),
# kubernetes types
({
'apiVersion': 'apps/v1',
'kind': 'Namespace'
}, 'V1Namespace',
client.V1Namespace(kind='Namespace', api_version='apps/v1')),
({
'fieldRef': {
'apiVersion': 'apps/v1beta1',
'fieldPath': 'foobar'
}
}, 'V1EnvVarSource',
client.V1EnvVarSource(field_ref=client.V1ObjectFieldSelector(
api_version='apps/v1beta1', field_path='foobar'))),
({
'finalizers': ['a', 'b', 'c']
}, 'V1ObjectMeta',
client.V1ObjectMeta(finalizers=['a', 'b', 'c'])),
])
def test_ok(self, value, t, expected):
"""Test casting values to the specified type successfully."""
actual = manifest.cast_value(value, t)
assert type(actual) == type(expected)
assert actual == expected
@pytest.mark.parametrize(
'value,t,error',
[
# builtin types
({
'foo': 'bar'
}, 'int', TypeError),
([1, 3, 5], 'float', TypeError),
(1.0, 'set', TypeError),
# kubernetes types
(11, 'V1Namespace', AttributeError),
('foo', 'V1Deployment', AttributeError),
(['a', 'b', 'c'], 'V1Service', AttributeError),
({1, 2, 3, 4}, 'V1Pod', AttributeError),
# unknown type
(11, 'NotARealType', ValueError),
])
def test_error(self, value, t, error):
"""Test casting values to the specified type unsuccessfully."""
with pytest.raises(error):
manifest.cast_value(value, t)
def __init__(
self,
component: tfx_base_component.BaseComponent,
component_launcher_class: Type[
base_component_launcher.BaseComponentLauncher],
depends_on: Set[dsl.ContainerOp],
pipeline: tfx_pipeline.Pipeline,
tfx_image: Text,
kubeflow_metadata_config: Optional[
kubeflow_pb2.KubeflowMetadataConfig],
):
"""Creates a new Kubeflow-based component.
This class essentially wraps a dsl.ContainerOp construct in Kubeflow
Pipelines.
Args:
component: The logical TFX component to wrap.
component_launcher_class: the class of the launcher to launch the
component.
depends_on: The set of upstream KFP ContainerOp components that this
component will depend on.
pipeline: The logical TFX pipeline to which this component belongs.
tfx_image: The container image to use for this component.
kubeflow_metadata_config: Configuration settings for
connecting to the MLMD store in a Kubeflow cluster.
"""
driver_class_path = '.'.join([
component.driver_class.__module__, component.driver_class.__name__
])
executor_spec = json_utils.dumps(component.executor_spec)
component_launcher_class_path = '.'.join([
component_launcher_class.__module__,
component_launcher_class.__name__
])
arguments = [
'--pipeline_name',
pipeline.pipeline_info.pipeline_name,
'--pipeline_root',
pipeline.pipeline_info.pipeline_root,
'--kubeflow_metadata_config',
json_format.MessageToJson(kubeflow_metadata_config),
'--additional_pipeline_args',
json.dumps(pipeline.additional_pipeline_args),
'--component_id',
component.component_id,
'--component_type',
component.component_type,
'--driver_class_path',
driver_class_path,
'--executor_spec',
executor_spec,
'--component_launcher_class_path',
component_launcher_class_path,
'--inputs',
artifact_utils.jsonify_artifact_dict(
_prepare_artifact_dict(component.inputs)),
'--outputs',
artifact_utils.jsonify_artifact_dict(
_prepare_artifact_dict(component.outputs)),
'--exec_properties',
json.dumps(component.exec_properties),
]
if pipeline.enable_cache:
arguments.append('--enable_cache')
self.container_op = dsl.ContainerOp(
name=component.component_id.replace('.', '_'),
command=_COMMAND,
image=tfx_image,
arguments=arguments,
)
tf.logging.info('Adding upstream dependencies for component {}'.format(
self.container_op.name))
for op in depends_on:
tf.logging.info(' -> Component: {}'.format(op.name))
self.container_op.after(op)
# TODO(b/140172100): Document the use of additional_pipeline_args.
if _WORKFLOW_ID_KEY in pipeline.additional_pipeline_args:
# Allow overriding pipeline's run_id externally, primarily for testing.
self.container_op.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value=pipeline.additional_pipeline_args[_WORKFLOW_ID_KEY]))
else:
# Add the Argo workflow ID to the container's environment variable so it
# can be used to uniquely place pipeline outputs under the pipeline_root.
field_path = "metadata.labels['workflows.argoproj.io/workflow']"
self.container_op.add_env_variable(
k8s_client.V1EnvVar(
name=_WORKFLOW_ID_KEY,
value_from=k8s_client.V1EnvVarSource(
field_ref=k8s_client.V1ObjectFieldSelector(
field_path=field_path))))
def test_sanitize_k8s_container_attribute(self):
# test cases for implicit type sanitization(conversion)
op = dsl.ContainerOp(name='echo', image='image', command=['sh', '-c'],
arguments=['echo test | tee /tmp/message.txt'],
file_outputs={'merged': '/tmp/message.txt'})
op.container \
.add_volume_mount(k8s_client.V1VolumeMount(
mount_path='/secret/gcp-credentials',
name='gcp-credentials')) \
.add_env_variable(k8s_client.V1EnvVar(
name=80,
value=80)) \
.add_env_variable(k8s_client.V1EnvVar(
name=80,
value_from=k8s_client.V1EnvVarSource(
config_map_key_ref=k8s_client.V1ConfigMapKeySelector(key=80, name=8080, optional='False'),
field_ref=k8s_client.V1ObjectFieldSelector(api_version=80, field_path=8080),
resource_field_ref=k8s_client.V1ResourceFieldSelector(container_name=80, divisor=8080, resource=8888),
secret_key_ref=k8s_client.V1SecretKeySelector(key=80, name=8080, optional='False')
)
)) \
.add_env_from(k8s_client.V1EnvFromSource(
config_map_ref=k8s_client.V1ConfigMapEnvSource(name=80, optional='True'),
prefix=999
)) \
.add_env_from(k8s_client.V1EnvFromSource(
secret_ref=k8s_client.V1SecretEnvSource(name=80, optional='True'),
prefix=888
)) \
.add_volume_mount(k8s_client.V1VolumeMount(
mount_path=111,
mount_propagation=222,
name=333,
read_only='False',
sub_path=444,
sub_path_expr=555
)) \
.add_volume_devices(k8s_client.V1VolumeDevice(
device_path=111,
name=222
)) \
.add_port(k8s_client.V1ContainerPort(
container_port='8080',
host_ip=111,
host_port='8888',
name=222,
protocol=333
)) \
.set_security_context(k8s_client.V1SecurityContext(
allow_privilege_escalation='True',
capabilities=k8s_client.V1Capabilities(add=[11, 22], drop=[33, 44]),
privileged='False',
proc_mount=111,
read_only_root_filesystem='False',
run_as_group='222',
run_as_non_root='True',
run_as_user='******',
se_linux_options=k8s_client.V1SELinuxOptions(level=11, role=22, type=33, user=44),
windows_options=k8s_client.V1WindowsSecurityContextOptions(
gmsa_credential_spec=11, gmsa_credential_spec_name=22)
)) \
.set_stdin(stdin='False') \
.set_stdin_once(stdin_once='False') \
.set_termination_message_path(termination_message_path=111) \
.set_tty(tty='False') \
.set_readiness_probe(readiness_probe=k8s_client.V1Probe(
_exec=k8s_client.V1ExecAction(command=[11, 22, 33]),
failure_threshold='111',
http_get=k8s_client.V1HTTPGetAction(
host=11,
http_headers=[k8s_client.V1HTTPHeader(name=22, value=33)],
path=44,
port='55',
scheme=66),
initial_delay_seconds='222',
period_seconds='333',
success_threshold='444',
tcp_socket=k8s_client.V1TCPSocketAction(host=555, port='666'),
timeout_seconds='777'
)) \
.set_liveness_probe(liveness_probe=k8s_client.V1Probe(
_exec=k8s_client.V1ExecAction(command=[11, 22, 33]),
failure_threshold='111',
http_get=k8s_client.V1HTTPGetAction(
host=11,
http_headers=[k8s_client.V1HTTPHeader(name=22, value=33)],
path=44,
port='55',
scheme=66),
initial_delay_seconds='222',
period_seconds='333',
success_threshold='444',
tcp_socket=k8s_client.V1TCPSocketAction(host=555, port='666'),
timeout_seconds='777'
)) \
.set_lifecycle(lifecycle=k8s_client.V1Lifecycle(
post_start=k8s_client.V1Handler(
_exec=k8s_client.V1ExecAction(command=[11, 22, 33]),
http_get=k8s_client.V1HTTPGetAction(
host=11,
http_headers=[k8s_client.V1HTTPHeader(name=22, value=33)],
path=44,
port='55',
scheme=66),
tcp_socket=k8s_client.V1TCPSocketAction(host=555, port='666')
),
pre_stop=k8s_client.V1Handler(
_exec=k8s_client.V1ExecAction(command=[11, 22, 33]),
http_get=k8s_client.V1HTTPGetAction(
host=11,
http_headers=[k8s_client.V1HTTPHeader(name=22, value=33)],
path=44,
port='55',
scheme=66),
tcp_socket=k8s_client.V1TCPSocketAction(host=555, port='666')
)
))
sanitize_k8s_object(op.container)
for e in op.container.env:
self.assertIsInstance(e.name, str)
if e.value:
self.assertIsInstance(e.value, str)
if e.value_from:
if e.value_from.config_map_key_ref:
self.assertIsInstance(e.value_from.config_map_key_ref.key, str)
if e.value_from.config_map_key_ref.name:
self.assertIsInstance(e.value_from.config_map_key_ref.name, str)
if e.value_from.config_map_key_ref.optional:
self.assertIsInstance(e.value_from.config_map_key_ref.optional, bool)
if e.value_from.field_ref:
self.assertIsInstance(e.value_from.field_ref.field_path, str)
if e.value_from.field_ref.api_version:
self.assertIsInstance(e.value_from.field_ref.api_version, str)
if e.value_from.resource_field_ref:
self.assertIsInstance(e.value_from.resource_field_ref.resource, str)
if e.value_from.resource_field_ref.container_name:
self.assertIsInstance(e.value_from.resource_field_ref.container_name, str)
if e.value_from.resource_field_ref.divisor:
self.assertIsInstance(e.value_from.resource_field_ref.divisor, str)
if e.value_from.secret_key_ref:
self.assertIsInstance(e.value_from.secret_key_ref.key, str)
if e.value_from.secret_key_ref.name:
self.assertIsInstance(e.value_from.secret_key_ref.name, str)
if e.value_from.secret_key_ref.optional:
self.assertIsInstance(e.value_from.secret_key_ref.optional, bool)
for e in op.container.env_from:
if e.prefix:
self.assertIsInstance(e.prefix, str)
if e.config_map_ref:
if e.config_map_ref.name:
self.assertIsInstance(e.config_map_ref.name, str)
if e.config_map_ref.optional:
self.assertIsInstance(e.config_map_ref.optional, bool)
if e.secret_ref:
if e.secret_ref.name:
self.assertIsInstance(e.secret_ref.name, str)
if e.secret_ref.optional:
self.assertIsInstance(e.secret_ref.optional, bool)
for e in op.container.volume_mounts:
if e.mount_path:
self.assertIsInstance(e.mount_path, str)
if e.mount_propagation:
self.assertIsInstance(e.mount_propagation, str)
if e.name:
self.assertIsInstance(e.name, str)
if e.read_only:
self.assertIsInstance(e.read_only, bool)
if e.sub_path:
self.assertIsInstance(e.sub_path, str)
if e.sub_path_expr:
self.assertIsInstance(e.sub_path_expr, str)
for e in op.container.volume_devices:
if e.device_path:
self.assertIsInstance(e.device_path, str)
if e.name:
self.assertIsInstance(e.name, str)
for e in op.container.ports:
if e.container_port:
self.assertIsInstance(e.container_port, int)
if e.host_ip:
self.assertIsInstance(e.host_ip, str)
if e.host_port:
self.assertIsInstance(e.host_port, int)
if e.name:
self.assertIsInstance(e.name, str)
if e.protocol:
self.assertIsInstance(e.protocol, str)
if op.container.security_context:
e = op.container.security_context
if e.allow_privilege_escalation:
self.assertIsInstance(e.allow_privilege_escalation, bool)
if e.capabilities:
for a in e.capabilities.add:
self.assertIsInstance(a, str)
for d in e.capabilities.drop:
self.assertIsInstance(d, str)
if e.privileged:
self.assertIsInstance(e.privileged, bool)
if e.proc_mount:
self.assertIsInstance(e.proc_mount, str)
if e.read_only_root_filesystem:
self.assertIsInstance(e.read_only_root_filesystem, bool)
if e.run_as_group:
self.assertIsInstance(e.run_as_group, int)
if e.run_as_non_root:
self.assertIsInstance(e.run_as_non_root, bool)
if e.run_as_user:
self.assertIsInstance(e.run_as_user, int)
if e.se_linux_options:
if e.se_linux_options.level:
self.assertIsInstance(e.se_linux_options.level, str)
if e.se_linux_options.role:
self.assertIsInstance(e.se_linux_options.role, str)
if e.se_linux_options.type:
self.assertIsInstance(e.se_linux_options.type, str)
if e.se_linux_options.user:
self.assertIsInstance(e.se_linux_options.user, str)
if e.windows_options:
if e.windows_options.gmsa_credential_spec:
self.assertIsInstance(e.windows_options.gmsa_credential_spec, str)
if e.windows_options.gmsa_credential_spec_name:
self.assertIsInstance(e.windows_options.gmsa_credential_spec_name, str)
if op.container.stdin:
self.assertIsInstance(op.container.stdin, bool)
if op.container.stdin_once:
self.assertIsInstance(op.container.stdin_once, bool)
if op.container.termination_message_path:
self.assertIsInstance(op.container.termination_message_path, str)
if op.container.tty:
self.assertIsInstance(op.container.tty, bool)
for e in [op.container.readiness_probe, op.container.liveness_probe]:
if e:
if e._exec:
for c in e._exec.command:
self.assertIsInstance(c, str)
if e.failure_threshold:
self.assertIsInstance(e.failure_threshold, int)
if e.http_get:
if e.http_get.host:
self.assertIsInstance(e.http_get.host, str)
if e.http_get.http_headers:
for h in e.http_get.http_headers:
if h.name:
self.assertIsInstance(h.name, str)
if h.value:
self.assertIsInstance(h.value, str)
if e.http_get.path:
self.assertIsInstance(e.http_get.path, str)
if e.http_get.port:
self.assertIsInstance(e.http_get.port, (str, int))
if e.http_get.scheme:
self.assertIsInstance(e.http_get.scheme, str)
if e.initial_delay_seconds:
self.assertIsInstance(e.initial_delay_seconds, int)
if e.period_seconds:
self.assertIsInstance(e.period_seconds, int)
if e.success_threshold:
self.assertIsInstance(e.success_threshold, int)
if e.tcp_socket:
if e.tcp_socket.host:
self.assertIsInstance(e.tcp_socket.host, str)
if e.tcp_socket.port:
self.assertIsInstance(e.tcp_socket.port, (str, int))
if e.timeout_seconds:
self.assertIsInstance(e.timeout_seconds, int)
if op.container.lifecycle:
for e in [op.container.lifecycle.post_start, op.container.lifecycle.pre_stop]:
if e:
if e._exec:
for c in e._exec.command:
self.assertIsInstance(c, str)
if e.http_get:
if e.http_get.host:
self.assertIsInstance(e.http_get.host, str)
if e.http_get.http_headers:
for h in e.http_get.http_headers:
if h.name:
self.assertIsInstance(h.name, str)
if h.value:
self.assertIsInstance(h.value, str)
if e.http_get.path:
self.assertIsInstance(e.http_get.path, str)
if e.http_get.port:
self.assertIsInstance(e.http_get.port, (str, int))
if e.http_get.scheme:
self.assertIsInstance(e.http_get.scheme, str)
if e.tcp_socket:
if e.tcp_socket.host:
self.assertIsInstance(e.tcp_socket.host, str)
if e.tcp_socket.port:
self.assertIsInstance(e.tcp_socket.port, (str, int))
# test cases for checking value after sanitization
check_value_op = dsl.ContainerOp(name='echo', image='image', command=['sh', '-c'],
arguments=['echo test | tee /tmp/message.txt'],
file_outputs={'merged': '/tmp/message.txt'})
check_value_op.container \
.add_env_variable(k8s_client.V1EnvVar(
name=80,
value=8080)) \
.set_security_context(k8s_client.V1SecurityContext(
allow_privilege_escalation='true',
capabilities=k8s_client.V1Capabilities(add=[11, 22], drop=[33, 44]),
privileged='false',
proc_mount=111,
read_only_root_filesystem='False',
run_as_group='222',
run_as_non_root='True',
run_as_user='******',
se_linux_options=k8s_client.V1SELinuxOptions(level=11, role=22, type=33, user=44),
windows_options=k8s_client.V1WindowsSecurityContextOptions(
gmsa_credential_spec=11, gmsa_credential_spec_name=22)
))
sanitize_k8s_object(check_value_op.container)
self.assertEqual(check_value_op.container.env[0].name, '80')
self.assertEqual(check_value_op.container.env[0].value, '8080')
self.assertEqual(check_value_op.container.security_context.allow_privilege_escalation, True)
self.assertEqual(check_value_op.container.security_context.capabilities.add[0], '11')
self.assertEqual(check_value_op.container.security_context.capabilities.add[1], '22')
self.assertEqual(check_value_op.container.security_context.capabilities.drop[0], '33')
self.assertEqual(check_value_op.container.security_context.capabilities.drop[1], '44')
self.assertEqual(check_value_op.container.security_context.privileged, False)
self.assertEqual(check_value_op.container.security_context.proc_mount, '111')
self.assertEqual(check_value_op.container.security_context.read_only_root_filesystem, False)
self.assertEqual(check_value_op.container.security_context.run_as_group, 222)
self.assertEqual(check_value_op.container.security_context.run_as_non_root, True)
self.assertEqual(check_value_op.container.security_context.run_as_user, 333)
self.assertEqual(check_value_op.container.security_context.se_linux_options.level, '11')
self.assertEqual(check_value_op.container.security_context.se_linux_options.role, '22')
self.assertEqual(check_value_op.container.security_context.se_linux_options.type, '33')
self.assertEqual(check_value_op.container.security_context.se_linux_options.user, '44')
self.assertEqual(check_value_op.container.security_context.windows_options.gmsa_credential_spec, '11')
self.assertEqual(check_value_op.container.security_context.windows_options.gmsa_credential_spec_name, '22')
# test cases for exception
with self.assertRaises(ValueError, msg='Invalid boolean string 2. Should be boolean.'):
exception_op = dsl.ContainerOp(name='echo', image='image')
exception_op.container \
.set_security_context(k8s_client.V1SecurityContext(
allow_privilege_escalation=1
))
sanitize_k8s_object(exception_op.container)
with self.assertRaises(ValueError, msg='Invalid boolean string Test. Should be "true" or "false".'):
exception_op = dsl.ContainerOp(name='echo', image='image')
exception_op.container \
.set_security_context(k8s_client.V1SecurityContext(
allow_privilege_escalation='Test'
))
sanitize_k8s_object(exception_op.container)
with self.assertRaises(ValueError, msg='Invalid test. Should be integer.'):
exception_op = dsl.ContainerOp(name='echo', image='image')
exception_op.container \
.set_security_context(k8s_client.V1SecurityContext(
run_as_group='test',
))
sanitize_k8s_object(exception_op.container)