pyflyte build: better container specification, build, and push UX

[cosmicBboy]

Objective

Flytekit today is completely unopinionated about how a user builds their container images. Even though this is amazing for large companies, for independent developers this might be cumbersome and hard. A solution is to keep the existing functionality but add additional way of automatically deducing the images and building them. This has the following advantages:

1. Smaller and leaner images
2. Faster and uniform builds, with better caching
3. Simplified management of multiple images
4. And finally improve the usage within hosted Flyte:
   1. We could limit the base images, pre-deploy and speed startup times
   2. we could also provide a build service that bypasses local builds and can be sped up tremendously

The proposed solution is to create custom image specs that result in a buildable docker specification. In case of using docker build, it could result in a dockerfile, in envd could result in build.envd file etc.

Task Update

@task(container_image="xyz/image:latest")
def foo():
  ...

@task(container_image="{{.image}}")
def foo2():
  ...
  
@task(container_image=ImageSpec(...))
def foo3():
  ...

ImageSpec

enum Capability(enum.Enum, int):
  CPU = 1
  GPU = 2
  Google_TPU = 3
  Amazon_Inf = 4
  ...
  
enum Versioneer(Enum, int):
   Git = 1
   SpechHash = 2
   Auto = 3
  
@dataclass
class BaseImage:
  os: str = "ubuntu"
  python_version: str = "3.9"
  capability: Capability

@dataclass
class ImageSpec:
  name: str
  dockerfile: Optional[str] = None
  versioneer: Versioneer = Versioneer.Git
  requirements: Optional[Union[List[Tuple[str, str]], str]] = None
  base_image: BaseImage = BaseImage.Default()
  registry: Optional[str] 
  
  def __post_init__(self):
    if self.dockerfile and self.requirements:
      raise ValueError("Only dockerfile or requirements can be used together")

pyflyte build

pyflyte build --plugin=[docker | envD | buildx] image.lock

pyflyte register / package will do the following things:

• Load all tasks and workflows
• For every task, if ImageSpec is specified then, it will collect all the imagespecs (should we do this in register as well?)
• Then build can launch the plugin and pass all image specs to build potentially in parallel
• Ideally when you pyflyte run the following happen:
  • Find recently built images that match the hash of the imagespec?
  • This could be done using git notes show, and find the first defined note and then use that to get all images
  • Question: When should we run pyflyte build? should this be run always unless --fast is specified. how is this backwards compatible?

Appendix

envD [link]

envD makes it simple to manage environments locally or on kubernetes. This is particularly useful if using gpus. It also eliminates the need to write the docker file, but ends up building the container images using a simplified spec. We should see if we can incorporate the image spec in ours directly.

[fg91]

At my previous company we built a hydra launcher plugin that allowed us to build images and execute Flyte workflows with one command. Back then, I gave a talk about this at the community sync: See here for how the launcher is used and here how the images are built. When we built this, we didn't use fast registration yet (which alleviates the need for something like this quite a bit) but we definitely felt that this was a missing feature when first starting out with Flyte.

In my current team, we build images in CICD exclusively.

Providing a way to pass a path to a Dockerfile in the task decorator and have it built automatically sounds great to me. Had this existed, we wouldn't have built the hydra launcher I mentioned above.

---

The proposed solution is to create custom image specs that result in a buildable docker specification. In case of using docker build, it could result in a dockerfile, in envd could result in build.envd file etc.

Does this mean creating a new "interface/abstraction" around Dockerfiles, e.g. in Python, and then generating a Dockerfile on the fly?
I personally am not a fan of this to be honest. I feel that most ML engineers are able to write or willing to learn how to write Dockerfiles and I fear that if we create a new abstraction in python to create Dockerfiles on the fly, there might be many follow up issues that Dockerfiles support a certain feature/syntax that the Flyte abstraction doesn't yet. At least I would limit it to e.g. additional pip installs etc.

On the other hand, ZenML for instance has such a feature where one can pass pip requirements or even clone the current environment by running pip-freeze behind the scenes and then using the package list to build an image.

should we do this in register as well?

I would say yes. The user could choose with pyflyte register --fast/--build.

[cameronraysmith]

@fg91 we saw the work you presented and prototyped a method of integrating with @rsokl 's https://github.com/mit-ll-responsible-ai/hydra-zen in https://github.com/sciexp/flytezen . Apologies it is far from perfect and almost totally undocumented save a few draft docstrings, but, perhaps after we have a chance to seed the documentation, some modified or pared down version could be upstreamed in some agreeable way, shape or form. It's analogous to what is constructed by pyflyte init but embeds a library of workflows in a python package and adds a hydra(-zen) cli to manage their execution. A very rough preview of intended usage is here. I am fairly confident this example could be improved by directly leveraging hydra-zen's dataclass construction and the associated consumption across task and workflow interfaces.

many thanks to both teams for these fantastic tools!

[cosmicBboy]

We should consider supporting packaging and registering workflows inside a docker image that contains the dependencies of the flyte project. The motivation behind this is that there are some cases when the build environment python version is different from the docker python version, which causes issues around python version dependencies @pingsutw @wild-endeavor @eapolinario

[davidmirror-ops]

2023-11-09 Contributor's meetup notes: the team agrees to close this entry as this is already implemented by ImageSpec. A related Docker plugin is being developed here