Kubernetes: add multi-node support (#3141)

* Discover and set instance's internal_ip (PodIP)
* Fix region mismatch
* Add `privileged: true` support
* [runner] Set RLIMIT_MEMLOCK to unlimited. Fixes issues with
  InfiniBand/RDMA

Part-of: #3126

Author: un-def

contributing/BACKENDS.md

@@ -27,15 +27,15 @@ git clone https://github.com/dstackai/gpuhunt.git
 
 - **Offline providers** offer static machine configurations that are not frequently updated.
   `gpuhunt` collects offline providers' instance offers on an hourly basis.
-  Examples: `aws`, `gcp`, `azure`, etc. 
+  Examples: `aws`, `gcp`, `azure`, etc.
 - **Online providers** offer dynamic machine configurations that are available at the very moment
   when you fetch configurations (e.g., GPU marketplaces).
   `gpuhunt` collects online providers' instance offers each time a `dstack` user provisions a new instance.
   Examples: `tensordock`, `vastai`, etc.
 
 ### 1.3. Create the provider class
 
-Create the provider class file under `src/gpuhunt/providers`. 
+Create the provider class file under `src/gpuhunt/providers`.
 
 Make sure your class extends the [`AbstractProvider`](https://github.com/dstackai/gpuhunt/blob/main/src/gpuhunt/providers/__init__.py)
 base class. See its docstrings for descriptions of the methods that your class should implement.
@@ -69,13 +69,13 @@ Refer to examples: [test_datacrunch.py](https://github.com/dstackai/gpuhunt/blob
 
 ### 1.6. Submit a pull request
 
-Once the cloud provider is added, submit a pull request. 
+Once the cloud provider is added, submit a pull request.
 
 > Anything unclear? Ask questions on the [Discord server](https://discord.gg/u8SmfwPpMd).
 
 ## 2. Integrate the cloud provider to dstackai/dstack
 
-Once the provider is added to `gpuhunt`, we can proceed with implementing 
+Once the provider is added to `gpuhunt`, we can proceed with implementing
 the corresponding backend with `dstack`. Follow the steps below.
 
 ### 2.1. Determine if you will implement a VM-based or a container-based backend
@@ -124,10 +124,10 @@ Then add these models to `AnyBackendConfig*` unions in [`src/dstack/_internal/co
 
 The script also generates `*BackendStoredConfig` that extends `*BackendConfig` to be able to store extra parameters in the DB. By the same logic, it generates `*Config` that extends `*BackendStoredConfig` with creds and uses it as the main `Backend` and `Compute` config instead of using `*BackendConfigWithCreds` directly.
 
-Refer to examples: 
-[datacrunch](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/datacrunch/models.py), 
-[aws](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/aws/models.py), 
-[gcp](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/gcp/models.py), 
+Refer to examples:
+[datacrunch](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/datacrunch/models.py),
+[aws](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/aws/models.py),
+[gcp](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/gcp/models.py),
 [azure](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/models.py), etc.
 
 ### 2.7. Implement the backend compute class
@@ -147,8 +147,8 @@ Go to `configurator.py` and implement custom `Configurator` logic. At minimum, y
 You may also need to validate other config parameters if there are any.
 
 Refer to examples: [datacrunch](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/datacrunch/configurator.py),
-[aws](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/aws/configurator.py), 
-[gcp](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/gcp/configurator.py), 
+[aws](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/aws/configurator.py),
+[gcp](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/gcp/configurator.py),
 [azure](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/azure/configurator.py), etc.
 
 Register configurator by appending it to `_CONFIGURATOR_CLASSES` in [`src/dstack/_internal/core/backends/configurators.py`](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/configurators.py).
@@ -181,6 +181,9 @@ The agent controls the VM and starts Docker containers for users' jobs.
 Since `dstack` controls the entire VM, VM-based backends can support more features,
 such as blocks, instance volumes, privileged containers, and reusable instances.
 
+Note, all VM-based backend `Compute`s should sublass the `ComputeWithPrivilegedSupport` mixin,
+as the `dstack-shim` agent provides this functionality OOTB.
+
 To support a VM-based backend, `dstack` expects the following:
 
 - An API for creating and terminating VMs

runner/internal/executor/executor.go

@@ -29,6 +29,7 @@ import (
 	"github.com/dstackai/dstack/runner/internal/schemas"
 	"github.com/dstackai/dstack/runner/internal/types"
 	"github.com/prometheus/procfs"
+	"golang.org/x/sys/unix"
 )
 
 // TODO: Tune these parameters for optimal experience/performance
@@ -518,6 +519,21 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error
 
 	cmd.Env = envMap.Render()
 
+	// Configure process resource limits
+	// TODO: Make rlimits customizable in the run configuration. Currently, we only set max locked memory
+	// to unlimited to fix the issue with InfiniBand/RDMA: "Cannot allocate memory".
+	// See: https://github.com/ofiwg/libfabric/issues/6437
+	// See: https://github.com/openucx/ucx/issues/8229
+	// Note: we already set RLIMIT_MEMLOCK to unlimited in the shim if we've detected IB devices
+	// (see configureHpcNetworkingIfAvailable() function), but, as it's on the shim side, it only works
+	// with VM-based backends.
+	rlimitMemlock := unix.Rlimit{Cur: unix.RLIM_INFINITY, Max: unix.RLIM_INFINITY}
+	// TODO: Check if we have CAP_SYS_RESOURCE. In container environments, even root usually doesn't have
+	// this capability.
+	if err := unix.Setrlimit(unix.RLIMIT_MEMLOCK, &rlimitMemlock); err != nil {
+		log.Error(ctx, "Failed to set resource limits", "err", err)
+	}
+
 	log.Trace(ctx, "Starting exec", "cmd", cmd.String(), "working_dir", cmd.Dir, "env", cmd.Env)
 
 	ptm, err := startCommand(cmd)

src/dstack/_internal/core/backends/aws/compute.py

@@ -24,6 +24,7 @@
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithReservationSupport,
     ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
@@ -90,6 +91,7 @@ def _ec2client_cache_methodkey(self, ec2_client, *args, **kwargs):
 class AWSCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
     ComputeWithReservationSupport,
     ComputeWithPlacementGroupSupport,

src/dstack/_internal/core/backends/azure/compute.py

@@ -43,6 +43,7 @@
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
     ComputeWithMultinodeSupport,
+    ComputeWithPrivilegedSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
     get_gateway_user_data,
@@ -78,6 +79,7 @@
 class AzureCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
     ComputeWithGatewaySupport,
     Compute,

src/dstack/_internal/core/backends/base/compute.py

@@ -320,6 +320,15 @@ def _restrict_instance_offer_az_to_volumes_az(
             ]
 
 
+class ComputeWithPrivilegedSupport:
+    """
+    Must be subclassed to support runs with `privileged: true`.
+    All VM-based Computes (that is, Computes that use the shim) should subclass this mixin.
+    """
+
+    pass
+
+
 class ComputeWithMultinodeSupport:
     """
     Must be subclassed to support multinode tasks and cluster fleets.

src/dstack/_internal/core/backends/cloudrift/compute.py

@@ -4,6 +4,7 @@
     Compute,
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     get_shim_commands,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
@@ -27,6 +28,7 @@
 class CloudRiftCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: CloudRiftConfig):

src/dstack/_internal/core/backends/cudo/compute.py

@@ -6,6 +6,7 @@
 from dstack._internal.core.backends.base.compute import (
     ComputeWithCreateInstanceSupport,
     ComputeWithFilteredOffersCached,
+    ComputeWithPrivilegedSupport,
     generate_unique_instance_name,
     get_shim_commands,
 )
@@ -32,6 +33,7 @@
 class CudoCompute(
     ComputeWithFilteredOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: CudoConfig):

src/dstack/_internal/core/backends/datacrunch/compute.py

@@ -8,6 +8,7 @@
 from dstack._internal.core.backends.base.compute import (
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     generate_unique_instance_name,
     get_shim_commands,
 )
@@ -39,6 +40,7 @@
 class DataCrunchCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: DataCrunchConfig):

src/dstack/_internal/core/backends/digitalocean_base/compute.py

@@ -7,6 +7,7 @@
 from dstack._internal.core.backends.base.compute import (
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     generate_unique_instance_name,
     get_user_data,
 )
@@ -40,6 +41,7 @@
 class BaseDigitalOceanCompute(
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
+    ComputeWithPrivilegedSupport,
     Compute,
 ):
     def __init__(self, config: BaseDigitalOceanConfig, api_url: str, type: BackendType):

src/dstack/_internal/core/backends/features.py

@@ -4,6 +4,7 @@
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
+    ComputeWithPrivilegedSupport,
     ComputeWithReservationSupport,
     ComputeWithVolumeSupport,
 )
@@ -38,6 +39,10 @@ def _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithCreateInstanceSupport,
 )
+BACKENDS_WITH_PRIVILEGED_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithPrivilegedSupport,
+)
 BACKENDS_WITH_MULTINODE_SUPPORT = [BackendType.REMOTE] + _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithMultinodeSupport,