class MephistoConfig:
blueprint: BlueprintArgs = BlueprintArgs()
provider: ProviderArgs = ProviderArgs()
architect: ArchitectArgs = ArchitectArgs()
task: TaskConfigArgs = TaskConfigArgs()
database: DatabaseArgs = DatabaseArgs()
log_level: str = "info"
class MephistoConfig:
blueprint: BlueprintArgs = BlueprintArgs()
provider: ProviderArgs = ProviderArgs()
architect: ArchitectArgs = ArchitectArgs()
task: TaskConfigArgs = TaskConfigArgs()
def test_run_jobs_with_restrictions(self):
"""Ensure allowed_concurrent and maximum_units_per_worker work"""
self.operator = Operator(self.db)
provider_args = MockProviderArgs(requester_name=self.requester_name)
architect_args = MockArchitectArgs(should_run_server=True)
config = MephistoConfig(
blueprint=MockBlueprintArgs(num_assignments=3, is_concurrent=True),
provider=provider_args,
architect=architect_args,
task=TaskConfigArgs(
task_title="title",
task_description="This is a description",
task_reward="0.3",
task_tags="1,2,3",
maximum_units_per_worker=2,
allowed_concurrent=1,
task_name="max-unit-test",
),
)
self.operator.validate_and_run_config(OmegaConf.structured(config))
tracked_runs = self.operator.get_running_task_runs()
self.assertEqual(len(tracked_runs), 1, "Run not launched")
task_run_id, tracked_run = list(tracked_runs.items())[0]
self.assertIsNotNone(tracked_run)
self.assertIsNotNone(tracked_run.task_launcher)
self.assertIsNotNone(tracked_run.task_runner)
self.assertIsNotNone(tracked_run.architect)
self.assertIsNotNone(tracked_run.task_run)
self.assertEqual(tracked_run.task_run.db_id, task_run_id)
self.await_server_start(tracked_run.architect)
# Create two agents to step through the task
architect = tracked_run.architect
self.assertIsInstance(architect, MockArchitect,
"Must use mock in testing")
# Register a worker
mock_worker_name = "MOCK_WORKER"
architect.server.register_mock_worker(mock_worker_name)
workers = self.db.find_workers(worker_name=mock_worker_name)
worker_id_1 = workers[0].db_id
self.assertEqual(len(tracked_run.task_runner.running_assignments), 0)
# Register an agent
mock_agent_details = "FAKE_ASSIGNMENT"
architect.server.register_mock_agent(worker_id_1, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 1, "Agent was not created properly")
agent = agents[0]
self.assertIsNotNone(agent)
# Try to register a second agent, which should fail due to concurrency
mock_agent_details = "FAKE_ASSIGNMENT_2"
architect.server.register_mock_agent(worker_id_1, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 1, "Second agent was created")
# Register another worker
mock_worker_name = "MOCK_WORKER_2"
architect.server.register_mock_worker(mock_worker_name)
workers = self.db.find_workers(worker_name=mock_worker_name)
worker_id_2 = workers[0].db_id
# Register an agent
mock_agent_details = "FAKE_ASSIGNMENT_2"
architect.server.register_mock_agent(worker_id_2, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 2, "Second agent was not created")
# wait for task to pass
self.wait_for_complete_assignment(
agents[1].get_unit().get_assignment(), 3)
# Pass a second task as well
mock_agent_details = "FAKE_ASSIGNMENT_3"
architect.server.register_mock_agent(worker_id_1, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 3, "Agent was not created properly")
mock_agent_details = "FAKE_ASSIGNMENT_4"
architect.server.register_mock_agent(worker_id_2, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 4, "Fourth agent was not created")
# wait for task to pass
self.wait_for_complete_assignment(
agents[3].get_unit().get_assignment(), 3)
# Both workers should have saturated their tasks, and not be granted agents
mock_agent_details = "FAKE_ASSIGNMENT_5"
architect.server.register_mock_agent(worker_id_1, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 4, "Additional agent was created")
architect.server.register_mock_agent(worker_id_2, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 4, "Additional agent was created")
# new workers should be able to work on these just fine though
mock_worker_name = "MOCK_WORKER_3"
architect.server.register_mock_worker(mock_worker_name)
workers = self.db.find_workers(worker_name=mock_worker_name)
worker_id_3 = workers[0].db_id
mock_worker_name = "MOCK_WORKER_4"
architect.server.register_mock_worker(mock_worker_name)
workers = self.db.find_workers(worker_name=mock_worker_name)
worker_id_4 = workers[0].db_id
# Register agents from new workers
mock_agent_details = "FAKE_ASSIGNMENT_5"
architect.server.register_mock_agent(worker_id_3, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 5, "Additional agent was not created")
mock_agent_details = "FAKE_ASSIGNMENT_6"
architect.server.register_mock_agent(worker_id_4, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 6, "Additional agent was not created")
# wait for task to pass
self.wait_for_complete_assignment(
agents[5].get_unit().get_assignment(), 3)
# Give up to 5 seconds for whole mock task to complete
start_time = time.time()
while time.time() - start_time < TIMEOUT_TIME:
if len(self.operator.get_running_task_runs()) == 0:
break
time.sleep(0.1)
self.assertLess(time.time() - start_time, TIMEOUT_TIME,
"Task not completed in time")
# Ensure all assignments are completed
task_run = tracked_run.task_run
assignments = task_run.get_assignments()
for assignment in assignments:
self.assertEqual(assignment.get_status(),
AssignmentState.COMPLETED)
# Create a new task
config = MephistoConfig(
blueprint=MockBlueprintArgs(num_assignments=1, is_concurrent=True),
provider=MockProviderArgs(requester_name=self.requester_name),
architect=MockArchitectArgs(should_run_server=True),
task=TaskConfigArgs(
task_title="title",
task_description="This is a description",
task_reward="0.3",
task_tags="1,2,3",
maximum_units_per_worker=2,
allowed_concurrent=1,
task_name="max-unit-test",
),
)
self.operator.validate_and_run_config(OmegaConf.structured(config))
tracked_runs = self.operator.get_running_task_runs()
self.assertEqual(len(tracked_runs), 1, "Run not launched")
task_run_id, tracked_run = list(tracked_runs.items())[0]
self.await_server_start(tracked_run.architect)
architect = tracked_run.architect
# Workers one and two still shouldn't be able to make agents
mock_agent_details = "FAKE_ASSIGNMENT_7"
architect.server.register_mock_agent(worker_id_1, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(
len(agents),
6,
"Additional agent was created for worker exceeding max units",
)
mock_agent_details = "FAKE_ASSIGNMENT_7"
architect.server.register_mock_agent(worker_id_2, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(
len(agents),
6,
"Additional agent was created for worker exceeding max units",
)
# Three and four should though
mock_agent_details = "FAKE_ASSIGNMENT_7"
architect.server.register_mock_agent(worker_id_3, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 7, "Additional agent was not created")
mock_agent_details = "FAKE_ASSIGNMENT_8"
architect.server.register_mock_agent(worker_id_4, mock_agent_details)
agents = self.db.find_agents()
self.assertEqual(len(agents), 8, "Additional agent was not created")
# Ensure the task run completed and that all assignments are done
start_time = time.time()
while time.time() - start_time < TIMEOUT_TIME:
if len(self.operator.get_running_task_runs()) == 0:
break
time.sleep(0.1)
self.assertLess(time.time() - start_time, TIMEOUT_TIME,
"Task not completed in time")
task_run = tracked_run.task_run
assignments = task_run.get_assignments()
for assignment in assignments:
self.assertEqual(assignment.get_status(),
AssignmentState.COMPLETED)
from mephisto.operations.operator import Operator
from mephisto.abstractions.architects.mock_architect import (
MockArchitect,
MockArchitectArgs,
)
from mephisto.operations.hydra_config import MephistoConfig
from mephisto.abstractions.providers.mock.mock_provider import MockProviderArgs
from mephisto.abstractions.blueprints.mock.mock_blueprint import MockBlueprintArgs
from mephisto.data_model.task_config import TaskConfigArgs
from omegaconf import OmegaConf
TIMEOUT_TIME = 10
MOCK_TASK_ARGS = TaskConfigArgs(
task_title="title",
task_description="This is a description",
task_reward="0.3",
task_tags="1,2,3",
)
class OperatorBaseTest(object):
"""
Unit testing for the Mephisto Operator
"""
DB_CLASS = None
def setUp(self):
self.data_dir = tempfile.mkdtemp()
database_path = os.path.join(self.data_dir, "mephisto.db")
assert self.DB_CLASS is not None, "Did not specify db to use"