Support setting container affinities

.changeset/real-bottles-report.md

@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+Support setting container affinities

packages/containers-shared/src/client/models/ApplicationAffinities.ts

@@ -3,11 +3,13 @@
 /* eslint-disable */
 
 import type { ApplicationAffinityColocation } from "./ApplicationAffinityColocation";
+import type { ApplicationAffinityHardwareGeneration } from "./ApplicationAffinityHardwareGeneration";
 
 /**
  * Defines affinity in application scheduling. (This still an experimental feature, some schedulers might not work with these affinities).
  *
  */
 export type ApplicationAffinities = {
 	colocation?: ApplicationAffinityColocation;
+	hardware_generation?: ApplicationAffinityHardwareGeneration;
 };

packages/containers-shared/src/client/models/ApplicationAffinityHardwareGeneration.ts

@@ -0,0 +1,11 @@
+/* istanbul ignore file */
+/* tslint:disable */
+/* eslint-disable */
+
+/**
+ * To the extend possible, prefer nodes with specified characteristics when placing application instances.
+ *
+ */
+export enum ApplicationAffinityHardwareGeneration {
+	HIGHEST_OVERALL_PERFORMANCE = "highest-overall-performance",
+}

packages/containers-shared/src/types.ts

@@ -1,4 +1,9 @@
-import type { InstanceType, SchedulingPolicy } from "./client";
+import type {
+	ApplicationAffinityColocation,
+	InstanceType,
+	SchedulingPolicy,
+} from "./client";
+import type { ApplicationAffinityHardwareGeneration } from "./client/models/ApplicationAffinityHardwareGeneration";
 
 export interface Logger {
 	debug: (...args: unknown[]) => void;
@@ -71,6 +76,10 @@ export type SharedContainerConfig = {
 		cities?: string[];
 		tier: number | undefined;
 	};
+	affinities?: {
+		colocation?: ApplicationAffinityColocation;
+		hardware_generation?: ApplicationAffinityHardwareGeneration;
+	};
 	observability: { logs_enabled: boolean };
 } & InstanceTypeOrLimits;
 

packages/wrangler/src/__tests__/cloudchamber/apply.test.ts

@@ -1,4 +1,5 @@
 import {
+	ApplicationAffinityColocation,
 	getCloudflareContainerRegistry,
 	SchedulingPolicy,
 	SecretAccessType,
@@ -2015,4 +2016,82 @@ describe("cloudchamber apply", () => {
 		const app = await applicationReqBodyPromise;
 		expect(app.configuration?.instance_type).toEqual("standard");
 	});
+
+	test("updates affinities", async () => {
+		setIsTTY(false);
+		const registry = getCloudflareContainerRegistry();
+		writeWranglerConfig({
+			name: "my-container",
+			containers: [
+				{
+					name: "my-container-app",
+					instances: 3,
+					class_name: "DurableObjectClass",
+					image: `${registry}/hello:1.0`,
+					instance_type: "dev",
+					constraints: {
+						tier: 1,
+					},
+					affinities: {
+						hardware_generation: "highest-overall-performance",
+					},
+				},
+			],
+		});
+
+		mockGetApplications([
+			{
+				id: "abc",
+				name: "my-container-app",
+				instances: 3,
+				created_at: new Date().toString(),
+				version: 1,
+				account_id: "1",
+				scheduling_policy: SchedulingPolicy.REGIONAL,
+				configuration: {
+					image: `${registry}/hello:1.0`,
+					disk: {
+						size: "2GB",
+						size_mb: 2000,
+					},
+					vcpu: 0.0625,
+					memory: "256MB",
+					memory_mib: 256,
+				},
+				constraints: {
+					tier: 1,
+				},
+				affinities: {
+					colocation: ApplicationAffinityColocation.DATACENTER,
+				},
+			},
+		]);
+
+		const applicationReqBodyPromise = mockModifyApplication();
+		await runWrangler("cloudchamber apply");
+		expect(std.stdout).toMatchInlineSnapshot(`
+			"╭ Deploy a container application deploy changes to your application
+			│
+			│ Container application changes
+			│
+			├ EDIT my-container-app
+			│
+			│   scheduling_policy = \\"regional\\"
+			│     [containers.affinities]
+			│ -   colocation = \\"datacenter\\"
+			│ +   hardware_generation = \\"highest-overall-performance\\"
+			│     [containers.configuration]
+			│     image = \\"registry.cloudflare.com/some-account-id/hello:1.0\\"
+			│
+			│
+			│  SUCCESS  Modified application my-container-app
+			│
+			╰ Applied changes
+
+			"
+		`);
+		expect(std.stderr).toMatchInlineSnapshot(`""`);
+		const app = await applicationReqBodyPromise;
+		expect(app.configuration?.instance_type).toEqual("dev");
+	});
 });

