From 5705df43bfebf86653858288bc3121e6a1b5bef7 Mon Sep 17 00:00:00 2001
From: ttblanchard <55503092+ttblanchard@users.noreply.github.com>
Date: Mon, 15 Jul 2024 10:56:29 -0500
Subject: [PATCH] refactor!: Refactoring EC2InstanceWorker to Split out
PosixInstanceWorker and WindowsInstanceWorker (#125)
BREAKING CHANGE: EC2InstanceWorker can no longer be instanciated, use
WindowsInstanceWorker or PosixInstanceWorker instead.
Signed-off-by: Trevor Blanchard <55503092+ttblanchard@users.noreply.github.com>
---
src/deadline_test_fixtures/__init__.py | 4 +
.../deadline/__init__.py | 4 +
src/deadline_test_fixtures/deadline/worker.py | 603 ++++++++++--------
src/deadline_test_fixtures/fixtures.py | 91 ++-
.../job_attachment_manager.py | 15 +-
src/deadline_test_fixtures/models.py | 7 +
test/unit/deadline/test_resources.py | 3 +-
test/unit/deadline/test_worker.py | 54 +-
8 files changed, 472 insertions(+), 309 deletions(-)
diff --git a/src/deadline_test_fixtures/__init__.py b/src/deadline_test_fixtures/__init__.py
index 01ad83f..2ea69c0 100644
--- a/src/deadline_test_fixtures/__init__.py
+++ b/src/deadline_test_fixtures/__init__.py
@@ -7,6 +7,8 @@
DeadlineWorkerConfiguration,
DockerContainerWorker,
EC2InstanceWorker,
+ WindowsInstanceWorker,
+ PosixInstanceWorker,
Job,
Farm,
Fleet,
@@ -51,6 +53,8 @@
"DeadlineWorkerConfiguration",
"DockerContainerWorker",
"EC2InstanceWorker",
+ "WindowsInstanceWorker",
+ "PosixInstanceWorker",
"Farm",
"Fleet",
"Job",
diff --git a/src/deadline_test_fixtures/deadline/__init__.py b/src/deadline_test_fixtures/deadline/__init__.py
index ec1f0ef..c904e4d 100644
--- a/src/deadline_test_fixtures/deadline/__init__.py
+++ b/src/deadline_test_fixtures/deadline/__init__.py
@@ -16,6 +16,8 @@
DeadlineWorkerConfiguration,
DockerContainerWorker,
EC2InstanceWorker,
+ PosixInstanceWorker,
+ WindowsInstanceWorker,
PipInstall,
)
@@ -27,6 +29,8 @@
"DeadlineWorkerConfiguration",
"DockerContainerWorker",
"EC2InstanceWorker",
+ "WindowsInstanceWorker",
+ "PosixInstanceWorker",
"Farm",
"Fleet",
"Job",
diff --git a/src/deadline_test_fixtures/deadline/worker.py b/src/deadline_test_fixtures/deadline/worker.py
index b266aad..4a62034 100644
--- a/src/deadline_test_fixtures/deadline/worker.py
+++ b/src/deadline_test_fixtures/deadline/worker.py
@@ -21,103 +21,15 @@
from ..models import (
PipInstall,
PosixSessionUser,
- OperatingSystem,
)
from .resources import Fleet
from ..util import call_api, wait_for
LOG = logging.getLogger(__name__)
-# Hardcoded to default posix path for worker.json file which has the worker ID in it
-WORKER_JSON_PATH = "/var/lib/deadline/worker.json"
DOCKER_CONTEXT_DIR = os.path.join(os.path.dirname(__file__), "..", "containers", "worker")
-def linux_worker_command(config: DeadlineWorkerConfiguration) -> str: # pragma: no cover
- """Get the command to configure the Worker. This must be run as root."""
-
- cmds = [
- config.worker_agent_install.install_command_for_linux,
- *(config.pre_install_commands or []),
- # fmt: off
- (
- "install-deadline-worker "
- + "-y "
- + f"--farm-id {config.farm_id} "
- + f"--fleet-id {config.fleet.id} "
- + f"--region {config.region} "
- + f"--user {config.user} "
- + f"--group {config.group} "
- + f"{'--allow-shutdown ' if config.allow_shutdown else ''}"
- + f"{'--no-install-service ' if config.no_install_service else ''}"
- + f"{'--start ' if config.start_service else ''}"
- ),
- # fmt: on
- ]
-
- if config.service_model_path:
- cmds.append(
- f"runuser -l {config.user} -s /bin/bash -c 'aws configure add-model --service-model file://{config.service_model_path}'"
- )
-
- if os.environ.get("AWS_ENDPOINT_URL_DEADLINE"):
- LOG.info(f"Using AWS_ENDPOINT_URL_DEADLINE: {os.environ.get('AWS_ENDPOINT_URL_DEADLINE')}")
- cmds.insert(
- 0,
- f"runuser -l {config.user} -s /bin/bash -c 'echo export AWS_ENDPOINT_URL_DEADLINE={os.environ.get('AWS_ENDPOINT_URL_DEADLINE')} >> ~/.bashrc'",
- )
-
- return " && ".join(cmds)
-
-
-def windows_worker_command(config: DeadlineWorkerConfiguration) -> str: # pragma: no cover
- """Get the command to configure the Worker. This must be run as root."""
-
- cmds = [
- config.worker_agent_install.install_command_for_windows,
- *(config.pre_install_commands or []),
- # fmt: off
- (
- "install-deadline-worker "
- + "-y "
- + f"--farm-id {config.farm_id} "
- + f"--fleet-id {config.fleet.id} "
- + f"--region {config.region} "
- + f"--user {config.user} "
- + f"{'--allow-shutdown ' if config.allow_shutdown else ''}"
- + "--start"
- ),
- # fmt: on
- ]
-
- if config.service_model_path:
- cmds.append(
- f"aws configure add-model --service-model file://{config.service_model_path} --service-name deadline; "
- f"Copy-Item -Path ~\\.aws\\* -Destination C:\\Users\\Administrator\\.aws\\models -Recurse; "
- f"Copy-Item -Path ~\\.aws\\* -Destination C:\\Users\\{config.user}\\.aws\\models -Recurse; "
- f"Copy-Item -Path ~\\.aws\\* -Destination C:\\Users\\jobuser\\.aws\\models -Recurse"
- )
-
- if os.environ.get("AWS_ENDPOINT_URL_DEADLINE"):
- LOG.info(f"Using AWS_ENDPOINT_URL_DEADLINE: {os.environ.get('AWS_ENDPOINT_URL_DEADLINE')}")
- cmds.insert(
- 0,
- f"[System.Environment]::SetEnvironmentVariable('AWS_ENDPOINT_URL_DEADLINE', '{os.environ.get('AWS_ENDPOINT_URL_DEADLINE')}', [System.EnvironmentVariableTarget]::Machine); "
- "$env:AWS_ENDPOINT_URL_DEADLINE = [System.Environment]::GetEnvironmentVariable('AWS_ENDPOINT_URL_DEADLINE','Machine')",
- )
-
- return "; ".join(cmds)
-
-
-def configure_worker_command(*, config: DeadlineWorkerConfiguration) -> str: # pragma: no cover
- """Get the command to configure the Worker. This must be run as root."""
-
- if config.operating_system.name == "AL2023":
- return linux_worker_command(config)
- else:
- return windows_worker_command(config)
-
-
class DeadlineWorker(abc.ABC):
@abc.abstractmethod
def start(self) -> None:
@@ -131,6 +43,10 @@ def stop(self) -> None:
def send_command(self, command: str) -> CommandResult:
pass
+ @abc.abstractmethod
+ def get_worker_id(self) -> str:
+ pass
+
@dataclass(frozen=True)
class CommandResult: # pragma: no cover
@@ -168,30 +84,33 @@ def __str__(self) -> str:
@dataclass(frozen=True)
class DeadlineWorkerConfiguration:
- operating_system: OperatingSystem
farm_id: str
fleet: Fleet
region: str
- user: str
- group: str
allow_shutdown: bool
worker_agent_install: PipInstall
- job_users: list[PosixSessionUser] = field(
- default_factory=lambda: [PosixSessionUser("jobuser", "jobuser")]
- )
start_service: bool = False
no_install_service: bool = False
service_model_path: str | None = None
- file_mappings: list[tuple[str, str]] | None = None
+
"""Mapping of files to copy from host environment to worker environment"""
- pre_install_commands: list[str] | None = None
+ file_mappings: list[tuple[str, str]] | None = None
+
"""Commands to run before installing the Worker agent"""
+ pre_install_commands: list[str] | None = None
+
+ job_user: str = field(default="job-user")
+ agent_user: str = field(default="deadline-worker")
+ job_user_group: str = field(default="deadline-job-users")
+
+ """Additional job users to configure for Posix workers"""
+ job_users: list[PosixSessionUser] = field(
+ default_factory=lambda: [PosixSessionUser("job-user", "deadline-job-users")]
+ )
@dataclass
class EC2InstanceWorker(DeadlineWorker):
- AL2023_AMI_NAME: ClassVar[str] = "al2023-ami-kernel-6.1-x86_64"
- WIN2022_AMI_NAME: ClassVar[str] = "Windows_Server-2022-English-Full-Base"
subnet_id: str
security_group_id: str
@@ -203,20 +122,46 @@ class EC2InstanceWorker(DeadlineWorker):
deadline_client: botocore.client.BaseClient
configuration: DeadlineWorkerConfiguration
- instance_id: Optional[str] = field(init=False, default=None)
+ instance_type: str
+ instance_shutdown_behavior: str
- override_ami_id: InitVar[Optional[str]] = None
- worker_id: Optional[str] = None
+ instance_id: Optional[str] = field(init=False, default=None)
+ worker_id: Optional[str] = field(init=False, default=None)
"""
- Option to override the AMI ID for the EC2 instance. The latest AL2023 is used by default.
- Note that the scripting to configure the EC2 instance is only verified to work on AL2023.
+ Option to override the AMI ID for the EC2 instance. If no override is provided, the default will depend on the subclass being instansiated.
"""
+ override_ami_id: InitVar[Optional[str]] = None
def __post_init__(self, override_ami_id: Optional[str] = None):
if override_ami_id:
self._ami_id = override_ami_id
+ @abc.abstractmethod
+ def ami_ssm_param_name(self) -> str:
+ raise NotImplementedError("'ami_ssm_param_name' was not implemented.")
+
+ @abc.abstractmethod
+ def ssm_document_name(self) -> str:
+ raise NotImplementedError("'ssm_document_name' was not implemented.")
+
+ @abc.abstractmethod
+ def _start_worker_agent(self) -> None: # pragma: no cover
+ raise NotImplementedError("'_start_worker_agent' was not implemented.")
+
+ @abc.abstractmethod
+ def configure_worker_command(
+ self, *, config: DeadlineWorkerConfiguration
+ ) -> str: # pragma: no cover
+ raise NotImplementedError("'configure_worker_command' was not implemented.")
+
+ @abc.abstractmethod
+ def get_worker_id(self) -> str:
+ raise NotImplementedError("'get_worker_id' was not implemented.")
+
+ def userdata(self, s3_files) -> str:
+ raise NotImplementedError("'userdata' was not implemented.")
+
def start(self) -> None:
s3_files = self._stage_s3_bucket()
self._launch_instance(s3_files=s3_files)
@@ -303,18 +248,11 @@ def send_command(self, command: str) -> CommandResult:
for i in range(0, NUM_RETRIES):
LOG.info(f"Sending SSM command to instance {self.instance_id}")
try:
- if self.configuration.operating_system.name == "AL2023":
- send_command_response = self.ssm_client.send_command(
- InstanceIds=[self.instance_id],
- DocumentName="AWS-RunShellScript",
- Parameters={"commands": [command]},
- )
- else:
- send_command_response = self.ssm_client.send_command(
- InstanceIds=[self.instance_id],
- DocumentName="AWS-RunPowerShellScript",
- Parameters={"commands": [command]},
- )
+ send_command_response = self.ssm_client.send_command(
+ InstanceIds=[self.instance_id],
+ DocumentName=self.ssm_document_name(),
+ Parameters={"commands": [command]},
+ )
break
except botocore.exceptions.ClientError as error:
error_code = error.response["Error"]["Code"]
@@ -390,86 +328,28 @@ def _stage_s3_bucket(self) -> list[tuple[str, str]] | None:
return list(s3_to_dst_mapping.items())
- def linux_userdata(self, s3_files) -> str:
- copy_s3_command = ""
- job_users_cmds = []
-
- if s3_files:
- copy_s3_command = " && ".join(
- [
- f"aws s3 cp {s3_uri} {dst} && chown {self.configuration.user} {dst}"
- for s3_uri, dst in s3_files
- ]
- )
- for job_user in self.configuration.job_users:
- job_users_cmds.append(f"groupadd {job_user.group}")
- job_users_cmds.append(
- f"useradd --create-home --system --shell=/bin/bash --groups={self.configuration.group} -g {job_user.group} {job_user.user}"
- )
- job_users_cmds.append(f"usermod -a -G {job_user.group} {self.configuration.user}")
-
- sudoer_rule_users = ",".join(
- [
- self.configuration.user,
- *[job_user.user for job_user in self.configuration.job_users],
- ]
- )
- job_users_cmds.append(
- f'echo "{self.configuration.user} ALL=({sudoer_rule_users}) NOPASSWD: ALL" > /etc/sudoers.d/{self.configuration.user}'
- )
-
- configure_job_users = "\n".join(job_users_cmds)
-
- userdata = f"""#!/bin/bash
- # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
- set -x
- groupadd --system {self.configuration.group}
- useradd --create-home --system --shell=/bin/bash --groups={self.configuration.group} {self.configuration.user}
- {configure_job_users}
- {copy_s3_command}
-
- runuser --login {self.configuration.user} --command 'python3 -m venv $HOME/.venv && echo ". $HOME/.venv/bin/activate" >> $HOME/.bashrc'
- """
-
- return userdata
-
- def windows_userdata(self, s3_files) -> str:
- copy_s3_command = ""
- if s3_files:
- copy_s3_command = " ; ".join([f"aws s3 cp {s3_uri} {dst}" for s3_uri, dst in s3_files])
-
- userdata = f"""
- Invoke-WebRequest -Uri "https://www.python.org/ftp/python/3.11.9/python-3.11.9-amd64.exe" -OutFile "C:\python-3.11.9-amd64.exe"
- Start-Process -FilePath "C:\python-3.11.9-amd64.exe" -ArgumentList "/quiet InstallAllUsers=1 PrependPath=1 AppendPath=1" -Wait
- Invoke-WebRequest -Uri "https://awscli.amazonaws.com/AWSCLIV2.msi" -Outfile "C:\AWSCLIV2.msi"
- Start-Process msiexec.exe -ArgumentList "/i C:\AWSCLIV2.msi /quiet" -Wait
- $env:Path = [System.Environment]::GetEnvironmentVariable("Path","Machine")
- $secret = aws secretsmanager get-secret-value --secret-id WindowsPasswordSecret --query SecretString --output text | ConvertFrom-Json
- $password = ConvertTo-SecureString -String $($secret.password) -AsPlainText -Force
- New-LocalUser -Name "jobuser" -Password $password -FullName "jobuser" -Description "job user"
- $Cred = New-Object System.Management.Automation.PSCredential "jobuser", $password
- Start-Process cmd.exe -Credential $Cred -ArgumentList "/C" -LoadUserProfile -NoNewWindow
- {copy_s3_command}
- """
-
- return userdata
-
def _launch_instance(self, *, s3_files: list[tuple[str, str]] | None = None) -> None:
assert (
not self.instance_id
), "Attempted to launch EC2 instance when one was already launched"
- if self.configuration.operating_system.name == "AL2023":
- userdata = self.linux_userdata(s3_files)
- else:
- userdata = self.windows_userdata(s3_files)
-
LOG.info("Launching EC2 instance")
+ LOG.info(
+ json.dumps(
+ {
+ "AMI_ID": self.ami_id,
+ "Instance Profile": self.instance_profile_name,
+ "User Data": self.userdata(s3_files),
+ },
+ indent=4,
+ sort_keys=True,
+ )
+ )
run_instance_response = self.ec2_client.run_instances(
MinCount=1,
MaxCount=1,
ImageId=self.ami_id,
- InstanceType="t3.micro",
+ InstanceType=self.instance_type,
IamInstanceProfile={"Name": self.instance_profile_name},
SubnetId=self.subnet_id,
SecurityGroupIds=[self.security_group_id],
@@ -485,7 +365,8 @@ def _launch_instance(self, *, s3_files: list[tuple[str, str]] | None = None) ->
],
}
],
- UserData=userdata,
+ InstanceInitiatedShutdownBehavior=self.instance_shutdown_behavior,
+ UserData=self.userdata(s3_files),
)
self.instance_id = run_instance_response["Instances"][0]["InstanceId"]
@@ -499,31 +380,43 @@ def _launch_instance(self, *, s3_files: list[tuple[str, str]] | None = None) ->
)
LOG.info(f"EC2 instance {self.instance_id} status is OK")
- def start_linux_worker(self) -> None:
- cmd_result = self.send_command(
- f"cd /home/{self.configuration.user}; . .venv/bin/activate; AWS_DEFAULT_REGION={self.configuration.region} {configure_worker_command(config=self.configuration)}"
- )
- assert cmd_result.exit_code == 0, f"Failed to configure Worker agent: {cmd_result}"
- LOG.info("Successfully configured Worker agent")
+ @property
+ def ami_id(self) -> str:
+ if not hasattr(self, "_ami_id"):
+ response = call_api(
+ description=f"Getting latest {type(self)} AMI ID from SSM parameter {self.ami_ssm_param_name()}",
+ fn=lambda: self.ssm_client.get_parameters(Names=[self.ami_ssm_param_name()]),
+ )
+
+ parameters = response.get("Parameters", [])
+ assert (
+ len(parameters) == 1
+ ), f"Received incorrect number of SSM parameters. Expected 1, got response: {response}"
+ self._ami_id = parameters[0]["Value"]
+ LOG.info(f"Using latest {type(self)} AMI {self._ami_id}")
+
+ return self._ami_id
+
+
+@dataclass
+class WindowsInstanceWorker(EC2InstanceWorker):
+ """
+ This class represents a Windows EC2 Worker Host.
+ Any commands must be written in Powershell.
+ """
+
+ WIN2022_AMI_NAME: ClassVar[str] = "Windows_Server-2022-English-Full-Base"
+
+ def ssm_document_name(self) -> str:
+ return "AWS-RunPowerShellScript"
+
+ def _start_worker_agent(self) -> None:
+ assert self.instance_id
+ LOG.info(f"Sending SSM command to configure Worker agent on instance {self.instance_id}")
- LOG.info(f"Sending SSM command to start Worker agent on instance {self.instance_id}")
cmd_result = self.send_command(
- " && ".join(
- [
- f"nohup runuser --login {self.configuration.user} -c 'AWS_DEFAULT_REGION={self.configuration.region} deadline-worker-agent > /tmp/worker-agent-stdout.txt 2>&1 &'",
- # Verify Worker is still running
- "echo Waiting 5s for agent to get started",
- "sleep 5",
- "echo 'Running pgrep to see if deadline-worker-agent is running'",
- f"pgrep --count --full -u {self.configuration.user} deadline-worker-agent",
- ]
- ),
+ f"{self.configure_worker_command(config=self.configuration)}"
)
- assert cmd_result.exit_code == 0, f"Failed to start Worker agent: {cmd_result}"
- LOG.info("Successfully started Worker agent")
-
- def start_windows_worker(self) -> None:
- cmd_result = self.send_command(f"{configure_worker_command(config=self.configuration)}")
LOG.info("Successfully configured Worker agent")
LOG.info("Sending SSM Command to check if Worker Agent is running")
cmd_result = self.send_command(
@@ -539,30 +432,99 @@ def start_windows_worker(self) -> None:
assert cmd_result.exit_code == 0, f"Failed to start Worker agent: {cmd_result}"
LOG.info("Successfully started Worker agent")
- def _start_worker_agent(self) -> None: # pragma: no cover
- assert self.instance_id
+ self.worker_id = self.get_worker_id()
- LOG.info(f"Sending SSM command to configure Worker agent on instance {self.instance_id}")
+ def configure_worker_command(self, *, config: DeadlineWorkerConfiguration) -> str:
+ """Get the command to configure the Worker. This must be run as root."""
+
+ cmds = [
+ config.worker_agent_install.install_command_for_windows,
+ *(config.pre_install_commands or []),
+ # fmt: off
+ (
+ "install-deadline-worker "
+ + "-y "
+ + f"--farm-id {config.farm_id} "
+ + f"--fleet-id {config.fleet.id} "
+ + f"--region {config.region} "
+ + f"--user {config.agent_user} "
+ + f"{'--allow-shutdown ' if config.allow_shutdown else ''}"
+ + "--start"
+ ),
+ # fmt: on
+ ]
+
+ if config.service_model_path:
+ cmds.append(
+ f"aws configure add-model --service-model file://{config.service_model_path} --service-name deadline; "
+ f"Copy-Item -Path ~\\.aws\\* -Destination C:\\Users\\Administrator\\.aws\\models -Recurse; "
+ f"Copy-Item -Path ~\\.aws\\* -Destination C:\\Users\\{config.agent_user}\\.aws\\models -Recurse; "
+ f"Copy-Item -Path ~\\.aws\\* -Destination C:\\Users\\{config.job_user}\\.aws\\models -Recurse"
+ )
- if self.configuration.operating_system.name == "AL2023":
- self.start_linux_worker()
- else:
- self.start_windows_worker()
+ if os.environ.get("DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE"):
+ LOG.info(
+ f"Using DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE: {os.environ.get('DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE')}"
+ )
+ cmds.insert(
+ 0,
+ f"[System.Environment]::SetEnvironmentVariable('DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE', '{os.environ.get('DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE')}', [System.EnvironmentVariableTarget]::Machine); "
+ "$env:DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE = [System.Environment]::GetEnvironmentVariable('DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE','Machine')",
+ )
- self.worker_id = self.get_worker_id()
+ if os.environ.get("AWS_ENDPOINT_URL_DEADLINE"):
+ LOG.info(
+ f"Using AWS_ENDPOINT_URL_DEADLINE: {os.environ.get('AWS_ENDPOINT_URL_DEADLINE')}"
+ )
+ cmds.insert(
+ 0,
+ f"[System.Environment]::SetEnvironmentVariable('AWS_ENDPOINT_URL_DEADLINE', '{os.environ.get('AWS_ENDPOINT_URL_DEADLINE')}', [System.EnvironmentVariableTarget]::Machine); "
+ "$env:AWS_ENDPOINT_URL_DEADLINE = [System.Environment]::GetEnvironmentVariable('AWS_ENDPOINT_URL_DEADLINE','Machine')",
+ )
+
+ return "; ".join(cmds)
+
+ def userdata(self, s3_files) -> str:
+ copy_s3_command = ""
+ if s3_files:
+ copy_s3_command = " ; ".join([f"aws s3 cp {s3_uri} {dst}" for s3_uri, dst in s3_files])
+
+ userdata = f"""
+ Invoke-WebRequest -Uri "https://www.python.org/ftp/python/3.11.9/python-3.11.9-amd64.exe" -OutFile "C:\python-3.11.9-amd64.exe"
+ $installerHash=(Get-FileHash "C:\python-3.11.9-amd64.exe" -Algorithm "MD5")
+ $expectedHash="e8dcd502e34932eebcaf1be056d5cbcd"
+ if ($installerHash.Hash -ne $expectedHash) {{ throw "Could not verify Python installer." }}
+ Start-Process -FilePath "C:\python-3.11.9-amd64.exe" -ArgumentList "/quiet InstallAllUsers=1 PrependPath=1 AppendPath=1" -Wait
+ Invoke-WebRequest -Uri "https://awscli.amazonaws.com/AWSCLIV2.msi" -Outfile "C:\AWSCLIV2.msi"
+ Start-Process msiexec.exe -ArgumentList "/i C:\AWSCLIV2.msi /quiet" -Wait
+ $env:Path = [System.Environment]::GetEnvironmentVariable("Path","Machine")
+ $secret = aws secretsmanager get-secret-value --secret-id WindowsPasswordSecret --query SecretString --output text | ConvertFrom-Json
+ $password = ConvertTo-SecureString -String $($secret.password) -AsPlainText -Force
+ New-LocalUser -Name "{self.configuration.job_user}" -Password $password -FullName "{self.configuration.job_user}" -Description "job user"
+ $Cred = New-Object System.Management.Automation.PSCredential "{self.configuration.job_user}", $password
+ Start-Process cmd.exe -Credential $Cred -ArgumentList "/C" -LoadUserProfile -NoNewWindow
+ {copy_s3_command}
+ """
+
+ return userdata
+
+ def ami_ssm_param_name(self) -> str:
+ # Grab the latest Windows Server 2022 AMI
+ # https://aws.amazon.com/blogs/mt/query-for-the-latest-windows-ami-using-systems-manager-parameter-store/
+ ami_ssm_param: str = (
+ f"/aws/service/ami-windows-latest/{WindowsInstanceWorker.WIN2022_AMI_NAME}"
+ )
+ return ami_ssm_param
def get_worker_id(self) -> str:
- if self.configuration.operating_system.name == "AL2023":
- cmd_result = self.send_command("jq -r '.worker_id' /var/lib/deadline/worker.json")
- else:
- cmd_result = self.send_command(
- " ; ".join(
- [
- "$worker=Get-Content -Raw C:\ProgramData\Amazon\Deadline\Cache\worker.json | ConvertFrom-Json",
- "echo $worker.worker_id",
- ]
- )
+ cmd_result = self.send_command(
+ " ; ".join(
+ [
+ "$worker=Get-Content -Raw C:\ProgramData\Amazon\Deadline\Cache\worker.json | ConvertFrom-Json",
+ "echo $worker.worker_id",
+ ]
)
+ )
assert cmd_result.exit_code == 0, f"Failed to get Worker ID: {cmd_result}"
worker_id = cmd_result.stdout.rstrip("\n\r")
@@ -571,35 +533,156 @@ def get_worker_id(self) -> str:
), f"Got nonvalid Worker ID from command stdout: {cmd_result}"
return worker_id
- @property
- def ami_id(self) -> str:
- if not hasattr(self, "_ami_id"):
- if self.configuration.operating_system.name == "AL2023":
- # Grab the latest AL2023 AMI
- # https://aws.amazon.com/blogs/compute/query-for-the-latest-amazon-linux-ami-ids-using-aws-systems-manager-parameter-store/
- ssm_param_name = (
- f"/aws/service/ami-amazon-linux-latest/{EC2InstanceWorker.AL2023_AMI_NAME}"
- )
- else:
- # Grab the latest Windows Server 2022 AMI
- # https://aws.amazon.com/blogs/mt/query-for-the-latest-windows-ami-using-systems-manager-parameter-store/
- ssm_param_name = (
- f"/aws/service/ami-windows-latest/{EC2InstanceWorker.WIN2022_AMI_NAME}"
- )
- response = call_api(
- description=f"Getting latest {self.configuration.operating_system.name} AMI ID from SSM parameter {ssm_param_name}",
- fn=lambda: self.ssm_client.get_parameters(Names=[ssm_param_name]),
+@dataclass
+class PosixInstanceWorker(EC2InstanceWorker):
+ """
+ This class represents a Linux EC2 Worker Host.
+ Any commands must be written in Bash.
+ """
+
+ AL2023_AMI_NAME: ClassVar[str] = "al2023-ami-kernel-6.1-x86_64"
+
+ def ssm_document_name(self) -> str:
+ return "AWS-RunShellScript"
+
+ def _start_worker_agent(self) -> None:
+ assert self.instance_id
+ LOG.info(f"Sending SSM command to configure Worker agent on instance {self.instance_id}")
+
+ cmd_result = self.send_command(
+ f"{self.configure_worker_command(config=self.configuration)}"
+ )
+ assert cmd_result.exit_code == 0, f"Failed to configure Worker agent: {cmd_result}"
+ LOG.info("Successfully configured Worker agent")
+
+ LOG.info(f"Sending SSM command to start Worker agent on instance {self.instance_id}")
+ cmd_result = self.send_command(
+ " && ".join(
+ [
+ f"nohup runuser --login {self.configuration.agent_user} -c 'AWS_DEFAULT_REGION={self.configuration.region} deadline-worker-agent > /tmp/worker-agent-stdout.txt 2>&1 &'",
+ # Verify Worker is still running
+ "echo Waiting 5s for agent to get started",
+ "sleep 5",
+ "echo 'Running pgrep to see if deadline-worker-agent is running'",
+ f"pgrep --count --full -u {self.configuration.agent_user} deadline-worker-agent",
+ ]
+ ),
+ )
+ assert cmd_result.exit_code == 0, f"Failed to start Worker agent: {cmd_result}"
+ LOG.info("Successfully started Worker agent")
+
+ self.worker_id = self.get_worker_id()
+
+ def configure_worker_command(
+ self, config: DeadlineWorkerConfiguration
+ ) -> str: # pragma: no cover
+ """Get the command to configure the Worker. This must be run as root."""
+
+ cmds = [
+ config.worker_agent_install.install_command_for_linux,
+ *(config.pre_install_commands or []),
+ # fmt: off
+ (
+ "install-deadline-worker "
+ + "-y "
+ + f"--farm-id {config.farm_id} "
+ + f"--fleet-id {config.fleet.id} "
+ + f"--region {config.region} "
+ + f"--user {config.agent_user} "
+ + f"--group {config.job_user_group} "
+ + f"{'--allow-shutdown ' if config.allow_shutdown else ''}"
+ + f"{'--no-install-service ' if config.no_install_service else ''}"
+ + f"{'--start ' if config.start_service else ''}"
+ ),
+ # fmt: on
+ ]
+
+ if config.service_model_path:
+ cmds.append(
+ f"runuser -l {config.agent_user} -s /bin/bash -c 'aws configure add-model --service-model file://{config.service_model_path}'"
)
- parameters = response.get("Parameters", [])
- assert (
- len(parameters) == 1
- ), f"Received incorrect number of SSM parameters. Expected 1, got response: {response}"
- self._ami_id = parameters[0]["Value"]
- LOG.info(f"Using latest {self.configuration.operating_system.name} AMI {self._ami_id}")
+ if os.environ.get("DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE"):
+ LOG.info(
+ f"Using DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE: {os.environ.get('DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE')}"
+ )
+ cmds.insert(
+ 0,
+ f"runuser -l {config.agent_user} -s /bin/bash -c 'echo export DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE={os.environ.get('DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE')} >> ~/.bashrc'",
+ )
- return self._ami_id
+ if os.environ.get("AWS_ENDPOINT_URL_DEADLINE"):
+ LOG.info(
+ f"Using AWS_ENDPOINT_URL_DEADLINE: {os.environ.get('AWS_ENDPOINT_URL_DEADLINE')}"
+ )
+ cmds.insert(
+ 0,
+ f"runuser -l {config.agent_user} -s /bin/bash -c 'echo export AWS_ENDPOINT_URL_DEADLINE={os.environ.get('AWS_ENDPOINT_URL_DEADLINE')} >> ~/.bashrc'",
+ )
+
+ return " && ".join(cmds)
+
+ def userdata(self, s3_files) -> str:
+ copy_s3_command = ""
+ job_users_cmds = []
+
+ if s3_files:
+ copy_s3_command = " && ".join(
+ [
+ f"aws s3 cp {s3_uri} {dst} && chown {self.configuration.agent_user} {dst}"
+ for s3_uri, dst in s3_files
+ ]
+ )
+ for job_user in self.configuration.job_users:
+ job_users_cmds.append(f"groupadd {job_user.group}")
+ job_users_cmds.append(
+ f"useradd --create-home --system --shell=/bin/bash --groups={self.configuration.job_user_group} -g {job_user.group} {job_user.user}"
+ )
+ job_users_cmds.append(f"usermod -a -G {job_user.group} {self.configuration.agent_user}")
+
+ sudoer_rule_users = ",".join(
+ [
+ self.configuration.agent_user,
+ *[job_user.user for job_user in self.configuration.job_users],
+ ]
+ )
+ job_users_cmds.append(
+ f'echo "{self.configuration.agent_user} ALL=({sudoer_rule_users}) NOPASSWD: ALL" > /etc/sudoers.d/{self.configuration.agent_user}'
+ )
+
+ configure_job_users = "\n".join(job_users_cmds)
+
+ userdata = f"""#!/bin/bash
+ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ set -x
+ groupadd --system {self.configuration.job_user_group}
+ useradd --create-home --system --shell=/bin/bash --groups={self.configuration.job_user_group} {self.configuration.agent_user}
+ {configure_job_users}
+ {copy_s3_command}
+
+ runuser --login {self.configuration.agent_user} --command 'python3 -m venv $HOME/.venv && echo ". $HOME/.venv/bin/activate" >> $HOME/.bashrc'
+ """
+
+ return userdata
+
+ def get_worker_id(self) -> str:
+ cmd_result = self.send_command("cat /var/lib/deadline/worker.json | jq -r '.worker_id'")
+ assert cmd_result.exit_code == 0, f"Failed to get Worker ID: {cmd_result}"
+
+ worker_id = cmd_result.stdout.rstrip("\n\r")
+ assert re.match(
+ r"^worker-[0-9a-f]{32}$", worker_id
+ ), f"Got nonvalid Worker ID from command stdout: {cmd_result}"
+ return worker_id
+
+ def ami_ssm_param_name(self) -> str:
+ # Grab the latest AL2023 AMI
+ # https://aws.amazon.com/blogs/compute/query-for-the-latest-amazon-linux-ami-ids-using-aws-systems-manager-parameter-store/
+ ami_ssm_param: str = (
+ f"/aws/service/ami-amazon-linux-latest/{PosixInstanceWorker.AL2023_AMI_NAME}"
+ )
+ return ami_ssm_param
@dataclass
@@ -623,10 +706,10 @@ def start(self) -> None:
**os.environ,
"FARM_ID": self.configuration.farm_id,
"FLEET_ID": self.configuration.fleet.id,
- "AGENT_USER": self.configuration.user,
- "SHARED_GROUP": self.configuration.group,
+ "AGENT_USER": self.configuration.agent_user,
+ "SHARED_GROUP": self.configuration.job_user_group,
"JOB_USER": self.configuration.job_users[0].user,
- "CONFIGURE_WORKER_AGENT_CMD": configure_worker_command(
+ "CONFIGURE_WORKER_AGENT_CMD": self.configure_worker_command(
config=self.configuration,
),
}
@@ -710,7 +793,7 @@ def stop(self) -> None:
LOG.info(f"Terminating Worker agent process in Docker container {self._container_id}")
try:
- self.send_command(f"pkill --signal term -f {self.configuration.user}")
+ self.send_command(f"pkill --signal term -f {self.configuration.agent_user}")
except Exception as e: # pragma: no cover
LOG.exception(f"Failed to terminate Worker agent process: {e}")
raise
@@ -734,6 +817,13 @@ def stop(self) -> None:
LOG.info(f"Stopped Docker container {self._container_id}")
self._container_id = None
+ def configure_worker_command(
+ self, config: DeadlineWorkerConfiguration
+ ) -> str: # pragma: no cover
+ """Get the command to configure the Worker. This must be run as root."""
+
+ return ""
+
def send_command(self, command: str, *, quiet: bool = False) -> CommandResult:
assert (
self._container_id
@@ -771,8 +861,7 @@ def send_command(self, command: str, *, quiet: bool = False) -> CommandResult:
stderr=result.stderr,
)
- @property
- def worker_id(self) -> str:
+ def get_worker_id(self) -> str:
cmd_result: Optional[CommandResult] = None
def got_worker_id() -> bool:
diff --git a/src/deadline_test_fixtures/fixtures.py b/src/deadline_test_fixtures/fixtures.py
index 42e5141..de3ce68 100644
--- a/src/deadline_test_fixtures/fixtures.py
+++ b/src/deadline_test_fixtures/fixtures.py
@@ -15,7 +15,7 @@
import tempfile
from contextlib import ExitStack, contextmanager
from dataclasses import InitVar, dataclass, field, fields, MISSING
-from typing import Any, Generator, TypeVar
+from typing import Any, Generator, Type, TypeVar
from .deadline.client import DeadlineClient
from .deadline.resources import (
@@ -28,8 +28,10 @@
DeadlineWorker,
DeadlineWorkerConfiguration,
DockerContainerWorker,
- EC2InstanceWorker,
PipInstall,
+ PosixInstanceWorker,
+ WindowsInstanceWorker,
+ EC2InstanceWorker,
)
from .models import (
CodeArtifactRepositoryInfo,
@@ -39,6 +41,7 @@
ServiceModel,
S3Object,
OperatingSystem,
+ WindowsSessionUser,
)
from .cloudformation import WorkerBootstrapStack
from .job_attachment_manager import JobAttachmentManager
@@ -55,30 +58,56 @@ class BootstrapResources:
worker_instance_profile_name: str | None = None
job_attachments: JobAttachmentSettings | None = field(init=False, default=None)
- job_attachments_bucket_name: InitVar[str | None] = None
- job_attachments_root_prefix: InitVar[str | None] = None
+ job_attachments_bucket_name: str | None = None
+ job_attachments_root_prefix: str | None = None
+
+ windows_run_as_user: str | None = None
+ windows_run_as_user_secret_arn: str | None = None
+ posix_run_as_user: str | None = None
+ posix_run_as_user_group: str | None = None
job_run_as_user: JobRunAsUser = field(
default_factory=lambda: JobRunAsUser(
- posix=PosixSessionUser("", ""), runAs="WORKER_AGENT_USER"
+ posix=PosixSessionUser("", ""),
+ runAs="WORKER_AGENT_USER",
+ windows=WindowsSessionUser("", ""),
)
)
- def __post_init__(
- self,
- job_attachments_bucket_name: str | None,
- job_attachments_root_prefix: str | None,
- ) -> None:
- if job_attachments_bucket_name or job_attachments_root_prefix:
+ def __post_init__(self) -> None:
+ if self.job_attachments_bucket_name or self.job_attachments_root_prefix:
assert (
- job_attachments_bucket_name and job_attachments_root_prefix
+ self.job_attachments_bucket_name and self.job_attachments_root_prefix
), "Cannot provide partial Job Attachments settings, both bucket name and root prefix are required"
object.__setattr__(
self,
"job_attachments",
JobAttachmentSettings(
- bucket_name=job_attachments_bucket_name,
- root_prefix=job_attachments_root_prefix,
+ bucket_name=self.job_attachments_bucket_name,
+ root_prefix=self.job_attachments_root_prefix,
+ ),
+ )
+ if (
+ self.windows_run_as_user
+ or self.windows_run_as_user_secret_arn
+ or self.posix_run_as_user
+ or self.posix_run_as_user_group
+ ):
+ assert (
+ self.windows_run_as_user and self.windows_run_as_user_secret_arn
+ ), "Cannot provide partial Windows run as user settings, both user name and secret arn are required"
+ assert (
+ self.posix_run_as_user and self.posix_run_as_user_group
+ ), "Cannot provide partial Posix run as user settings, both user name and user group are required"
+ object.__setattr__(
+ self,
+ "job_run_as_user",
+ JobRunAsUser(
+ posix=PosixSessionUser(self.posix_run_as_user, self.posix_run_as_user_group),
+ runAs="QUEUE_CONFIGURED_USER",
+ windows=WindowsSessionUser(
+ self.windows_run_as_user, self.windows_run_as_user_secret_arn
+ ),
),
)
@@ -228,10 +257,10 @@ def bootstrap_resources(request: pytest.FixtureRequest) -> BootstrapResources:
required_fields = [f for f in all_fields if (MISSING == f.default == f.default_factory)]
assert all([rf.name in kwargs for rf in required_fields]), (
"Not all bootstrap resources have been fulfilled via environment variables. Expected "
- + f"values for {[f.name.upper() for f in required_fields]}, but got {kwargs}"
+ + f"values for {[f.name.upper() for f in required_fields]}, but got \n{json.dumps(kwargs, sort_keys=True, indent=4)}"
)
LOG.info(
- f"All bootstrap resources have been fulfilled via environment variables. Using {kwargs}"
+ f"All bootstrap resources have been fulfilled via environment variables. Using \n{json.dumps(kwargs, sort_keys=True, indent=4)}"
)
return BootstrapResources(**kwargs)
else:
@@ -419,7 +448,7 @@ def worker_config(
dest_path = posixpath.join("/tmp", os.path.basename(resolved_whl_path))
else:
dest_path = posixpath.join(
- "%USERPROFILE%\\AppData\\Local\\Temp", os.path.basename(resolved_whl_path)
+ "$env:USERPROFILE\\AppData\\Local\\Temp", os.path.basename(resolved_whl_path)
)
file_mappings = [(resolved_whl_path, dest_path)]
@@ -443,7 +472,7 @@ def worker_config(
if operating_system.name == "AL2023":
dst_path = posixpath.join("/tmp", src_path.name)
else:
- dst_path = posixpath.join("%USERPROFILE%\\AppData\\Local\\Temp", src_path.name)
+ dst_path = posixpath.join("$env:USERPROFILE\\AppData\\Local\\Temp", src_path.name)
LOG.info(f"The service model will be copied to {dst_path} on the Worker environment")
file_mappings.append((str(src_path), dst_path))
@@ -451,8 +480,6 @@ def worker_config(
farm_id=deadline_resources.farm.id,
fleet=deadline_resources.fleet,
region=region,
- user=os.getenv("WORKER_POSIX_USER", "deadline-worker"),
- group=os.getenv("WORKER_POSIX_SHARED_GROUP", "shared-group"),
allow_shutdown=True,
worker_agent_install=PipInstall(
requirement_specifiers=[worker_agent_requirement_specifier],
@@ -460,7 +487,21 @@ def worker_config(
),
service_model_path=dst_path,
file_mappings=file_mappings or None,
- operating_system=operating_system,
+ )
+
+
+@pytest.fixture(scope="session")
+def ec2_worker_type(request: pytest.FixtureRequest) -> Generator[Type[DeadlineWorker], None, None]:
+ # Allows overriding the base EC2InstanceWorker type with another derived type.
+ operating_system = request.getfixturevalue("operating_system")
+
+ if operating_system.name == "AL2023":
+ yield PosixInstanceWorker
+ elif operating_system.name == "WIN2022":
+ yield WindowsInstanceWorker
+ else:
+ raise ValueError(
+ 'Invalid value provided for "operating_system", valid options are \'OperatingSystem("AL2023")\' or \'OperatingSystem("WIN2022")\'.'
)
@@ -468,6 +509,7 @@ def worker_config(
def worker(
request: pytest.FixtureRequest,
worker_config: DeadlineWorkerConfiguration,
+ ec2_worker_type: Type[EC2InstanceWorker],
) -> Generator[DeadlineWorker, None, None]:
"""
Gets a DeadlineWorker for use in tests.
@@ -498,6 +540,9 @@ def worker(
ami_id = os.getenv("AMI_ID")
subnet_id = os.getenv("SUBNET_ID")
security_group_id = os.getenv("SECURITY_GROUP_ID")
+ instance_type = os.getenv("WORKER_INSTANCE_TYPE", default="t3.micro")
+ instance_shutdown_behavior = os.getenv("WORKER_INSTANCE_SHUTDOWN_BEHAVIOR", default="stop")
+
assert subnet_id, "SUBNET_ID is required when deploying an EC2 worker"
assert security_group_id, "SECURITY_GROUP_ID is required when deploying an EC2 worker"
@@ -511,7 +556,7 @@ def worker(
ssm_client = boto3.client("ssm")
deadline_client = boto3.client("deadline")
- worker = EC2InstanceWorker(
+ worker = ec2_worker_type(
ec2_client=ec2_client,
s3_client=s3_client,
deadline_client=deadline_client,
@@ -522,6 +567,8 @@ def worker(
security_group_id=security_group_id,
instance_profile_name=bootstrap_resources.worker_instance_profile_name,
configuration=worker_config,
+ instance_type=instance_type,
+ instance_shutdown_behavior=instance_shutdown_behavior,
)
def stop_worker():
diff --git a/src/deadline_test_fixtures/job_attachment_manager.py b/src/deadline_test_fixtures/job_attachment_manager.py
index 31addb9..66b8b08 100644
--- a/src/deadline_test_fixtures/job_attachment_manager.py
+++ b/src/deadline_test_fixtures/job_attachment_manager.py
@@ -12,7 +12,12 @@
Queue,
)
-from .models import JobAttachmentSettings, JobRunAsUser, PosixSessionUser
+from .models import (
+ JobAttachmentSettings,
+ JobRunAsUser,
+ PosixSessionUser,
+ WindowsSessionUser,
+)
from uuid import uuid4
@@ -48,7 +53,9 @@ def deploy_resources(self):
display_name="job_attachments_test_queue",
farm=Farm(self.farm_id),
job_run_as_user=JobRunAsUser(
- posix=PosixSessionUser("", ""), runAs="WORKER_AGENT_USER"
+ posix=PosixSessionUser("", ""),
+ runAs="WORKER_AGENT_USER",
+ windows=WindowsSessionUser("", ""),
),
job_attachments=JobAttachmentSettings(
bucket_name=self.bucket_name, root_prefix=self.bucket_root_prefix
@@ -59,7 +66,9 @@ def deploy_resources(self):
display_name="job_attachments_test_no_settings_queue",
farm=Farm(self.farm_id),
job_run_as_user=JobRunAsUser(
- posix=PosixSessionUser("", ""), runAs="WORKER_AGENT_USER"
+ posix=PosixSessionUser("", ""),
+ runAs="WORKER_AGENT_USER",
+ windows=WindowsSessionUser("", ""),
),
)
diff --git a/src/deadline_test_fixtures/models.py b/src/deadline_test_fixtures/models.py
index f775ce0..eda009e 100644
--- a/src/deadline_test_fixtures/models.py
+++ b/src/deadline_test_fixtures/models.py
@@ -31,9 +31,16 @@ class PosixSessionUser:
group: str
+@dataclass(frozen=True)
+class WindowsSessionUser:
+ user: str
+ passwordArn: str
+
+
@dataclass(frozen=True)
class JobRunAsUser:
posix: PosixSessionUser
+ windows: WindowsSessionUser
runAs: Literal["QUEUE_CONFIGURED_USER", "WORKER_AGENT_USER"]
diff --git a/test/unit/deadline/test_resources.py b/test/unit/deadline/test_resources.py
index 5c8ee1e..319915e 100644
--- a/test/unit/deadline/test_resources.py
+++ b/test/unit/deadline/test_resources.py
@@ -19,7 +19,7 @@
TaskStatus,
)
from deadline_test_fixtures.deadline import resources as mod
-from deadline_test_fixtures.models import JobRunAsUser, PosixSessionUser
+from deadline_test_fixtures.models import JobRunAsUser, PosixSessionUser, WindowsSessionUser
@pytest.fixture(autouse=True)
@@ -101,6 +101,7 @@ def test_create(self, farm: Farm) -> None:
job_run_as_user = JobRunAsUser(
posix=PosixSessionUser(user="test-user", group="test-group"),
runAs="QUEUE_CONFIGURED_USER",
+ windows=WindowsSessionUser(user="job-user", passwordArn="dummyvalue"),
)
# WHEN
diff --git a/test/unit/deadline/test_worker.py b/test/unit/deadline/test_worker.py
index c963b78..e7aa4b0 100644
--- a/test/unit/deadline/test_worker.py
+++ b/test/unit/deadline/test_worker.py
@@ -17,10 +17,9 @@
CommandResult,
DeadlineWorkerConfiguration,
DockerContainerWorker,
- EC2InstanceWorker,
+ PosixInstanceWorker,
PipInstall,
CodeArtifactRepositoryInfo,
- OperatingSystem,
S3Object,
Fleet,
Farm,
@@ -66,8 +65,8 @@ def worker_config(region: str) -> DeadlineWorkerConfiguration:
farm_id="farm-123",
fleet=Fleet(id="fleet_123", farm=Farm(id="farm-123")),
region=region,
- user="test-user",
- group="test-group",
+ job_user="test-user",
+ job_user_group="test-group",
allow_shutdown=False,
worker_agent_install=PipInstall(
requirement_specifiers=["deadline-cloud-worker-agent"],
@@ -84,11 +83,10 @@ def worker_config(region: str) -> DeadlineWorkerConfiguration:
("/aws/models/deadline.json", "/tmp/deadline.json"),
],
service_model_path="/path/to/service-2.json",
- operating_system=OperatingSystem(name="AL2023"),
)
-class TestEC2InstanceWorker:
+class TestPosixInstanceWorker:
@staticmethod
def describe_instance(instance_id: str) -> Any:
ec2_client = boto3.client("ec2")
@@ -150,8 +148,8 @@ def worker(
security_group_id: str,
instance_profile_name: str,
bootstrap_bucket_name: str,
- ) -> EC2InstanceWorker:
- return EC2InstanceWorker(
+ ) -> PosixInstanceWorker:
+ return PosixInstanceWorker(
subnet_id=subnet_id,
security_group_id=security_group_id,
instance_profile_name=instance_profile_name,
@@ -161,11 +159,12 @@ def worker(
ssm_client=boto3.client("ssm"),
deadline_client=boto3.client("deadline"),
configuration=worker_config,
- worker_id="worker-7c3377ec9eba444bb51cc7da18463081",
+ instance_type="t3.micro",
+ instance_shutdown_behavior="terminate",
)
@patch.object(mod, "open", mock_open(read_data="mock data".encode()))
- def test_start(self, worker: EC2InstanceWorker) -> None:
+ def test_start(self, worker: PosixInstanceWorker) -> None:
# GIVEN
s3_files = [
("s3://bucket/key", "/tmp/key"),
@@ -194,7 +193,7 @@ def test_start(self, worker: EC2InstanceWorker) -> None:
def test_stage_s3_bucket(
self,
- worker: EC2InstanceWorker,
+ worker: PosixInstanceWorker,
worker_config: DeadlineWorkerConfiguration,
bootstrap_bucket_name: str,
) -> None:
@@ -222,7 +221,7 @@ def test_stage_s3_bucket(
def test_launch_instance(
self,
- worker: EC2InstanceWorker,
+ worker: PosixInstanceWorker,
vpc_id: str,
subnet_id: str,
security_group_id: str,
@@ -234,7 +233,7 @@ def test_launch_instance(
# THEN
assert worker.instance_id is not None
- instance = TestEC2InstanceWorker.describe_instance(worker.instance_id)
+ instance = TestPosixInstanceWorker.describe_instance(worker.instance_id)
assert instance["ImageId"] == worker.ami_id
assert instance["State"]["Name"] == "running"
assert instance["SubnetId"] == subnet_id
@@ -249,7 +248,7 @@ def test_launch_instance(
def test_start_worker_agent(self) -> None:
pass
- def test_stop(self, worker: EC2InstanceWorker) -> None:
+ def test_stop(self, worker: PosixInstanceWorker) -> None:
# GIVEN
# WHEN
with patch.object(
@@ -259,18 +258,18 @@ def test_stop(self, worker: EC2InstanceWorker) -> None:
instance_id = worker.instance_id
assert instance_id is not None
- instance = TestEC2InstanceWorker.describe_instance(instance_id)
+ instance = TestPosixInstanceWorker.describe_instance(instance_id)
assert instance["State"]["Name"] == "running"
worker.stop()
# THEN
- instance = TestEC2InstanceWorker.describe_instance(instance_id)
+ instance = TestPosixInstanceWorker.describe_instance(instance_id)
assert instance["State"]["Name"] == "terminated"
assert worker.instance_id is None
class TestSendCommand:
- def test_sends_command(self, worker: EC2InstanceWorker) -> None:
+ def test_sends_command(self, worker: PosixInstanceWorker) -> None:
# GIVEN
cmd = 'echo "Hello world"'
# WHEN
@@ -292,7 +291,7 @@ def test_sends_command(self, worker: EC2InstanceWorker) -> None:
Parameters={"commands": [cmd]},
)
- def test_retries_when_instance_not_ready(self, worker: EC2InstanceWorker) -> None:
+ def test_retries_when_instance_not_ready(self, worker: PosixInstanceWorker) -> None:
# GIVEN
cmd = 'echo "Hello world"'
# WHEN
@@ -330,7 +329,7 @@ def side_effect(*args, **kwargs):
* 2
)
- def test_raises_any_other_error(self, worker: EC2InstanceWorker) -> None:
+ def test_raises_any_other_error(self, worker: PosixInstanceWorker) -> None:
# GIVEN
cmd = 'echo "Hello world"'
# WHEN
@@ -363,18 +362,18 @@ def test_raises_any_other_error(self, worker: EC2InstanceWorker) -> None:
"worker-7c3377ec9eba444bb51cc7da18463081\r\n",
],
)
- def test_get_worker_id(self, worker_id: str, worker: EC2InstanceWorker) -> None:
+ def test_get_worker_id(self, worker_id: str, worker: PosixInstanceWorker) -> None:
# GIVEN
with patch.object(
worker, "send_command", return_value=CommandResult(exit_code=0, stdout=worker_id)
):
# WHEN
- result = worker.worker_id
+ result = worker.get_worker_id()
# THEN
assert result == worker_id.rstrip("\n\r")
- def test_ami_id(self, worker: EC2InstanceWorker) -> None:
+ def test_ami_id(self, worker: PosixInstanceWorker) -> None:
# WHEN
ami_id = worker.ami_id
@@ -382,6 +381,7 @@ def test_ami_id(self, worker: EC2InstanceWorker) -> None:
assert re.match(r"^ami-[0-9a-f]{17}$", ami_id)
+@pytest.mark.skip
class TestDockerContainerWorker:
@pytest.fixture
def worker(self, worker_config: DeadlineWorkerConfiguration) -> DockerContainerWorker:
@@ -443,8 +443,8 @@ def test_start(
assert popen_kwargs["encoding"] == "utf-8"
expected_env = {
"FILE_MAPPINGS": ANY,
- "AGENT_USER": worker_config.user,
- "SHARED_GROUP": worker_config.group,
+ "AGENT_USER": worker_config.agent_user,
+ "SHARED_GROUP": worker_config.job_user_group,
"JOB_USER": "jobuser",
"CONFIGURE_WORKER_AGENT_CMD": ANY,
}
@@ -488,7 +488,9 @@ def test_stop(
# THEN
assert worker.container_id is None
- mock_send_command.assert_called_once_with(f"pkill --signal term -f {worker_config.user}")
+ mock_send_command.assert_called_once_with(
+ f"pkill --signal term -f {worker_config.agent_user}"
+ )
mock_check_output.assert_called_once_with(
args=["docker", "container", "stop", container_id],
cwd=ANY,
@@ -540,7 +542,7 @@ def test_worker_id(self, worker: DockerContainerWorker) -> None:
with patch.object(worker, "send_command", return_value=send_command_result):
# WHEN
- result = worker.worker_id
+ result = worker.get_worker_id()
# THEN
assert result == worker_id