Skip to content
Snippets Groups Projects
Commit 2e3e9c88 authored by ml bonhomme's avatar ml bonhomme :bee: Committed by Bastien Abadie
Browse files

Add 'mode' field on ponos agent model, and make hardware requirements nullable

parent ecf2c46b
No related branches found
No related tags found
1 merge request!2305Add 'mode' field on ponos agent model, and make hardware requirements nullable
# Generated by Django 4.1.7 on 2024-04-29 08:59
import django.core.validators
import enumfields.fields
from django.db import migrations, models
import arkindex.ponos.models
class Migration(migrations.Migration):
dependencies = [
("ponos", "0007_remove_task_has_docker_socket"),
]
operations = [
migrations.AddField(
model_name="agent",
name="mode",
field=enumfields.fields.EnumField(default="docker", enum=arkindex.ponos.models.AgentMode, max_length=20),
),
migrations.AlterField(
model_name="agent",
name="cpu_cores",
field=models.PositiveSmallIntegerField(blank=True, null=True, validators=[django.core.validators.MinValueValidator(1)]),
),
migrations.AlterField(
model_name="agent",
name="cpu_frequency",
field=models.BigIntegerField(blank=True, null=True, validators=[django.core.validators.MinValueValidator(1)]),
),
migrations.AlterField(
model_name="agent",
name="ram_total",
field=models.BigIntegerField(blank=True, null=True, validators=[django.core.validators.MinValueValidator(1)]),
),
migrations.AddConstraint(
model_name="agent",
constraint=models.CheckConstraint(check=models.Q(("mode", arkindex.ponos.models.AgentMode["Slurm"]), models.Q(("cpu_cores__isnull", False), ("cpu_frequency__isnull", False), ("ram_total__isnull", False)), _connector="OR"), name="slurm_or_hardware_requirements"),
),
]
......@@ -66,6 +66,11 @@ class Farm(models.Model):
return True
class AgentMode(Enum):
Docker = "docker"
Slurm = "slurm"
class Agent(models.Model):
"""
A remote host that can run tasks.
......@@ -76,18 +81,27 @@ class Agent(models.Model):
updated = models.DateTimeField(auto_now=True)
farm = models.ForeignKey(Farm, on_delete=models.PROTECT)
public_key = models.TextField()
mode = EnumField(AgentMode, default=AgentMode.Docker, max_length=20)
hostname = models.SlugField(max_length=64, db_index=False)
cpu_cores = models.PositiveSmallIntegerField(validators=[MinValueValidator(1)])
cpu_frequency = models.BigIntegerField(validators=[MinValueValidator(1)])
cpu_cores = models.PositiveSmallIntegerField(null=True, blank=True, validators=[MinValueValidator(1)])
cpu_frequency = models.BigIntegerField(null=True, blank=True, validators=[MinValueValidator(1)])
# Total amount of RAM on this agent in bytes
ram_total = models.BigIntegerField(validators=[MinValueValidator(1)])
ram_total = models.BigIntegerField(null=True, blank=True, validators=[MinValueValidator(1)])
# Last minute average CPU load measure on this agent
cpu_load = models.FloatField(null=True, blank=True)
# Last RAM load measure expressed as a percentage (0 ≤ ram_load ≤ 1)
ram_load = models.FloatField(null=True, blank=True)
last_ping = models.DateTimeField(editable=False)
class Meta:
constraints = [
models.CheckConstraint(
check=Q(mode=AgentMode.Slurm) | Q(cpu_cores__isnull=False, cpu_frequency__isnull=False, ram_total__isnull=False),
name="slurm_or_hardware_requirements",
),
]
def __str__(self) -> str:
return self.hostname
......
from unittest.mock import patch
from django.core.exceptions import ValidationError
from django.db import transaction
from django.db.models import prefetch_related_objects
from django.db.utils import IntegrityError
from django.utils import timezone
from arkindex.ponos.models import FINAL_STATES, State
from arkindex.ponos.models import FINAL_STATES, Agent, AgentMode, Farm, State
from arkindex.process.models import ProcessMode
from arkindex.project.tests import FixtureAPITestCase
......@@ -12,6 +15,7 @@ class TestModels(FixtureAPITestCase):
@classmethod
def setUpTestData(cls):
super().setUpTestData()
cls.farm = Farm.objects.create(name="Invisible corn farm")
cls.process = cls.corpus.processes.create(
creator=cls.user,
mode=ProcessMode.Workers,
......@@ -132,3 +136,158 @@ class TestModels(FixtureAPITestCase):
self.assertEqual(self.process.get_state(0), State.Running)
self.assertEqual(self.process.get_state(1), State.Unscheduled)
self.assertEqual(self.process.get_state(-1), State.Unscheduled)
def test_agent_default_mode_docker(self):
Agent.objects.create(
hostname="agent_smith",
cpu_cores=2,
cpu_frequency=4.2e9,
public_key="",
farm=self.farm,
ram_total=2e9,
last_ping=timezone.now(),
ram_load=0.49,
cpu_load=0.99
)
test_agent=Agent.objects.get(hostname="agent_smith")
self.assertEqual(test_agent.mode, AgentMode.Docker)
def test_agent_slurm_mode(self):
Agent.objects.create(
hostname="agent_smith",
cpu_cores=2,
cpu_frequency=4.2e9,
public_key="",
farm=self.farm,
ram_total=2e9,
last_ping=timezone.now(),
ram_load=0.49,
cpu_load=0.99,
mode=AgentMode.Slurm.value
)
test_agent=Agent.objects.get(hostname="agent_smith")
self.assertEqual(test_agent.mode, AgentMode.Slurm)
def test_agent_invalid_mode(self):
with self.assertRaisesRegex(ValidationError, "bad_mode is not a valid value for enum"):
Agent.objects.create(
hostname="agent_smith",
cpu_cores=2,
cpu_frequency=4.2e9,
public_key="",
farm=self.farm,
ram_total=2e9,
last_ping=timezone.now(),
ram_load=0.49,
cpu_load=0.99,
mode="bad_mode"
)
def test_agent_mode_hardware_constraint(self):
"""
When the agent's mode is not AgentMode.Slurm, the hardware requirement parameters
cannot be null. When the agent's mode is AgentMode.Slurm, they can be null or not.
"""
cases = [
{
# no cpu_cores
"mode": AgentMode.Docker,
"params": {
"hostname": "agent_smith",
"cpu_frequency": 4.2e9,
"public_key": "",
"farm": self.farm,
"ram_total": 2e9,
"last_ping": timezone.now(),
"ram_load": 0.49,
"cpu_load": 0.99,
},
"failure": True
},
{
# no cpu_frequency
"mode": AgentMode.Docker,
"params": {
"hostname": "agent_smith",
"cpu_cores": 2,
"public_key": "",
"farm": self.farm,
"ram_total": 2e9,
"last_ping": timezone.now(),
"ram_load": 0.49,
"cpu_load": 0.99,
},
"failure": True
},
{
# no ram_total
"mode": AgentMode.Docker,
"params": {
"hostname": "agent_smith",
"cpu_cores": 2,
"cpu_frequency": 4.2e9,
"public_key": "",
"farm": self.farm,
"last_ping": timezone.now(),
"ram_load": 0.49,
"cpu_load": 0.99,
},
"failure": True
},
{
# hardware parameters set to None
"mode": AgentMode.Docker,
"params": {
"hostname": "agent_smith",
"cpu_cores": None,
"cpu_frequency": None,
"public_key": "",
"farm": self.farm,
"ram_total": None,
"last_ping": timezone.now(),
"ram_load": 0.49,
"cpu_load": 0.99,
},
"failure": True
},
{
# slurm mode with hardware requirements
"mode": AgentMode.Slurm,
"params": {
"hostname": "agent_smith",
"cpu_cores": 2,
"cpu_frequency": 4.2e9,
"public_key": "",
"farm": self.farm,
"ram_total": 2e9,
"last_ping": timezone.now(),
"ram_load": 0.49,
"cpu_load": 0.99,
},
"failure": False
},
{
# slurm mode without hardware requirements
"mode": AgentMode.Slurm,
"params": {
"hostname": "agent_smith",
"public_key": "",
"farm": self.farm,
"last_ping": timezone.now(),
"ram_load": 0.49,
"cpu_load": 0.99,
},
"failure": False
},
]
for case in cases:
with self.subTest(case=case):
if case["failure"]:
# Each create must be run in an atomic block to avoid TransactionManagementError
with transaction.atomic():
with self.assertRaisesRegex(IntegrityError, "slurm_or_hardware_requirements"):
Agent.objects.create(**case["params"], mode=case["mode"])
else:
Agent.objects.create(**case["params"], mode=case["mode"])
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment