feat(deadline): add custom user data commands to Worker instance star…

…tup (#239)

examples/deadline/All-In-AWS-Infrastructure-Basic/python/package/lib/compute_tier.py

@@ -18,17 +18,22 @@
     IVpc,
     Port
 )
+from aws_cdk.aws_s3_assets import (
+  Asset
+)
 
 from aws_rfdk import (
     HealthMonitor,
 )
 from aws_rfdk.deadline import (
+    InstanceUserDataProvider,
     IRenderQueue,
     UsageBasedLicense,
     UsageBasedLicensing,
     WorkerInstanceFleet,
 )
 
+import os
 
 @dataclass
 class ComputeTierProps(StackProps):
@@ -50,6 +55,30 @@ class ComputeTierProps(StackProps):
     # List of the usage-based liceses that the worker nodes will be served.
     licenses: Optional[List[UsageBasedLicense]] = None
 
+class UserDataProvider(InstanceUserDataProvider):
+    def __init__(self, scope: Construct, stack_id: str):
+        super().__init__(scope, stack_id)
+        self.test_script=Asset(scope, "SampleAsset",
+            path=os.path.join(os.getcwd(), "..", "scripts", "configure_worker.sh")
+        )
+
+    def pre_cloud_watch_agent(self, host) -> None:
+        host.user_data.add_commands("echo preCloudWatchAgent")
+
+    def pre_render_queue_configuration(self, host) -> None:
+        host.user_data.add_commands("echo preRenderQueueConfiguration")
+
+    def pre_worker_configuration(self, host) -> None:
+        host.user_data.add_commands("echo preWorkerConfiguration")
+
+    def post_worker_launch(self, host) -> None:
+        host.user_data.add_commands("echo postWorkerLaunch")
+        self.test_script.grant_read(host)
+        local_path = host.user_data.add_s3_download_command(
+            bucket=self.test_script.bucket,
+            bucket_key=self.test_script.s3_object_key
+        )
+        host.user_data.add_execute_file_command(file_path=local_path)
 
 class ComputeTier(Stack):
     """
@@ -85,6 +114,7 @@ def __init__(self, scope: Construct, stack_id: str, *, props: ComputeTierProps,
             worker_machine_image=props.worker_machine_image,
             health_monitor=self.health_monitor,
             key_name=props.key_pair_name,
+            user_data_provider=UserDataProvider(self, 'UserDataProvider')
         )
 
         if props.usage_based_licensing and props.licenses:

examples/deadline/All-In-AWS-Infrastructure-Basic/scripts/configure_worker.sh

@@ -0,0 +1,4 @@
+#!/bin/sh
+
+# This script is only here to illustrate how to add custom user data when using our WorkerInstanceFleet and WorkerInstanceConfiguration constructs
+mkdir test_dir 

examples/deadline/All-In-AWS-Infrastructure-Basic/ts/lib/compute-tier.ts

@@ -11,6 +11,8 @@ import {
 } from '@aws-cdk/aws-ec2';
 import * as cdk from '@aws-cdk/core';
 import {
+  IHost,
+  InstanceUserDataProvider,  
   IRenderQueue,
   IWorkerFleet,
   UsageBasedLicense,
@@ -21,6 +23,8 @@ import {
   HealthMonitor,
   IHealthMonitor,
 } from 'aws-rfdk';
+import { Asset } from '@aws-cdk/aws-s3-assets';
+import * as path from 'path'
 
 /**
  * Properties for {@link ComputeTier}.
@@ -62,6 +66,36 @@ export interface ComputeTierProps extends cdk.StackProps {
   readonly licenses?: UsageBasedLicense[];
 }
 
+class UserDataProvider extends InstanceUserDataProvider {
+  preCloudWatchAgent(host: IHost): void {
+    host.userData.addCommands('echo preCloudWatchAgent');
+  }
+  preRenderQueueConfiguration(host: IHost): void {
+    host.userData.addCommands('echo preRenderQueueConfiguration');
+  }  
+  preWorkerConfiguration(host: IHost): void {
+    host.userData.addCommands('echo preWorkerConfiguration');
+  }
+  postWorkerLaunch(host: IHost): void {
+    host.userData.addCommands('echo postWorkerLaunch');
+    if (host.node.scope != undefined) {
+      const testScript = new Asset(
+        host.node.scope as cdk.Construct, 
+        'SampleAsset',
+        {path: path.join(__dirname, '..', '..', 'scripts', 'configure_worker.sh')},
+      );
+      testScript.grantRead(host);
+      const localPath = host.userData.addS3DownloadCommand({
+        bucket: testScript.bucket,
+        bucketKey: testScript.s3ObjectKey,
+      });
+      host.userData.addExecuteFileCommand({
+        filePath: localPath,
+      })
+    }
+  }
+}
+
 /**
  * The computer tier consists of raw compute power. For a Deadline render farm,
  * this will be the fleet of Worker nodes that render Deadline jobs.
@@ -100,6 +134,7 @@ export class ComputeTier extends cdk.Stack {
       workerMachineImage: props.workerMachineImage,
       healthMonitor: this.healthMonitor,
       keyName: props.keyPairName,
+      userDataProvider: new UserDataProvider(this, 'UserDataProvider'),
     });
 
     if (props.usageBasedLicensing && props.licenses) {

packages/aws-rfdk/lib/deadline/README.md

@@ -319,3 +319,23 @@ const workerFleet = new WorkerInstanceFleet(stack, 'WorkerFleet', {
   }
 });
 ```
+
+### User data scripts for the Worker configuration
+
+You have possibility to run user data scripts at various points during the Worker configuration lifecycle.
+
+To do this, subclass `InstanceUserDataProvider` and override desired methods:
+```ts
+class UserDataProvider extends InstanceUserDataProvider {
+  preCloudWatchAgent(host: IHost): void {
+    host.userData.addCommands('echo preCloudWatchAgent');
+  }
+}
+const fleet = new WorkerInstanceFleet(stack, 'WorkerFleet', {
+  vpc,
+  renderQueue,
+  workerMachineImage: /* ... */,
+  userDataProvider: new UserDataProvider(stack, 'UserDataProvider'),
+});
+
+```

packages/aws-rfdk/lib/deadline/lib/worker-configuration.ts

@@ -29,6 +29,67 @@ import {
   Version,
 } from './version';
 
+/**
+ * Provider for adding user data scripts
+ * Methods of this interface will be invoked in WorkerInstanceConfiguration
+ * on different stages of worker configuration
+ */
+export interface IInstanceUserDataProvider {
+  /**
+   * Method that is invoked before configuring the Cloud Watch Agent.
+   */
+  preCloudWatchAgent(host: IHost): void;
+
+  /**
+   * Method that is invoked before the render queue configuration.
+   */
+  preRenderQueueConfiguration(host: IHost): void;
+
+  /**
+   * Method that is invoked after configuring the connection to the render queue and before configuring the Deadline Worker.
+   */
+  preWorkerConfiguration(host: IHost): void;
+
+  /**
+   * Method that is invoked after all configuration is done and worker started.
+   */
+  postWorkerLaunch(host: IHost): void;
+}
+
+/**
+ * Implementation of {@link IInstanceUserDataProvider}.
+ * Can be used as sub-class with override the desired methods
+ * to add custom user data commands for WorkerInstanceFleet or WorkerInstanceConfiguration.
+ */
+export class InstanceUserDataProvider extends Construct  implements IInstanceUserDataProvider{
+  constructor(scope: Construct, id: string) {
+    super(scope, id);
+  }
+
+  /**
+   * @inheritdoc
+   */
+  preCloudWatchAgent(_host: IHost): void {
+  }
+
+  /**
+   * @inheritdoc
+   */
+  preRenderQueueConfiguration(_host: IHost): void {
+  }
+
+  /**
+   * @inheritdoc
+   */
+  preWorkerConfiguration(_host: IHost): void {
+  }
+
+  /**
+   * @inheritdoc
+   */
+  postWorkerLaunch(_host: IHost): void {
+  }
+}
 /**
  * Configuration settings for Deadline Workers
  */
@@ -86,6 +147,12 @@ export interface WorkerInstanceConfigurationProps {
    * @default The Worker is assigned the default settings as outlined in the WorkerSettings interface.
    */
   readonly workerSettings?: WorkerSettings;
+
+  /**
+   * An optional provider of user data commands to be injected at various points during the Worker configuration lifecycle.
+   * You can provide a subclass of InstanceUserDataProvider with the methods overridden as desired.
+   */
+  readonly userDataProvider?: IInstanceUserDataProvider;
 }
 
 /**
@@ -94,6 +161,14 @@ export interface WorkerInstanceConfigurationProps {
  *
  * The configuration happens on instance start-up using user data scripting.
  *
+ * This configuration performs the following steps in order:
+ * 1) Configure Cloud Watch Agent
+ * 2) Configure Deadline Worker RenderQueue connection
+ * 3) Configure Deadline Worker settings
+ *
+ * A `userDataProvider` can be specified that defines callback functions.
+ * These callbacks can be used to inject user data commands at different points during the Worker instance configuration.
+ *
  * Security Considerations
  * ------------------------
  * - The instances configured by this construct will download and run scripts from your CDK bootstrap bucket when that instance
@@ -105,12 +180,15 @@ export interface WorkerInstanceConfigurationProps {
 export class WorkerInstanceConfiguration extends Construct {
   constructor(scope: Construct, id: string, props: WorkerInstanceConfigurationProps) {
     super(scope, id);
-
+    props.userDataProvider?.preCloudWatchAgent(props.worker);
     if (props.cloudwatchLogSettings) {
       this.configureCloudWatchLogStream(props.worker, id, props.cloudwatchLogSettings);
     }
+    props.userDataProvider?.preRenderQueueConfiguration(props.worker);
     props.renderQueue?.configureClientInstance({ host: props.worker});
+    props.userDataProvider?.preWorkerConfiguration(props.worker);
     this.configureWorkerSettings(props.worker, id, props.workerSettings);
+    props.userDataProvider?.postWorkerLaunch(props.worker);
   }
 
   /**

packages/aws-rfdk/lib/deadline/lib/worker-fleet.ts

@@ -55,6 +55,7 @@ import {
 } from './render-queue';
 import { Version } from './version';
 import {
+  IInstanceUserDataProvider,
   WorkerInstanceConfiguration,
   WorkerSettings,
 } from './worker-configuration';
@@ -189,6 +190,12 @@ export interface WorkerInstanceFleetProps extends WorkerSettings {
    * @default The default devices of the provided ami will be used.
    */
   readonly blockDevices?: BlockDevice[];
+
+  /**
+   * An optional provider of user data commands to be injected at various points during the Worker configuration lifecycle.
+   * You can provide a subclass of InstanceUserDataProvider with the methods overridden as desired.
+   */
+  readonly userDataProvider?: IInstanceUserDataProvider;
 }
 
 /**
@@ -449,6 +456,7 @@ export class WorkerInstanceFleet extends WorkerInstanceFleetBase {
       },
       renderQueue: props.renderQueue,
       workerSettings: props,
+      userDataProvider: props.userDataProvider,
     });
 
     // Updating the user data with successful cfn-signal commands.

packages/aws-rfdk/lib/deadline/test/worker-fleet.test.ts

@@ -44,6 +44,8 @@ import {
   escapeTokenRegex,
 } from '../../core/test/token-regex-helpers';
 import {
+  IHost,
+  InstanceUserDataProvider,
   IRenderQueue,
   RenderQueue,
   Repository,
@@ -1461,6 +1463,40 @@ describe('HealthMonitor Tests', () => {
       TargetType: 'instance',
     }));
   });
+
+  test('UserData is added', () => {
+    // WHEN
+    class UserDataProvider extends InstanceUserDataProvider {
+      preCloudWatchAgent(host: IHost): void {
+        host.userData.addCommands('echo preCloudWatchAgent');
+      }
+      preRenderQueueConfiguration(host: IHost): void {
+        host.userData.addCommands('echo preRenderQueueConfiguration');
+      }
+      preWorkerConfiguration(host: IHost): void {
+        host.userData.addCommands('echo preWorkerConfiguration');
+      }
+      postWorkerLaunch(host: IHost): void {
+        host.userData.addCommands('echo postWorkerLaunch');
+      }
+    }
+    const fleet = new WorkerInstanceFleet(wfstack, 'workerFleet', {
+      vpc,
+      workerMachineImage: new GenericLinuxImage({
+        'us-east-1': 'ami-any',
+      }),
+      renderQueue,
+      healthMonitor,
+      userDataProvider: new UserDataProvider(wfstack, 'UserDataProvider'),
+    });
+    const userData = fleet.fleet.userData.render();
+
+    // THEN
+    expect(userData).toContain('echo preCloudWatchAgent');
+    expect(userData).toContain('echo preRenderQueueConfiguration');
+    expect(userData).toContain('echo preWorkerConfiguration');
+    expect(userData).toContain('echo postWorkerLaunch');
+  });
 });
 
 describe('tagging', () => {