Implement GPU `blocks` property #2253

un-def · 2025-01-31T14:14:26Z

Add blocks property to cloud and SSH fleet
Reverse instance->job relation, switch from 1-to-1 to many-to-1 (many jobs to one instance)
Add InstanceModel.shared_info structure to hold total and currenty used number of blocks
Generate virtual shared offers (fraction of a real offer) on the fly
Update CLI and API to display shared resources and instances
Keep track of volumes used by jobs

Part-of: #1780

* Add `blocks` property to cloud and SSH fleet * Reverse instance->job relation, switch from 1-to-1 to many-to-1 (many jobs to one instance) * Add InstanceModel.shared_info structure to hold total and currenty used number of blocks * Generate virtual shared offers (fraction of a real offer) on the fly * Update CLI and API to display shared resources and instances * Keep track of volumes used by jobs Part-of: #1780

r4victor · 2025-02-03T07:12:38Z

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

+        Union[Literal["auto"], int],
+        Field(
+            description=(
+                "The amount of blocks to split the instance into, a number or `auto` (e.g., `4`)."


Maybe elaborate what auto means?

r4victor · 2025-02-03T07:46:02Z

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

@@ -110,6 +110,11 @@ def get_public_keys(self) -> List[str]:
        return [ssh_key.public.strip() for ssh_key in self.ssh_keys]


+class InstanceSharedInfo(CoreModel):
+    total_blocks: Union[Literal["auto"], int]
+    busy_blocks: int = 0


Not sure about storing busy_blocks in json vs in a separate column. Any particular reasons to do so? I'm generally inclined to keep static data in json and changing fields like this as separate columns.

Maybe using json is more flexible if we take a complete different approach instead of relying on blocks, Idk.

Replaced with two integer columns 8b555f7

r4victor · 2025-02-03T07:48:08Z

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

+    return shareable_offers
+
+
+def generate_shared_offer(offer: InstanceOfferWithAvailability, blocks: int, total_blocks: int):


missing return type

r4victor · 2025-02-03T08:00:02Z

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

@@ -416,7 +421,7 @@ def _error(cls, values) -> Dict:

 class JobPlan(CoreModel):
    job_spec: JobSpec
-    offers: List[InstanceOfferWithAvailability]
+    offers: List[Union[InstanceSharedOffer, InstanceOfferWithAvailability]]


Do we really need to introduce different types of offers to the API? Can't every InstanceOfferWithAvailability be InstanceSharedOffer with total_blocks/blocks = 1?

r4victor · 2025-02-03T08:35:33Z

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

+    if isinstance(blocks, int) and blocks < 2:
+        blocks = None


Setting shared_info to None when blocks: 1 is very confusing since you need to know that total_blocks cannot be 1 (e.g. when shared offers are generated). I'd reverse it: if blocks is None, then assume blocks: 1.

Again, with my solution, it would be more appropriate to normalize shared_info to keep total/used blocks in separate columns so that it's easier to migrate.

Refactored as part of 8b555f7

r4victor · 2025-02-03T08:46:02Z

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

+    # List of volumes used by the job
+    volume_names: Optional[list[str]] = None  # None for backward compalibility
+    # Virtual shared offer. None if the instance is not shared.
+    offer: Optional[InstanceSharedOffer] = None


Maybe it's more useful to always store the offer used by the job, not just when it was shared?

Changed in 8b555f7

r4victor · 2025-02-03T08:51:20Z

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

    # None if data is not yet available (on vm-based backends and ssh instances)
    # or not applicable (container-based backends)
    ports: Optional[dict[int, int]] = None
+    # List of volumes used by the job
+    volume_names: Optional[list[str]] = None  # None for backward compalibility
+    # Virtual shared offer. None if the instance is not shared.
+    offer: Optional[InstanceSharedOffer] = None


I think it's crucial now to add description of JobRuntimeData and what should go there so that it does not become a catch-all struct for everything. Probably makes sense to add a similar description to JobProvisioningData to differentiate it with JobRuntimeData.

r4victor · 2025-02-03T09:04:28Z

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

+        if shared_info is not None:
+            if blocks == "auto":
+                blocks = len(instance_offer.instance.resources.gpus)
+                if blocks > 1:
+                    shared_info.total_blocks = blocks
+                else:
+                    shared_info = None
+            if shared_info is not None:
+                instance.shared_info = shared_info.json()
+            else:
+                instance.shared_info = None


The shared_info updating is convoluted. I believe it can be created with total_blocks: "auto" and then it's overridden. I don't see much value in that. I'd keep user's configured blocks as a static property somewhere and wound't touch it. total_blocks can then be resolved when instance resources are known. No need to allow for auto.

Replaced with NULL -> calculated auto blocks in 8b555f7

NULL as "info is not yet available" seems to fit there

r4victor · 2025-02-03T09:26:56Z

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

+    for backend, offer in offers:
+        resources = offer.instance.resources
+        gpus = resources.gpus
+        if not gpus:
+            continue
+        if gpus[0].vendor == gpuhunt.AcceleratorVendor.GOOGLE:
+            # TPU instances cannot be shared
+            continue
+        if blocks == "auto":
+            divisor = len(gpus)
+        else:
+            divisor = blocks
+        if len(gpus) % divisor or resources.cpus % divisor:
+            # gpus or cpus are not divisible by blocks
+            continue
+        shareable_offers.append((backend, offer))


I think this code can be nicely refactored into a separate func.

Refactored in 8b555f7

r4victor · 2025-02-03T09:38:15Z

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

+                    shared_offer.availability = InstanceAvailability.BUSY
+                if shared_offer.availability == InstanceAvailability.IDLE or not idle_only:
+                    instances_with_offers.append((instance, shared_offer))
+                break


So we show only the smallest matching offer? Is it intended. Asking because the doc showed all possible offers.

Yes, IMO there is no reason to show all virual offers as they won't be picked anyway

r4victor

Tested with cloud fleets. Everything works as expected.

I'd reconsider how shared_info is stored/updated (mainly to make the code more comprehendible) and if InstanceSharedOffer/InstanceOfferWithAvailability should be differentiated in the API. Otherwise, looks good.

r4victor · 2025-02-03T09:59:07Z

...stack/_internal/server/migrations/versions/1338b788b612_reverse_job_instance_relationship.py

+
+def downgrade() -> None:
+    with op.batch_alter_table("instances", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("job_id", sa.UUID(), autoincrement=False, nullable=True))


sa.UUID() fails with sqlite

change to:

batch_op.add_column(sa.Column("job_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=True))

Suitable for single-node setups

un-def requested review from jvstme and r4victor January 31, 2025 14:23

r4victor reviewed Feb 3, 2025

View reviewed changes

r4victor approved these changes Feb 3, 2025

View reviewed changes

r4victor reviewed Feb 3, 2025

View reviewed changes

un-def added 5 commits February 4, 2025 03:03

Get rid of InstanceSharedOffer

3342a34

Fix downgrade migration

fa57f44

Update blocks description

5d5434b

Replace shared_info with two int columns

8b555f7

Add JobRuntimeData docstring

a6dcde2

un-def force-pushed the issue_1780_gpu_blocks branch from 0767ad7 to a6dcde2 Compare February 5, 2025 13:32

un-def added 2 commits February 5, 2025 14:07

Add DSTACK_FORCE_BRIDGE_NETWORK env var

16de87d

Suitable for single-node setups

Update JobRuntimeData.offer description

73f4ba7

r4victor approved these changes Feb 6, 2025

View reviewed changes

un-def merged commit c1206c5 into master Feb 6, 2025
24 checks passed

un-def deleted the issue_1780_gpu_blocks branch February 6, 2025 08:52

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement GPU `blocks` property #2253

Implement GPU `blocks` property #2253

un-def commented Jan 31, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor Feb 3, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025 •

edited

Loading

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025 •

edited

Loading

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor Feb 3, 2025

un-def Feb 5, 2025

r4victor left a comment

r4victor Feb 3, 2025

un-def Feb 5, 2025

		return shareable_offers


		def generate_shared_offer(offer: InstanceOfferWithAvailability, blocks: int, total_blocks: int):

Implement GPU blocks property #2253

Implement GPU blocks property #2253

Conversation

un-def commented Jan 31, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

un-def Feb 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

un-def Feb 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

r4victor left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Implement GPU `blocks` property #2253

Implement GPU `blocks` property #2253

un-def Feb 5, 2025 •

edited

Loading

un-def Feb 5, 2025 •

edited

Loading