#!/usr/bin/env python3.11
# Copyright 2024, Gurobi Optimization, LLC
# Assign workers to shifts; each worker may or may not be available on a
# particular day. The optimization problem is solved as a batch, and
# the schedule constructed only from the meta data available in the solution
# JSON.
#
# NOTE: You'll need a license file configured to use a Cluster Manager
# for this example to run.
import time
import json
import sys
import gurobipy as gp
from gurobipy import GRB
from collections import OrderedDict, defaultdict
# For later pretty printing names for the shifts
shiftname = OrderedDict(
[
("Mon1", "Monday 8:00"),
("Mon8", "Monday 14:00"),
("Tue2", "Tuesday 8:00"),
("Tue9", "Tuesday 14:00"),
("Wed3", "Wednesday 8:00"),
("Wed10", "Wednesday 14:00"),
("Thu4", "Thursday 8:00"),
("Thu11", "Thursday 14:00"),
("Fri5", "Friday 8:00"),
("Fri12", "Friday 14:00"),
("Sat6", "Saturday 8:00"),
("Sat13", "Saturday 14:00"),
("Sun7", "Sunday 9:00"),
("Sun14", "Sunday 12:00"),
]
)
# Build the assignment problem in a Model, and submit it for batch optimization
#
# Required input: A Cluster Manager environment setup for batch optimization
def submit_assigment_problem(env):
# Number of workers required for each shift
shifts, shiftRequirements = gp.multidict(
{
"Mon1": 3,
"Tue2": 2,
"Wed3": 4,
"Thu4": 4,
"Fri5": 5,
"Sat6": 5,
"Sun7": 3,
"Mon8": 2,
"Tue9": 2,
"Wed10": 3,
"Thu11": 4,
"Fri12": 5,
"Sat13": 7,
"Sun14": 5,
}
)
# Amount each worker is paid to work one shift
workers, pay = gp.multidict(
{
"Amy": 10,
"Bob": 12,
"Cathy": 10,
"Dan": 8,
"Ed": 8,
"Fred": 9,
"Gu": 11,
}
)
# Worker availability
availability = gp.tuplelist(
[
("Amy", "Tue2"),
("Amy", "Wed3"),
("Amy", "Thu4"),
("Amy", "Sun7"),
("Amy", "Tue9"),
("Amy", "Wed10"),
("Amy", "Thu11"),
("Amy", "Fri12"),
("Amy", "Sat13"),
("Amy", "Sun14"),
("Bob", "Mon1"),
("Bob", "Tue2"),
("Bob", "Fri5"),
("Bob", "Sat6"),
("Bob", "Mon8"),
("Bob", "Thu11"),
("Bob", "Sat13"),
("Cathy", "Wed3"),
("Cathy", "Thu4"),
("Cathy", "Fri5"),
("Cathy", "Sun7"),
("Cathy", "Mon8"),
("Cathy", "Tue9"),
("Cathy", "Wed10"),
("Cathy", "Thu11"),
("Cathy", "Fri12"),
("Cathy", "Sat13"),
("Cathy", "Sun14"),
("Dan", "Tue2"),
("Dan", "Thu4"),
("Dan", "Fri5"),
("Dan", "Sat6"),
("Dan", "Mon8"),
("Dan", "Tue9"),
("Dan", "Wed10"),
("Dan", "Thu11"),
("Dan", "Fri12"),
("Dan", "Sat13"),
("Dan", "Sun14"),
("Ed", "Mon1"),
("Ed", "Tue2"),
("Ed", "Wed3"),
("Ed", "Thu4"),
("Ed", "Fri5"),
("Ed", "Sat6"),
("Ed", "Mon8"),
("Ed", "Tue9"),
("Ed", "Thu11"),
("Ed", "Sat13"),
("Ed", "Sun14"),
("Fred", "Mon1"),
("Fred", "Tue2"),
("Fred", "Wed3"),
("Fred", "Sat6"),
("Fred", "Mon8"),
("Fred", "Tue9"),
("Fred", "Fri12"),
("Fred", "Sat13"),
("Fred", "Sun14"),
("Gu", "Mon1"),
("Gu", "Tue2"),
("Gu", "Wed3"),
("Gu", "Fri5"),
("Gu", "Sat6"),
("Gu", "Sun7"),
("Gu", "Mon8"),
("Gu", "Tue9"),
("Gu", "Wed10"),
("Gu", "Thu11"),
("Gu", "Fri12"),
("Gu", "Sat13"),
("Gu", "Sun14"),
]
)
# Start environment, get model in this environment
with gp.Model("assignment", env=env) as m:
# Assignment variables: x[w,s] == 1 if worker w is assigned to shift s.
# Since an assignment model always produces integer solutions, we use
# continuous variables and solve as an LP.
x = m.addVars(availability, ub=1, name="x")
# Set tags encoding the assignments for later retrieval of the schedule.
# Each tag is a JSON string of the format
# {
# "Worker": "<Name of the worker>",
# "Shift": "String representation of the shift"
# }
#
for k, v in x.items():
name, timeslot = k
d = {"Worker": name, "Shift": shiftname[timeslot]}
v.VTag = json.dumps(d)
# The objective is to minimize the total pay costs
m.setObjective(
gp.quicksum(pay[w] * x[w, s] for w, s in availability), GRB.MINIMIZE
)
# Constraints: assign exactly shiftRequirements[s] workers to each shift
reqCts = m.addConstrs(
(x.sum("*", s) == shiftRequirements[s] for s in shifts), "_"
)
# Submit this model for batch optimization to the cluster manager
# and return its batch ID for later querying the solution
batchID = m.optimizeBatch()
return batchID
# Wait for the final status of the batch.
# Initially the status of a batch is "submitted"; the status will change
# once the batch has been processed (by a compute server).
def waitforfinalbatchstatus(batch):
# Wait no longer than ten seconds
maxwaittime = 10
starttime = time.time()
while batch.BatchStatus == GRB.BATCH_SUBMITTED:
# Abort this batch if it is taking too long
curtime = time.time()
if curtime - starttime > maxwaittime:
batch.abort()
break
# Wait for one second
time.sleep(1)
# Update the resident attribute cache of the Batch object with the
# latest values from the cluster manager.
batch.update()
# Print the schedule according to the solution in the given dict
def print_shift_schedule(soldict):
schedule = defaultdict(list)
# Iterate over the variables that take a non-zero value (i.e.,
# an assignment), and collect them per day
for v in soldict["Vars"]:
# There is only one VTag, the JSON dict of an assignment we passed
# in as the VTag
assignment = json.loads(v["VTag"][0])
schedule[assignment["Shift"]].append(assignment["Worker"])
# Print the schedule
for k in shiftname.values():
day, time = k.split()
workers = ", ".join(schedule[k])
print(f" - {day:10} {time:>5}: {workers}")
if __name__ == "__main__":
# Create Cluster Manager environment in batch mode.
env = gp.Env(empty=True)
env.setParam("CSBatchMode", 1)
# env is a context manager; upon leaving, Env.dispose() is called
with env.start():
# Submit the assignment problem to the cluster manager, get batch ID
batchID = submit_assigment_problem(env)
# Create a batch object, wait for batch to complete, query solution JSON
with gp.Batch(batchID, env) as batch:
waitforfinalbatchstatus(batch)
if batch.BatchStatus != GRB.BATCH_COMPLETED:
print("Batch request couldn't be completed")
sys.exit(0)
jsonsol = batch.getJSONSolution()
# Dump JSON solution string into a dict
soldict = json.loads(jsonsol)
# Has the assignment problem been solved as expected?
if soldict["SolutionInfo"]["Status"] != GRB.OPTIMAL:
# Shouldn't happen...
print("Assignment problem could not be solved to optimality")
sys.exit(0)
# Print shift schedule from solution JSON
print_shift_schedule(soldict)