Add fleets property to run configurations and CLI (#2488)

Closes: #1447

Author: un-def

src/dstack/_internal/cli/services/profile.py

@@ -65,6 +65,13 @@ def register_profile_args(parser: argparse.ArgumentParser):
     )
 
     fleets_group = parser.add_argument_group("Fleets")
+    fleets_group.add_argument(
+        "--fleet",
+        action="append",
+        metavar="NAME",
+        dest="fleets",
+        help="Consider only instances from the specified fleet(s) for reuse",
+    )
     fleets_group_exc = fleets_group.add_mutually_exclusive_group()
     fleets_group_exc.add_argument(
         "-R",
@@ -147,6 +154,8 @@ def apply_profile_args(
     if args.max_duration is not None:
         profile_settings.max_duration = args.max_duration
 
+    if args.fleets:
+        profile_settings.fleets = args.fleets
     if args.idle_duration is not None:
         profile_settings.idle_duration = args.idle_duration
     elif args.dont_destroy:

src/dstack/_internal/core/models/profiles.py

@@ -240,6 +240,9 @@ class ProfileParams(CoreModel):
         Optional[UtilizationPolicy],
         Field(description="Run termination policy based on utilization"),
     ] = None
+    fleets: Annotated[
+        Optional[list[str]], Field(description="The fleets considered for reuse")
+    ] = None
 
     # Deprecated and unused. Left for compatibility with 0.18 clients.
     pool_name: Annotated[Optional[str], Field(exclude=True)] = None

src/dstack/_internal/server/background/tasks/process_submitted_jobs.py

@@ -212,6 +212,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
                     InstanceModel.deleted == False,
                     InstanceModel.total_blocks > InstanceModel.busy_blocks,
                 )
+                .options(joinedload(InstanceModel.fleet))
                 .execution_options(populate_existing=True)
             )
             pool_instances = list(res.unique().scalars().all())

src/dstack/_internal/server/services/instances.py

@@ -181,6 +181,9 @@ def filter_pool_instances(
             continue
         if instance.unreachable:
             continue
+        fleet = instance.fleet
+        if profile.fleets is not None and (fleet is None or fleet.name not in profile.fleets):
+            continue
         if status is not None and instance.status != status:
             continue
         jpd = get_instance_provisioning_data(instance)
@@ -268,10 +271,12 @@ async def get_pool_instances(
     project: ProjectModel,
 ) -> List[InstanceModel]:
     res = await session.execute(
-        select(InstanceModel).where(
+        select(InstanceModel)
+        .where(
             InstanceModel.project_id == project.id,
             InstanceModel.deleted == False,
         )
+        .options(joinedload(InstanceModel.fleet))
     )
     instance_models = list(res.unique().scalars().all())
     return instance_models

src/dstack/api/server/_fleets.py

@@ -64,11 +64,9 @@ def _get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[Dict]:
     spec_excludes: Dict[str, Any] = {}
     configuration_excludes: Dict[str, Any] = {}
     profile_excludes: set[str] = set()
-    # Fields can be excluded like this:
-    # if fleet_spec.configuration.availability_zones is None:
-    #     configuration_excludes["availability_zones"] = True
-    # if fleet_spec.profile is not None and fleet_spec.profile.availability_zones is None:
-    #     profile_excludes.add("availability_zones")
+    profile = fleet_spec.profile
+    if profile.fleets is None:
+        profile_excludes.add("fleet")
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:

src/dstack/api/server/_runs.py

@@ -104,13 +104,12 @@ def _get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
     spec_excludes: dict[str, Any] = {}
     configuration_excludes: dict[str, Any] = {}
     profile_excludes: set[str] = set()
-    # configuration = run_spec.configuration
-    # profile = run_spec.profile
-    # Fields can be excluded like this:
-    # if configuration.availability_zones is None:
-    #     configuration_excludes["availability_zones"] = True
-    # if profile is not None and profile.availability_zones is None:
-    #     profile_excludes.add("availability_zones")
+    configuration = run_spec.configuration
+    profile = run_spec.profile
+    if configuration.fleets is None:
+        configuration_excludes["fleet"] = True
+    if profile is not None and profile.fleets is None:
+        profile_excludes.add("fleet")
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:

src/tests/_internal/server/background/tasks/test_process_submitted_jobs.py

@@ -536,6 +536,35 @@ async def test_assigns_job_to_shared_instance(self, test_db, session: AsyncSessi
         assert instance.total_blocks == 4
         assert instance.busy_blocks == 2
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_assigns_job_to_specific_fleet(self, test_db, session: AsyncSession):
+        project = await create_project(session)
+        user = await create_user(session)
+        repo = await create_repo(session=session, project_id=project.id)
+        fleet_a = await create_fleet(session=session, project=project, name="a")
+        await create_instance(session=session, project=project, fleet=fleet_a, price=1.0)
+        fleet_b = await create_fleet(session=session, project=project, name="b")
+        await create_instance(session=session, project=project, fleet=fleet_b, price=2.0)
+        fleet_c = await create_fleet(session=session, project=project, name="c")
+        await create_instance(session=session, project=project, fleet=fleet_c, price=3.0)
+        run_spec = get_run_spec(run_name="test-run", repo_id=repo.name)
+        # When more than one fleet is requested, the cheapest one is selected
+        run_spec.configuration.fleets = ["c", "b"]
+        run = await create_run(
+            session=session, project=project, repo=repo, user=user, run_spec=run_spec
+        )
+        job = await create_job(session=session, run=run)
+
+        await process_submitted_jobs()
+
+        await session.refresh(job)
+        res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
+        job = res.unique().scalar_one()
+        assert job.status == JobStatus.SUBMITTED
+        assert job.instance is not None
+        assert job.instance.fleet == fleet_b
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_creates_new_instance_in_existing_fleet(self, test_db, session: AsyncSession):

src/tests/_internal/server/routers/test_fleets.py

@@ -366,6 +366,7 @@ async def test_creates_fleet(self, test_db, session: AsyncSession, client: Async
                     "name": "",
                     "default": False,
                     "reservation": None,
+                    "fleets": None,
                 },
                 "autocreated": False,
             },
@@ -484,6 +485,7 @@ async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: A
                     "name": "",
                     "default": False,
                     "reservation": None,
+                    "fleets": None,
                 },
                 "autocreated": False,
             },

src/tests/_internal/server/routers/test_runs.py

@@ -124,6 +124,7 @@ def get_dev_env_run_plan_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "configuration_path": "dstack.yaml",
             "profile": {
@@ -142,6 +143,7 @@ def get_dev_env_run_plan_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "repo_code_hash": None,
             "repo_data": {"repo_dir": "/repo", "repo_type": "local"},
@@ -274,6 +276,7 @@ def get_dev_env_run_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "configuration_path": "dstack.yaml",
             "profile": {
@@ -292,6 +295,7 @@ def get_dev_env_run_dict(
                 "idle_duration": None,
                 "utilization_policy": None,
                 "reservation": None,
+                "fleets": None,
             },
             "repo_code_hash": None,
             "repo_data": {"repo_dir": "/repo", "repo_type": "local"},