Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
shell.nix to enable training ML models remotely with the GCP AI Platform

shell.nix for Training in GCP AI Platform

This gist contains a shell.nix file that can be used to create a Python environment for running training jobs in the GCP AI Platform.

This is specifically for the following tutorial:

This uses code from in the census/tf-keras directory.

However, this shell.nix file should be able to be modified to work for almost training job.

How To Use

First run nix-shell to get into the shell.

Then, install all required python Packages:

$ pip install -r requirements.txt

If you add additional Python packages to the buildInputs line in the shell.nix file, you should be able to use these system-level packages, instead of having to download them.

Now you should be able to actually run training:

$ python3 -m trainer.task --job-dir local-training-output

The tutorial linked above recommends using the following command, however, this fails when using gcloud from Nixpkgs :-\

$ gcloud ai-platform local train --package-path trainer --module-name trainer.task --job-dir local-training-output

As long as training directly with python works, you can launch a training job on the GCP AI Platform.

First, you need to login with gcloud using OAuth:

$ gcloud auth login

Set your default project name so you don't have to specify it in each command below:

$ gcloud config set project inner-melody-274800

Next, you need to create a bucket to store the trained models:

$ BUCKET_NAME="my-training-example-task-3"
$ REGION="us-central1"
$ gsutil mb -l $REGION gs://$BUCKET_NAME

Finally, actually launch the training job:

$ gcloud ai-platform jobs submit training "my_first_keras_job" --package-path trainer/ --module-name trainer.task --region $REGION --python-version 3.7 --runtime-version 1.15 --job-dir "gs://$BUCKET_NAME/keras-job-dir" --stream-logs
nixpkgs-src = builtins.fetchTarball {
# Recent version of nixpkgs master as of 2020-03-30.
url = "";
sha256 = "sha256:0aw6rw4r13jij8hn27z2pbilvwzcpvaicc59agqznmr2bd2742xl";
nixpkgs = import nixpkgs-src { config = { allowUnfree = true; }; };
with nixpkgs;
lib-path = lib.makeLibraryPath [
mkShell {
name = "pip-env";
buildInputs = with python37Packages; [
# System requirements.
# needed for buildilng certain packages
src = null;
shellHook = ''
# Allow the use of wheels.
# Augment the dynamic linker path
export "LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${lib-path}"
if test ! -d $VENV; then
virtualenv $VENV
source ./$VENV/bin/activate
export PYTHONPATH=`pwd`/$VENV/${python.sitePackages}/:$PYTHONPATH
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment