Skip to main content

Client tutorial

This tutorial walks through the key concepts used by the ExaDeploy client. Make sure that you've installed the Python client before following along.

Setting up the cluster​

In a terminal, start the end-to-end image. This starts up a full temporary ExaDeploy cluster on your local machine.

python3 -m exa local_e2e gcr.io/exafunction/local_e2e:...
note

If you don't have a local_e2e image, contact Exafunction to get one.

tip

If you've set up a cluster in the cloud from following the quick start guides, you can feel free to use those as well.

Start the interpreter​

# These variables assume you're using the local_e2e image.
# Change them if using a real cluster.
EXA_SCHEDULER_ADDRESS=localhost:50050 \
EXA_MODULE_REPOSITORY_ADDRESS=localhost:50051 \
python3

Defining remote code​

We'll create a simple function to add two arrays.

import exa
import numpy as np

class AdditionAutoModule(exa.AutoModule):
    @exa.export
    def add(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return a + b

Subclasses of exa.AutoModule contain related groups of pure Python functions which may share resources, such as a deep learning model and weights. The exa.export decorator marks a function as a remote function. Commonly used types such as np.ndarray are automatically recognized and serialized in exa.AutoModule. When relying on autoconversion of types, the type annotations are required.

Uploading the code​

from exa import repo
repo.register_py_interpreter_module("InterpreterModule")

The repo module offers us shortcuts to access a singleton of the exa.ModuleRepository class, which is a client to the ExaDeploy module repository. The register_py_interpreter_module function defines a set of requirements and a file tree to be uploaded to the repository. In this case, neither is necessary.

module = exa.new_module_from_cls(AdditionAutoModule, "InterpreterModule")

The exa module contains some shortcuts to access a singleton of the exa.Session class, which is a client to the ExaDeploy scheduler and runners. The new_module_from_cls function uploads the code for AdditionAutoModule to a runner that has already loaded the InterpreterModule that was just defined and uploaded to the module repository.

tip

There are three places where we can declare code:

  • In exported functions: These are uploaded after the session is created on a per-client basis.
  • In the InterpreterModule or some other plugin: These files and requirements are set up once per runner when it gets assigned. They can be dynamically loaded and unloaded without replacing the runner entirely, but require registration in the module repository.
  • In the image of the runner: (Not shown in this example.) These are baked into the image of the runner.

You have the flexibility of deciding where to put your code between these options. Exafunction has implemented a large number of common plugins and functions to conveniently wrap common frameworks.

If possible, we recommend putting code in the plugin. However, adding them to the runner image is most convenient.

Executing the code​

result = module.add(np.array([1, 2, 3]), np.array([4, 5, 6]))
assert np.array_equal(exa.get(result), np.array([5, 7, 9]))

The module.add function is now a remote function that can be called with local arguments. The exa.get function retrieves the result of a remote function call and converts it back to local types. If the remote function is not yet complete, it will block until the result is available.

Remote pipelining​

The result object contains exa.Value objects and tries to behave asynchronously — the function will usually return immediately, and the exa.Value objects will contain handles to the remote values which may not yet be computed. If a value isn't needed locally, the exa.Value object can be passed to another remote function call, without using exa.get.

result2 = module.add(np.array([1, 2, 3]), np.array([4, 5, 6]))
result3 = module.add(result, result2)
assert np.array_equal(exa.get(result3), np.array([10, 14, 18]))

In this example, the data for result2 never needs to be downloaded locally to the client. The data for result has already been fetched in the previous code block, but it doesn't need to be reuploaded. This helps save network bandwidth.

Freeing GPU resources​

While the exa.Value object exists, the runner will retain the value, so it must be explicitly deleted to avoid growing memory usage over time. This convention was chosen to allow for simple pipelining of remote functions.

del result, result2, result3
note

Usually it's not necessary to del the exa.Value objects. The Python garbage collector will automatically free them when they are no longer referenced, such as when they go out of scope from a function.

Putting it all together​

demo_simple.py
import exa
import numpy as np
from exa import repo

class AdditionAutoModule(exa.AutoModule):
    @exa.export
    def add(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
        return a + b

repo.register_py_interpreter_module("InterpreterModule")
module = exa.new_module_from_cls(AdditionAutoModule, "InterpreterModule")

result = module.add(np.array([1, 2, 3]), np.array([4, 5, 6]))
assert np.array_equal(exa.get(result), np.array([5, 7, 9]))

result2 = module.add(np.array([1, 2, 3]), np.array([4, 5, 6]))
result3 = module.add(result, result2)
assert np.array_equal(exa.get(result3), np.array([10, 14, 18]))

del result, result2, result3

The code can be run with:

# These variables assume you're using the local_e2e image.
# Change them if using a real cluster.
EXA_SCHEDULER_ADDRESS=localhost:50050 \
EXA_MODULE_REPOSITORY_ADDRESS=localhost:50051 \
python3 demo_simple.py