feat: Support running configuration scripts after worker reaches STARTED state (#601)

Signed-off-by: David Leong <[email protected]>

Author: leongdl

docs/logging.md

@@ -53,17 +53,18 @@ Log events may also contain a `type`, `subtype`, icon (`ti`), and additional fie
 | AgentInfo | None | None | platform; python[interpreter,version]; agent[version,installedAt,runningAs]; depenencies | Information about the running Agent software. |
 | API | Req | 📤 | operation; request_url; params; resource (optional) | A request to an AWS API. Only requests to AWS Deadline Cloud APIs contain a resource field. |
 | API | Resp | 📥 | operation; params; status_code, request_id; error (optional) | A response from an AWS API request. |
-| FileSystem | Read/Write/Create/Delete | 💾 | filepath; message | A filesystem operation. |
 | AWSCreds | Load/Install/Delete | 🔑 | resource; message; role_arn (optional) | Related to an operation for AWS Credentials. |
 | AWSCreds | Query | 🔑 | resource; message; role_arn (optional); expiry (optional) | Related to an operation for AWS Credentials. |
 | AWSCreds | Refresh | 🔑 | resource; message; role_arn (optional); expiry (optional); scheduled_time (optional) | Related to an operation for AWS Credentials. |
+| FileSystem | Read/Write/Create/Delete | 💾 | filepath; message | A filesystem operation. |
 | Metrics | System | 📊 | many | System metrics. |
 | Session | Starting/Failed/AWSCreds/Complete/Info | 🔷 | queue_id; job_id; session_id | An update or information related to a Session. |
 | Session | Add/Remove | 🔷 | queue_id; job_id; session_id; action_ids; queued_actions | Adding or removing SessionActions in a Session. |
 | Session | Logs | 🔷 | queue_id; job_id; session_id; log_dest | Information regarding where the Session logs are located. |
 | Session | User | 🔷 | queue_id; job_id; session_id; user | The user that a Session is running Actions as. |
 | Session | Runtime | 🔷 | queue_id; job_id; session_id | Information related to the running Session. This includes information about the host, process control, and encountered Exceptions which could contain information like filepaths. |
 | Worker | Create/Load/ID/Status/Delete | 💻 | farm_id; fleet_id; worker_id (optional); message | A notification related to a Worker resource within AWS Deadline Cloud. |
+| Worker | HostConfiguration | 📜 | farm_id; fleet_id; worker_id (optional); message; status; exit_code (optional); success (optional) | Worker Host configuration event. |
 
 If you prefer structured logs to be emited on your host, then you can configure your Worker Agent to emit structured logs instead. Please see the
 `structured_logs` option in the [`worker.toml.example`](../src/deadline_worker_agent/installer/worker.toml.example)

docs/state.md

@@ -45,7 +45,8 @@ The worker maintains a worker state which is a JSON representation of the follow
 ```jsonc
 {
     "worker_id": "<WORKER_ID>",
-    "instance_id": "<INSTANCE_ID>"  # OPTIONAL
+    "instance_id": "<INSTANCE_ID>",  # OPTIONAL
+    "host_configuration_succeeded": bool # OPTINAL
 }
 ```
 
@@ -58,6 +59,9 @@ This contains a unique identifier that represents the worker in your AWS Deadlin
 
 This is an optional parameter which contains an identifier of the EC2 instance running the worker, if applicable. 
 
+#### `host_configuration_succeeded`
+
+This is an optional parameter which indicates host configuration has been executed on the worker. Workers will only run once per worker.
 
 #### How the State File Affects Startup Behavior
 

pyproject.toml

@@ -14,7 +14,8 @@ dependencies = [
     "requests ~= 2.31",
     "boto3 >= 1.34.75",
     "deadline == 0.49.*",
-    "openjd-sessions == 0.10.*",
+    # Pinned to patch version due to Host Config Script runner usage of private OpenJD Sessions API.
+    "openjd-sessions == 0.10.1",
     # tomli became tomllib in standard library in Python 3.11
     "tomli == 2.0.* ; python_version<'3.11'",
     "tomlkit == 0.13.*",

src/deadline_worker_agent/api_models.py

@@ -103,6 +103,14 @@ class AttachmentUploadAction(TypedDict):
     taskId: str
 
 
+class HostConfiguration(TypedDict):
+    scriptBody: str
+    """Host Configuration Script Body."""
+
+    scriptTimeoutSeconds: int
+    """Customer Supplied timeout."""
+
+
 class LogConfiguration(TypedDict):
     error: NotRequired[str]
     logDriver: str
@@ -414,6 +422,7 @@ class UpdateWorkerScheduleRequest(TypedDict):
 
 class UpdateWorkerResponse(TypedDict):
     log: NotRequired[LogConfiguration]
+    hostConfiguration: NotRequired[HostConfiguration]
 
 
 class IpAddresses(TypedDict):

src/deadline_worker_agent/aws/deadline/__init__.py

@@ -112,6 +112,17 @@ class DeadlineRequestInterrupted(Exception):
     pass
 
 
+@dataclass(frozen=True)
+class WorkerHostConfiguration:
+    """Host Configuration after a worker has started."""
+
+    script_body: str
+    """Host Configuration Script Body."""
+
+    script_timeout_seconds: int
+    """Customer Supplied timeout."""
+
+
 @dataclass(frozen=True)
 class WorkerLogConfig:
     """The destination where the Worker Agent should synchronize its logs to"""

src/deadline_worker_agent/feature_flag.py

@@ -7,3 +7,5 @@
 ASSET_SYNC_JOB_USER_FEATURE = (
     os.environ.get("ASSET_SYNC_JOB_USER_FEATURE", "false").lower() == "true"
 )
+
+HOST_CONFIGURATION_FEATURE = os.environ.get("HOST_CONFIGURATION_FEATURE", "false").lower() == "true"

src/deadline_worker_agent/log_messages.py

@@ -164,6 +164,7 @@ class WorkerLogEventOp(str, Enum):
     ID = "ID"  # The ID that the Agent is running as
     STATUS = "Status"
     DELETE = "Delete"
+    HOST_CONFIGURATION = "HostConfiguration"
 
 
 class WorkerLogEvent(BaseLogEvent):
@@ -201,6 +202,53 @@ def asdict(self) -> dict[str, Any]:
         return self.add_exception_to_dict(dd)
 
 
+class WorkerHostConfigurationStatus(str, Enum):
+    RUNNING = "Running"
+    SUCCEEDED = "Succeeded"
+    FAILED = "Failed"
+    SKIPPED = "Skipped"
+
+
+class WorkerHostConfigurationLogEvent(WorkerLogEvent):
+    ti = "📜"
+    type = "Worker"
+    status: WorkerHostConfigurationStatus
+    exit_code: Optional[int]
+    success: Optional[bool]
+
+    def __init__(
+        self,
+        *,
+        farm_id: str,
+        fleet_id: str,
+        message: str,
+        status: WorkerHostConfigurationStatus,
+        worker_id: Optional[str] = None,
+        exit_code: Optional[int] = None,
+        success: Optional[bool] = None,
+    ) -> None:
+        self.exit_code = exit_code
+        self.success = success
+        self.status = status
+
+        super().__init__(
+            op=WorkerLogEventOp.HOST_CONFIGURATION,
+            farm_id=farm_id,
+            fleet_id=fleet_id,
+            worker_id=worker_id,
+            message=message,
+        )
+
+    def asdict(self) -> dict[str, str]:
+        dd = super().asdict()
+        dd.update(status=self.status)
+        if self.exit_code:
+            dd.update(exit_code=self.exit_code)
+        if self.success:
+            dd.update(success=self.success)
+        return self.add_exception_to_dict(dd)
+
+
 class FilesystemLogEventOp(str, Enum):
     READ = "Read"
     WRITE = "Write"

src/deadline_worker_agent/startup/bootstrap.py

@@ -4,7 +4,7 @@
 
 from dataclasses import dataclass, asdict, fields, field
 from time import sleep
-from typing import Optional
+from typing import Any, Dict, Optional, Tuple, cast
 import json
 import logging as _logging
 import stat
@@ -16,6 +16,7 @@
 from ..aws.deadline import (
     DeadlineRequestConditionallyRecoverableError,
     DeadlineRequestUnrecoverableError,
+    WorkerHostConfiguration,
     WorkerLogConfig,
     construct_worker_log_config,
     create_worker,
@@ -100,6 +101,8 @@ class WorkerPersistenceInfo:
     )
     """The EC2 instance ID of the Worker, if applicable"""
 
+    host_configuration_succeeded: bool | None = field(default=None)
+
     @classmethod
     def load(cls, *, config: Configuration) -> Optional[WorkerPersistenceInfo]:
         """Load the Worker Bootstrap from the Worker Agent state persistence file"""
@@ -115,7 +118,7 @@ def load(cls, *, config: Configuration) -> Optional[WorkerPersistenceInfo]:
         )
 
         with config.worker_state_file.open("r", encoding="utf8") as fh:
-            data: dict[str, str] = json.load(fh)
+            data: dict[str, str | bool] = json.load(fh)
 
         own_fields = set(f.name for f in fields(class_or_instance=WorkerPersistenceInfo))
         selected_data = {key: value for key, value in data.items() if key in own_fields}
@@ -129,7 +132,7 @@ def load(cls, *, config: Configuration) -> Optional[WorkerPersistenceInfo]:
                 )
             )
 
