|
26 | 26 | from dstack._internal.core.models.duration import Duration |
27 | 27 | from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus |
28 | 28 | from dstack._internal.core.models.instances import InstanceStatus |
29 | | -from dstack._internal.core.models.profiles import StartupOrder, UtilizationPolicy |
| 29 | +from dstack._internal.core.models.profiles import Profile, StartupOrder, UtilizationPolicy |
30 | 30 | from dstack._internal.core.models.runs import ( |
31 | 31 | ClusterInfo, |
32 | 32 | ImagePullProgress, |
|
44 | 44 | from dstack._internal.server import settings as server_settings |
45 | 45 | from dstack._internal.server.background.pipeline_tasks.jobs_running import ( |
46 | 46 | JOB_DISCONNECTED_RETRY_TIMEOUT, |
| 47 | + MAX_DURATION_ENFORCEMENT_GRACE, |
47 | 48 | ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS, |
48 | 49 | JobRunningFetcher, |
49 | 50 | JobRunningPipeline, |
@@ -1752,6 +1753,128 @@ async def test_inactivity_duration( |
1752 | 1753 | assert job.termination_reason == expected_termination_reason |
1753 | 1754 | assert job.inactivity_secs == expected_inactivity_secs |
1754 | 1755 |
|
| 1756 | + @pytest.mark.parametrize( |
| 1757 | + ( |
| 1758 | + "max_duration", |
| 1759 | + "running_for", |
| 1760 | + "stamp_running_at", |
| 1761 | + "expected_status", |
| 1762 | + "expected_termination_reason", |
| 1763 | + "expect_pull", |
| 1764 | + ), |
| 1765 | + [ |
| 1766 | + pytest.param( |
| 1767 | + 600, |
| 1768 | + timedelta(seconds=600), |
| 1769 | + True, |
| 1770 | + JobStatus.RUNNING, |
| 1771 | + None, |
| 1772 | + True, |
| 1773 | + id="deadline-reached-but-runner-still-within-grace", |
| 1774 | + ), |
| 1775 | + pytest.param( |
| 1776 | + 600, |
| 1777 | + timedelta(seconds=600) + MAX_DURATION_ENFORCEMENT_GRACE - timedelta(seconds=1), |
| 1778 | + True, |
| 1779 | + JobStatus.RUNNING, |
| 1780 | + None, |
| 1781 | + True, |
| 1782 | + id="grace-not-elapsed", |
| 1783 | + ), |
| 1784 | + pytest.param( |
| 1785 | + 600, |
| 1786 | + timedelta(seconds=600) + MAX_DURATION_ENFORCEMENT_GRACE, |
| 1787 | + True, |
| 1788 | + JobStatus.TERMINATING, |
| 1789 | + JobTerminationReason.MAX_DURATION_EXCEEDED, |
| 1790 | + False, |
| 1791 | + id="grace-elapsed", |
| 1792 | + ), |
| 1793 | + pytest.param( |
| 1794 | + 600, |
| 1795 | + timedelta(days=1), |
| 1796 | + False, |
| 1797 | + JobStatus.RUNNING, |
| 1798 | + None, |
| 1799 | + True, |
| 1800 | + id="job-started-before-upgrade", |
| 1801 | + ), |
| 1802 | + pytest.param( |
| 1803 | + "off", |
| 1804 | + timedelta(days=1), |
| 1805 | + True, |
| 1806 | + JobStatus.RUNNING, |
| 1807 | + None, |
| 1808 | + True, |
| 1809 | + id="max-duration-off", |
| 1810 | + ), |
| 1811 | + ], |
| 1812 | + ) |
| 1813 | + async def test_max_duration_enforced_by_server( |
| 1814 | + self, |
| 1815 | + test_db, |
| 1816 | + session: AsyncSession, |
| 1817 | + worker: JobRunningWorker, |
| 1818 | + max_duration, |
| 1819 | + running_for: timedelta, |
| 1820 | + stamp_running_at: bool, |
| 1821 | + expected_status: JobStatus, |
| 1822 | + expected_termination_reason: Optional[JobTerminationReason], |
| 1823 | + expect_pull: bool, |
| 1824 | + ) -> None: |
| 1825 | + now = datetime(2023, 1, 2, 5, 0, tzinfo=timezone.utc) |
| 1826 | + project = await create_project(session=session) |
| 1827 | + user = await create_user(session=session) |
| 1828 | + repo = await create_repo(session=session, project_id=project.id) |
| 1829 | + run = await create_run( |
| 1830 | + session=session, |
| 1831 | + project=project, |
| 1832 | + repo=repo, |
| 1833 | + user=user, |
| 1834 | + status=RunStatus.RUNNING, |
| 1835 | + run_name="test-run", |
| 1836 | + run_spec=get_run_spec( |
| 1837 | + run_name="test-run", |
| 1838 | + repo_id=repo.name, |
| 1839 | + profile=Profile(name="default", max_duration=max_duration), |
| 1840 | + ), |
| 1841 | + ) |
| 1842 | + instance = await create_instance( |
| 1843 | + session=session, |
| 1844 | + project=project, |
| 1845 | + status=InstanceStatus.BUSY, |
| 1846 | + ) |
| 1847 | + job = await create_job( |
| 1848 | + session=session, |
| 1849 | + run=run, |
| 1850 | + status=JobStatus.RUNNING, |
| 1851 | + running_at=(now - running_for) if stamp_running_at else None, |
| 1852 | + job_provisioning_data=get_job_provisioning_data(), |
| 1853 | + instance=instance, |
| 1854 | + instance_assigned=True, |
| 1855 | + ) |
| 1856 | + with ( |
| 1857 | + patch("dstack._internal.server.services.runner.pool.SSHTunnel"), |
| 1858 | + patch( |
| 1859 | + "dstack._internal.server.services.runner.client.RunnerClient.from_address" |
| 1860 | + ) as runner_client_cls, |
| 1861 | + freeze_time(now), |
| 1862 | + ): |
| 1863 | + runner_client_mock = runner_client_cls.return_value |
| 1864 | + runner_client_mock.pull.return_value = PullResponse( |
| 1865 | + job_states=[], |
| 1866 | + job_logs=[], |
| 1867 | + runner_logs=[], |
| 1868 | + last_updated=0, |
| 1869 | + no_connections_secs=0, |
| 1870 | + ) |
| 1871 | + await _process_job(session, worker, job) |
| 1872 | + assert runner_client_mock.pull.called == expect_pull |
| 1873 | + |
| 1874 | + await session.refresh(job) |
| 1875 | + assert job.status == expected_status |
| 1876 | + assert job.termination_reason == expected_termination_reason |
| 1877 | + |
1755 | 1878 | @pytest.mark.parametrize( |
1756 | 1879 | ["samples", "expected_status"], |
1757 | 1880 | [ |
|
0 commit comments