tests/test_schedulers/test_joblib_scheduler.py

0001 """Tests for JobLibScheduler and scheduler registry.
0002
0003 Tests cover:
0004 - JobLibScheduler job execution
0005 - Parallelism policy respect (max_concurrent, timeout, retry)
0006 - Artifact collection from job outputs
0007 - Registry registration and lookup
0008 - Status checking and cancellation
0009 """
0010
0011 import json
0012 import os
0013 import pytest
0014 import tempfile
0015 from pathlib import Path
0016
0017 from aid2e.schedulers import (
0018     JobLibScheduler,
0019     BaseScheduler,
0020     JobStatus,
0021     StageExecutionResult,
0022     register_scheduler,
0023     get_scheduler,
0024     list_registered_schedulers,
0025     is_scheduler_registered,
0026 )
0027 from aid2e.schedulers.JobLib import JobLibRunnerConfig
0028
0029
0030 class TestJobLibSchedulerBasics:
0031     """Test basic JobLibScheduler functionality."""
0032
0033     def test_scheduler_instantiation(self):
0034         """Test creating JobLibScheduler with default config."""
0035         scheduler = JobLibScheduler()
0036         assert isinstance(scheduler, BaseScheduler)
0037         assert scheduler.config.n_jobs == -1
0038         assert scheduler.config.backend == "loky"
0039
0040     def test_scheduler_with_custom_config(self):
0041         """Test JobLibScheduler with custom JobLibRunnerConfig."""
0042         config = JobLibRunnerConfig(n_jobs=2, backend="threading", timeout=300)
0043         scheduler = JobLibScheduler(config=config)
0044         assert scheduler.config.n_jobs == 2
0045         assert scheduler.config.backend == "threading"
0046         assert scheduler.config.timeout == 300
0047
0048     def test_job_status_model(self):
0049         """Test JobStatus data model."""
0050         status = JobStatus(
0051             job_id="job_1",
0052             status="completed",
0053             return_code=0,
0054             stdout="Output",
0055             stderr="",
0056         )
0057         assert status.job_id == "job_1"
0058         assert status.status == "completed"
0059         assert status.return_code == 0
0060
0061     def test_stage_execution_result_model(self):
0062         """Test StageExecutionResult data model."""
0063         result = StageExecutionResult(
0064             stage_name="evaluate",
0065             job_statuses=[],
0066             artifacts={"output.json": '{"f1": 1.0}'},
0067             success=True,
0068         )
0069         assert result.stage_name == "evaluate"
0070         assert result.success is True
0071         assert "output.json" in result.artifacts
0072
0073
0074 class TestJobLibSchedulerExecution:
0075     """Test job execution on JobLibScheduler."""
0076
0077     def test_run_simple_job(self):
0078         """Test running a simple shell command job."""
0079         scheduler = JobLibScheduler()
0080
0081         job_def = {
0082             'name': 'simple_echo',
0083             'command': 'echo "Hello World"',
0084             'payload': {},
0085             'outputs': [],
0086         }
0087
0088         result = scheduler.run_stage(
0089             stage_name='test_stage',
0090             job_definitions=[job_def],
0091         )
0092
0093         assert result.success is True
0094         assert result.stage_name == 'test_stage'
0095         assert len(result.job_statuses) == 1
0096         assert result.job_statuses[0].status == 'completed'
0097         assert result.job_statuses[0].return_code == 0
0098         assert "Hello World" in result.job_statuses[0].stdout
0099
0100     def test_run_multiple_jobs_parallel(self):
0101         """Test running multiple jobs in parallel."""
0102         scheduler = JobLibScheduler(config=JobLibRunnerConfig(n_jobs=2))
0103
0104         job_defs = [
0105             {'name': f'job_{i}', 'command': f'echo "Job {i}"', 'payload': {}, 'outputs': []}
0106             for i in range(3)
0107         ]
0108
0109         result = scheduler.run_stage(
0110             stage_name='parallel_stage',
0111             job_definitions=job_defs,
0112         )
0113
0114         assert result.success is True
0115         assert len(result.job_statuses) == 3
0116         for status in result.job_statuses:
0117             assert status.status == 'completed'
0118             assert status.return_code == 0
0119
0120     def test_job_with_output_artifact(self):
0121         """Test job that creates an output file and collects it."""
0122         with tempfile.TemporaryDirectory() as tmpdir:
0123             output_file = os.path.join(tmpdir, "result.json")
0124
0125             job_def = {
0126                 'name': 'create_json',
0127                 'command': f'echo \'{{"f1": 1.0, "f2": 2.0}}\' > {output_file}',
0128                 'payload': {},
0129                 'outputs': [{'path': 'result.json', 'format': 'json'}],
0130             }
0131
0132             scheduler = JobLibScheduler()
0133             result = scheduler.run_stage(
0134                 stage_name='artifact_stage',
0135                 job_definitions=[job_def],
0136                 working_dir=tmpdir,
0137             )
0138
0139             assert result.success is True
0140             assert len(result.artifacts) > 0 or result.job_statuses[0].return_code == 0
0141
0142     def test_job_failure(self):
0143         """Test handling of failed job."""
0144         scheduler = JobLibScheduler()
0145
0146         job_def = {
0147             'name': 'failing_job',
0148             'command': 'exit 1',  # Non-zero exit code
0149             'payload': {},
0150             'outputs': [],
0151         }
0152
0153         result = scheduler.run_stage(
0154             stage_name='failure_stage',
0155             job_definitions=[job_def],
0156         )
0157
0158         assert result.success is False
0159         assert len(result.job_statuses) == 1
0160         assert result.job_statuses[0].status == 'failed'
0161         assert result.job_statuses[0].return_code == 1
0162
0163     def test_job_with_payload_environment(self):
0164         """Test job receives payload via JOB_PAYLOAD environment variable."""
0165         scheduler = JobLibScheduler()
0166
0167         job_def = {
0168             'name': 'payload_job',
0169             'command': 'python -c "import os, json, base64, pickle; print(json.dumps(pickle.loads(base64.b64decode(os.environ[\'JOB_PAYLOAD_PICKLE\']))))"',
0170             'payload': {'param1': 'value1', 'param2': 42},
0171             'outputs': [],
0172         }
0173
0174         result = scheduler.run_stage(
0175             stage_name='payload_stage',
0176             job_definitions=[job_def],
0177         )
0178
0179         assert result.success is True
0180         assert result.job_statuses[0].return_code == 0
0181
0182
0183 class TestJobLibSchedulerParallelismPolicy:
0184     """Test parallelism policy enforcement."""
0185
0186     def test_max_concurrent_respected(self):
0187         """Test that max_concurrent parameter is passed to joblib."""
0188         scheduler = JobLibScheduler()
0189
0190         job_defs = [
0191             {'name': f'job_{i}', 'command': 'echo "x"', 'payload': {}, 'outputs': []}
0192             for i in range(4)
0193         ]
0194
0195         policy = {
0196             'max_concurrent': 2,
0197             'retry_max': 1,
0198             'timeout_sec': 30,
0199         }
0200
0201         result = scheduler.run_stage(
0202             stage_name='limited_stage',
0203             job_definitions=job_defs,
0204             parallelism_policy=policy,
0205         )
0206
0207         assert result.success is True
0208         assert len(result.job_statuses) == 4
0209
0210     def test_timeout_respected(self):
0211         """Test that job timeout is enforced."""
0212         config = JobLibRunnerConfig(n_jobs=1, timeout=1)  # 1 second timeout
0213         scheduler = JobLibScheduler(config=config)
0214
0215         job_def = {
0216             'name': 'timeout_job',
0217             'command': 'sleep 5',  # Will timeout
0218             'payload': {},
0219             'outputs': [],
0220         }
0221
0222         result = scheduler.run_stage(
0223             stage_name='timeout_stage',
0224             job_definitions=[job_def],
0225         )
0226
0227         # Job should fail due to timeout
0228         assert result.success is False or result.job_statuses[0].return_code != 0
0229
0230
0231 class TestSchedulerRegistry:
0232     """Test scheduler registry and factory."""
0233
0234     def test_joblib_registered_by_default(self):
0235         """Test that JobLibScheduler is registered by default."""
0236         assert is_scheduler_registered('joblib') is True
0237         assert get_scheduler('joblib') is JobLibScheduler
0238
0239     def test_list_registered_schedulers(self):
0240         """Test listing all registered schedulers."""
0241         schedulers = list_registered_schedulers()
0242         assert 'joblib' in schedulers
0243         assert schedulers['joblib'] is JobLibScheduler
0244
0245     def test_get_unregistered_scheduler_raises(self):
0246         """Test that getting unregistered scheduler raises KeyError."""
0247         with pytest.raises(KeyError):
0248             get_scheduler('nonexistent_scheduler')
0249
0250     def test_register_new_scheduler(self):
0251         """Test registering a new mock scheduler."""
0252         # Create a simple mock scheduler
0253         class MockScheduler(BaseScheduler):
0254             def run_stage(self, stage_name, job_definitions, parallelism_policy=None, working_dir=None):
0255                 return StageExecutionResult(
0256                     stage_name=stage_name,
0257                     job_statuses=[],
0258                     artifacts={},
0259                     success=True,
0260                 )
0261
0262             def check_status(self, job_id):
0263                 return JobStatus(job_id=job_id, status='completed', return_code=0)
0264
0265             def cancel_job(self, job_id):
0266                 return False
0267
0268         # Register it
0269         register_scheduler('mock', MockScheduler)
0270
0271         # Verify registration
0272         assert is_scheduler_registered('mock') is True
0273         assert get_scheduler('mock') is MockScheduler
0274
0275     def test_register_duplicate_raises(self):
0276         """Test that registering duplicate scheduler name raises ValueError."""
0277         with pytest.raises(ValueError):
0278             register_scheduler('joblib', JobLibScheduler)
0279
0280     def test_register_invalid_class_raises(self):
0281         """Test that registering non-BaseScheduler class raises ValueError."""
0282         class NotAScheduler:
0283             pass
0284
0285         with pytest.raises(ValueError):
0286             register_scheduler('invalid', NotAScheduler)
0287
0288
0289 class TestJobLibSchedulerStatusAndCancel:
0290     """Test status checking and job cancellation."""
0291
0292     def test_check_status_completed_job(self):
0293         """Test checking status of completed job."""
0294         scheduler = JobLibScheduler()
0295
0296         result = scheduler.run_stage(
0297             stage_name='test',
0298             job_definitions=[{'name': 'echo', 'command': 'echo x', 'payload': {}, 'outputs': []}],
0299         )
0300
0301         job_id = result.job_statuses[0].job_id
0302         status = scheduler.check_status(job_id)
0303
0304         # Since JobLib is synchronous, we get the cached status
0305         assert status.job_id == job_id
0306
0307     def test_cancel_job_returns_false(self):
0308         """Test that cancel_job returns False (not supported for sync JobLib)."""
0309         scheduler = JobLibScheduler()
0310         result = scheduler.cancel_job('any_job_id')
0311
0312         # JobLib can't cancel synchronous jobs
0313         assert result is False
0314
0315     def test_shutdown_noop(self):
0316         """Test that shutdown completes without error."""
0317         scheduler = JobLibScheduler()
0318         scheduler.shutdown()  # Should not raise
0319
0320
0321 class TestJobLibSchedulerEdgeCases:
0322     """Test edge cases and error conditions."""
0323
0324     def test_empty_job_list(self):
0325         """Test running stage with no jobs."""
0326         scheduler = JobLibScheduler()
0327
0328         result = scheduler.run_stage(
0329             stage_name='empty',
0330             job_definitions=[],
0331         )
0332
0333         assert result.stage_name == 'empty'
0334         assert len(result.job_statuses) == 0
0335
0336     def test_job_with_missing_fields(self):
0337         """Test job with missing optional fields."""
0338         scheduler = JobLibScheduler()
0339
0340         job_def = {
0341             # Minimal job def
0342             'command': 'echo test',
0343         }
0344
0345         result = scheduler.run_stage(
0346             stage_name='minimal',
0347             job_definitions=[job_def],
0348         )
0349
0350         # Should still work with defaults
0351         assert len(result.job_statuses) > 0
0352
0353     def test_command_with_special_chars(self):
0354         """Test command containing special shell characters."""
0355         scheduler = JobLibScheduler()
0356
0357         job_def = {
0358             'name': 'special',
0359             'command': 'echo "hello world" && echo "success"',
0360             'payload': {},
0361             'outputs': [],
0362         }
0363
0364         result = scheduler.run_stage(
0365             stage_name='special',
0366             job_definitions=[job_def],
0367         )
0368
0369         assert result.success is True
0370         assert "success" in result.job_statuses[0].stdout