Claus

13 minute read

Welcome to the wonderful world of Rust … or to be more specific it’s about Rust and

  • WASM
  • Azure Functions
  • Genetic Algorithms
  • Solving a complex optimization problem, fast

This is part four in a series of Rust on Azure Functions. The other parts are about performance comparisons (part 1 and part 2), and explaining the PMX algorithm (part 3). If you want to learn more about Genetic Algorithms be sure to read part 3 first).

Rust and JavaScript on Azure Functions

This is the finished thing: a traveling salesman problem solver built in Rust and JavaScript each running on their own Azure Function (a Function as a Service offering). Click on the buttons to generate a solution the routing problem and see the fitness history and time elapsed below. (Isn’t it curious how chaning )

Want to know how it’s built? Scroll down!

Optimized Traveling

The traveling salesman problem (TSP) is an NP-complete problem which is defined as follows:

“Given a list of cities and the distances between each pair of cities, what is the shortest possible route that visits each city and returns to the origin city?” (from Wikipedia).

So, mathematically speaking the goal is find the minimal, distinct combination of all distances between cities. This can be extended to include one-way streets and higher-ranking as well as lower-ranking connections, but for now the naive version will do just fine. In order to solve these kinds of problems, a special branch of machine learning is required: (Meta) Heuristic algorithms.

These algorithms trade speed and ability to solve at all, for less optimal solutions using a repeatable strategy. This strategy can either be greedy (i.e. always taking the lowest cost solution, even if it leads to a local optimum) or more accepting, by picking a more expensive solution every once in a while. Genetic Algorithms are a problem agnostic approach that always use the lowest-cost solution but introduce variation by randomly changing each solution candidate.

Solution Architecture

This is going to be a deep dive into some aspects of the solution, for the full code please check out the GitHub repository.

A solution candidate (or genome or citizen) is a potential solution to the problem, which in the case of the TSP is an ordering of the cities that have a minimal total distance. If you are interested in the representation and its implications, more vocabulary and strategies, part 3 goes into that.

Functions

FaaS type services like Azure Functions have a huge potential to simplify operationalizing machine learning on many levels and by using WASM (and in my case Rust) the necessary performance can be achieved without the vendors having to change too many things.

Operationalizing? This term involves everything that is required to get a machine learning model or operation into production - meaning tracking inputs, outputs, and model version at the very least and FaaS services can do most of that already!

This is why the project is implemented on Azure Functions and, to compare performance, each language has their own instance within a Function app. Both instances are running Azure’s Function Runtime v2, supporting node.js versions > 6 and therefore WASM targets.

Note: The Rust target is deployed as a pre-compiled WASM file that is pushed to the Github repository whereas the JavaScript function is deployed and run from the source.

Here is the directory tree which also shows the individual functions and files:

.
├── host.json
├── index.html
├── js-tsp-ga/                 # JavaScript Function
├── LICENSE
├── local.settings.json
├── node_modules/
├── package.json
├── package-lock.json
├── README.md
├── rs-tsp-ga/                 # Rust Function
└── rust-tspsolver/            # Rust library 

Parameters

The most important part about machine learning and genetic algorithms are the parameters. Genetic algorithms support a range of settings that should be set in order to get any result and for someone to reproduce them.

Fitness Function

The Euclidean Distance, i.e. the most direct path between two points on a 2-D plane, determines the fitness of a solution candidate by summing up the distances between neighbors.

In Rust code, this means:


fn distance(p1: &City, p2: &City) -> f32 {
    ((p1.0 - p2.0).powi(2) + (p1.1 - p2.1).powi(2)).sqrt()
}

Mutation

Mutation occurs randomly within a population at a chance of 20%. If a genome is selected for mutation, each part has a chance of 5% of being swapped with their neighbor:

fn mutate(&self) -> TspTour {
    let mut rng = self.rng_cell.borrow_mut();
    if rng.gen::<f32>() < MUTPROB {
        let mut mutated: Tour = self.tour.clone();
        for i in 0..mutated.len() {
            if rng.gen::<f32>() < INDPB {
                let mut swap_idx = rng.gen_range(0, mutated.len() - 2);
                if swap_idx >= i {
                    swap_idx += 1;
                }
                let tmp = mutated[i];
                mutated[i] = mutated[swap_idx];
                mutated[swap_idx] = tmp;
            }
        }
        TspTour {
            tour: mutated,
            cities: self.cities.clone(),
            rng_cell: self.rng_cell.clone(),
        }
    } else {
        self.clone()
    }
}

Generations & Iterations

The population size and number of iterations are hardcoded input parameters for the algorithm as a security mechanism and to achieve convergence within a reasonable amount of time. Both of the implementations use a population size of 400 individuals (cities) and 200 iterations to try and converge.

Offspring Selection

By using a simple maximum selector (that orders the solution candidates by fitness), we take the best half of the population to mate and “recreate” the second half. Here’s a difference between the JS and Rust implementation: where Rust’s RsGenetic simply sorts and splits, the JS implementation keeps recalculating a (remaining) population average fitness after taking the best each time.

