Dependency handling for electrons

[kessler-frost]

Deps Design Document

Terminology

• task - electron decorated function by the user
• backend environment - the environment where the task function is actually executed, e.g an aws instance or inside a slurm cluster
• call_before and call_after - functions to be called before task and after task respectively in the same backend environment
• partial - a function decorated with the functools.partial decorator (link) which basically “freezes” the function

Problems

Dependency features that are not available right now in a simpler more direct way:

• call_before and call_after function executions. It’ll basically mean packing all three types of these functions into 1 function and calling that instead.
• pip, conda, module package dependency installation on the backend environment.
• Passing environment variables to the backend environment.
• Having certain other dependencies already present on the backend environment like a file.
• Error handling for all of the above points.
• Should be able to do all of the above without the necessity of having covalent installed on the backend environment.
• Either order of execution should be maintained or an order of precedence be established beforehand.

Proposal

UX - Explicit

UX - Implicit conversion to objects

Deps Class

This will be the parent class for any kind of dependency. For implementing any new type of dependency we’ll have to subclass this and override the __init__() and apply() methods.

• __init__():

  • Function to initialize the Deps object with given variables.
  • This is important since once the values of the internal variables have been set it should not be changed and this is the only place to provide any values to the Deps object.

• apply():

  • Exact function that’ll be executed on the backend environment.
  • Does not take any argument and only uses self.* kind of variables, i.e, internal variables local to the object assigned at initialization.
  • Reasons for not taking any arguments:
    • Assumption of complete knowledge of dependency at build_graph time by the electron
    • Any function that needs to be dynamically assigned a value should be an electron
    • Gives a uniform way of applying any dependency without worrying about what kind of dependency it is; so even if we decide to expand our list of supported Deps subclasses, or the user wants to have a custom Deps, they can easily do it without having to worry about argument management
    • Need to get transformed into a partial function runnable without having this Deps object available on the backend environment

Initial Deps SubClasses

• PipDeps:

  • Class to install pip packages
  • args/kwargs:
    • packages: List of strings which will be combined to do a pip install hence even version number can be provided as “numpy==0.23”
    • requirements: requirements.txt file path can also be given
  • Both of the kwargs can also be combined into one list if desired and run in a single command inside the apply() method
  • apply():
    • Contains package installation code using pip

• CondaDeps:

  • Class to install conda packages
  • args/kwargs:
    • channels: List of strings of channel names in order of priority when installing conda packages
    • packages: List of strings of package names with version number if wanted to be installed using conda
    • yaml: File path for environment.yml file for creating a new conda environment
  • apply():
    • Contains packages installation code using conda
  • If there’s a yaml then creation of the environment with given configuration is also done

• ModuleDeps:

  • Class to install module packages
  • args/kwargs:
    • packages: List of strings of package names to be installed using module package manager
  • apply():
    • Contains package installation code using module

• EnvDeps:

  • Class to assign environment variables
  • args/kwargs:
    • vars: Dictionary of strings (both keys and values) to assign a environment variables to the backend environment
  • apply():
    • Contains env variable assignment code

• BashDeps:

  • Class to execute bash command
  • args/kwargs:
    • commands: List of strings of commands to be executed inside the shell in the same order as the list
  • apply():
    • Contains code to execute the list of commands given in a shell

• CallDeps:

  • Class to execute the given function with given arguments
  • args/kwargs:
    • func: Function callable to be run
    • args: Tuple of args for func
    • kwargs: Dictionary of kwargs for func
  • Will be used in both call_before list and call_after list of ordered executable Deps objects, e.g:
    • call_before=[CallDeps(a, args=(1, 2))], CallDeps(a, args=(3, 4))]
    • This list is considered ordered so a(1, 2) function will be run before a(3, 4)
  • apply():
    • Contains code to execute the function callable with given args/kwargs

• Special Case of ImportDeps or any such composite “Deps”:

  • Not a subclass of Deps but just a proxy for UX to handle multiple package installation dependencies together
  • args/kwargs:
    • pip: Initialized PipDeps object
    • conda: Initialized CondaDeps object
    • module: Initialized ModuleDeps object
  • Purpose of this is to enclose all kinds importing dependencies into 1 object and also eases the implementation of short-handed dependency assignment
  • Each package dependency based electron metadata field mentioned in this will be assigned the respective initialized Deps object, e.g:
    • pip metadata field will be assigned the import_deps_object.pip value, etc.
  • Resolution of this will be done at the time of metadata field assignment and through overriding the __new__() method of this class