-        return cls(**selected_data)
+        return cls(**cast(Dict[str, Any], selected_data))
 
     def save(self, *, config: Configuration) -> None:
         """Save the Worker Bootstrap to the Worker Agent state persistence file"""
@@ -185,6 +188,9 @@ class WorkerBootstrap:
     log_config: Optional[WorkerLogConfig] = None
     """The log configuration for the Worker"""
 
+    host_config: Optional[WorkerHostConfiguration] = None
+    """The host configuration for the Worker"""
+
 
 def bootstrap_worker(config: Configuration, *, use_existing_worker: bool = True) -> WorkerBootstrap:
     """Contains startup logic to ensure that the Worker is created and started"""
@@ -217,7 +223,7 @@ def bootstrap_worker(config: Configuration, *, use_existing_worker: bool = True)
 
     try:
         # raises: BootstrapWithoutWorkerLoad, SystemExit
-        log_config = _start_worker(
+        log_config, host_config = _start_worker(
             deadline_client=deadline_client,
             config=config,
             worker_id=worker_info.worker_id,
@@ -246,6 +252,7 @@ def bootstrap_worker(config: Configuration, *, use_existing_worker: bool = True)
         worker_info=worker_info,
         session=worker_session,
         log_config=log_config,
+        host_config=host_config,
     )
 
 
@@ -437,15 +444,16 @@ def _start_worker(
     config: Configuration,
     worker_id: str,
     has_existing_worker: bool,
-) -> Optional[WorkerLogConfig]:
+) -> Tuple[Optional[WorkerLogConfig], Optional[WorkerHostConfiguration]]:
     """Updates the Worker in the service to the STARTED state.
 
     Returns:
         Optional[WorkerLogConfig] -- Non-None only if the UpdateWorker request
             contained a log configuration for the Worker Agent to use for writing
             its own logs. The returned WorkerLogConfig is the configuration that it
             should use.
-
+        Optional[WorkerHostConfiguration] -- Non-None if UpdateWorker request contains a
+            host configuration for the Worker Agent to configure the host.
     Raises:
         BootstrapWithoutWorkerLoad
         SystemExit
@@ -498,9 +506,19 @@ def _start_worker(
         )
     )
 
+    logging_config = None
     if log_config := response.get("log"):
-        return construct_worker_log_config(log_config=log_config)
-    return None
+        logging_config = construct_worker_log_config(log_config=log_config)
+
+    host_config = None
+    if response_host_config := response.get("hostConfiguration"):
+        # scriptTimeoutSeconds is technically always there from the backend.
+        host_config = WorkerHostConfiguration(
+            script_body=response_host_config.get("scriptBody", ""),
+            script_timeout_seconds=response_host_config["scriptTimeoutSeconds"],
+        )
+
+    return logging_config, host_config
 
 
 def _enforce_no_instance_profile_or_stop_worker(

src/deadline_worker_agent/startup/entrypoint.py

@@ -15,6 +15,8 @@
 from typing import Optional
 from pathlib import Path
 
+from ..feature_flag import HOST_CONFIGURATION_FEATURE
+
 from ..api_models import WorkerStatus
 from ..aws.deadline import (
     update_worker,
@@ -30,13 +32,16 @@
 from ..log_messages import (
     AgentInfoLogEvent,
     LogRecordStringTranslationFilter,
+    WorkerHostConfigurationStatus,
     WorkerLogEvent,
     WorkerLogEventOp,
+    WorkerHostConfigurationLogEvent,
 )
 from ..log_sync.cloudwatch import stream_cloudwatch_logs
 from ..log_sync.loggers import ROOT_LOGGER, logger as log_sync_logger
 from ..worker import Worker
 from .bootstrap import bootstrap_worker
+from .host_configuration_script import HostConfigurationScriptRunner
 
 __all__ = ["entrypoint"]
 _logger = logging.getLogger(__name__)
@@ -163,6 +168,78 @@ def filter(self, record: logging.LogRecord) -> bool:
             # logs that we forward to CloudWatch.
             _log_agent_info()
 
+            # If there was a host config, and it was bootstrapped before, only log a message.
+            if worker_bootstrap.host_config and not HOST_CONFIGURATION_FEATURE:
+                _logger.info(
+                    WorkerHostConfigurationLogEvent(
+                        farm_id=config.farm_id,
+                        fleet_id=config.fleet_id,
+                        worker_id=worker_id,
+                        message="Host Configuration Feature is not enabled.",
+                        status=WorkerHostConfigurationStatus.SKIPPED,
+                    )
+                )
+            elif (
+                worker_bootstrap.host_config
+                and worker_bootstrap.worker_info.host_configuration_succeeded
+            ):
+                _logger.info(
+                    WorkerHostConfigurationLogEvent(
+                        farm_id=config.farm_id,
+                        fleet_id=config.fleet_id,
+                        worker_id=worker_id,
+                        message="Host Configuration has been setup before. Not running config scripts.",
+                        status=WorkerHostConfigurationStatus.SKIPPED,
+                    )
+                )
+            # Before the run looop starts, run the Host Configuration script.
+            elif worker_bootstrap.host_config:
+                _logger.info(
+                    WorkerHostConfigurationLogEvent(
+                        farm_id=config.farm_id,
+                        fleet_id=config.fleet_id,
+                        worker_id=worker_id,
+                        message="Running host configuration script.",
+                        status=WorkerHostConfigurationStatus.RUNNING,
+                    )
+                )
+                host_config_runner = HostConfigurationScriptRunner(
+                    logger=_logger,
+                    configuration=config,
+                    worker_id=worker_id,
+                    session_directory=config.worker_persistence_dir,
+                    worker_boto3_session=session,
+                    host_configuration_script=worker_bootstrap.host_config.script_body,
+                    host_configuration_timeout_seconds=worker_bootstrap.host_config.script_timeout_seconds,
+                )
+                exit_code = host_config_runner.run()
+                if exit_code == 0:
+                    _logger.info(
+                        WorkerHostConfigurationLogEvent(
+                            farm_id=config.farm_id,
+                            fleet_id=config.fleet_id,
+                            worker_id=worker_id,
+                            message="Worker Agent host configuration succeeded. Starting worker session loop.",
+                            status=WorkerHostConfigurationStatus.SUCCEEDED,
+                            exit_code=exit_code,
+                        )
+                    )
+                    # Persist host configuration has been performed.
+                    worker_bootstrap.worker_info.host_configuration_succeeded = True
+                    worker_bootstrap.worker_info.save(config=config)
+                else:
+                    _logger.critical(
+                        WorkerHostConfigurationLogEvent(
+                            farm_id=config.farm_id,
+                            fleet_id=config.fleet_id,
+                            worker_id=worker_id,
+                            message=f"Worker Agent host configuration failed with exit code {exit_code}. Cannot run jobs, exiting.",
+                            status=WorkerHostConfigurationStatus.FAILED,
+                            exit_code=exit_code,
+                        )
+                    )
+                    sys.exit(1)
+
             worker_sessions = Worker(
                 farm_id=config.farm_id,
                 fleet_id=config.fleet_id,