Skip to content

Commit

Permalink
feat: Add stop/start worker agent service method
Browse files Browse the repository at this point in the history
Signed-off-by: Charles Moore <moorec@amazon.com>
  • Loading branch information
moorec-aws committed Jul 16, 2024
1 parent cb76921 commit cc1006b
Showing 1 changed file with 40 additions and 3 deletions.
43 changes: 40 additions & 3 deletions src/deadline_test_fixtures/deadline/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,14 @@ def configure_worker_command(
) -> str: # pragma: no cover
raise NotImplementedError("'configure_worker_command' was not implemented.")

@abc.abstractmethod
def start_worker_service(self) -> None: # pragma: no cover
raise NotImplementedError("'_start_worker_service' was not implemented.")

@abc.abstractmethod
def stop_worker_service(self) -> None: # pragma: no cover
raise NotImplementedError("'_stop_worker_service' was not implemented.")

@abc.abstractmethod
def get_worker_id(self) -> str:
raise NotImplementedError("'get_worker_id' was not implemented.")
Expand Down Expand Up @@ -422,8 +430,8 @@ def _start_worker_agent(self) -> None:
cmd_result = self.send_command(
" ; ".join(
[
"echo Waiting 20s for the agent service to get started",
"sleep 20",
"echo Waiting 30s for the agent service to get started",
"sleep 30",
"echo 'Running running Get-Process to check if the agent is running'",
"IF(Get-Process pythonservice){echo 'service is running'}ELSE{exit 1}",
]
Expand All @@ -432,6 +440,9 @@ def _start_worker_agent(self) -> None:
assert cmd_result.exit_code == 0, f"Failed to start Worker agent: {cmd_result}"
LOG.info("Successfully started Worker agent")

# Wait 20 seconds for the worker to start before attempting to get ID
time.sleep(20)

self.worker_id = self.get_worker_id()

def configure_worker_command(self, *, config: DeadlineWorkerConfiguration) -> str:
Expand Down Expand Up @@ -508,6 +519,19 @@ def userdata(self, s3_files) -> str:

return userdata

def start_worker_service(self):
LOG.info("Sending command to start the Worker Agent service")

cmd_result = self.send_command('Start-Service -Name "DeadlineWorker"')

assert cmd_result.exit_code == 0, f"Failed to start Worker Agent service: : {cmd_result}"

def stop_worker_service(self):
LOG.info("Sending command to stop the Worker Agent service")
cmd_result = self.send_command('Stop-Service -Name "DeadlineWorker"')

assert cmd_result.exit_code == 0, f"Failed to stop Worker Agent service: : {cmd_result}"

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/
Expand Down Expand Up @@ -551,7 +575,7 @@ def _start_worker_agent(self) -> None:
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)}"
f"cd /home/{self.configuration.agent_user}; . .venv/bin/activate; AWS_DEFAULT_REGION={self.configuration.region} {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")
Expand Down Expand Up @@ -666,6 +690,19 @@ def userdata(self, s3_files) -> str:

return userdata

def start_worker_service(self):
LOG.info("Sending command to start the Worker Agent service")

cmd_result = self.send_command("systemctl start deadline-worker")

assert cmd_result.exit_code == 0, f"Failed to start Worker Agent service: {cmd_result}"

def stop_worker_service(self):
LOG.info("Sending command to stop the Worker Agent service")
cmd_result = self.send_command("systemctl stop deadline-worker")

assert cmd_result.exit_code == 0, f"Failed to start Worker Agent service: {cmd_result}"

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}"
Expand Down

0 comments on commit cc1006b

Please sign in to comment.