Always require upgrade budget for OS upgrades in controller

14 files changed, 51 insertions, 76 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java
index c6dddc0f223..5b2c2d74d20 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java
@@ -217,7 +217,7 @@ public class Controller extends AbstractComponent {
     }
 
     /** Set the target OS version for infrastructure on cloud in this system */
-    public void upgradeOsIn(CloudName cloudName, Version version, Optional<Duration> upgradeBudget, boolean force) {
+    public void upgradeOsIn(CloudName cloudName, Version version, Duration upgradeBudget, boolean force) {
         if (version.isEmpty()) {
             throw new IllegalArgumentException("Invalid version '" + version.toFullString() + "'");
         }
@@ -225,9 +225,6 @@ public class Controller extends AbstractComponent {
         if (!clouds.contains(cloudName)) {
             throw new IllegalArgumentException("Cloud '" + cloudName + "' does not exist in this system");
         }
-        if (!zoneRegistry.zones().ofCloud(cloudName).reprovisionToUpgradeOs().ids().isEmpty() && upgradeBudget.isEmpty()) {
-            throw new IllegalArgumentException("Cloud '" + cloudName.value() + "' requires a time budget for OS upgrades");
-        }
         try (Lock lock = curator.lockOsVersions()) {
             Set<OsVersionTarget> targets = new TreeSet<>(curator.readOsVersionTargets());
             if (!force && targets.stream().anyMatch(target -> target.osVersion().cloud().equals(cloudName) &&
@@ -239,7 +236,7 @@ public class Controller extends AbstractComponent {
             targets.add(new OsVersionTarget(new OsVersion(version, cloudName), upgradeBudget));
             curator.writeOsVersionTargets(targets);
             log.info("Triggered OS upgrade to " + version.toFullString() + " in cloud " +
-                     cloudName.value() + upgradeBudget.map(b -> ", with upgrade budget " + b).orElse(""));
+                     cloudName.value() + ", with upgrade budget " + upgradeBudget);
         }
     }
 
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeScheduler.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeScheduler.java
index 9fabe9f0e49..04b06b3e1f6 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeScheduler.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeScheduler.java
@@ -46,7 +46,7 @@ public class OsUpgradeScheduler extends ControllerMaintainer {
         for (var cloud : supportedClouds()) {
             Optional<Version> newTarget = newTargetIn(cloud);
             if (newTarget.isEmpty()) continue;
-            controller().upgradeOsIn(cloud, newTarget.get(), Optional.of(upgradeBudget()), false);
+            controller().upgradeOsIn(cloud, newTarget.get(), upgradeBudget(), false);
         }
         return true;
     }
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
index 79bea294472..e71fcf12b23 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
@@ -43,15 +43,14 @@ public class OsUpgrader extends InfrastructureUpgrader<OsVersionTarget> {
 
     @Override
     protected void upgrade(OsVersionTarget target, SystemApplication application, ZoneApi zone) {
-        Optional<Duration> zoneUpgradeBudget = target.upgradeBudget()
-                                                     .map(totalBudget -> zoneBudgetOf(totalBudget, zone));
+        Duration zoneUpgradeBudget = zoneBudgetOf(target.upgradeBudget(), zone);
         log.info(String.format("Upgrading OS of %s to version %s in %s in cloud %s%s", application.id(),
                                target.osVersion().version().toFullString(),
                                zone.getVirtualId(), zone.getCloudName(),
-                               zoneUpgradeBudget.map(d -> " with time budget " + d).orElse("")));
+                               " with time budget " + zoneUpgradeBudget));
         controller().serviceRegistry().configServer().nodeRepository().upgradeOs(zone.getVirtualId(), application.nodeType(),
                                                                                  target.osVersion().version(),
-                                                                                 zoneUpgradeBudget);
+                                                                                 Optional.of(zoneUpgradeBudget));
     }
 
     @Override
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializer.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializer.java
index 627ad9f709c..6942fc0549f 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializer.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializer.java
@@ -10,7 +10,6 @@ import com.yahoo.vespa.hosted.controller.versions.OsVersionTarget;
 
 import java.time.Duration;
 import java.util.Collections;
-import java.util.Optional;
 import java.util.Set;
 import java.util.TreeSet;
 
@@ -43,7 +42,9 @@ public class OsVersionTargetSerializer {
         Set<OsVersionTarget> osVersionTargets = new TreeSet<>();
         array.traverse((ArrayTraverser) (i, inspector) -> {
             OsVersion osVersion = osVersionSerializer.fromSlime(inspector);
-            Optional<Duration> upgradeBudget = Serializers.optionalDuration(inspector.field(upgradeBudgetField));
+            // TODO(mpolden): Require this field after 2021-05-01
+            Duration upgradeBudget = Serializers.optionalDuration(inspector.field(upgradeBudgetField))
+                                                .orElse(Duration.ZERO);
             osVersionTargets.add(new OsVersionTarget(osVersion, upgradeBudget));
         });
         return Collections.unmodifiableSet(osVersionTargets);
@@ -51,7 +52,7 @@ public class OsVersionTargetSerializer {
 
     private void toSlime(OsVersionTarget target, Cursor object) {
         osVersionSerializer.toSlime(target.osVersion(), object);
-        target.upgradeBudget().ifPresent(d -> object.setLong(upgradeBudgetField, d.toMillis()));
+        object.setLong(upgradeBudgetField, target.upgradeBudget().toMillis());
     }
 
 }
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java
index ac63f0cbcaf..b34c6685693 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java
@@ -137,8 +137,8 @@ public class OsApiHandler extends AuditLoggingRequestHandler {
         Inspector cloudField = root.field("cloud");
         Inspector upgradeBudgetField = root.field("upgradeBudget");
         boolean force = root.field("force").asBool();
-        if (!versionField.valid() || !cloudField.valid()) {
-            throw new IllegalArgumentException("Fields 'version' and 'cloud' are required");
+        if (!versionField.valid() || !cloudField.valid() || !upgradeBudgetField.valid()) {
+            throw new IllegalArgumentException("Fields 'version', 'cloud' and 'upgradeBudget' are required");
         }
 
         CloudName cloud = CloudName.from(cloudField.asString());
@@ -148,22 +148,17 @@ public class OsApiHandler extends AuditLoggingRequestHandler {
         } catch (IllegalArgumentException e) {
             throw new IllegalArgumentException("Invalid version '" + versionField.asString() + "'", e);
         }
-        Optional<Duration> upgradeBudget = Optional.of(upgradeBudgetField)
-                                                   .filter(Inspector::valid)
-                                                   .map(Inspector::asString).map(s -> {
-                    try {
-                        return Duration.parse(s);
-                    } catch (Exception e2) {
-                        throw new IllegalArgumentException("Invalid duration '" + s + "'", e2);
-                    }
-                });
-
+        Duration upgradeBudget;
+        try {
+            upgradeBudget = Duration.parse(upgradeBudgetField.asString());
+        } catch (Exception e) {
+            throw new IllegalArgumentException("Invalid duration '" + upgradeBudgetField.asString() + "'", e);
+        }
         controller.upgradeOsIn(cloud, target, upgradeBudget, force);
         Slime response = new Slime();
         Cursor cursor = response.setObject();
         cursor.setString("message", "Set target OS version for cloud '" + cloud.value() + "' to " +
-                                    target.toFullString() + upgradeBudget.map(d -> " with upgrade budget " + d)
-                                                                         .orElse(""));
+                                    target.toFullString() + " with upgrade budget " + upgradeBudget);
         return response;
     }
 
@@ -178,7 +173,7 @@ public class OsApiHandler extends AuditLoggingRequestHandler {
             currentVersionObject.setString("version", osVersion.version().toFullString());
             Optional<OsVersionTarget> target = targets.stream().filter(t -> t.osVersion().equals(osVersion)).findFirst();
             currentVersionObject.setBool("targetVersion", target.isPresent());
-            target.flatMap(OsVersionTarget::upgradeBudget).ifPresent(budget -> currentVersionObject.setString("upgradeBudget", budget.toString()));
+            target.ifPresent(t -> currentVersionObject.setString("upgradeBudget", t.upgradeBudget().toString()));
             currentVersionObject.setString("cloud", osVersion.cloud().value());
             Cursor nodesArray = currentVersionObject.setArray("nodes");
             nodeVersions.asMap().values().forEach(nodeVersion -> {
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionTarget.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionTarget.java
index bacd2ada298..35b6ac7a1a9 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionTarget.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionTarget.java
@@ -5,7 +5,6 @@ import org.jetbrains.annotations.NotNull;
 
 import java.time.Duration;
 import java.util.Objects;
-import java.util.Optional;
 
 /**
  * An {@link OsVersion} and its upgrade budget.
@@ -22,11 +21,12 @@ public class OsVersionTarget implements Comparable<OsVersionTarget> {
     //          - CHANGING THE FORMAT OF A FIELD: Don't do it bro.
 
     private final OsVersion osVersion;
-    private final Optional<Duration> upgradeBudget;
+    private final Duration upgradeBudget;
 
-    public OsVersionTarget(OsVersion osVersion, Optional<Duration> upgradeBudget) {
+    public OsVersionTarget(OsVersion osVersion, Duration upgradeBudget) {
         this.osVersion = Objects.requireNonNull(osVersion);
-        this.upgradeBudget = requireNotNegative(upgradeBudget);
+        this.upgradeBudget = Objects.requireNonNull(upgradeBudget);
+        if (upgradeBudget.isNegative()) throw new IllegalArgumentException("upgradeBudget cannot be negative");
     }
 
     /** The OS version contained in this target */
@@ -35,7 +35,7 @@ public class OsVersionTarget implements Comparable<OsVersionTarget> {
     }
 
     /** The total time budget across all zones for applying target, if any */
-    public Optional<Duration> upgradeBudget() {
+    public Duration upgradeBudget() {
         return upgradeBudget;
     }
 
@@ -53,13 +53,6 @@ public class OsVersionTarget implements Comparable<OsVersionTarget> {
         return Objects.hash(osVersion, upgradeBudget);
     }
 
-    private static Optional<Duration> requireNotNegative(Optional<Duration> duration) {
-        Objects.requireNonNull(duration);
-        if (duration.isEmpty()) return duration;
-        if (duration.get().isNegative()) throw new IllegalArgumentException("Duration cannot be negative");
-        return duration;
-    }
-
     @Override
     public int compareTo(@NotNull OsVersionTarget o) {
         return osVersion.compareTo(o.osVersion);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java
index 19a84ec93f0..1b43fe2aaa5 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java
@@ -24,7 +24,6 @@ import org.junit.Test;
 import java.time.Duration;
 import java.util.Comparator;
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
@@ -353,7 +352,7 @@ public class MetricsReporterTest {
 
         // All nodes upgrade to initial OS version
         var version0 = Version.fromString("8.0");
-        tester.controller().upgradeOsIn(cloud, version0, Optional.empty(), false);
+        tester.controller().upgradeOsIn(cloud, version0, Duration.ZERO, false);
         osUpgrader.maintain();
         tester.configServer().setOsVersion(version0, SystemApplication.tenantHost.id(), zone);
         tester.configServer().setOsVersion(version0, SystemApplication.configServerHost.id(), zone);
@@ -367,7 +366,7 @@ public class MetricsReporterTest {
             var currentVersion = i == 0 ? version0 : targets.get(i - 1);
             var nextVersion = targets.get(i);
             // System starts upgrading to next OS version
-            tester.controller().upgradeOsIn(cloud, nextVersion, Optional.empty(), false);
+            tester.controller().upgradeOsIn(cloud, nextVersion, Duration.ZERO, false);
             runAll(osUpgrader, statusUpdater, reporter);
             assertOsChangeDuration(Duration.ZERO, hosts);
             assertOsNodeCount(hosts.size(), currentVersion);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeSchedulerTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeSchedulerTest.java
index 4aed9b0bffe..76781d964a1 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeSchedulerTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgradeSchedulerTest.java
@@ -11,7 +11,6 @@ import org.junit.Test;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
-import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -38,7 +37,7 @@ public class OsUpgradeSchedulerTest {
 
         // Target is set
         Version version0 = Version.fromString("7.0.0.20210123190005");
-        tester.controller().upgradeOsIn(cloud, version0, Optional.of(Duration.ofDays(1)), false);
+        tester.controller().upgradeOsIn(cloud, version0, Duration.ofDays(1), false);
 
         // Target remains unchanged as it hasn't expired yet
         for (var interval : List.of(Duration.ZERO, Duration.ofDays(15))) {
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
index 42057e7880c..3e2fd4ec0b9 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
@@ -17,13 +17,11 @@ import org.junit.Test;
 import java.time.Duration;
 import java.util.Collection;
 import java.util.List;
-import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 
 /**
  * @author mpolden
@@ -67,8 +65,8 @@ public class OsUpgraderTest {
 
         // New OS version released
         Version version1 = Version.fromString("7.1");
-        tester.controller().upgradeOsIn(cloud1, Version.fromString("7.0"), Optional.empty(), false);
-        tester.controller().upgradeOsIn(cloud1, version1, Optional.empty(), false);
+        tester.controller().upgradeOsIn(cloud1, Version.fromString("7.0"), Duration.ZERO, false);
+        tester.controller().upgradeOsIn(cloud1, version1, Duration.ZERO, false);
         assertEquals(1, tester.controller().osVersionTargets().size()); // Only allows one version per cloud
         statusUpdater.maintain();
 
@@ -136,16 +134,10 @@ public class OsUpgraderTest {
         tester.configServer().bootstrap(List.of(zone1.getId(), zone2.getId(), zone3.getId(), zone4.getId()),
                                         nodeTypes);
 
-        // Upgrade without budget fails
-        Version version = Version.fromString("7.1");
-        try {
-            tester.controller().upgradeOsIn(cloud, version, Optional.empty(), false);
-            fail("Expected exception");
-        } catch (IllegalArgumentException ignored) {}
-
         // Upgrade with budget
-        tester.controller().upgradeOsIn(cloud, version, Optional.of(Duration.ofHours(12)), false);
-        assertEquals(Duration.ofHours(12), tester.controller().osVersionTarget(cloud).get().upgradeBudget().get());
+        Version version = Version.fromString("7.1");
+        tester.controller().upgradeOsIn(cloud, version, Duration.ofHours(12), false);
+        assertEquals(Duration.ofHours(12), tester.controller().osVersionTarget(cloud).get().upgradeBudget());
         statusUpdater.maintain();
         osUpgrader.maintain();
 
@@ -195,8 +187,8 @@ public class OsUpgraderTest {
 
         // New OS version released
         Version version = Version.fromString("7.1");
-        tester.controller().upgradeOsIn(cloud, Version.fromString("7.0"), Optional.empty(), false);
-        tester.controller().upgradeOsIn(cloud, version, Optional.empty(), false);
+        tester.controller().upgradeOsIn(cloud, Version.fromString("7.0"), Duration.ZERO, false);
+        tester.controller().upgradeOsIn(cloud, version, Duration.ZERO, false);
         statusUpdater.maintain();
 
         // zone 1 upgrades
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
index e9a1adcfe88..95f6af06230 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
@@ -11,7 +11,6 @@ import com.yahoo.vespa.hosted.controller.versions.OsVersionStatus;
 import org.junit.Test;
 
 import java.time.Duration;
-import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -41,7 +40,7 @@ public class OsVersionStatusUpdaterTest {
         // Setting a new target adds it to current status
         Version version1 = Version.fromString("7.1");
         CloudName cloud = CloudName.defaultName();
-        tester.controller().upgradeOsIn(cloud, version1, Optional.empty(), false);
+        tester.controller().upgradeOsIn(cloud, version1, Duration.ZERO, false);
         statusUpdater.maintain();
 
         var osVersions = tester.controller().osVersionStatus().versions();
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializerTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializerTest.java
index dea09998952..6a3632ec2c9 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializerTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/persistence/OsVersionTargetSerializerTest.java
@@ -9,7 +9,6 @@ import com.yahoo.vespa.hosted.controller.versions.OsVersionTarget;
 import org.junit.Test;
 
 import java.time.Duration;
-import java.util.Optional;
 import java.util.Set;
 
 import static org.junit.Assert.assertEquals;
@@ -23,8 +22,8 @@ public class OsVersionTargetSerializerTest {
     public void serialization() {
         OsVersionTargetSerializer serializer = new OsVersionTargetSerializer(new OsVersionSerializer());
         Set<OsVersionTarget> targets = ImmutableSet.of(
-                new OsVersionTarget(new OsVersion(Version.fromString("7.1"), CloudName.defaultName()), Optional.empty()),
-                new OsVersionTarget(new OsVersion(Version.fromString("7.1"), CloudName.from("foo")), Optional.of(Duration.ofDays(1)))
+                new OsVersionTarget(new OsVersion(Version.fromString("7.1"), CloudName.defaultName()), Duration.ZERO),
+                new OsVersionTarget(new OsVersion(Version.fromString("7.1"), CloudName.from("foo")), Duration.ofDays(1))
         );
         Set<OsVersionTarget> serialized = serializer.fromSlime(serializer.toSlime(targets));
         assertEquals(targets, serialized);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
index 857530df1fe..7364723f5f0 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
@@ -69,8 +69,8 @@ public class OsApiTest extends ControllerContainerTest {
         // All nodes are initially on empty version
         upgradeAndUpdateStatus();
         // Upgrade OS to a different version in each cloud
-        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.5.2\", \"cloud\": \"cloud1\"}", Request.Method.PATCH),
-                       "{\"message\":\"Set target OS version for cloud 'cloud1' to 7.5.2\"}", 200);
+        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.5.2\", \"cloud\": \"cloud1\", \"upgradeBudget\": \"PT0S\"}", Request.Method.PATCH),
+                       "{\"message\":\"Set target OS version for cloud 'cloud1' to 7.5.2 with upgrade budget PT0S\"}", 200);
         assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"8.2.1\", \"cloud\": \"cloud2\", \"upgradeBudget\": \"PT24H\"}", Request.Method.PATCH),
                        "{\"message\":\"Set target OS version for cloud 'cloud2' to 8.2.1 with upgrade budget PT24H\"}", 200);
 
@@ -85,27 +85,27 @@ public class OsApiTest extends ControllerContainerTest {
         assertFile(new Request("http://localhost:8080/os/v1/"), "versions-all-upgraded.json");
 
         // Downgrade with force is permitted
-        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.5.1\", \"cloud\": \"cloud1\", \"force\": true}", Request.Method.PATCH),
-                       "{\"message\":\"Set target OS version for cloud 'cloud1' to 7.5.1\"}", 200);
+        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.5.1\", \"cloud\": \"cloud1\", \"force\": true, \"upgradeBudget\": \"PT0S\"}", Request.Method.PATCH),
+                       "{\"message\":\"Set target OS version for cloud 'cloud1' to 7.5.1 with upgrade budget PT0S\"}", 200);
 
-        // Error: Missing field 'cloud' or 'version'
+        // Error: Missing fields
         assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.6\"}", Request.Method.PATCH),
-                       "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Fields 'version' and 'cloud' are required\"}", 400);
+                       "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Fields 'version', 'cloud' and 'upgradeBudget' are required\"}", 400);
         assertResponse(new Request("http://localhost:8080/os/v1/", "{\"cloud\": \"cloud1\"}", Request.Method.PATCH),
-                       "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Fields 'version' and 'cloud' are required\"}", 400);
+                       "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Fields 'version', 'cloud' and 'upgradeBudget' are required\"}", 400);
 
         // Error: Invalid versions
-        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": null, \"cloud\": \"cloud1\"}", Request.Method.PATCH),
+        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": null, \"cloud\": \"cloud1\", \"upgradeBudget\": \"PT0S\"}", Request.Method.PATCH),
                        "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Invalid version '0.0.0'\"}", 400);
-        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"foo\", \"cloud\": \"cloud1\"}", Request.Method.PATCH),
+        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"foo\", \"cloud\": \"cloud1\", \"upgradeBudget\": \"PT0S\"}", Request.Method.PATCH),
                        "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Invalid version 'foo': For input string: \\\"foo\\\"\"}", 400);
 
         // Error: Invalid cloud
-        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.6\", \"cloud\": \"foo\"}", Request.Method.PATCH),
+        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.6\", \"cloud\": \"foo\", \"upgradeBudget\": \"PT0S\"}", Request.Method.PATCH),
                        "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Cloud 'foo' does not exist in this system\"}", 400);
 
         // Error: Downgrade OS
-        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.4.1\", \"cloud\": \"cloud1\"}", Request.Method.PATCH),
+        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.4.1\", \"cloud\": \"cloud1\", \"upgradeBudget\": \"PT0S\"}", Request.Method.PATCH),
                        "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Cannot downgrade cloud 'cloud1' to version 7.4.1\"}", 400);
 
         // Request firmware checks in all zones.
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json
index 4350b70f09f..e1b70c780c5 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json
@@ -3,6 +3,7 @@
     {
       "version": "7.5.2",
       "targetVersion": true,
+      "upgradeBudget": "PT0S",
       "cloud": "cloud1",
       "nodes": [
         {
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json
index dcdecaaa2cc..a9186a9d733 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json
@@ -55,6 +55,7 @@
     {
       "version": "7.5.2",
       "targetVersion": true,
+      "upgradeBudget": "PT0S",
       "cloud": "cloud1",
       "nodes": [
         {