diff --git a/tardis/adapters/sites/slurm.py b/tardis/adapters/sites/slurm.py index 049147aa..ffb47421 100644 --- a/tardis/adapters/sites/slurm.py +++ b/tardis/adapters/sites/slurm.py @@ -6,7 +6,6 @@ from ...interfaces.siteadapter import SiteAdapter from ...utilities.staticmapping import StaticMapping from ...utilities.attributedict import AttributeDict -from ...utilities.attributedict import convert_to_attribute_dict from ...utilities.executors.shellexecutor import ShellExecutor from ...utilities.asynccachemap import AsyncCacheMap from ...utilities.utils import ( @@ -19,7 +18,6 @@ from asyncio import TimeoutError from contextlib import contextmanager from functools import partial -from datetime import datetime import logging import re @@ -127,13 +125,11 @@ async def deploy_resource( logger.debug(f"{self.site_name} sbatch returned {result}") pattern = re.compile(r"^Submitted batch job (\d*)", flags=re.MULTILINE) remote_resource_uuid = int(pattern.findall(result.stdout)[0]) - resource_attributes.update( + + return AttributeDict( remote_resource_uuid=remote_resource_uuid, - created=datetime.now(), - updated=datetime.now(), resource_status=ResourceStatus.Booting, ) - return resource_attributes async def resource_status( self, resource_attributes: AttributeDict @@ -156,20 +152,12 @@ async def resource_status( "State": "COMPLETED", } logger.debug(f"{self.site_name} has status {resource_status}.") - resource_attributes.update(updated=datetime.now()) - return convert_to_attribute_dict( - {**resource_attributes, **self.handle_response(resource_status)} - ) - async def terminate_resource(self, resource_attributes: AttributeDict): + return self.handle_response(resource_status) + + async def terminate_resource(self, resource_attributes: AttributeDict) -> None: request_command = f"scancel {resource_attributes.remote_resource_uuid}" await self._executor.run_command(request_command) - resource_attributes.update( - resource_status=ResourceStatus.Stopped, updated=datetime.now() - ) - return self.handle_response( - {"JobId": resource_attributes.remote_resource_uuid}, **resource_attributes - ) def sbatch_cmdline_options(self, drone_uuid, machine_meta_data_translation_mapping): sbatch_options = self.machine_type_configuration.get( @@ -205,9 +193,9 @@ def sbatch_cmdline_options(self, drone_uuid, machine_meta_data_translation_mappi ), ) - async def stop_resource(self, resource_attributes: AttributeDict): + async def stop_resource(self, resource_attributes: AttributeDict) -> None: logger.debug("Slurm jobs cannot be stopped gracefully. Terminating instead.") - return await self.terminate_resource(resource_attributes) + await self.terminate_resource(resource_attributes) @contextmanager def handle_exceptions(self): diff --git a/tests/adapters_t/sites_t/test_slurm.py b/tests/adapters_t/sites_t/test_slurm.py index a372abb0..fa1abd6b 100644 --- a/tests/adapters_t/sites_t/test_slurm.py +++ b/tests/adapters_t/sites_t/test_slurm.py @@ -5,8 +5,10 @@ from tardis.exceptions.executorexceptions import CommandExecutionFailure from tardis.interfaces.siteadapter import ResourceStatus from tardis.utilities.attributedict import AttributeDict -from ...utilities.utilities import mock_executor_run_command -from ...utilities.utilities import run_async +from ...utilities.utilities import ( + mock_executor_run_command, + run_async, +) from unittest import TestCase from unittest.mock import MagicMock, patch @@ -153,11 +155,6 @@ def test_start_up_command_deprecation_warning(self): @mock_executor_run_command(TEST_DEPLOY_RESOURCE_RESPONSE) def test_deploy_resource(self): - expected_resource_attributes = self.resource_attributes - expected_resource_attributes.update( - created=datetime.now(), updated=datetime.now() - ) - resource_attributes = AttributeDict( machine_type="test2large", site_name="TestSite", @@ -169,19 +166,11 @@ def test_deploy_resource(self): drone_uuid="testsite-1390065", ) - returned_resource_attributes = run_async( - self.slurm_adapter.deploy_resource, resource_attributes - ) - - self.assertLess( - returned_resource_attributes.created - expected_resource_attributes.created, - timedelta(seconds=1), - ) - - self.check_attribute_dicts( - expected_resource_attributes, - returned_resource_attributes, - exclude=("created", "updated"), + self.assertDictEqual( + AttributeDict( + remote_resource_uuid=1390065, resource_status=ResourceStatus.Booting + ), + run_async(self.slurm_adapter.deploy_resource, resource_attributes), ) self.mock_executor.return_value.run_command.assert_called_with( @@ -247,26 +236,14 @@ def test_site_name(self): @mock_executor_run_command(TEST_RESOURCE_STATUS_RESPONSE) def test_resource_status(self): - expected_resource_attributes = self.resource_attributes - expected_resource_attributes.update(updated=datetime.now()) - - returned_resource_attributes = run_async( - self.slurm_adapter.resource_status, - resource_attributes=self.resource_attributes, - ) - - self.assertLess( - ( - returned_resource_attributes.updated - - expected_resource_attributes.updated + self.assertDictEqual( + AttributeDict( + resource_status=ResourceStatus.Booting, remote_resource_uuid=1390065 + ), + run_async( + self.slurm_adapter.resource_status, + resource_attributes=self.resource_attributes, ), - timedelta(seconds=1), - ) - - self.check_attribute_dicts( - expected_resource_attributes, - returned_resource_attributes, - exclude=("created", "updated"), ) self.mock_executor.return_value.run_command.assert_called_with( @@ -279,18 +256,14 @@ def test_update_resource_status(self): self.resource_attributes["resource_status"], ResourceStatus.Booting ) - return_resource_attributes = run_async( - self.slurm_adapter.resource_status, - resource_attributes=self.resource_attributes, - ) - - self.assertEqual( - return_resource_attributes["resource_status"], ResourceStatus.Running - ) - - self.assertEqual( - return_resource_attributes["drone_uuid"], - self.resource_attributes["drone_uuid"], + self.assertDictEqual( + AttributeDict( + resource_status=ResourceStatus.Running, remote_resource_uuid=1390065 + ), + run_async( + self.slurm_adapter.resource_status, + resource_attributes=self.resource_attributes, + ), ) self.mock_executor.return_value.run_command.assert_called_with(