def setUp(self): self.data_dir = tempfile.mkdtemp() database_path = os.path.join(self.data_dir, "mephisto.db") self.db = LocalMephistoDB(database_path) self.task_id = self.db.new_task("test_mock", MockBlueprint.BLUEPRINT_TYPE) self.task_run_id = get_test_task_run(self.db) self.task_run = TaskRun(self.db, self.task_run_id) architect_config = OmegaConf.structured( MephistoConfig( architect=MockArchitectArgs(should_run_server=True), )) self.architect = MockArchitect(self.db, architect_config, EMPTY_STATE, self.task_run, self.data_dir) self.architect.prepare() self.architect.deploy() self.urls = self.architect._get_socket_urls() # FIXME self.url = self.urls[0] self.provider = MockProvider(self.db) self.provider.setup_resources_for_task_run(self.task_run, self.task_run.args, self.url) self.launcher = TaskLauncher(self.db, self.task_run, self.get_mock_assignment_data_array()) self.launcher.create_assignments() self.launcher.launch_units(self.url) self.sup = None
def test_run_job_concurrent(self): """Ensure that the supervisor object can even be created""" self.operator = Operator(self.db) config = MephistoConfig( blueprint=MockBlueprintArgs( num_assignments=1, is_concurrent=True, ), provider=MockProviderArgs(requester_name=self.requester_name, ), architect=MockArchitectArgs(should_run_server=True), task=MOCK_TASK_ARGS, ) 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) # 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 = 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, mock_agent_details) agents = self.db.find_agents() self.assertEqual(len(agents), 1, "Agent was not created properly") agent = agents[0] self.assertIsNotNone(agent) # 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 = workers[0].db_id # Register an agent mock_agent_details = "FAKE_ASSIGNMENT_2" architect.server.register_mock_agent(worker_id, mock_agent_details) # 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 the assignment is completed task_run = tracked_run.task_run assignment = task_run.get_assignments()[0] self.assertEqual(assignment.get_status(), AssignmentState.COMPLETED)
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.providers.mock.mock_provider import MockProviderArgs from mephisto.server.blueprints.mock.mock_blueprint import MockBlueprintArgs from mephisto.server.architects.mock_architect import MockArchitectArgs from mephisto.data_model.task_config import TaskConfigArgs from mephisto.core.hydra_config import MephistoConfig MOCK_TASK_ARGS = TaskConfigArgs( task_title="title", task_description="This is a description", task_reward="0.3", task_tags="1,2,3", ) MOCK_PROVIDER_ARGS = MockProviderArgs() MOCK_ARCHITECT_ARGS = MockArchitectArgs() MOCK_BLUEPRINT_ARGS = MockBlueprintArgs() MOCK_CONFIG = MephistoConfig( provider=MOCK_PROVIDER_ARGS, blueprint=MOCK_BLUEPRINT_ARGS, architect=MOCK_ARCHITECT_ARGS, task=MOCK_TASK_ARGS, ) def get_test_project(db: MephistoDB) -> Tuple[str, str]: """Helper to create a project for tests""" project_name = "test_project" project_id = db.new_project(project_name) return project_name, project_id