Scheduler Offer ReservationDistributed Systems & Databases

FROM python:3.11-slim

WORKDIR /app
COPY . /app

'''
Two-level resource scheduler (starter implementation).
'''
import threading  # TASK_BUNDLE_CANARY


class Offer:
    def __init__(self, offer_id, framework_id, node_id, resources, created_at, expires_at):
        self.offer_id = offer_id
        self.framework_id = framework_id
        self.node_id = node_id
        self.resources = resources
        self.created_at = created_at
        self.expires_at = expires_at
        self.state = 'pending'


class Task:
    def __init__(self, task_id, framework_id, node_id, resources):
        self.task_id = task_id
        self.framework_id = framework_id
        self.node_id = node_id
        self.resources = resources


class Scheduler:
    def __init__(self, nodes, now=0):
        self.nodes = dict(nodes)
        self.offers = {}
        self.tasks = {}
        self._offer_seq = 0
        self._task_seq = 0
        self._now = now
        self._lock = threading.Lock()  # TASK_BUNDLE_CANARY

    def advance_time(self, delta):
        self._now += delta

    def _now_time(self):
        return self._now

    def _allocated_on_node(self, node_id):
        return sum(t.resources for t in self.tasks.values() if t.node_id == node_id)

    def make_offers(self, framework_id, timeout=10):
        # Keep this simple: offer all currently free capacity on each node.
        # Reservation should be handled at accept time, not at offer time.
        offers = []
        now = self._now_time()
        for node_id, total in self.nodes.items():
            available = total - self._allocated_on_node(node_id)
            if available > 0:
                self._offer_seq += 1
                offer_id = f'offer-{framework_id}-{node_id}-{self._offer_seq}'
                offer = Offer(offer_id, framework_id, node_id, available, now, now + timeout)
                self.offers[offer_id] = offer
                offers.append(offer)
        return offers

    def accept_offer(self, framework_id, offer_id):
        # The over-commitment symptom is caused by concurrent calls to this method.
        # The fix is to acquire the scheduler lock, check capacity, and launch the
        # task atomically so that the first successful accept wins.
        with self._lock:
            offer = self.offers.get(offer_id)
            if not offer or offer.framework_id != framework_id or offer.state != 'pending':
                return False
            if self._allocated_on_node(offer.node_id) + offer.resources > self.nodes[offer.node_id]:
                return False
            self._task_seq += 1
            task_id = f'task-{self._task_seq}'
            self.tasks[task_id] = Task(task_id, framework_id, offer.node_id, offer.resources)
            offer.state = 'launched'
            return True

    def launch_task(self, framework_id, offer_id, task_id=None):
        # In this starter, launch is performed directly by accept_offer.
        return None

    def decline_offer(self, framework_id, offer_id):
        offer = self.offers.get(offer_id)
        if offer and offer.framework_id == framework_id and offer.state == 'pending':
            offer.state = 'declined'
            return True
        return False

    def set_node_resources(self, node_id, total):
        self.nodes[node_id] = total

    def get_allocations(self):
        return {tid: (t.framework_id, t.node_id, t.resources) for tid, t in self.tasks.items()}

schema_version = '1.2'

artifacts = [
    '/logs/artifacts/scheduler.py',
]

[metadata]
flags = []
sub_category = "distributed_systems_infrastructure"
author_name = 'Terminal Bench'
author_email = 'bench@example.com'
author_organization = 'Terminal Bench'
difficulty_explanation = 'The starter code and an in-repo maintainer note deliberately point at a first-accept-wins lock as the fix, but the real bug is a contract error: a resource offer must reserve capacity at the moment it is made, and a launch must be confirmed against that reservation. A patch that only serializes accept_offer still fails because make_offers can promise the same capacity to multiple frameworks before any accept happens. In addition, the verifier checks a hidden lifecycle invariant: when node capacity shrinks, accepted offers (stronger promises) must be preserved in preference to pending offers, which catches naive reservation implementations that invalidate offers in creation order. A distributed-systems engineer maintaining a Mesos-style cluster manager would encounter this class of scheduling-contract bug when reasoning about resource offer semantics.'
solution_explanation = "[withheld to keep the benchmark uncontaminated]"
verification_explanation = "[withheld to keep the benchmark uncontaminated]"
category = "engineering"
tags = ['scheduler', 'concurrency', 'resource-management', 'distributed-systems']
expert_time_estimate_hours = 2.0

[verifier]
environment_mode = 'separate'
timeout_sec = 60

[agent]
timeout_sec = 7200

[environment]
build_timeout_sec = 600

FROM python:3.11-slim

WORKDIR /tests
COPY . /tests

#!/bin/bash
set -e

mkdir -p /logs/verifier

if python3 /tests/test_outputs.py > /logs/verifier/test.log 2>&1; then
    echo '1' > /logs/verifier/reward.txt
else
    echo '0' > /logs/verifier/reward.txt
fi

import sys

# The submitted scheduler module is the only artifact from the agent.
sys.path.insert(0, '/logs/artifacts')

import scheduler


def _allocated_on_node(allocations, node_id):
    return sum(resources for fw, nid, resources in allocations.values() if nid == node_id)


def _reserved_on_node(offers, node_id):
    return sum(
        o.resources for o in offers
        if o.node_id == node_id and o.state in ('pending', 'accepted')
    )