• Some things worth mentioning:

  • Since the apply() methods will be run every time for every electron, so if some dependencies are already present in the environment then we shouldn’t try to install/download them in there - the creator of the Deps subclass should keep this in mind when writing the apply() method
  • It is yet to figure out, although I feel I’m close, how to implement apply() in a way where the backend environment does not need the Deps object to be there hence no covalent dependency should be there
  • This method is essentially doing something like this:

def new_func(*args, **kwargs):
	deps_object_1.apply()
	deps_object_2.apply()
	
	task(*args, **kwargs)

	deps_object_3.apply()
….

executor.execute(func=new_func, args=inputs[args], kwargs=inputs[kwargs])

• Passing return/input values between apply() functions is intentionally avoided but can be implemented if need be, although some thought should be given when deciding that as to why exactly is that needed.

[cjao]

Just noting here that this will involve changes to the DB schema. We'll need to have a broader discussion about how we're storing electron metadata in the DB and object store.

[AlejandroEsquivel]

Discussion about this moved to #686

[santoshkumarradha]

@kessler-frost / one quick thing thats missing in this class is the dependency on local import files. This is majorly allowed by simply passing the imported module which can be passed by value to cloudpickle. (refer cloudpickle Readme)

CC @AlejandroEsquivel

[santoshkumarradha]

To elaborate a bit more, If you are using a locally written python module, and using it inside an electron, cloudpickle would not pickle them. One has to explicitly tell cloud pickle to pickle that module. We also need a same interface where we pass on these local non-pip-installable modules to the Deps that can be pickled by value by cloudpickle.

[AlejandroEsquivel]

Ah yes, thanks for elaborating @santoshkumarradha, I will add an issue to the Import Deps epic to cover the UX proposal and implementation.

[cjao]

This design doc is really well written. I just have a few questions.

0. Would the new UX significantly beat manually writing the config commands in the electron? What if we were to provide various the Deps classes like PipDeps or CondaDeps and ask users to explicitly write

def new_func(*args, **kwargs):
	deps_object_1.apply()
	deps_object_2.apply()
	
	task(*args, **kwargs)

	deps_object_3.apply()

