Skip to content

geyang/zaku

Repository files navigation

Zaku, a fast Task Queue for ML Workloads

To get a quick overview of what you can do with zaku, check out the following:

  • take a look at the basic tutorial or the tutorial for robotics:
  • or try to take a look at the example gallery here

Install zaku --- the latest version is {VERSION} on pypi.

pip install -U 'zaku[all]=={VERSION}'

Learning Usage Patterns by Running the Tests (specs)

Running tests is the best way to learn how to use a library. To run the tests in the ./specs folder, first download and setup redis-stack-server. Then start a zaku task queue server at local port 9000.

  1. install zaku and look at the options:
    pip install -U 'zaku[all]=={VERSION}'
    zaku -h
  2. install and run the redis-stack-server:
    brew install redis-stack-server
    redis-stack-server
  3. run the zaku task queue at port 9000:
    zaku --port 9000 --verbose

Now you can run the tests by running

make test

In pycharm, you should see the following:

Example Usage

Adding Jobs:

Supposed you have a TaskServer running at localhost:9000.

from zaku import TaskQ

queue = TaskQ(name="my-test-queue", uri="http://localhost:9000")

for i in range(100):
    queue.add_job({"job_id": i, "seed": i * 100})

Retrieving Jobs:

from zaku import TaskQ

queue = TaskQ(name="my-test-queue", uri="http://localhost:9000")

job_id, job = queue.take()

Now, after you have finished the job, you need to mark the job for completion. The way we do so is by calling

queue.mark_done(job_id)

Sometimes when you worker responsible for completeing the job encounters a failure, you need to also put the job back into the queue so that other workers can retry. You can do so by calling

queue.mark_reset()

Now, we offer a context manager TaskQ.pop, which automatically catches exceptions and resets the job (or marks it complete).

from zaku import TaskQ

queue = TaskQ(name="my-test-queue", uri="http://localhost:9000")

with queue.pop() as job:
  if job is None:
    print("No job available")
  
  print("Retrieved job:", job)

Developing Zaku (Optional)

First, git clone this repo, and install it in editable mode plus dependencies relevant for building the documentations:

git clone https://github.com/geyang/zaku.git
cd zaku
pip install -e '.[dev]'

Improving Documentation

We use sphinx to generate the documentation. To get started, make sure you install the dev dependencies. This should already be included if you installed zaku via the zaku[dev] flag. Take a look at ./docs/requirements.txt to see what packages are required for development.

When editing the documentation, you can use a preview server that automatically refreshes the browser upon changes in source. Simply do

make preview

Note: restarts are required, when the page index is changes. This includes adding and removing new pages.

To run the full build, call

make docs

License

Built with ❤️ by @episodeyang. Distributed under the MIT license. See LICENSE for more information.