Pyflyte --env Option for Register and Serialize

[MinuraPunchihewa]

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

This PR adds the --envs option for the register and serialize commands to set environment variables in the images.

Tracking Issue

Fixes flyteorg/flyte#4092

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).

[MinuraPunchihewa]

Hey @kumare3, @wild-endeavor,
I have just created this draft PR by adding the --envs option to the register and launchplan commands. Can you please give me some insight into how these envs are to be used and passed into the image?

[PudgyPigeon]

Hey @kumare3, @wild-endeavor, I have just created this draft PR by adding the --envs option to the register and launchplan commands. Can you please give me some insight into how these envs are to be used and passed into the image?

When we used our fork in production, we didn't need to pass in ENVs into launchplan. It's only for register as that is the CLI command that compiles the tasks/workflows/launchplans and sends it to the remote backend. Our use case is dynamically defining a variable amount of tasks/config settings/workflows/etc at compile time and running it in the remote backend through a singular workflow - abstracted for data scientist ease of use. We can share some of the workflow code (barring the confidential stuff) - if maintainers/developers want to see what our use case was specifically.

pyflyte launchplan as far as we know only takes in the version, launchplan name, configs to activate or deactivate a certain version of a launchplan. So no need to pass in envs into that command - I don't know a scenario where you could or should pass envs into pyflyte launchplan.

Also, feel free to use our code snippets in your PR, we only ask for some due consideration in crediting.

[MinuraPunchihewa]

Hey @PudgyPigeon,
Thank you very much for letting me know. I will remove the option from the launchplan command.

[MinuraPunchihewa]

I have removed the option from launchplan. Can I be given some insight as to how the envs can be used within register, please?

[wild-endeavor]

@MinuraPunchihewa - thank you for working on this. The purpose of this issue is to make the behavior of pyflyte run --envs ... similar to that of pyflyte register --envs .... The run part is defined here, but for register it's a bit different.

For the run command, the envvars are fed into the FlyteRemote object's execute function and these ultimately land here. This fills in the ExecutionSpec message, which is ultimately how they're passed into task containers.

For register, rather than filling in an execution spec object, which is only used one time, we want to fill them into the task templates themselves. This means that the compilation step for each task will need to take in the dictionary of env vars. Fortunately, the compilation process already respects the env vars in the SerializationSettings object. So if you can get them in there, the rest should just follow.

Let's also add this option to pyflyte serialize... just follow the code in serialize.py and you'll see that it's the same thing. It just needs to be added to that object and things should work.

Not sure about the input syntax... i know the original issue uses json - but pyflyte run currently uses something else. I don't have much of a preference but i think they should match. And we shouldn't change the existing format (cuz that can break people) but we can add support for a new format, but that's outside the scope of this ticket. i think probably best to just make it match for now.

[MinuraPunchihewa]

Hey @wild-endeavor,
Thank you very much for your comments. I have updated how the envs are passed and the option has also been added to the serialize command.
Can you let me know if it looks OK?

I will then work on using this env variables to configure the image.

[MinuraPunchihewa]

@wild-endeavor With the amazing explanation that you have given, I believe I was able to implement this successfully. Please let me know if there is anything else that I need to do; run tests, documentation etc.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (c413570) 55.04% compared to head (52558bd) 54.70%.
Report is 29 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1880      +/-   ##
==========================================
- Coverage   55.04%   54.70%   -0.35%     
==========================================
  Files         296      306      +10     
  Lines       22242    22788     +546     
  Branches     3357     2255    -1102     
==========================================
+ Hits        12244    12466     +222     
- Misses       9835    10166     +331     
+ Partials      163      156       -7     

Files	Coverage Δ
flytekit/tools/repo.py	0.00% <ø> (ø)

... and 71 files with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[wild-endeavor]

@MinuraPunchihewa thanks for the changes! Could you do me a favor though? Since we don't have time to manually test everything - could you please run a test against sandbox if you haven't already. take a screenshot of the command line, and then take a screenshot of the Task pane in the console? (If you click on a task, in a running or completed workflow on the UI, you should see a bunch of json in the task pane). Take a screenshot of the definition there. it should have the env vars from the command line.

Also could you just add one unit test please?

Thanks so much.

[wild-endeavor]

oh also, can you run make lint and fix please? It looks like that test is failing. Thank you. Much appreciated.

[brian-odonovan]

Just wanted to comment that I am anxiously awaiting this feature! Thanks for putting this together.

[MinuraPunchihewa]

Hey @wild-endeavor, @brian-odonovan,
Sorry for the delay, I will try and finish this up as soon I can.

[brian-odonovan]

Hey @wild-endeavor, @brian-odonovan, Sorry for the delay, I will try and finish this up as soon I can.

No worries! We have a totally sufficient workaround by setting the FLYTE_INTERNAL_DOMAIN as an env variable in our github actions workflows and then ingesting that during registration.

[MinuraPunchihewa]

Hey @wild-endeavor,
Sorry for the delay in my response. So, I have run a register command and it seems to execute successfully,

I am not entirely sure how to check if this worked though. I have to check this in the console, is that right?

[MinuraPunchihewa]

Oh, actually I see it here,

This is after launching the workflow.

Does it look correct? If so, what more do I have to do?

[wild-endeavor]

Thanks!

lint is failing unf... could you run make lint and take a look please?

flake8...................................................................Failed
- hook id: flake8
- exit code: 1

flytekit/clis/sdk_in_container/helpers.py:1:1: F401 'json' imported but unused
flytekit/clis/sdk_in_container/register.py:30:1: E302 expected 2 blank lines, found 1

black....................................................................Failed
- hook id: black
- files were modified by this hook

reformatted flytekit/clis/sdk_in_container/register.py
reformatted flytekit/clis/sdk_in_container/serialize.py

All done! ✨ 🍰 ✨
2 files reformatted, 6[51](https://github.com/flyteorg/flytekit/actions/runs/6460241412/job/17542075205?pr=1880#step:6:52) files left unchanged.

isort....................................................................Failed
- hook id: isort
- files were modified by this hook

Fixing /home/runner/work/flytekit/flytekit/flytekit/clis/sdk_in_container/register.py
Fixing /home/runner/work/flytekit/flytekit/flytekit/clis/sdk_in_container/serialize.py

should be pretty easy. make fmt will format if lint doesn't do that for you.

[MinuraPunchihewa]

@wild-endeavor I think I was able to fix the lint issues. The make commands did not actually work for me, but I applied fixes based on the report above.

[wild-endeavor]

lint still failing

[INFO] This may take a few minutes...
flake8...................................................................Passed
black....................................................................Failed
- hook id: black
- files were modified by this hook

reformatted flytekit/clis/sdk_in_container/register.py
reformatted flytekit/clis/sdk_in_container/serialize.py

All done! ✨ 🍰 ✨
2 files reformatted, 668 files left unchanged.

isort....................................................................Failed
- hook id: isort
- files were modified by this hook

Fixing /home/runner/work/flytekit/flytekit/flytekit/clis/sdk_in_container/register.py
Fixing /home/runner/work/flytekit/flytekit/flytekit/clis/sdk_in_container/serialize.py

if you want, just give me write access to your fork and i can push a commit to this branch.

running make setup might help with the commands failing, there are things you need to install.

[MinuraPunchihewa]

@wild-endeavor I think I got the make commands running now. I've made a commit after make fmt. Please let me know if this has solved the problem.
I've sent you an invite to my fork as well, just in case.

[welcome]

Congrats on merging your first pull request! 🎉