Пример #1
0
def step_repr(
    step_name=None,
    tmpl_name=None,
    image=None,
    command=None,
    source=None,
    env=None,
    script_output=None,
    args=None,
    input=None,
    output=None,
    manifest=None,
    success_cond=None,
    failure_cond=None,
    canned_step_name=None,
    canned_step_args=None,
    resources=None,
    secret=None,
    action=None,
    volume_mounts=None,
    cache=None,
):
    assert step_name is not None
    assert tmpl_name is not None
    # generate protobuf step representation
    pb_step = couler_pb2.Step()
    pb_step.id = get_uniq_step_id()
    pb_step.name = step_name
    pb_step.tmpl_name = tmpl_name

    if env is not None:
        _add_env_to_step(pb_step, env)

    if secret is not None:
        for k, _ in secret.data.items():
            pb_secret = pb_step.secrets.add()
            pb_secret.key = k
            pb_secret.name = secret.name

    if cache is not None:
        pb_step.cache.name = cache.name
        pb_step.cache.key = cache.key
        pb_step.cache.max_age = cache.max_age

    # image can be None if manifest specified.
    if image is not None:
        pb_step.container_spec.image = image
    if manifest is not None:
        pb_step.resource_spec.manifest = manifest
        if success_cond is not None and failure_cond is not None:
            pb_step.resource_spec.success_condition = success_cond
            pb_step.resource_spec.failure_condition = failure_cond
        if action is not None:
            pb_step.resource_spec.action = action
    else:
        if command is None:
            pb_step.container_spec.command.append("python")
        elif isinstance(command, list):
            pb_step.container_spec.command.extend(command)
        elif isinstance(command, str):
            pb_step.container_spec.command.append(command)
        else:
            raise ValueError("command must be str or list")
    if source is not None:
        if isinstance(source, str):
            pb_step.script = source
        else:
            pb_step.script = utils.body(source)
    if canned_step_name is not None and canned_step_args is not None:
        pb_step.canned_step_spec.name = canned_step_name
        if isinstance(canned_step_args, dict):
            for k, v in canned_step_args.items():
                pb_step.canned_step_spec.args[k] = v
        else:
            raise ValueError("canned_step_spec.args must be a dictionary")

    # setup resources
    if resources is not None:
        if not isinstance(resources, dict):
            raise ValueError("resources must be type dict")
        for k, v in resources.items():
            # key: cpu, memory, gpu
            # value: "1", "8", "500m", "1Gi" etc.
            pb_step.container_spec.resources[k] = str(v)

    # Attach volume mounts
    if volume_mounts is not None:
        for vm in volume_mounts:
            vol_mount = pb_step.container_spec.volume_mounts.add()
            vol_mount.name = vm.name
            vol_mount.path = vm.mount_path

    if states._when_prefix is not None:
        pb_step.when = states._when_prefix

    # add template to proto workflow
    wf = get_default_proto_workflow()
    if tmpl_name not in wf.templates:
        proto_step_tmpl = couler_pb2.StepTemplate()
        proto_step_tmpl.name = tmpl_name
        _add_io_to_template(proto_step_tmpl, pb_step.id, input, output,
                            script_output)
        wf.templates[tmpl_name].CopyFrom(proto_step_tmpl)

    # add step arguments
    if args is not None:
        for i, arg in enumerate(args):
            if isinstance(arg, OutputArtifact):
                pb_art = pb_step.args.add()
                pb_art.artifact.name = arg.artifact["name"]
                pb_art.artifact.value = arg.value
            else:
                pb_param = pb_step.args.add()
                pb_param.parameter.name = utils.input_parameter_name(
                    pb_step.name, i)
                pb_param.parameter.value = (
                    # This is casted to string for protobuf and Argo
                    # will automatically convert it to a correct type.
                    str(arg) if isinstance(arg, (str, float, bool,
                                                 int)) else arg.value)

    if states._exit_handler_enable:
        # add exit handler steps
        eh_step = wf.exit_handler_steps.add()
        eh_step.CopyFrom(pb_step)
    else:
        # add step to proto workflow
        concurrent_step = wf.steps.add()
        inner_step = concurrent_step.steps.add()
        inner_step.CopyFrom(pb_step)
    return pb_step
