From e5862aede714e47c3eaa165d636808f025eb1317 Mon Sep 17 00:00:00 2001 From: Graeme McHale Date: Wed, 20 Sep 2023 16:21:46 -0700 Subject: [PATCH] feat: basic Windows support --- .../aws_credentials/aws_configs.py | 11 +- .../aws_credentials/queue_boto3_session.py | 19 +- .../aws_credentials/worker_boto3_session.py | 3 - .../installer/__init__.py | 5 +- .../installer/install.ps1 | 272 ++++++++++++++++++ .../installer/worker.toml.windows.example | 256 +++++++++++++++++ .../scheduler/scheduler.py | 13 +- .../sessions/job_entities/job_details.py | 2 +- .../set_windows_permissions.py | 40 +++ src/deadline_worker_agent/startup/config.py | 3 +- .../startup/config_file.py | 2 + .../startup/entrypoint.py | 7 +- src/deadline_worker_agent/startup/settings.py | 16 +- src/deadline_worker_agent/worker.py | 19 +- test/unit/aws_credentials/test_aws_configs.py | 24 +- .../test_queue_boto3_session.py | 21 +- test/unit/conftest.py | 36 ++- test/unit/install/test_install.py | 3 +- test/unit/log_sync/test_cloudwatch.py | 10 +- test/unit/scheduler/test_scheduler.py | 13 +- test/unit/scheduler/test_session_cleanup.py | 62 ++-- test/unit/sessions/test_job_entities.py | 2 + test/unit/sessions/test_session.py | 12 +- test/unit/startup/test_config.py | 43 ++- test/unit/startup/test_settings.py | 9 +- test/unit/test_set_windows_permissions.py | 50 ++++ 26 files changed, 869 insertions(+), 84 deletions(-) create mode 100644 src/deadline_worker_agent/installer/install.ps1 create mode 100644 src/deadline_worker_agent/installer/worker.toml.windows.example create mode 100644 src/deadline_worker_agent/set_windows_permissions.py create mode 100644 test/unit/test_set_windows_permissions.py diff --git a/src/deadline_worker_agent/aws_credentials/aws_configs.py b/src/deadline_worker_agent/aws_credentials/aws_configs.py index a9f3cd5d..0f039db0 100644 --- a/src/deadline_worker_agent/aws_credentials/aws_configs.py +++ b/src/deadline_worker_agent/aws_credentials/aws_configs.py @@ -3,7 +3,7 @@ from __future__ import annotations import stat - +import os import logging from abc import ABC, abstractmethod from configparser import ConfigParser @@ -11,6 +11,7 @@ from typing import Optional from openjd.sessions import PosixSessionUser, SessionUser from subprocess import run, DEVNULL, PIPE, STDOUT +from ..set_windows_permissions import grant_full_control __all__ = [ "AWSConfig", @@ -28,8 +29,12 @@ def _run_cmd_as(*, user: PosixSessionUser, cmd: list[str]) -> None: def _setup_parent_dir(*, dir_path: Path, owner: SessionUser | None = None) -> None: if owner is None: - create_perms: int = stat.S_IRWXU - dir_path.mkdir(mode=create_perms, exist_ok=True) + if os.name == "posix": + create_perms: int = stat.S_IRWXU + dir_path.mkdir(mode=create_perms, exist_ok=True) + else: + dir_path.mkdir(exist_ok=True) + grant_full_control(dir_path) else: assert isinstance(owner, PosixSessionUser) _run_cmd_as(user=owner, cmd=["mkdir", "-p", str(dir_path)]) diff --git a/src/deadline_worker_agent/aws_credentials/queue_boto3_session.py b/src/deadline_worker_agent/aws_credentials/queue_boto3_session.py index e80eb79d..7439b740 100644 --- a/src/deadline_worker_agent/aws_credentials/queue_boto3_session.py +++ b/src/deadline_worker_agent/aws_credentials/queue_boto3_session.py @@ -91,8 +91,6 @@ def __init__( interrupt_event: Event, worker_persistence_dir: Path, ) -> None: - if os.name != "posix": - raise NotImplementedError("Windows not supported.") super().__init__() self._deadline_client = deadline_client @@ -110,7 +108,11 @@ def __init__( self._credentials_filename = ( "aws_credentials" # note: .json extension added by JSONFileCache ) - self._credentials_process_script_path = self._credential_dir / "get_aws_credentials.sh" + + if os.name == "posix": + self._credentials_process_script_path = self._credential_dir / "get_aws_credentials.sh" + else: + self._credentials_process_script_path = self._credential_dir / "get_aws_credentials.cmd" self._aws_config = AWSConfig(self._os_user) self._aws_credentials = AWSCredentials(self._os_user) @@ -321,9 +323,14 @@ def _generate_credential_process_script(self) -> str: Generates the bash script which generates the credentials as JSON output on STDOUT. This script will be used by the installed credential process. """ - return ("#!/bin/bash\nset -eu\ncat {0}\n").format( - (self._credential_dir / self._credentials_filename).with_suffix(".json") - ) + if os.name == "posix": + return ("#!/bin/bash\nset -eu\ncat {0}\n").format( + (self._credential_dir / self._credentials_filename).with_suffix(".json") + ) + else: + return ("@echo off\ntype {0}\n").format( + (self._credential_dir / self._credentials_filename).with_suffix(".json") + ) def _uninstall_credential_process(self) -> None: """ diff --git a/src/deadline_worker_agent/aws_credentials/worker_boto3_session.py b/src/deadline_worker_agent/aws_credentials/worker_boto3_session.py index df452785..f1d20e57 100644 --- a/src/deadline_worker_agent/aws_credentials/worker_boto3_session.py +++ b/src/deadline_worker_agent/aws_credentials/worker_boto3_session.py @@ -1,7 +1,6 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. from __future__ import annotations -import os import logging from typing import Any, cast @@ -34,8 +33,6 @@ def __init__( config: Configuration, worker_id: str, ) -> None: - if os.name != "posix": - raise NotImplementedError("Windows not supported.") super().__init__() self._bootstrap_session = bootstrap_session diff --git a/src/deadline_worker_agent/installer/__init__.py b/src/deadline_worker_agent/installer/__init__.py index 3d0efbfa..45246362 100644 --- a/src/deadline_worker_agent/installer/__init__.py +++ b/src/deadline_worker_agent/installer/__init__.py @@ -11,13 +11,14 @@ INSTALLER_PATH = { "linux": Path(__file__).parent / "install.sh", + "win32": Path(__file__).parent / "install.ps1", } def install() -> None: """Installer entrypoint for the Amazon Deadline Cloud Worker Agent""" - if sys.platform != "linux": + if sys.platform not in ["linux", "win32"]: print(f"ERROR: Unsupported platform {sys.platform}") sys.exit(1) @@ -26,7 +27,7 @@ def install() -> None: worker_agent_program = Path(sysconfig.get_path("scripts")) / "deadline-worker-agent" cmd = [ - "sudo", + "sudo" if sys.platform == "linux" else "", str(INSTALLER_PATH[sys.platform]), "--farm-id", args.farm_id, diff --git a/src/deadline_worker_agent/installer/install.ps1 b/src/deadline_worker_agent/installer/install.ps1 new file mode 100644 index 00000000..adb43922 --- /dev/null +++ b/src/deadline_worker_agent/installer/install.ps1 @@ -0,0 +1,272 @@ +# Amazon Deadline Cloud Worker Agent Installer + +# Parse command-line arguments +param( + [string]$FarmId, + [string]$FleetId, + [string]$Region = "us-west-2", + [string]$User = "deadline-worker", + [string]$Group = "deadline-job-users", + [string]$WorkerAgentProgram, + [switch]$NoInstallService = $false, + [switch]$Start = $false, + [switch]$Confirm = $false +) + +# Defaults +$default_wa_user = "deadline-worker" +$default_job_group = "deadline-job-users" + +function Usage { + Write-Host "Usage: install.ps1 -FarmId -FleetId [-Region ] [-User ] [-Group ] [-WorkerAgentProgram ] [-NoInstallService] [-Start] [-Confirm]" + Write-Host "" + Write-Host "Arguments" + Write-Host "---------" + Write-Host " -FarmId " + Write-Host " The Amazon Deadline Cloud Farm ID that the Worker belongs to." + Write-Host " -FleetId " + Write-Host " The Amazon Deadline Cloud Fleet ID that the Worker belongs to." + Write-Host " -Region " + Write-Host " The AWS region of the Amazon Deadline Cloud farm. Defaults to $Region." + Write-Host " -User " + Write-Host " A user name that the Amazon Deadline Cloud Worker Agent will run as. Defaults to $default_wa_user." + Write-Host " -Group " + Write-Host " A group name that the Worker Agent shares with the user(s) that Jobs will be running as." + Write-Host " Do not use the primary/effective group of the Worker Agent user specified in -User as" + Write-Host " this is not a secure configuration. Defaults to $default_job_group." + Write-Host " -WorkerAgentProgram " + Write-Host " An optional path to the Worker Agent program. This is used as the program path" + Write-Host " when creating the service. If not specified, the first program named" + Write-Host " deadline-worker-agent found in the PATH will be used." + Write-Host " -NoInstallService" + Write-Host " Skips the worker agent service installation." + Write-Host " -Start" + Write-Host " Starts the service as part of the installation. By default, the service" + Write-Host " is configured to start on system boot but not started immediately." + Write-Host " This option is ignored if -NoInstallService is used." + Write-Host " -Confirm" + Write-Host " Skips a confirmation prompt before performing the installation." + exit 2 +} + +function Banner { + Write-Host "===========================================================" + Write-Host "| Amazon Deadline Cloud Worker Agent Installer |" + Write-Host "===========================================================" +} + +function UserExists { + param([string]$User) + $UserObject = Get-WmiObject -Class Win32_UserAccount -Filter "Name='$User'" + return $UserObject -ne $null +} + +function GroupExists { + param([string]$group) + $groupObject = Get-WmiObject -Class Win32_Group -Filter "Name='$group'" + return $groupObject -ne $null +} + +function ValidateDeadlineId { + param([string]$prefix, [string]$text) + $pattern = "^$prefix-[a-f0-9]{32}$" + return $text -match $pattern +} + +function Set-DirectoryPermissions { + param ( + [string]$Path, + [string]$User, + [string]$Permission + ) + + # Get the current ACL of the directory + $acl = Get-Acl -Path $Path + + # Remove existing inheritance (if any) this makes the permissions predictable + $acl.SetAccessRuleProtection($true, $false) + + # Create a new rule for the specified user and permission level + $userRule = New-Object System.Security.AccessControl.FileSystemAccessRule( + $User, $Permission, "ContainerInherit,ObjectInherit", "None", "Allow" + ) + + # Create a rule for administrators with Full Control + $adminRule = New-Object System.Security.AccessControl.FileSystemAccessRule( + "Administrators", $Permission, "ContainerInherit,ObjectInherit", "None", "Allow" + ) + + # Add the rules to the ACL + $acl.AddAccessRule($adminRule) + $acl.SetAccessRule($userRule) + + # Apply the modified ACL to the directory + Set-Acl -Path $Path -AclObject $acl +} + + +# Validate required command-line arguments +if (-not $FarmId) { + Write-Host "ERROR: -FarmId not specified" + usage +} +elseif (-not (ValidateDeadlineId -prefix "farm" -text $FarmId)) { + Write-Host "ERROR: Not a valid value for -FarmId: ${FarmId}" + usage +} + +if (-not $FleetId) { + Write-Host "ERROR: -FleetId not specified" + usage +} +elseif (-not (ValidateDeadlineId -prefix "fleet" -text $FleetId)) { + Write-Host "ERROR: Not a valid value for -FleetId: ${FleetId}" + usage +} + +if (-not $WorkerAgentProgram) { + $WorkerAgentProgram=(Get-Command deadline-worker-agent -ErrorAction SilentlyContinue).Path + + if (-not $WorkerAgentProgram) { + Write-Host "ERROR: Could not find deadline-worker-agent in search path" + exit 1 + } +} +elseif (-not (Test-Path -Path $WorkerAgentProgram -PathType Leaf)) { + Write-Host "ERROR: The specified Worker Agent path is not found: ${worker_agent_program}" + usage +} + + +# Output configuration +Banner +Write-Host "" +Write-Host "Farm ID: $FarmId" +Write-Host "Fleet ID: $FleetId" +Write-Host "Region: $Region" +Write-Host "Worker agent user: $User" +Write-Host "Worker job group: $Group" +Write-Host "Worker agent program path: $WorkerAgentProgram" +Write-Host "Start service: $Start" + +# Confirmation prompt +if (!$Confirm) { + while ($true) { + $choice = Read-Host "Confirm install with the above settings (y/n):" + if ($choice -eq "y") { + $Confirm = $true + break + } + elseif ($choice -eq "n") { + Write-Host "Installation aborted" + exit 1 + } + else { + Write-Host "Not a valid choice ($choice). Please try again." + } + } +} + +Write-Host "" + +# Check if the worker agent user exists, and create it if not +if (!(UserExists $User)) { + Write-Host "Creating worker agent user ($User)" + $null = New-LocalUser -Name $User -NoPassword # -UserMayNotChangePassword -PasswordNeverExpires -AccountNeverExpires + Write-Host "Done creating worker agent user ($User)" +} +else { + Write-Host "Worker agent user $User already exists" +} + +# Check if the job group exists, and create it if not +if (!(GroupExists $Group)) { + Write-Host "Creating job group ($Group)" + $null = New-LocalGroup -Name $Group + Write-Host "Done creating job group ($Group)" +} +else { + Write-Host "Job group $Group already exists" +} + +# Add the worker agent user to the job group +$groupMembers = Get-LocalGroupMember -Group $Group -ErrorAction SilentlyContinue +if ($groupMembers -eq $null -or ($groupMembers | Where-Object { $_.Name -eq $User }) -eq $null) { + # User is not a member of the group, so add them + Add-LocalGroupMember -Group $Group -Member $User -ErrorAction SilentlyContinue + Write-Host "User $User added to group $Group." +} else { + Write-Host "User $User is already a member of group $Group." +} + + +$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path + +# Provision directories +$deadLineDirectory = "C:\ProgramData\Amazon\Deadline" +Write-Host "Provisioning root directory ($deadlineDirectory)" +$null = New-Item -Path $deadLineDirectory -ItemType Directory -Force +# Setting the permissions correctly on this directory will set them correctly on everything +# within it +Set-DirectoryPermissions -Path $deadLineDirectory -User $User -Permission FullControl +Write-Host "Done provisioning root directory ($deadlineDirectory)" + +$deadlineLogSubDirectory = Join-Path $deadLineDirectory "Logs" +Write-Host "Provisioning log directory ($deadlineLogSubDirectory)" +$null = New-Item -Path $deadlineLogSubDirectory -ItemType Directory -Force +Write-Host "Done provisioning log directory ($deadlineLogSubDirectory)" + +$deadlinePersistenceSubDirectory = Join-Path $deadLineDirectory "Cache" +Write-Host "Provisioning persistence directory ($deadlinePersistenceSubDirectory)" +$null = New-Item -Path $deadlinePersistenceSubDirectory -ItemType Directory -Force +Write-Host "Done provisioning persistence directory ($deadlinePersistenceSubDirectory)" + +$deadlineConfigSubDirectory = Join-Path $deadLineDirectory "Config" +Write-Host "Provisioning config directory ($deadlineConfigSubDirectory)" +$null = New-Item -Path $deadlineConfigSubDirectory -ItemType Directory -Force +Write-Host "Done provisioning config directory ($deadlineConfigSubDirectory)" + + +Write-Host "Configuring farm and fleet" +$workerConfigFile = Join-Path $deadlineConfigSubDirectory "worker.toml" +if (-not (Test-Path -Path $$workerConfigFile -PathType Leaf)) { + Copy-Item -Path "$ScriptDir/worker.toml.windows.example" -Destination $workerConfigFile +} +$backupWorkerConfig = "$workerConfigFile.bak" +Copy-Item -Path $workerConfigFile -Destination $backupWorkerConfig +$content = Get-Content -Path $workerConfigFile +$content = $content -replace '^# farm_id\s*=\s*("REPLACE-WTIH-WORKER-FARM-ID")$', "farm_id = `"$FarmId`"" +$content = $content -replace '^# fleet_id\s*=\s*("REPLACE-WITH-WORKER-FLEET-ID")$', "fleet_id = `"$FleetId`"" +$content | Set-Content -Path $workerConfigFile +Write-Host "Done configuring farm and fleet" + + +if (!$no_install_service) { + # Set up the service + Write-Host "Installing Windows service" + $deadlineServiceName = "DeadlineCloudWorkerAgent" + $deadlineServiceDisplayName = "Amazon Deadline Cloud Worker Agent" + $deadlineServiceDescription = "Amazon Deadline Cloud Worker Agent" + $deadlineServiceExecutable = $WorkerAgentProgram + $credentials = New-Object System.Management.Automation.PSCredential ($User, (New-Object System.Security.SecureString)) + + # Check if the service exists + if (-not (Get-Service -Name $deadlineServiceName -ErrorAction SilentlyContinue)) { + # Service does not exist, so create it + New-Service -Name $deadlineServiceName -Credential $credentials -DisplayName $deadlineServiceDisplayName -Description $deadlineServiceDescription -BinaryPathName $deadlineServiceExecutable -StartupType "Automatic" + Write-Host "Service created." + } else { + Write-Host "Service already exists." + } + + Write-Host "Done installing Windows service" + + # Start the service + if ($Start) { + Write-Host "Starting the service" + Start-Service -Name $deadlineServiceName + Write-Host "Done starting the service" + } +} + +Write-Host "Done" \ No newline at end of file diff --git a/src/deadline_worker_agent/installer/worker.toml.windows.example b/src/deadline_worker_agent/installer/worker.toml.windows.example new file mode 100644 index 00000000..44e78975 --- /dev/null +++ b/src/deadline_worker_agent/installer/worker.toml.windows.example @@ -0,0 +1,256 @@ +[worker] + +# The unique identifier of the Amazon Deadline Cloud farm that the Worker belongs to. This value is overridden +# when the DEADLINE_WORKER_FARM_ID environment variable is set or the --farm-id command-line argument +# is specified. +# +# The following is an example for setting the farm ID in this configuration file: +# +# farm_id = "farm-aabbccddeeff11223344556677889900" +# +# Uncomment the line below and replace the value with your Amazon Deadline Cloud farm ID: +# +# farm_id = "REPLACE-WTIH-WORKER-FARM-ID" + + +# The unique identifier of the Amazon Deadline Cloud fleet that the Worker belongs to. This value is overridden +# when the DEADLINE_WORKER_FLEET_ID environment variable is set or the --fleet-id command-line +# argument is specified. +# +# The following is an example for setting the fleet ID in this configuration file: +# +# fleet_id = "fleet-aabbccddeeff11223344556677889900" +# +# Uncomment the line below and replace the value with your Amazon Deadline Cloud fleet ID: +# +# fleet_id = "REPLACE-WITH-WORKER-FLEET-ID" + + +# The directory where the worker agent persists its settings and credentials. The persistence directory +# is expected to be located on a file-system that is local to the Worker Host. The Worker's ID +# and credentials are persisted and these should not be accessible by other Workers. This value is overridden +# when the DEADLINE_WORKER_PERSISTENCE_DIR environment variable is set or the --persistence-dir command-line argument +# is specified. +# +# The following is the default worker persistence dir on POSIX systems. +# worker_persistence_dir = "/var/lib/deadline" + +[aws] + +# The worker agent requires initial AWS credentials in order to bootstrap the worker. Bootstrapping +# the worker consists of two steps: +# +# 1. Creating a worker resource if one has not been previously created and persisted +# to /var/lib/deadline/worker.json +# 2. Obtaining credentials for the worker resource +# +# Once the worker credentials are obtained, those credentials are used and rotated throughout the +# worker agent's lifecycle +# +# By default, the worker agent uses boto3's default credential resolution process to obtain the +# bootstrapping credentials. For details, on this, see +# https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html. +# +# Optionally, the name of an AWS profile can be specified to obtain bootstrapping credentials. For +# an overview of the AWS config file and AWS profiles, see +# https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html +# +# To use an AWS profile, uncomment the line below and replace with your desired AWS profile name. +# This value is overridden when the DEADLINE_WORKER_PROFILE environment variable is set or if a +# --profile command-line argument is specified. +# +# profile = "my_aws_profile_name" + + +# The "allow_ec2_instance_profile" setting controls whether the Worker will run with an EC2 instance +# profile associated to the host instance. This value is overridden when the +# DEADLINE_WORKER_ALLOW_INSTANCE_PROFILE environment variable is set using one of the following +# case-insensitive values: +# +# '0', 'off', 'f', 'false', 'n', 'no', '1', 'on', 't', 'true', 'y', 'yes'. +# +# or if the --allow-instance-profile command-line flag is specified. +# +# By default, this value is false and the worker agent makes requests to the EC2 instance meta-data +# service (IMDS) to check for an instance profile. If an instance profile is detected, the worker +# agent will stop and exit. +# +# ***************************************** WARNING ***************************************** +# * * +# * IF THIS IS TRUE, THEN ANY SESSIONS RUNNING ON THE WORKER CAN ASSUME THE INSTANCE * +# * PROFILE AND USE ITS ASSOCIATED PRIVILEGES * +# * * +# ******************************************************************************************* +# +# To accept this risk and allow the worker agent to run with an EC2 instance profile, uncomment +# the line below: +# +# allow_ec2_instance_profile = true + +[logging] + +# Setting "verbose" to true causes more detailed log output which may be useful for troubleshooting. +# This value is overridden when the DEADLINE_WORKER_VERBOSE environment variable is set using one +# of the following case-insensitive values: +# +# '0', 'off', 'f', 'false', 'n', 'no', '1', 'on', 't', 'true', 'y', 'yes'. +# +# or if the --verbose command-line flag is specified. +# +# ***************************************** WARNING ***************************************** +# * * +# * IT IS RECOMMENDED TO NOT TURN ON VERBOSE LOGGING FOR PRODUCTION USE * +# * * +# ******************************************************************************************* +# +# To turn on verbose logging, uncomment the line below: +# +# verbose = true + + +# The directory where the worker agent writes log files. This value is overridden +# when the DEADLINE_WORKER_LOGS_DIR environment variable is set or the --logs-dir command-line argument +# is specified. +# +# The following is the default worker logs dir on POSIX systems. To change the log directory, +# uncomment and modify the line below: +# +# worker_logs_dir = "/var/log/amazon/deadline" + + +# Whether the worker agent should write session logs to the local file-system. +# +# ******************************************* NOTE ****************************************** +# * * +# # The Amazon Deadline Cloud Worker Agent does not have any feature to delete/clean up session logs. # +# # It is up to the worker host admin to manage the deletion of session logs written to # +# # local file-system. Otherwise, the worker could eventually fill up all file-system # +# # capacity. # +# * * +# ******************************************************************************************* +# +# The logs are written to: +# +# //.log +# +# This value is overridden when the DEADLINE_WORKER_LOCAL_SESSION_LOGS environment variable is set using +# one of the following case-insensitive values: +# +# '0', 'off', 'f', 'false', 'n', 'no', '1', 'on', 't', 'true', 'y', 'yes'. +# +# or if the --no-local-session-logs command-line flag is specified. +# +# By default local session logging is turned on. To turn off local session logging, uncomment the +# line below: +# +# local_session_logs = false + + +[os] + +# Amazon Deadline Cloud may specify an OS user to impersonate when running session actions. By setting +# "impersonation" to false, the OS user is ignored and the session actions are run as the same user +# that the worker agent process runs as. +# +# This value is overridden when the DEADLINE_WORKER_IMPERSONATION environment variable is set using +# one of the following case-insensitive values: +# +# '0', 'off', 'f', 'false', 'n', 'no', '1', 'on', 't', 'true', 'y', 'yes'. +# +# or if the --no-impersonation command-line flag is specified. +# +# To turn off OS user impersonation, uncomment the line below +# +impersonation = false + + +# Amazon Deadline Cloud may tell the worker to stop. If the "shutdown_on_stop" setting below is true, then the +# Worker will attempt to shutdown the host system after the Worker has been stopped. +# +# This value is overridden when the DEADLINE_WORKER_NO_SHUTDOWN environment variable is set using +# one of the following case-insensitive values: +# +# '0', 'off', 'f', 'false', 'n', 'no', '1', 'on', 't', 'true', 'y', 'yes'. +# +# or if the --no-shutdown command-line flag is specified. +# +# To prevent the worker agent from shutting down the host when being told to stop, uncomment the +# line below: +# +# shutdown_on_stop = false + + +[capabilities] + +# Capabilities for the Worker can be declared in this config section. There are two types of +# capabilities - amounts and attributes. +# +# The values in the following configuration sections are overridden if the +# DEADLINE_WORKER_CAPABILITIES environment variable is set using a JSON syntax, for example: +# +# {"amounts": { "amount.apples": 12 }, "attributes": { "attr.license_servers": ["green.internal"] }} +# +# For further documentation, see the sections below. + +[capabilities.amounts] + +# Amount Capabilities +# ------------------- +# +# Amounts represent a capability with a numeric quantity. Amount capabilities are expressed in the +# config file as: +# +# "[:]amount." = +# +# Where: +# +# +# An optional namespace that indicates the capabilitiy is specific to a particular +# Open Job Description-compatable render farm management system. +# +# A name for the amount capability +# +# A positive floating-point value which indicates the amount of capability available to +# the worker. +# +# Examples +# -------- +# +# "amount.slots" = 20 +# "deadline:amount.pets" = 99 + +[capabilities.attributes] + +# +# Attributes Capabilities +# ------------------- +# +# Attributes represent a capability with a one or more string labels, e.g. major version of a +# software component. Attribute capabilities are expressed in the config file as: +# +# "[:]attr." = [ "", ... ] +# [...] +# +# Where: +# +# +# An optional namespace that indicates the capabilitiy is specific to a particular +# Open Job Description-compatable render farm component. +# +# A name for the attribute capability +# +# One or more string labels +# +# Examples +# -------- +# +# "attr.groups" = [ +# "simulation", +# "maya", +# "nuke" +# ] +# +# "acmewidgetsco:attr.admins" = [ +# "bob", +# "alice" +# ] diff --git a/src/deadline_worker_agent/scheduler/scheduler.py b/src/deadline_worker_agent/scheduler/scheduler.py index 5a2b4761..6594828c 100644 --- a/src/deadline_worker_agent/scheduler/scheduler.py +++ b/src/deadline_worker_agent/scheduler/scheduler.py @@ -53,7 +53,8 @@ from .session_queue import SessionActionQueue, SessionActionStatus from ..startup.config import ImpersonationOverrides from ..utils import MappingWithCallbacks - +from ..set_windows_permissions import grant_full_control +import subprocess logger = LOGGER @@ -219,7 +220,7 @@ def run(self) -> None: The Worker begins by hydrating its assigned work using the UpdateWorkerSchedule API. The scheduler then enters a loop of processing assigned actions - creating and deleting - Worker sessions as required. If no actions are assigned, the Worke idles for 5 seconds. + Worker sessions as required. If no actions are assigned, the Worker idles for 5 seconds. If any action completes, finishes cancelation, or if the Worker is done idling, an UpdateWorkerSchedule API request is made with any relevant changes specified in the request. @@ -636,8 +637,12 @@ def _create_new_sessions( if self._worker_logs_dir: queue_log_dir = self._queue_log_dir_path(queue_id=session_spec["queueId"]) try: - queue_log_dir.mkdir(mode=stat.S_IRWXU, exist_ok=True) - except OSError: + if os.name == "posix": + queue_log_dir.mkdir(mode=stat.S_IRWXU, exist_ok=True) + else: + queue_log_dir.mkdir(exist_ok=True) + grant_full_control(queue_log_dir.name) + except (OSError, subprocess.CalledProcessError): error_msg = ( f"Failed to create local session log directory on worker: {queue_log_dir}" ) diff --git a/src/deadline_worker_agent/sessions/job_entities/job_details.py b/src/deadline_worker_agent/sessions/job_entities/job_details.py index 0195b46b..e9f80171 100644 --- a/src/deadline_worker_agent/sessions/job_entities/job_details.py +++ b/src/deadline_worker_agent/sessions/job_entities/job_details.py @@ -106,7 +106,7 @@ def jobs_runs_as_api_model_to_worker_agent( ) else: # TODO: windows support - raise NotImplementedError(f"{os.name} is not supported") + return None return jobs_run_as diff --git a/src/deadline_worker_agent/set_windows_permissions.py b/src/deadline_worker_agent/set_windows_permissions.py new file mode 100644 index 00000000..29a7f1ce --- /dev/null +++ b/src/deadline_worker_agent/set_windows_permissions.py @@ -0,0 +1,40 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +from typing import Optional +import subprocess +import getpass + + +def grant_full_control(path: str, username: Optional[str] = None): + """ + Set permissions for a specified file or directory (and any child objects) + to give full control only to the specified user. + + Args: + path (str): The path of the file or directory for which permissions will be set. + username (str, optional): The username for whom permissions will be granted. If none is + provided the current username will be used. + + Example: + path = "C:\\example_directory_or_file" + username = "a_username" + grant_full_control(path, username) + """ + + if not username: + username = getpass.getuser() + + subprocess.run( + [ + "icacls", + path, + # Remove any existing permissions + "/inheritance:r", + # OI - Contained objects will inherit + # CI - Sub-directories will inherit + # F - Full control + "/grant", + ("{0}:(OI)(CI)(F)").format(username), + "/T", # Apply recursively for directories + ] + ) diff --git a/src/deadline_worker_agent/startup/config.py b/src/deadline_worker_agent/startup/config.py index 124e23c4..694ed29c 100644 --- a/src/deadline_worker_agent/startup/config.py +++ b/src/deadline_worker_agent/startup/config.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os import logging as _logging from dataclasses import dataclass from pathlib import Path @@ -119,7 +120,7 @@ def __init__( settings = WorkerSettings(**settings_kwargs) - if settings.posix_job_user is not None: + if os.name == "posix" and settings.posix_job_user is not None: user, group = self._get_user_and_group_from_posix_job_user(settings.posix_job_user) self.impersonation = ImpersonationOverrides( inactive=not settings.impersonation, diff --git a/src/deadline_worker_agent/startup/config_file.py b/src/deadline_worker_agent/startup/config_file.py index 3ca74a39..7075cd12 100644 --- a/src/deadline_worker_agent/startup/config_file.py +++ b/src/deadline_worker_agent/startup/config_file.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any, Optional import sys +import os from pydantic import BaseModel, BaseSettings, Field @@ -20,6 +21,7 @@ DEFAULT_CONFIG_PATH: dict[str, Path] = { "darwin": Path("/etc/amazon/deadline/worker.toml"), "linux": Path("/etc/amazon/deadline/worker.toml"), + "win32": Path(os.path.expandvars(r"%PROGRAMDATA%/Amazon/Deadline/Config/worker.toml")), } diff --git a/src/deadline_worker_agent/startup/entrypoint.py b/src/deadline_worker_agent/startup/entrypoint.py index 79335d9c..c99c8cb6 100644 --- a/src/deadline_worker_agent/startup/entrypoint.py +++ b/src/deadline_worker_agent/startup/entrypoint.py @@ -44,9 +44,14 @@ def detect_system_capabilities() -> Capabilities: "linux": "linux", "windows": "windows", } + platform_machine = platform.machine().lower() + python_machine_to_openjd_cpu_arch = {"x86_64": "x86_64", "amd64": "x86_64"} if openjd_os_family := python_system_to_openjd_os_family.get(platform_system): attributes[AttributeCapabilityName("attr.worker.os.family")] = [openjd_os_family] - attributes[AttributeCapabilityName("attr.worker.cpu.arch")] = [platform.machine()] + if openjd_cpu_arch := python_machine_to_openjd_cpu_arch.get(platform_machine): + attributes[AttributeCapabilityName("attr.worker.cpu.arch")] = [openjd_cpu_arch] + else: + raise NotImplementedError(f"{platform_machine} not supported") amounts[AmountCapabilityName("amount.worker.vcpu")] = float(psutil.cpu_count()) amounts[AmountCapabilityName("amount.worker.memory")] = float(psutil.virtual_memory().total) / ( 1024.0**2 diff --git a/src/deadline_worker_agent/startup/settings.py b/src/deadline_worker_agent/startup/settings.py index 98b5087b..f473ede3 100644 --- a/src/deadline_worker_agent/startup/settings.py +++ b/src/deadline_worker_agent/startup/settings.py @@ -11,14 +11,20 @@ from .capabilities import Capabilities from .config_file import ConfigFile +import os + # Default path for the worker's logs. DEFAULT_POSIX_WORKER_LOGS_DIR = Path("/var/log/amazon/deadline") +DEFAULT_WINDOWS_WORKER_LOGS_DIR = Path(os.path.expandvars(r"%PROGRAMDATA%/Amazon/Deadline/Logs")) # Default path for the worker persistence directory. # The persistence directory is expected to be located on a file-system that is local to the Worker # Node. The Worker's ID and credentials are persisted and these should not be accessible by other # Worker Nodes. DEFAULT_POSIX_WORKER_PERSISTENCE_DIR = Path("/var/lib/deadline") +DEFAULT_WINDOWS_WORKER_PERSISTENCE_DIR = Path( + os.path.expandvars(r"%PROGRAMDATA%/Amazon/Deadline/Cache") +) class WorkerSettings(BaseSettings): @@ -80,8 +86,14 @@ class WorkerSettings(BaseSettings): capabilities: Capabilities = Field( default_factory=lambda: Capabilities(amounts={}, attributes={}) ) - worker_logs_dir: Path = DEFAULT_POSIX_WORKER_LOGS_DIR - worker_persistence_dir: Path = DEFAULT_POSIX_WORKER_PERSISTENCE_DIR + worker_logs_dir: Path = ( + DEFAULT_WINDOWS_WORKER_LOGS_DIR if os.name == "nt" else DEFAULT_POSIX_WORKER_LOGS_DIR + ) + worker_persistence_dir: Path = ( + DEFAULT_WINDOWS_WORKER_PERSISTENCE_DIR + if os.name == "nt" + else DEFAULT_POSIX_WORKER_PERSISTENCE_DIR + ) local_session_logs: bool = True class Config: diff --git a/src/deadline_worker_agent/worker.py b/src/deadline_worker_agent/worker.py index acfd3b32..7491b0a3 100644 --- a/src/deadline_worker_agent/worker.py +++ b/src/deadline_worker_agent/worker.py @@ -4,6 +4,7 @@ import json import signal +import os import sys import traceback from concurrent.futures import Executor, Future, ThreadPoolExecutor, wait @@ -108,8 +109,10 @@ def __init__( signal.signal(signal.SIGTERM, self._signal_handler) signal.signal(signal.SIGINT, self._signal_handler) - # TODO: Remove this once WA is stable or put behind a debug flag - signal.signal(signal.SIGUSR1, self._output_thread_stacks) + + if os.name == "posix": + # TODO: Remove this once WA is stable or put behind a debug flag + signal.signal(signal.SIGUSR1, self._output_thread_stacks) def _signal_handler(self, signum: int, frame: FrameType | None = None) -> None: """ @@ -156,7 +159,7 @@ def id(self) -> str: @property def sessions(self) -> WorkerSessionCollection: - raise NotImplementedError("Worker.sessions property not implemeneted") + raise NotImplementedError("Worker.sessions property not implemented") def run(self) -> None: """Runs the main Worker loop for processing sessions.""" @@ -373,7 +376,15 @@ def _get_spot_instance_shutdown_action_timeout(self, *, imdsv2_token: str) -> ti logger.info(f"Spot {action} happening at {shutdown_time}") # Spot gives the time in UTC with a trailing Z, but Python can't handle # the Z so we strip it - shutdown_time = datetime.fromisoformat(shutdown_time[:-1]).astimezone(timezone.utc) + if os.name == "posix": + shutdown_time = datetime.fromisoformat(shutdown_time[:-1]).astimezone( + timezone.utc + ) + else: + # astimezone() appears to behave differently on Windows, it will make a timestamp that + # is already utc incorrect + shutdown_time = datetime.fromisoformat(shutdown_time[:-1]) + shutdown_time = shutdown_time.replace(tzinfo=timezone.utc) current_time = datetime.now(timezone.utc) time_delta = shutdown_time - current_time time_delta_seconds = int(time_delta.total_seconds()) diff --git a/test/unit/aws_credentials/test_aws_configs.py b/test/unit/aws_credentials/test_aws_configs.py index 2323fa55..c8eefc60 100644 --- a/test/unit/aws_credentials/test_aws_configs.py +++ b/test/unit/aws_credentials/test_aws_configs.py @@ -14,6 +14,7 @@ _setup_parent_dir, ) from openjd.sessions import PosixSessionUser, SessionUser +import os @pytest.fixture @@ -27,9 +28,13 @@ def mock_run_cmd_as() -> Generator[MagicMock, None, None]: yield mock_run_cmd_as -@pytest.fixture(params=(PosixSessionUser(user="some-user", group="some-group"), None)) -def os_user(request: pytest.FixtureRequest) -> Optional[SessionUser]: - return request.param +@pytest.fixture +def os_user() -> Optional[SessionUser]: + if os.name == "posix": + return PosixSessionUser(user="user", group="group") + else: + # TODO: Revisit when Windows impersonation is added + return None class TestSetupParentDir: @@ -62,10 +67,15 @@ def test_creates_dir( ) mock_run_cmd_as.assert_any_call(user=os_user, cmd=["chmod", "770", str(dir_path)]) else: - mkdir.assert_called_once_with( - mode=0o700, - exist_ok=True, - ) + if os.name == "posix": + mkdir.assert_called_once_with( + mode=0o700, + exist_ok=True, + ) + else: + mkdir.assert_called_once_with( + exist_ok=True, + ) def test_sets_group_ownership( self, diff --git a/test/unit/aws_credentials/test_queue_boto3_session.py b/test/unit/aws_credentials/test_queue_boto3_session.py index 8748efad..42fc65ea 100644 --- a/test/unit/aws_credentials/test_queue_boto3_session.py +++ b/test/unit/aws_credentials/test_queue_boto3_session.py @@ -51,9 +51,13 @@ def deadline_client() -> MagicMock: return MagicMock() -@pytest.fixture(params=(PosixSessionUser(user="some-user", group="some-group"), None)) -def os_user(request: pytest.FixtureRequest) -> Optional[SessionUser]: - return request.param +@pytest.fixture +def os_user() -> Optional[SessionUser]: + if os.name == "posix": + return PosixSessionUser(user="user", group="group") + else: + # TODO: Revisit when Windows impersonation is added + return None class TestInit: @@ -614,9 +618,14 @@ def test_success( session._install_credential_process() # THEN - credentials_process_script_path = ( - Path(tmpdir) / "queues" / queue_id / "get_aws_credentials.sh" - ) + if os.name == "posix": + credentials_process_script_path = ( + Path(tmpdir) / "queues" / queue_id / "get_aws_credentials.sh" + ) + else: + credentials_process_script_path = ( + Path(tmpdir) / "queues" / queue_id / "get_aws_credentials.cmd" + ) mock_os_open.assert_called_once_with( path=str(credentials_process_script_path), flags=os.O_WRONLY | os.O_CREAT | os.O_TRUNC, diff --git a/test/unit/conftest.py b/test/unit/conftest.py index 6f6eb290..2e924374 100644 --- a/test/unit/conftest.py +++ b/test/unit/conftest.py @@ -51,16 +51,34 @@ def logs_client() -> MagicMock: return MagicMock() -@pytest.fixture(params=(PosixSessionUser(user="some-user", group="some-group"),)) -def posix_job_user(request: pytest.FixtureRequest) -> Optional[SessionUser]: - return request.param +if os.name == "posix": + @pytest.fixture(params=(PosixSessionUser(user="some-user", group="some-group"),)) + def os_user(request: pytest.FixtureRequest) -> Optional[SessionUser]: + return request.param -@pytest.fixture(params=(False,)) -def impersonation( - request: pytest.FixtureRequest, posix_job_user: Optional[SessionUser] -) -> ImpersonationOverrides: - return ImpersonationOverrides(inactive=request.param, posix_job_user=posix_job_user) +else: + + @pytest.fixture() + def os_user() -> Optional[SessionUser]: + return None + + +if os.name == "posix": + + @pytest.fixture(params=(False,)) + def impersonation( + request: pytest.FixtureRequest, os_user: Optional[SessionUser] + ) -> ImpersonationOverrides: + return ImpersonationOverrides(inactive=request.param, posix_job_user=os_user) + +else: + + @pytest.fixture(params=(True,)) + def impersonation( + request: pytest.FixtureRequest, + ) -> ImpersonationOverrides: + return ImpersonationOverrides(inactive=request.param) @pytest.fixture @@ -242,7 +260,7 @@ def jobs_run_as() -> JobsRunAs | None: """The OS user/group associated with the job's queue""" # TODO: windows support if os.name != "posix": - raise NotImplementedError(f"{os.name} is not supported") + return None return JobsRunAs(posix=PosixSessionUser(user="job-user", group="job-user")) diff --git a/test/unit/install/test_install.py b/test/unit/install/test_install.py index c1bd78b0..6dacbdd5 100644 --- a/test/unit/install/test_install.py +++ b/test/unit/install/test_install.py @@ -258,7 +258,6 @@ def test_runs_expected_subprocess( "aix", "emscripten", "wasi", - "win32", "cygwin", "darwin", ), @@ -273,4 +272,4 @@ def test_unsupported_platform_raises(platform: str, capsys: pytest.CaptureFixtur assert raise_ctx.value.code == 1 capture = capsys.readouterr() - assert capture.out == f"ERROR: Unsupported platform {platform}{os.linesep}" + assert capture.out == f"ERROR: Unsupported platform {platform}\n" diff --git a/test/unit/log_sync/test_cloudwatch.py b/test/unit/log_sync/test_cloudwatch.py index 9740ed01..6290edb1 100644 --- a/test/unit/log_sync/test_cloudwatch.py +++ b/test/unit/log_sync/test_cloudwatch.py @@ -36,7 +36,7 @@ def mock_module_logger() -> Generator[MagicMock, None, None]: class TestCloudWatchLogEventBatch: @fixture(autouse=True) def now(self) -> datetime: - return datetime.fromtimestamp(123) + return datetime(2000, 1, 1) @fixture(autouse=True) def datetime_mock(self, now: datetime) -> Generator[MagicMock, None, None]: @@ -51,7 +51,7 @@ def datetime_mock(self, now: datetime) -> Generator[MagicMock, None, None]: @fixture def event(self, now: datetime) -> PartitionedCloudWatchLogEvent: return PartitionedCloudWatchLogEvent( - log_event=CloudWatchLogEvent(timestamp=int(now.timestamp()), message="abc"), + log_event=CloudWatchLogEvent(timestamp=int(now.timestamp() * 1000), message="abc"), size=len("abc".encode("utf-8")), ) @@ -168,10 +168,12 @@ def test_valid_log_event( datetime_mock: MagicMock, ): # GIVEN - now = datetime.fromtimestamp(1) + now = datetime(2000, 1, 1) datetime_mock.now.return_value = now event = PartitionedCloudWatchLogEvent( - log_event=CloudWatchLogEvent(message="abc", timestamp=int(now.timestamp())), + log_event=CloudWatchLogEvent( + message="abc", timestamp=(int(now.timestamp()) * 1000) + ), size=3, ) batch = CloudWatchLogEventBatch() diff --git a/test/unit/scheduler/test_scheduler.py b/test/unit/scheduler/test_scheduler.py index cd54898f..a6ec057d 100644 --- a/test/unit/scheduler/test_scheduler.py +++ b/test/unit/scheduler/test_scheduler.py @@ -11,6 +11,7 @@ from openjd.sessions import ActionState, ActionStatus from botocore.exceptions import ClientError import pytest +import os from deadline_worker_agent.api_models import ( AssignedSession, @@ -574,6 +575,7 @@ def test_local_logging( session_log_file_path = MagicMock() with ( + patch.object(scheduler_mod, "grant_full_control") as mock_grant_full_control, patch.object(scheduler, "_executor"), patch.object(scheduler_mod.LogConfiguration, "from_boto") as mock_log_config_from_boto, patch.object( @@ -588,7 +590,11 @@ def test_local_logging( # THEN mock_queue_log_dir.assert_called_once_with(queue_id=queue_id) - queue_log_dir_path.mkdir.assert_called_once_with(mode=0o700, exist_ok=True) + if os.name == "posix": + queue_log_dir_path.mkdir.assert_called_once_with(mode=0o700, exist_ok=True) + else: + queue_log_dir_path.mkdir.assert_called_once_with(exist_ok=True) + mock_grant_full_control.assert_called_once() mock_queue_session_log_file_path.assert_called_once_with( session_id=session_id, queue_log_dir=queue_log_dir_path ) @@ -672,7 +678,10 @@ def test_local_logging_os_error( # THEN mock_queue_log_dir.assert_called_once_with(queue_id=queue_id) - queue_log_dir_path.mkdir.assert_called_once_with(mode=0o700, exist_ok=True) + if os.name == "posix": + queue_log_dir_path.mkdir.assert_called_once_with(mode=0o700, exist_ok=True) + else: + queue_log_dir_path.mkdir.assert_called_once_with(exist_ok=True) if mkdir_side_effect: mock_queue_session_log_file_path.assert_not_called() else: diff --git a/test/unit/scheduler/test_session_cleanup.py b/test/unit/scheduler/test_session_cleanup.py index b37da445..2f57d36f 100644 --- a/test/unit/scheduler/test_session_cleanup.py +++ b/test/unit/scheduler/test_session_cleanup.py @@ -8,6 +8,7 @@ from openjd.sessions import SessionUser, PosixSessionUser import pytest +import os from deadline_worker_agent.scheduler.session_cleanup import ( SessionUserCleanupManager, @@ -20,6 +21,11 @@ def __init__(self, user: str): self.user = user +class WindowsSessionUser(SessionUser): + def __init__(self, user: str): + self.user = user + + class TestSessionUserCleanupManager: @pytest.fixture def manager(self) -> SessionUserCleanupManager: @@ -34,8 +40,11 @@ def user_session_map_lock_mock( yield mock @pytest.fixture - def os_user(self) -> PosixSessionUser: - return PosixSessionUser(user="user", group="group") + def os_user(self) -> SessionUser: + if os.name == "posix": + return PosixSessionUser(user="user", group="group") + else: + return WindowsSessionUser(user="user") @pytest.fixture def session(self, os_user: PosixSessionUser) -> MagicMock: @@ -45,11 +54,12 @@ def session(self, os_user: PosixSessionUser) -> MagicMock: return session_stub class TestRegister: + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_registers_session( self, manager: SessionUserCleanupManager, session: MagicMock, - os_user: PosixSessionUser, + os_user: SessionUser, user_session_map_lock_mock: MagicMock, ): # WHEN @@ -99,11 +109,12 @@ def test_register_raises_windows_not_supported( user_session_map_lock_mock.__exit__.assert_not_called() class TestDeregister: + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_deregisters_session( self, manager: SessionUserCleanupManager, session: MagicMock, - os_user: PosixSessionUser, + os_user: SessionUser, user_session_map_lock_mock: MagicMock, ): # GIVEN @@ -119,11 +130,12 @@ def test_deregisters_session( user_session_map_lock_mock.__enter__.assert_called_once() user_session_map_lock_mock.__exit__.assert_called_once() + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_deregister_skipped_no_user( self, manager: SessionUserCleanupManager, session: MagicMock, - os_user: PosixSessionUser, + os_user: SessionUser, user_session_map_lock_mock: MagicMock, ): # GIVEN @@ -166,9 +178,10 @@ def cleanup_session_user_processes_mock(self) -> Generator[MagicMock, None, None with patch.object(SessionUserCleanupManager, "cleanup_session_user_processes") as mock: yield mock + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_calls_cleanup_session_user_processes( self, - os_user: PosixSessionUser, + os_user: SessionUser, manager: SessionUserCleanupManager, cleanup_session_user_processes_mock: MagicMock, ): @@ -178,9 +191,10 @@ def test_calls_cleanup_session_user_processes( # THEN cleanup_session_user_processes_mock.assert_called_once_with(os_user) + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_skips_cleanup_when_configured_to( self, - os_user: PosixSessionUser, + os_user: SessionUser, cleanup_session_user_processes_mock: MagicMock, ): # GIVEN @@ -193,17 +207,28 @@ def test_skips_cleanup_when_configured_to( cleanup_session_user_processes_mock.assert_not_called() class TestCleanupSessionUserProcesses: - @pytest.fixture - def agent_user( - self, - os_user: PosixSessionUser, - ) -> PosixSessionUser: - return PosixSessionUser(user=f"agent_{os_user.user}", group=f"agent_{os_user.group}") + if os.name == "posix": + + @pytest.fixture + def agent_user( + self, + os_user: PosixSessionUser, + ) -> PosixSessionUser: + return PosixSessionUser( + user=f"agent_{os_user.user}", group=f"agent_{os_user.group}" + ) + + else: + @pytest.fixture + def agent_user(self) -> SessionUser: + return WindowsSessionUser(user="user") + + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") @pytest.fixture(autouse=True) def subprocess_check_output_mock( self, - agent_user: PosixSessionUser, + agent_user: SessionUser, ) -> Generator[MagicMock, None, None]: with patch.object( session_cleanup_mod.subprocess, @@ -212,14 +237,16 @@ def subprocess_check_output_mock( ) as mock: yield mock + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") @pytest.fixture(autouse=True) def subprocess_run_mock(self) -> Generator[MagicMock, None, None]: with patch.object(session_cleanup_mod.subprocess, "run") as mock: yield mock + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_cleans_up_processes( self, - os_user: PosixSessionUser, + os_user: SessionUser, subprocess_run_mock: MagicMock, caplog: pytest.LogCaptureFixture, ): @@ -258,6 +285,7 @@ def test_not_posix_user( assert str(raised_err.value) == "Windows not supported" subprocess_run_mock.assert_not_called() + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_no_processes_to_clean_up( self, os_user: PosixSessionUser, @@ -278,9 +306,10 @@ def test_no_processes_to_clean_up( in caplog.text ) + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_fails_to_clean_up_processes( self, - os_user: PosixSessionUser, + os_user: SessionUser, subprocess_run_mock: MagicMock, caplog: pytest.LogCaptureFixture, ): @@ -298,6 +327,7 @@ def test_fails_to_clean_up_processes( assert f"Failed to stop processes running as '{os_user.user}': {err}" in caplog.text assert raised_err.value is err + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_skips_if_session_user_is_agent_user( self, subprocess_run_mock: MagicMock, diff --git a/test/unit/sessions/test_job_entities.py b/test/unit/sessions/test_job_entities.py index c07fefa5..c74ea345 100644 --- a/test/unit/sessions/test_job_entities.py +++ b/test/unit/sessions/test_job_entities.py @@ -18,6 +18,7 @@ import pytest +import os from deadline_worker_agent.api_models import ( Attachments, @@ -174,6 +175,7 @@ def test_has_path_mapping_rules( assert job_details.path_mapping_rules not in (None, []) assert len(job_details.path_mapping_rules) == len(path_mapping_rules) + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_jobs_run_as(self) -> None: """Ensures that if we receive a jobs_run_as field in the response, that the created entity has a (Posix) SessionUser created with the diff --git a/test/unit/sessions/test_session.py b/test/unit/sessions/test_session.py index 78904756..a19fd902 100644 --- a/test/unit/sessions/test_session.py +++ b/test/unit/sessions/test_session.py @@ -10,6 +10,8 @@ from unittest.mock import patch, MagicMock, ANY import pytest +import os + from openjd.model.v2023_09 import ( Action, Environment, @@ -51,9 +53,12 @@ import deadline_worker_agent.sessions.session as session_mod -@pytest.fixture(params=(PosixSessionUser(user="some-user", group="some-group"),)) -def os_user(request: pytest.FixtureRequest) -> Optional[SessionUser]: - return request.param +@pytest.fixture +def os_user() -> Optional[SessionUser]: + if os.name == "posix": + return PosixSessionUser(user="user", group="group") + else: + return None @pytest.fixture @@ -485,6 +490,7 @@ def mock_asset_sync(self, session: Session) -> Generator[MagicMock, None, None]: # This overrides the asset_loading_method fixture in tests/unit/conftest.py which feeds into # the job_attachment_details fixture @pytest.mark.parametrize("asset_loading_method", [e.value for e in AssetLoadingMethod]) + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") def test_asset_loading_method( self, session: Session, diff --git a/test/unit/startup/test_config.py b/test/unit/startup/test_config.py index cd3a775d..24ff4dcd 100644 --- a/test/unit/startup/test_config.py +++ b/test/unit/startup/test_config.py @@ -9,8 +9,9 @@ from typing import Any, Generator, List, Optional import logging import pytest +import os -from openjd.sessions import PosixSessionUser +from openjd.sessions import SessionUser, PosixSessionUser from deadline_worker_agent.startup.cli_args import ParsedCommandLineArguments from deadline_worker_agent.startup import config as config_mod @@ -73,6 +74,14 @@ def arg_parser( yield arg_parser +@pytest.fixture +def user(self) -> Optional[SessionUser]: + if os.name == "posix": + return PosixSessionUser(user="user", group="group") + else: + return None + + class TestLoad: """Tests for Configuration.load()""" @@ -264,8 +273,18 @@ def test_uses_parsed_cleanup_session_user_processes( @pytest.mark.parametrize( argnames="worker_logs_dir", argvalues=( - Path("/foo"), - Path("/bar"), + pytest.param( + Path("/foo"), marks=pytest.mark.skipif(os.name != "posix", reason="Not posix") + ), + pytest.param( + Path("/bar"), marks=pytest.mark.skipif(os.name != "posix", reason="Not posix") + ), + pytest.param( + Path("D:\\foo"), marks=pytest.mark.skipif(os.name != "nt", reason="Not windows") + ), + pytest.param( + Path("D:\\bar"), marks=pytest.mark.skipif(os.name != "nt", reason="Not windows") + ), ), ) def test_uses_worker_logs_dir( @@ -287,8 +306,18 @@ def test_uses_worker_logs_dir( @pytest.mark.parametrize( argnames="persistence_dir", argvalues=( - Path("/foo"), - Path("/bar"), + pytest.param( + Path("/foo"), marks=pytest.mark.skipif(os.name != "posix", reason="Not posix") + ), + pytest.param( + Path("/bar"), marks=pytest.mark.skipif(os.name != "posix", reason="Not posix") + ), + pytest.param( + Path("D:\\foo"), marks=pytest.mark.skipif(os.name != "nt", reason="Not windows") + ), + pytest.param( + Path("D:\\bar"), marks=pytest.mark.skipif(os.name != "nt", reason="Not windows") + ), ), ) def test_uses_worker_persistence_dir( @@ -596,6 +625,7 @@ def test_impersonation_passed_to_settings_initializer( else: assert "impersonation" not in call.kwargs + @pytest.mark.skipif(os.name != "posix", reason="Posix-only test.") @pytest.mark.parametrize( argnames="posix_job_user", argvalues=("user:group", None), @@ -774,8 +804,9 @@ def test_local_session_logs_passed_to_settings_initializer( argvalues=( pytest.param( "user:group", - PosixSessionUser(group="group", user="user"), + user, id="has-posix-job-user-setting", + marks=pytest.mark.skipif(os.name != "posix", reason="Posix-only test."), ), pytest.param( None, diff --git a/test/unit/startup/test_settings.py b/test/unit/startup/test_settings.py index 3bff5847..245f01e8 100644 --- a/test/unit/startup/test_settings.py +++ b/test/unit/startup/test_settings.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, Mock, patch from typing import Any, Generator, NamedTuple, Type import pytest +import os from pathlib import Path from pydantic import ConstrainedStr @@ -105,14 +106,18 @@ class FieldTestCaseParams(NamedTuple): field_name="worker_logs_dir", expected_type=Path, expected_required=False, - expected_default=Path("/var/log/amazon/deadline"), + expected_default=Path("/var/log/amazon/deadline") + if os.name == "posix" + else Path(os.path.expandvars(r"%PROGRAMDATA%/Amazon/Deadline/Logs")), expected_default_factory_return_value=None, ), FieldTestCaseParams( field_name="worker_persistence_dir", expected_type=Path, expected_required=False, - expected_default=Path("/var/lib/deadline"), + expected_default=Path("/var/lib/deadline") + if os.name == "posix" + else Path(os.path.expandvars(r"%PROGRAMDATA%/Amazon/Deadline/Cache")), expected_default_factory_return_value=None, ), FieldTestCaseParams( diff --git a/test/unit/test_set_windows_permissions.py b/test/unit/test_set_windows_permissions.py new file mode 100644 index 00000000..5adcddef --- /dev/null +++ b/test/unit/test_set_windows_permissions.py @@ -0,0 +1,50 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + +import unittest +from unittest.mock import patch +from deadline_worker_agent.set_windows_permissions import grant_full_control + + +class TestGrantFullControl(unittest.TestCase): + @patch("subprocess.run") + @patch("getpass.getuser", return_value="testuser") + def test_grant_full_control_with_default_username(self, mock_getuser, mock_subprocess_run): + path = "C:\\example_directory_or_file" + grant_full_control(path) + + expected_command = [ + "icacls", + path, + "/inheritance:r", + "/grant", + "{0}:(OI)(CI)(F)".format("testuser"), + "/T", + ] + + mock_subprocess_run.assert_called_once_with(expected_command) + + mock_getuser.assert_called_once() + + @patch("subprocess.run") + @patch("getpass.getuser") + def test_grant_full_control_with_custom_username(self, mock_getuser, mock_subprocess_run): + path = "C:\\example_directory_or_file" + custom_username = "customuser" + grant_full_control(path, username=custom_username) + + expected_command = [ + "icacls", + path, + "/inheritance:r", + "/grant", + "{0}:(OI)(CI)(F)".format(custom_username), + "/T", + ] + + mock_subprocess_run.assert_called_once_with(expected_command) + + mock_getuser.assert_not_called() + + +if __name__ == "__main__": + unittest.main()