Idea for RFC to extend EnvironmentTemplate to support containers by wrapping the task onRun

[mwiebe]

Background

Environment templates can accept parameters for a job, and use them to prepare the environment variables, file system contents, or other aspects of the context in which a job runs. They can be included inside job templates at the job and step level, or can be added externally to a job with the openjd-cli --environment option or scheduler support like with a Deadline Cloud queue environment.

Some existing example use cases for context:

• These sample Conda and Rez queue environments accept software requirements of a job as a parameter, and use it to provide a virtual environment for the job to run within.
• This job_env_daemon_process sample runs a background daemon process, that each task of the job sends a command to for running the task.
• This bash-in-docker sample uses a step environment to run a docker container, and its task execs a script in the container.

To follow the Open Job Description Design Tenets, a solution for containers would interoperate with the example Conda and Rez use cases. Job templates should be portable in a way to run them, unmodified, with either a Conda, Rez, Docker, or Apptainer environment template that provides the software environment to run in.

The RFC Idea

The bash-in-docker example, and the example that runs a background daemon process and then communicating it in each task are not far from starting a docker container in the background and then running each task in that container. The missing ingredient is a way to let the environment template say how to run the task inside of the container without having to modify the job template it's running.

The idea for this RFC is to extend the <Environment> with a third session action, onWrapTaskRun. The runtime would provide this action with enough context to run the task's onRun action itself. That context includes all the properties of the <Action>: The command, args, timeout, and cancelation method. Because all the context for running the task lives inside the session working directory, the container commands would need to create a bind mount that maps the path {{Session.WorkingDirectory}} to the identical path within the container.

Credit for inspiration of this idea goes to Pydantic's WrapValidator, while implementing OpenJobDescription/openjd-model-for-python#164.

A sketch of how a docker environment template could look with this feature:

specificationVersion: "environment-2023-09"
parameterDefinitions:
- name: ContainerImage
  type: STRING
  description: >
    The container image to run jobs in.
  default: "ubuntu:latest"
  userInterface:
    control: LINE_EDIT
    label: Container Image

environment:
  name: Docker
  script:
    actions:
      onEnter:
        command: "bash"
        args: ["{{Env.File.Enter}}"]
      # This is the new session action of this idea
      onWrapTaskRun:
        command: "bash"
        args: ["{{Env.File.WrapTaskRun}}"]
      onExit:
        command: "bash"
        args: ["{{Env.File.Exit}}"]
    embeddedFiles:
    - name: Enter
      filename: docker-env-enter.sh
      type: TEXT
      data: |
        #!/bin/env bash
        set -euo pipefail

        # Pull the requested image
        docker image pull '{{Param.ContainerImage}}'

        # Run the container in the background, and get its container ID
        DOCKER_CONTAINER_ID=$(docker container run --rm \
            --pull never \
            --detach \
            --mount 'type=bind,src={{Session.WorkingDirectory}},dst={{Session.WorkingDirectory}}' \
            '{{Param.ContainerImage}}' \
            bash -c 'sleep infinity')
        
        # Save the container ID to an environment variable for the onWrapTaskRun action to use
        echo "openjd_env: DOCKER_CONTAINER_ID=$DOCKER_CONTAINER_ID"

    - name: WrapTaskRun
      filename: docker-wrap-task-run.sh
      type: TEXT
      data: |
        #!/bin/env bash
        set -euo pipefail

        # Execute the task within the container that `onEnter` ran
        # This relies on the following potential new things:
        #    {{Task.Environment}} contains a list ["ENVVAR1=value1", ...] with every
        #        environment variable defined by `openjd_env`
        #    {{Task.Command}} contains the command from the `onRun` task session action
        #    {{Task.Args}} contains the args from the `onRun` task session action
        #    {{ shlex_join() }} is a function joining items of a list into a shell-quoted string,
        #        like Python's shlex.join https://docs.python.org/3/library/shlex.html#shlex.join
        #    {{ list_join() }} is a function equivalent to Python [v for item in lst for v in item]
        docker container exec \
            $DOCKER_CONTAINER_ID \
            {{ shlex_join(list_join([["--env", envvar] for envvar in Task.Environment]) }} \
            {{ shlex_quote(Task.Command) }} \
            {{ shlex_join(Task.Args) }}

    - name: Exit
      filename: docker-env-exit.sh
      type: TEXT
      data: |
        #!/bin/env bash
        set -euo pipefail

        # Stop and clean up the container that `onEnter` ran
        # This relies on the following potential new things:
        #    {{Env.Action.Timeout}} contains the value specified for the action's timeout, in seconds
        docker container stop \
            $DOCKER_CONTAINER_ID \
            --timeout {{Env.Action.Timeout}}

[leongdl]

How would Conda inter-operate with a docker container in this instance? If there was a Queue Environment using Conda, then a Docker Environment.

Would it install to the docker container? (I think not)

[mwiebe]

It could interoperate, or it could be an alternative. You could use a container image that already has all the software installed in it, or you could nest a Conda virtual environment inside the docker container. I think usually you'd use it as an alternative, either set up your queue to use container images to provide software, or use conda.

[leongdl]

Right, I think the simpler approach would be "use as alternative" as suggested. Docker has other extensions like opening ports and sharing local file paths, which is why I was thinking if we want to go down the path of "opening up the conda env to docker", or "running conda in the docker".

[leongdl]

One more thought on this is what if the user wants more than 1 instance of a container? Given DOCKER_CONTAINER_ID, is there any risk the system maybe compromisable?

[mwiebe]

Can you give some more ideas of what the user is trying to do? Are they running a multi-container application like with docker compose? I expect that would be implemented as a different environment template.

[leongdl]

I was thinking also along the lines of sidecar containers, so yes to docker compose. The main container performs computation, the side car container pumps metrics / logs or other functionality.

[crowecawcaw]

One thing worth considering: how would credentials, especially queue creds, be made available inside the container? The current sketch only bind-mounts Session.WorkingDirectory and forwards Task.Environment vars, but host-level credentials from instance metadata or credential files wouldn't automatically be accessible to the containerized task. It might be worth defining a pattern or extension point for securely propagating credentials into the container.

[leongdl]

For this - I was actually thinking to not re-invent the wheel. The credential rotation and its processes are well documented in the worker agent doc. I would re-use the same queue (aka session) credential refresh, bind mount it into the container.

As an example, our default environment wrapper would do the following:

1. Mount the session folder
2. Share all license environment variables
3. Share all worker agent session environment variables (everything from path mapping onwards)
4. Queue shared files such as credentials.

For long running jobs, the credentials must rotate with the session, so I think we should not re-invent what works well currently.

[mwiebe]

Here's another older request from deadline-cloud about this feature: aws-deadline/deadline-cloud#321

[leongdl]

For my container experiments, I've been directly mapping the session folder, along with the creds files that worker agent generates. This has 2 positive side effects. 1) they are auto rotated by the worker agent, 2) paths are mapped the same, so the same command to setup credentials is shared with the container.

That also goes the same for sharing job attachments too! Would be a good extension into the OpenJD Sessions for Python I think?

[mwiebe]

Is there any part of that needs special handling? Or is that all reasonable to do from the environment actions currently as defined by this idea?

[leongdl]

Generally works, just needs good expamples for discoverability.

Also interesting things like mounted folders on the worker (EG over VPC lattice)