Пример #2
0
def step_repr(
    step_name=None,
    tmpl_name=None,
    image=None,
    command=None,
    source=None,
    env=None,
    script_output=None,
    args=None,
    input=None,
    output=None,
    manifest=None,
    success_cond=None,
    failure_cond=None,
    canned_step_name=None,
    canned_step_args=None,
    resources=None,
):
    assert step_name is not None
    assert tmpl_name is not None
    # generate protobuf step representation
    pb_step = couler_pb2.Step()
    pb_step.id = get_uniq_step_id()
    pb_step.name = step_name
    pb_step.tmpl_name = tmpl_name
    if env is not None:
        for k, v in env.items():
            if isinstance(v, str):
                pb_step.container_spec.env[k] = v
            else:
                pb_step.container_spec.env[k] = json.dumps(v)

    # image can be None if manifest specified.
    if image is not None:
        pb_step.container_spec.image = image
    if manifest is not None:
        pb_step.resource_spec.manifest = manifest
        if success_cond is not None and failure_cond is not None:
            pb_step.resource_spec.success_condition = success_cond
            pb_step.resource_spec.failure_condition = failure_cond
    else:
        if command is None:
            pb_step.container_spec.command.append("python")
        elif isinstance(command, list):
            pb_step.container_spec.command.extend(command)
        elif isinstance(command, str):
            pb_step.container_spec.command.append(command)
        else:
            raise ValueError("command must be str or list")
    if source is not None:
        if isinstance(source, str):
            pb_step.script = source
        else:
            pb_step.script = utils.body(source)
    if canned_step_name is not None and canned_step_args is not None:
        pb_step.canned_step_spec.name = canned_step_name
        if isinstance(canned_step_args, dict):
            for k, v in canned_step_args.items():
                pb_step.canned_step_spec.args[k] = v
        else:
            raise ValueError("canned_step_spec.args must be a dictionary")

    # setup resources
    if resources is not None:
        if not isinstance(resources, dict):
            raise ValueError("resources must be type dict")
        for k, v in resources.items():
            # key: cpu, memory, gpu
            # value: "1", "8", "500m", "1Gi" etc.
            pb_step.container_spec.resources[k] = str(v)

    if states._when_prefix is not None:
        pb_step.when = states._when_prefix

    # add template to proto workflow
    wf = get_default_proto_workflow()
    if tmpl_name not in wf.templates:
        proto_step_tmpl = couler_pb2.StepTemplate()
        proto_step_tmpl.name = tmpl_name
        _add_io_to_template(proto_step_tmpl, pb_step.id, input, output,
                            script_output)
        wf.templates[tmpl_name].CopyFrom(proto_step_tmpl)

    # add step arguments
    if args is not None:
        for i, arg in enumerate(args):
            if isinstance(arg, OutputArtifact):
                pb_art = pb_step.args.add()
                pb_art.artifact.name = arg.artifact["name"]
                pb_art.artifact.value = ('"{{inputs.artifacts.%s}}"' %
                                         pb_art.artifact.name)
            else:
                pb_param = pb_step.args.add()
                pb_param.parameter.name = utils.input_parameter_name(
                    pb_step.name, i)
                pb_param.parameter.value = ('"{{inputs.parameters.%s}}"' %
                                            pb_param.parameter.name)

    if states._exit_handler_enable:
        # add exit handler steps
        eh_step = wf.exit_handler_steps.add()
        eh_step.CopyFrom(pb_step)
    else:
        # add step to proto workflow
        concurrent_step = wf.steps.add()
        inner_step = concurrent_step.steps.add()
        inner_step.CopyFrom(pb_step)
    return pb_step
Пример #3
0
def step_repr(
    step_name=None,
    tmpl_name=None,
    image=None,
    command=None,
    source=None,
    script_output=None,
    args=None,
    input=None,
    output=None,
    manifest=None,
    success_cond=None,
    failure_cond=None,
):
    assert step_name is not None
    assert tmpl_name is not None
    # generate protobuf step representation
    pb_step = couler_pb2.Step()
    pb_step.id = get_uniq_step_id()
    pb_step.name = step_name
    pb_step.tmpl_name = tmpl_name
    # image can be None if manifest specified.
    if image is not None:
        pb_step.container_spec.image = image
    if manifest is not None:
        pb_step.resource_spec.manifest = manifest
    if success_cond is not None:
        pb_step.resource_spec.success_condition = success_cond
        pb_step.resource_spec.failure_condition = failure_cond
    if command is None:
        pb_step.container_spec.command.append("python")
    elif isinstance(command, list):
        pb_step.container_spec.command.extend(command)
    elif isinstance(command, str):
        pb_step.container_spec.command.append(command)
    else:
        raise ValueError("command must be str or list")
    if source is not None:
        if isinstance(source, str):
            pb_step.script = source
        else:
            pb_step.script = utils.body(source)

    # add template to proto workflow
    wf = get_default_proto_workflow()
    if tmpl_name not in wf.templates:
        proto_step_tmpl = couler_pb2.StepTemplate()
        proto_step_tmpl.name = tmpl_name
        _add_io_to_template(proto_step_tmpl, pb_step.id, input, output,
                            script_output)
        wf.templates[tmpl_name].CopyFrom(proto_step_tmpl)

    # add step arguments
    if args is not None:
        for i, arg in enumerate(args):
            if isinstance(arg, OutputArtifact):
                pb_art = couler_pb2.StepIO()
                pb_art.artifact.name = arg.artifact["name"]
                pb_art.artifact.value = ('"{{inputs.artifacts.%s}}"' %
                                         pb_art.artifact.name)
                pb_step.args.append(pb_art)
            else:
                pb_param = couler_pb2.StepIO()
                pb_param.parameter.name = utils.input_parameter_name(
                    pb_step.name, i)
                pb_param.parameter.value = ('"{{inputs.parameters.%s}}"' %
                                            pb_param.parameter.name)
                pb_step.args.append(pb_param)

    # add step to proto workflow
    concurrent_step = couler_pb2.ConcurrentSteps()
    concurrent_step.steps.append(pb_step)
    wf.steps.append(concurrent_step)
    return pb_step