(This implementation wouldn't actually work if the deps need to be installed before unpickling the task, but one might imagine introducing "dep" electrons, which would fit into the Transport Graph framework.)

There a few considerations here: readability, ease of use, and ease of reasoning about the code. The latter is especially important when code doesn't run the way one expects, and error handling would be most straightforward if the user issues the setup commands explicitly.

1. Are the different classes of deps assumed to be independent? What would happen if a pip package requires some system package to be installed using shell commands?

2. Where should the PipDeps be installed? In a venv? Or a condaenv?

3. Should the CondaDeps automatically activate the condaenv? If so, would that overlap with the executor's execute_in_conda_env method?

[santoshkumarradha]

@cjao , great qs,

0. I would agree with you and this is more on the line of the following

deps_object_1.apply()
deps_object_2.apply()

def new_func(*args, **kwargs):	
	task(*args, **kwargs)

deps_object_3.apply()

Rather than the former. It is meant to have a unified set of call back/call forward methods before and after python function execution ( not inside of it) good catch.

Are the different classes of deps assumed to be independent? What would happen if a pip package requires some system package to be installed using shell commands?
Where should the PipDeps be installed? In a venv? Or a condaenv?

Great questions, although we see pip and conda in the same variable there, what we need is to differentiate pip with pipenv and conda with conda env and through error while defining ImportDeps to make sure only one is passed.

Should the CondaDeps automatically activate the condaenv? If so, would that overlap with the executor's execute_in_conda_env method?

Again great Q, essentially the execute_in_conda_env was a lean patch to support multiple conda env which should be depreciated once this is release. The goal with Deps is to run the given function in the required environment by installing the unavailable packages (or creating the env if it is not available in conda) in the given list.

Hope this helps

[cjao]

Thanks for the clarifications. Could you elaborate on the order in which the dependencies are to be processed? Conda first, and then Pip? What happens if users define their own Deps subclasses? And how do these dependencies interact with the File deps? Which ones are to be handled first?

We must take care to not compromise the clarity of the program logic whenever we perform magic behind the scenes (in this case, translating dependency declarations into Python statements). Unless the deps are truly independent -- any order which they are processed yields the same results -- it will be hard to reason about the code unless there is a deterministic processing order.

[cjao]

For now, we will process the different classes of deps in a canonical order.

[santoshkumarradha]

I agree, the order I would choose for default ( after which we can have a backlog feature addition to create UX for modifying this behaviour) is

Env->Bash->Module->Conda->Pip->Call_before
*function*
Call_after

typical cluster modules contain conda as well as python, so having module load conda and then doing stuff of instaling conda env or switching conda env would be useful. Same for pip where module load python is often used to invoke cluster's default python.

[cjao]

Sounds good, I see @AlejandroEsquivel had already written down this order in the MVP planning issue.

[cjao]

Do we expect the Deps classes to be cloud-picklable?

[santoshkumarradha]

again, good q. Ideally lets' design it in a way that we dont include Deps object in the executer but rather just the list of commands/functions to be executed. We want to keep the transferring to executer as lean as possible. Having said that designing objects to be picklable is extremely important for future proofing and transporting (especially since we are dealing with remote execution)

tldr - lets make it picklable but design deps in a way that we dont need to ship the object.

[cjao]

The bulk of a Deps object's footprint would be its associated data (e.g. a list of pip packages). If the Deps object can in fact be serialized, it should be easy to ship the object and call apply() in the remote environment.

I don't see much point in introducing a Deps object to store the list of bash commands or pip dependencies only to deconstruct the object right before execution.

[santoshkumarradha]

I don't see much point in introducing a Deps object to store the list of bash commands or pip dependencies only to deconstruct the object right before execution.

I guess the reasoning to send it via non pickled sense is that we might want to restrict pickle to only objects that definitely need to be pickled (only functions) to the executer while others are done via REST and such. This might maybe let us be more clever on how we handle larger units ( by breaking the message etc..) instead of being tied to one unit blob. Also, since pickle is not cross compatable, it might be worth while to keep them separate and just execute them which reduces one head ache off us taking care of making sure of pickle compatibility.

The reason to have Deps is not just for our convenience to store, but rather for user to have more structure while defining. Passing requirements as second image in sankalp's proposal although we allow it, ideally is not how we want user to interact as one can very much make mistakes and not even have type hints on what all is allowed. Hence passing more structured data gives the user the ease to exactly know what all they can pass and how to pass them.

[cjao]

I see, thanks for clarifying.

[cjao]

One way to avoid pickling any class or function is to have Deps.apply() simply emit the sequence of bash commands and leave it to the executor to figure out how to apply them. functools.partial doesn't seem to avoid pickling; it just fills in some of the parameters of a function.

[cjao]

For the implementation, we need to distinguish between two types of deps. Some classes of Deps can be applied in a separate task before the actual task; these can be packaged in "Dep" electrons and marked as a dependency of the core electron. But others only have an effect if applied in the same session as the actual task and need special handling by the executor.

To illustrate, recall that the SSH and Slurm executors runs each task through a series of noninteractive SSH sessions:

1. Upload the python/sbatch script
2. Run the python/sbatch script
3. Download the results

Installing pip packages into a venv/condaenv can be done in a separate SSH session before Step 1. Other setup commands, however, only persist for one shell session. These must be run in the same session as Step 2. For example, any environment variables must be set at the beginning of Step 2 since they will be unset in the next SSH session.

To apply the latter kind of deps, the executor.execute method should grow hooks for call_before and call_after, and any session-specific environment setup should occur in call_before.

[cjao]

One constraint here is that we really don't want to keep SSH connections open for very long; thus either call_before and call_after must be extremely lightweight, or we pack the hooks into the same task as the main task; but this would prevent environment setup (e.g. conda activate) before the main task is deserialized.

[santoshkumarradha]

Can you elaborate on why we would do this rather than ship all the deps execution statement along with the function to happen before and after ? for instance, one can easily define install and module load statements in the SLURM job while in case of SLURM and similarly for others. Down side is that we cant get back communication of whether the said Deps state have been finished or not, but thats an okay state to be in for MVP and then add it on as feature.

[cjao]

I think we're talking about the same thing. module load or conda activate are the sort of commands that would go in call_before.

[cjao]

I was thinking initially that putting deps (especially complicated deps involving many shell commands) in their own electrons might be easier for error handling -- failure of any dep would be handled like failure of any other electron, and one can easily see exactly what step failed.

[cjao]

Mutating the state of the executor's environment outside of the user's explicitly defined tasks can introduce some new and potentially subtle failure modes:

• Errors from Deps.apply() (e.g. pip dependency resolver fails)
• Errors from Deps.apply() not being idempotent
• Errors from two electrons having potentially conflicting deps; could one electron's Deps.apply() interfere with other electron's?

We need to figure out how to handle these in a transparent way.

[santoshkumarradha]

Finally closed by #1876