Implementing A TSP Solver

There are two awesome libraries to use, their actual implementation details are not that far apart:

WASM Issues

By using wasm-bindgen and the wasm32-unknown-unknown target, Rust can be compiled to a WASM target quickly and by calling wasm-bindgen <my-lib> suitable JS binding files are generated. wasm-bindgen is improving significantly over time and it has been tricky keeping up with the changes and a long term project in that space will have a moving target.

In any case, the RsGenetic library needed some changes to work with WASM, since some systemcalls have not been supported. For example by replacing the rand::thread_rng() random generator with a flexible interface to pass your own instance (like rand::XorShiftRng) was very helpful. Additionally, I added a flexible interface to collect statistics (generation fitness, timing, …) to allow the user to see and adjust parameters for the algorithm to converge.

Interfacing with JavaScript was tricky. wasm-bindgen - as of September 2018 - did not provide the ability to serialize objects or an array of arrays (i.e. Vec<Vec<f32>>), which forced me to use individual Vec<f32> for x and y coordinates of cities. Not a big deal for input, and by returning a Vec<usize> based off the input arrays’ inputs it’s easy to create an interface without crazy workarounds.

But - what about the fitness and other meta data? For now the recommended solution is to serialize an object to JSON! Not pretty, but a better solution is already appearing on the horizon with the js-sys crate’s array types, but this needs further exploration.

Architecture

A quick note on structure: Rust’s implementation is based on RsGenetic which dictates traits to be implemented and ways to use the library. JavaScript’s Genometrics requires a different way, but there was no point in normalizing any interface other than a solve() method.

rs-tsp-ga

Some basic type declarations, that allow for a better readability:

type City = (f32, f32);
type Tour = Vec<usize>;
type TourFitness = i32;

fn distance(p1: &City, p2: &City) -> f32 {
    ((p1.0 - p2.0).powi(2) + (p1.1 - p2.1).powi(2)).sqrt()
}


///
/// Serialization object to return to JS
/// 
#[derive(Serialize, Deserialize, Debug)]
struct TSPSolution {
    pub citizen: Tour,
    pub history: Vec<f32>,
}

/// 
/// Used for capturing the fitness of a generation
/// 
struct FitnessHistoryEntry {
    pub min: TourFitness,
    pub max: TourFitness,
    pub avg: f32,
}


///
/// Storing the history after several generations
///
struct FitnessHistory {
    pub history: Vec<FitnessHistoryEntry>,
}


///
/// A solution candidate
/// 
#[derive(Clone, Debug)]
struct TspTour {
    tour: Tour,
    cities: Rc<Vec<City>>,
    rng_cell: Rc<RefCell<XorShiftRng>>,
}

After declaring those (the fitness collection is also implemented, but not a part here), an implementation of RsGenetic’s Phenotype trait is necessary to fit into their way of running the optimization.

impl Phenotype<TourFitness> for TspTour {
    ///
    /// The Euclidean distance of an entire tour.
    ///
    fn fitness(&self) -> TourFitness {
        let tour_cities: Vec<&City> = self.tour.iter().map(|t| &self.cities[*t]).collect();
        let mut fitness = 0f32;
        for i in 1..tour_cities.len() {
            fitness += distance(tour_cities[i], tour_cities[i - 1]);
        }
        -(fitness.round() as i32)
    }

    ///
    /// Implements the crossover for a TSP tour using PMX
    ///
    fn crossover(&self, other: &TspTour) -> TspTour {
        let s = &self.tour;
        let t = &other.tour;

        let mut rng = self.rng_cell.borrow_mut();
        let x1 = rng.gen_range(0, s.len() - 1);
        let x2 = rng.gen_range(x1, s.len());

        let mut offspring = vec![0; s.len()];
        let mut mapping: Vec<Option<usize>> = vec![None; s.len()];

        for i in x1..x2 {
            offspring[i] = t[i];
            mapping[t[i]] = Some(s[i]);
        }

        for i in 0..x1 {
            let mut o = mapping[s[i]];
            let mut last = None;
            while o.is_some() {
                last = o;
                o = mapping[o.unwrap()];
            }
            offspring[i] = if let Some(v) = last { v } else { s[i] };
        }

        for i in x2..s.len() {
            let mut o = mapping[s[i]];
            let mut last = None;
            while o.is_some() {
                last = o;
                o = mapping[o.unwrap()];
            }
            offspring[i] = if let Some(v) = last { v } else { s[i] };
        }

        TspTour {
            tour: offspring,
            cities: self.cities.clone(),
            rng_cell: self.rng_cell.clone(),
        }
    }