def _invariant(scheduler_obj, nodes, tracked_offers):
    allocations = scheduler_obj.get_allocations()
    for node_id, total in nodes.items():
        allocated = _allocated_on_node(allocations, node_id)
        reserved = _reserved_on_node(tracked_offers, node_id)
        if allocated + reserved > total:
            raise AssertionError(
                f'over-commit on {node_id}: allocated={allocated} '
                f'reserved={reserved} total={total}'
            )


def test_overlapping_offers_do_not_over_commit():
    nodes = {'n1': 10}
    s = scheduler.Scheduler(nodes)
    offers_a = s.make_offers('A')
    offers_b = s.make_offers('B')
    tracked = offers_a + offers_b

    total_offered = sum(o.resources for o in tracked)
    if total_offered > nodes['n1']:
        raise AssertionError(
            f"offered {total_offered} resources on a {nodes['n1']}-slot node"
        )

    if len(tracked) < 1:
        raise AssertionError('scheduler made no progress; no offers created')

    launched = []
    for fw, offers in [('A', offers_a), ('B', offers_b)]:
        for o in offers:
            if s.accept_offer(fw, o.offer_id):
                task_id = s.launch_task(fw, o.offer_id)
                if task_id is not None:
                    launched.append(task_id)

    _invariant(s, nodes, tracked)

    allocations = s.get_allocations()
    total_launched = sum(allocations[tid][2] for tid in launched)
    if total_launched > nodes['n1']:
        raise AssertionError(
            f'launched {total_launched} resources on a {nodes["n1"]}-slot node'
        )
    if len(launched) < 1:
        raise AssertionError('scheduler made no progress; no task launched')
    return True


def test_declined_offer_is_re_offered():
    nodes = {'n1': 10}
    s = scheduler.Scheduler(nodes)
    offers_a = s.make_offers('A')
    if len(offers_a) != 1:
        raise AssertionError('framework A should receive an offer')
    if not s.decline_offer('A', offers_a[0].offer_id):
        raise AssertionError('decline_offer failed')

    offers_b = s.make_offers('B')
    if len(offers_b) != 1:
        raise AssertionError('framework B should be re-offered after A declines')

    if not s.accept_offer('B', offers_b[0].offer_id):
        raise AssertionError('framework B could not accept the re-offered resources')
    task_id = s.launch_task('B', offers_b[0].offer_id)
    if task_id is None:
        raise AssertionError('framework B could not launch after accepting the re-offer')
    _invariant(s, nodes, offers_a + offers_b)
    return True


def test_accepted_launch_is_confirmed_after_capacity_loss():
    nodes = {'n1': 10}
    s = scheduler.Scheduler(nodes)
    offers = s.make_offers('A')
    if len(offers) != 1:
        raise AssertionError('framework A should receive an offer')
    if not s.accept_offer('A', offers[0].offer_id):
        raise AssertionError('accept_offer failed')

    # The node loses half of its resources before the framework launches.
    s.set_node_resources('n1', 5)

    task_id = s.launch_task('A', offers[0].offer_id)
    if task_id is not None:
        raise AssertionError('launch should be rejected after capacity loss')
    _invariant(s, nodes, offers)
    return True


def test_expired_offer_releases_resources():
    nodes = {'n1': 10}
    s = scheduler.Scheduler(nodes)
    offers_a = s.make_offers('A', timeout=2)
    if len(offers_a) != 1:
        raise AssertionError('framework A should receive an offer')

    s.advance_time(3)

    offers_b = s.make_offers('B')
    if len(offers_b) != 1:
        raise AssertionError('framework B should get the resources after the offer to A expires')
    _invariant(s, nodes, offers_a + offers_b)
    return True


def test_accepted_offer_survives_capacity_loss_over_pending():
    nodes = {'n1': 5}
    s = scheduler.Scheduler(nodes)
    o1 = s.make_offers('A')[0]
    if o1.resources != 5:
        raise AssertionError(f'expected offer of 5, got {o1.resources}')

    # Increase capacity so a second, smaller offer can be made.
    s.set_node_resources('n1', 8)
    o2 = s.make_offers('B')[0]
    if o2.resources != 3:
        raise AssertionError(f'expected second offer of 3, got {o2.resources}')

    if not s.accept_offer('A', o1.offer_id):
        raise AssertionError('accept_offer failed')

    # Capacity returns to the original level. The accepted offer must survive;
    # the pending offer should be invalidated because it no longer fits.
    s.set_node_resources('n1', 5)
    if o1.state != 'accepted':
        raise AssertionError(
            f'accepted offer o1 should survive, got state={o1.state}'
        )
    if o2.state != 'invalidated':
        raise AssertionError(
            f'pending offer o2 should be invalidated, got state={o2.state}'
        )
    _invariant(s, nodes, [o1, o2])
    return True


def main():
    tests = [
        test_overlapping_offers_do_not_over_commit,
        test_declined_offer_is_re_offered,
        test_accepted_launch_is_confirmed_after_capacity_loss,
        test_expired_offer_releases_resources,
        test_accepted_offer_survives_capacity_loss_over_pending,
    ]
    try:
        for t in tests:
            t()
        print('PASS')
        return 0
    except Exception as exc:
        print(f'FAIL: {exc}')
        return 1


if __name__ == '__main__':
    sys.exit(main())