packages/wrangler/src/__tests__/containers/config.test.ts

@@ -400,6 +400,9 @@ describe("getNormalizedContainerOptions", () => {
 						regions: ["us-east-1", "us-west-2"],
 						cities: ["NYC", "SF"],
 					},
+					affinities: {
+						hardware_generation: "highest-overall-performance",
+					},
 				},
 			],
 			durable_objects: {
@@ -429,6 +432,9 @@ describe("getNormalizedContainerOptions", () => {
 				regions: ["US-EAST-1", "US-WEST-2"],
 				cities: ["nyc", "sf"],
 			},
+			affinities: {
+				hardware_generation: "highest-overall-performance",
+			},
 			observability: {
 				logs_enabled: true,
 			},

packages/wrangler/src/__tests__/containers/deploy.test.ts

@@ -6,6 +6,7 @@ import {
 	InstanceType,
 	SchedulingPolicy,
 } from "@cloudflare/containers-shared";
+import { ApplicationAffinityHardwareGeneration } from "@cloudflare/containers-shared/src/client/models/ApplicationAffinityHardwareGeneration";
 import { http, HttpResponse } from "msw";
 import { clearCachedAccount } from "../../cloudchamber/locations";
 import { mockAccountV4 as mockContainersAccount } from "../cloudchamber/utils";
@@ -1497,6 +1498,143 @@ describe("wrangler deploy with containers", () => {
 			"
 		`);
 	});
+
+	describe("affinities", () => {
+		it("may be specified on creation", async () => {
+			mockGetVersion("Galaxy-Class");
+			writeWranglerConfig({
+				...DEFAULT_DURABLE_OBJECTS,
+				containers: [
+					{
+						...DEFAULT_CONTAINER_FROM_REGISTRY,
+						affinities: {
+							hardware_generation: "highest-overall-performance",
+						},
+					},
+				],
+			});
+
+			mockGetApplications([]);
+
+			mockCreateApplication();
+
+			await runWrangler("deploy index.js");
+
+			expect(cliStd.stdout).toMatchInlineSnapshot(`
+				"╭ Deploy a container application deploy changes to your application
+				│
+				│ Container application changes
+				│
+				├ NEW my-container
+				│
+				│   [[containers]]
+				│   name = \\"my-container\\"
+				│   scheduling_policy = \\"default\\"
+				│   instances = 0
+				│   max_instances = 10
+				│   rollout_active_grace_period = 0
+				│
+				│     [containers.configuration]
+				│     image = \\"docker.io/hello:world\\"
+				│     instance_type = \\"dev\\"
+				│
+				│     [containers.constraints]
+				│     tier = 1
+				│
+				│     [containers.affinities]
+				│     hardware_generation = \\"highest-overall-performance\\"
+				│
+				│     [containers.durable_objects]
+				│     namespace_id = \\"1\\"
+				│
+				│
+				│  SUCCESS  Created application my-container (Application ID: undefined)
+				│
+				╰ Applied changes
+
+				"
+			`);
+		});
+
+		it("may be specified on modification", async () => {
+			mockGetVersion("Galaxy-Class");
+			writeWranglerConfig({
+				...DEFAULT_DURABLE_OBJECTS,
+				containers: [
+					{
+						...DEFAULT_CONTAINER_FROM_REGISTRY,
+						affinities: {
+							hardware_generation: "highest-overall-performance",
+						},
+					},
+				],
+			});
+
+			mockGetApplications([
+				{
+					id: "abc",
+					name: "my-container",
+					instances: 0,
+					max_instances: 10,
+					created_at: new Date().toString(),
+					version: 1,
+					account_id: "1",
+					scheduling_policy: SchedulingPolicy.DEFAULT,
+					rollout_active_grace_period: 0,
+					configuration: {
+						image: "docker.io/hello:world",
+						disk: {
+							size: "2GB",
+							size_mb: 2000,
+						},
+						vcpu: 0.0625,
+						memory: "256MB",
+						memory_mib: 256,
+					},
+					constraints: {
+						tier: 1,
+					},
+					durable_objects: {
+						namespace_id: "1",
+					},
+				},
+			]);
+
+			mockModifyApplication({
+				affinities: {
+					hardware_generation:
+						ApplicationAffinityHardwareGeneration.HIGHEST_OVERALL_PERFORMANCE,
+				},
+			});
+			mockCreateApplicationRollout({
+				description: "Progressive update",
+				strategy: "rolling",
+				kind: "full_auto",
+			});
+
+			await runWrangler("deploy index.js");
+
+			expect(cliStd.stdout).toMatchInlineSnapshot(`
+				"╭ Deploy a container application deploy changes to your application
+				│
+				│ Container application changes
+				│
+				├ EDIT my-container
+				│
+				│     [containers.constraints]
+				│     tier = 1
+				│ +   [containers.affinities]
+				│ +   hardware_generation = \\"highest-overall-performance\\"
+				│
+				│
+				│  SUCCESS  Modified application my-container (Application ID: abc)
+				│
+				╰ Applied changes
+
+				"
+			`);
+		});
+	});
 });
 
 // This is a separate describe block because we intentionally do not mock any

packages/wrangler/src/cloudchamber/apply.ts

@@ -42,6 +42,8 @@ import type {
 } from "../yargs-types";
 import type {
 	Application,
+	ApplicationAffinities,
+	ApplicationAffinityColocation,
 	ApplicationID,
 	ApplicationName,
 	CreateApplicationRequest,
@@ -50,6 +52,7 @@ import type {
 	Observability as ObservabilityConfiguration,
 	UserDeploymentConfiguration,
 } from "@cloudflare/containers-shared";
+import type { ApplicationAffinityHardwareGeneration } from "@cloudflare/containers-shared/src/client/models/ApplicationAffinityHardwareGeneration";
 import type { JsonMap } from "@iarna/toml";
 
 function mergeDeep<T>(target: T, source: Partial<T>): T {
@@ -220,6 +223,28 @@ function containerAppToInstanceType(
 	return configuration;
 }
 
+/**
+ * Perform type conversion of affinities so that they can be fed to the API.
+ */
+function convertContainerAffinitiesForApi(
+	container: ContainerApp
+): ApplicationAffinities | undefined {
+	if (container.affinities === undefined) {
+		return undefined;
+	}
+
+	const affinities: ApplicationAffinities = {
+		colocation: container.affinities?.colocation as
+			| ApplicationAffinityColocation
+			| undefined,
+		hardware_generation: container.affinities?.hardware_generation as
+			| ApplicationAffinityHardwareGeneration
+			| undefined,
+	};
+
+	return affinities;
+}
+
 function containerAppToCreateApplication(
 	accountId: string,
 	containerApp: ContainerApp,
@@ -267,6 +292,7 @@ function containerAppToCreateApplication(
 				region.toUpperCase()
 			),
 		},
+		affinities: convertContainerAffinitiesForApi(containerApp),
 	};
 
 	// delete the fields that should not be sent to API

packages/wrangler/src/config/environment.ts

@@ -150,6 +150,15 @@ export type ContainerApp = {
 		tier?: number;
 	};
 
+	/**
+	 * Scheduling affinities
+	 * @hidden
+	 */
+	affinities?: {
+		colocation?: "datacenter";
+		hardware_generation?: "highest-overall-performance";
+	};
+
 	// not used when deploying container with wrangler deploy
 	/**
 	 * @deprecated use the `class_name` field instead.

packages/wrangler/src/config/validation.ts

@@ -2709,6 +2709,7 @@ function validateContainerApp(
 					"instance_type",
 					"configuration",
 					"constraints",
+					"affinities",
 					"rollout_step_percentage",
 					"rollout_kind",
 					"durable_objects",