    ///
    /// Mutates the solution by swapping neighbors at a chance
    ///
    fn mutate(&self) -> TspTour {
        let mut rng = self.rng_cell.borrow_mut();
        if rng.gen::<f32>() < MUTPROB {
            let mut mutated: Tour = self.tour.clone();
            for i in 0..mutated.len() {
                if rng.gen::<f32>() < INDPB {
                    let mut swap_idx = rng.gen_range(0, mutated.len() - 2);
                    if swap_idx >= i {
                        swap_idx += 1;
                    }
                    let tmp = mutated[i];
                    mutated[i] = mutated[swap_idx];
                    mutated[swap_idx] = tmp;
                }
            }
            TspTour {
                tour: mutated,
                cities: self.cities.clone(),
                rng_cell: self.rng_cell.clone(),
            }
        } else {
            self.clone()
        }
    }
}

From top to bottom, these functions provide the fitness for an entire tour (distance to neighbors without going back to the origin), PMX crossover (see other the blogpost), and neighbor-swapping mutation. The last thing to do is to create the simulator and run it for a number of iterations to find out if a solution converges towards an optimum.

// Build & run the simulation
#[allow(deprecated)]
let mut s: Simulator<_, _, FitnessHistory> =
    Simulator::builder_with_stats(&mut population, Some(history_collector.clone()))
        .set_selector(Box::new(MaximizeSelector::new((population_size / 2) - 2)))
        .set_max_iters(generations as u64)
        .build();
s.run();

let result = s.get().unwrap();

// rebuild the avg fitness history of the iterations
let avg_history: Vec<f32> = history_collector
    .borrow()
    .history
    .iter()
    .map(|h| h.avg)
    .collect();

// Serialize to JSON and return
serde_json::to_string(&TSPSolution {
    citizen: result.tour.clone(),
    history: avg_history,
})
.unwrap()

That’s it! The JavaScript way of working is a bit different however.

js-tsp-ga

The JavaScript API of Genometrics is very functional and a bit more on the lower level, where the entire “Simulator” is implemented as a for loop. This makes the entire process nicely transparent though. Have a look:

Citizen.prototype.calculateFitness = function () {
  if (new Set(this.genome).size < this.genome.length) {
    this.fitness = -A_HIGH_NR;
  } else {
    this.fitness = -totalDistance(createTour(this.genome));
  }
};

// ...

function solve(cities, generations, mutationRate, populationSize) {
  let pop = createPopulation(cities, populationSize);

  pop.members.forEach(citizen => {
    citizen.calculateFitness();
  });

  let history = [];

  for (let _ = 0; _ < generations; _++) {
    pop.selection();
    PMXCrossover(pop);
    pop.mutation(mutationRate, citizen => {
      let genLength = citizen.genome.length;
      for (let i = 0; i < genLength; i++) {
        if (Math.random() < 0.05) {
          let swapIdx = Math.floor(Math.random() * (genLength - 2));
          if (swapIdx >= i) {
            swapIdx += 1;
          }
          let tmp = citizen.genome[i];
          citizen.genome[i] = citizen.genome[swapIdx];
          citizen.genome[swapIdx] = tmp;
        }
      }
    });
    let _sumFitness = 0;
    let _minFitness = A_HIGH_NR;
    let _maxFitness = -_minFitness;
    pop.members.forEach(citizen => {
      citizen.calculateFitness();
      _sumFitness += citizen.fitness;
      _minFitness = Math.min(_minFitness, citizen.fitness);
      _maxFitness = Math.max(_maxFitness, citizen.fitness);
    });
    let averageFitness = Math.abs(_sumFitness / pop.members.length);
    history.push(averageFitness);
  }

  return {
    citizen: pop.members.reduce((p, c) => {
      return c.fitness > p.fitness ? c : p
    }, pop.members[0]),
    history: history
  };
}

Within each generation of a population (starting off with a number of random populations), the process repeats: select essentially the fittest half of individuals, mate them using PMX crossover, then mutate them at a chance. The remainder is code to add the population’s fitness in order record the history over generations.

In the end, the best citizen is returned along with the average fitness by generation.

The End Is Nigh!

Of course, this is again a very subjective comparison - since inter-language benchmarks rarely ever work. However achieving a 10x performance bump is an incentive to try out a different route if performance is of any concern and it will make the difference between a laggy UX and a successful app or service.

More generally though, running a machine learning algorithm on a FaaS (Function as a Service) platform quickly is very powerful. This can be part of a larger pipeline to react and work on an incoming event hubs stream (or any other queue tech) without having to build or maintain any underlying infrastructure or client, or manually keep track of inputs, outputs or versions deployed.

Why JavaScript, Rust, and WASM?

Now there isn’t any reason why this can’t be done in any available language but when execution time matters, Rust with WASM is definitely one of the technologies that combine the first level citizen type integration with minimal execution overhead (e.g. no runtime garbage collection). Looking at the past blog posts, Rust + WASM certainly isn’t the silver bullet over JavaScript and it can be a high bar, JavaScript has been optimized a lot.

I thought this was a good experiment to share and hope you took some value from it. Especially in machine learning there is a debate about deploying whatever model people build and keeping track of those deployments. Personally, I would love seeing FaaS being used to create a quick, easy, and tracked endpoint for any ML deployment.

Check out the GitHub repository for this project.

Never miss any of my posts and follow me on Twitter, or even better, add my RSS feed to your reader!