fkat 0.1.2__py3-none-any.whl

fkat/utils/cuda/preflight/health_check/aws_instance_config.py

@@ -0,0 +1,82 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class InstanceBenchmarkConfig:
+    name: str
+    gpu_memory_gb: int
+    config_dim: int
+    num_loops_single: int
+    num_loops_multi: int
+    baseline_num_loops: int
+    baseline_single_node_latency: float
+    baseline_pair_nodes_latency: float
+
+
+P4D_24XLARGE = InstanceBenchmarkConfig(
+    name="p4d.24xlarge",
+    gpu_memory_gb=18,
+    config_dim=4_000_000_000,
+    num_loops_single=200,
+    num_loops_multi=200,
+    baseline_num_loops=43,
+    baseline_single_node_latency=261,
+    baseline_pair_nodes_latency=772,
+)
+
+P4DE_24XLARGE = InstanceBenchmarkConfig(
+    name="p4de.24xlarge",
+    gpu_memory_gb=18,
+    config_dim=4_000_000_000,
+    num_loops_single=200,
+    num_loops_multi=200,
+    baseline_num_loops=43,
+    baseline_single_node_latency=261,
+    baseline_pair_nodes_latency=772,
+)
+
+P5_48XLARGE = InstanceBenchmarkConfig(
+    name="p5.48xlarge",
+    gpu_memory_gb=40,
+    config_dim=8_000_000_000,
+    num_loops_single=200,
+    num_loops_multi=200,
+    baseline_num_loops=24,
+    baseline_single_node_latency=324,
+    baseline_pair_nodes_latency=419,
+)
+
+P5EN_48XLARGE = InstanceBenchmarkConfig(
+    name="p5en.48xlarge",
+    gpu_memory_gb=75,
+    config_dim=16_700_000_000,
+    num_loops_single=100,
+    num_loops_multi=100,
+    baseline_num_loops=9,
+    baseline_single_node_latency=660,
+    baseline_pair_nodes_latency=780,
+)
+
+DEFAULT_INSTANCE = InstanceBenchmarkConfig(
+    name="default",
+    gpu_memory_gb=10,
+    config_dim=1_000,
+    num_loops_single=0,
+    num_loops_multi=0,
+    baseline_num_loops=0,
+    baseline_single_node_latency=0,
+    baseline_pair_nodes_latency=0,
+)
+
+INSTANCE_BENCHMARK_CONFIGS: dict[str, InstanceBenchmarkConfig] = {
+    cfg.name: cfg
+    for cfg in [
+        P4D_24XLARGE,
+        P4DE_24XLARGE,
+        P5_48XLARGE,
+        P5EN_48XLARGE,
+        DEFAULT_INSTANCE,
+    ]
+}

fkat/utils/cuda/preflight/health_check/constants.py

@@ -0,0 +1,23 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+INSTANCE_HEALTH_STATUS_DDB_TABLE_NAME = "NodeHealthStatus"
+
+# GPU stress test configs:
+STRESS_TEST_MAX_RUNTIME_IN_SEC = 120
+HEALTH_CHECK_TIMEOUT_SECS = 150
+
+# Good/bad node decision factor
+NUM_LOOPS_RANGE = 1
+SINGLE_NODE_LATENCY_THRESHOLD_FACTOR = 1.1
+PAIR_NODES_LATENCY_THRESHOLD_FACTOR = 1.1
+PREFLIGHT_MLFLOW_METRIC_PREFIX = "preflight"
+
+# Max len of a node IPv4 address
+MAX_ADDR_LENGTH = 100
+PASS = "pass"
+FAIL = "fail"
+
+MLFLOW_EXPERIMENT_NAME = "bad_node_detection_{region}"
+
+AWS_BATCH_JOB_ID = "AWS_BATCH_JOB_ID"
+AWS_BATCH_LINK = "https://{region}.console.aws.amazon.com/batch/home?region={region}#jobs/ec2/detail/{batch_id}"

fkat/utils/cuda/preflight/health_check/ddb_client.py

@@ -0,0 +1,82 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import boto3
+from boto3.dynamodb.conditions import Key
+from botocore.exceptions import ClientError
+import logging
+from typing import TYPE_CHECKING, Any
+from fkat.utils.cuda.preflight.health_check.constants import INSTANCE_HEALTH_STATUS_DDB_TABLE_NAME
+import datetime
+
+if TYPE_CHECKING:
+    from types_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table
+
+
+class HealthStatusDDBClient:
+    _PARTITION_KEY = "instance_gpu_hash_id"
+    _SORT_KEY = "time_checked"
+    _is_initialized = False
+    _instance = None
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> HealthStatusDDBClient:
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, region: str = "us-east-1") -> None:
+        if not self._is_initialized:
+            session = boto3.Session()
+            self.ddb_resource: DynamoDBServiceResource = session.resource("dynamodb", region_name=region)  # type: ignore[assignment]
+            self.table: Table = self.ddb_resource.Table(INSTANCE_HEALTH_STATUS_DDB_TABLE_NAME)
+
+            self._is_initialized = True
+
+    def generate_ddb_item(
+        self,
+        instance_gpu_hash_id: str,
+        instance_health: bool,
+        gpu_stats: dict[str | int, dict[str, Any]],
+        batch_job_id: str,
+        instance_id: str,
+        instance_type: str,
+        test_result: dict[str, Any],
+    ) -> dict[str, Any]:
+        return {
+            "instance_gpu_hash_id": instance_gpu_hash_id,
+            "time_checked": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+            "batch_job_id": batch_job_id,
+            "instance_type": instance_type,
+            "gpu_info": {str(key): value for key, value in gpu_stats.items()},
+            "instance_id": instance_id,
+            "healthy": instance_health,
+            "test_result": test_result,
+        }
+
+    def put_item(self, item: dict[str, Any]) -> None:
+        try:
+            respones = self.table.put_item(Item=item, ReturnValues="ALL_OLD")
+            logging.info(f"Item {respones} successfully added to ddb.")
+        except ClientError as e:
+            # If the item already exists, an exception will be raised
+            if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
+                logging.error("Item already exists. Duplicate insertion prevented.")
+            else:
+                logging.error(f"Error inserting item: {e.response['Error']['Message']}")
+
+            raise e
+
+    def get_item(self, partition_key: str) -> dict[str, Any] | None:
+        try:
+            response = self.table.query(
+                KeyConditionExpression=Key(self._PARTITION_KEY).eq(partition_key),
+                ScanIndexForward=False,
+                Limit=1,
+            )
+
+            logging.info(f"successfully get response {response} from table with key {partition_key}")
+            return response["Items"][0] if "Items" in response and response["Items"] else None
+        except Exception as e:
+            logging.error("An unexpected error occurred:", e)
+            raise e

fkat/utils/cuda/preflight/health_check/gpu_connection_test.py

@@ -0,0 +1,104 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import torch
+import mlflow
+from datetime import timedelta
+import torch.distributed as dist
+from torch.multiprocessing import Queue
+
+import logging
+from fkat.utils.cuda.preflight.health_check.timer import Timer
+from fkat.utils.cuda.preflight.health_check.helpers import (
+    destroy_process_group_if_initialized,
+)
+from fkat.utils.cuda.preflight.health_check.logger import end_all_mlflow_active_runs
+from fkat.utils.cuda.preflight.health_check.constants import (
+    MLFLOW_EXPERIMENT_NAME,
+    PREFLIGHT_MLFLOW_METRIC_PREFIX,
+)
+
+PG_TIMEOUT_MIN = float(os.environ.get("PG_TIMEOUT", 5))
+
+
+def run_gpu_connection_test(
+    mlflow_run_id: str,
+    result_queue: Queue,
+    dim_items: int,
+    loops: int,
+    master_addr: str,
+    master_port: str,
+    world_size: int,
+    rank: int,
+    device_id: int | None = None,
+    mode: str = "single",
+) -> None:
+    """
+    Runs a GPU connectivity and communication benchmark test using NCCL and logs performance metrics to MLflow.
+
+    This function initializes a distributed process group with NCCL, performs a warm-up `all_reduce`, and
+    repeatedly performs `all_reduce` operations to test GPU communication latency. It records timing statistics
+    for each iteration, logs them to MLflow, and places the results in the provided queue.
+
+    Args:
+        mlflow_run_id (str): The ID of the MLflow run to log metrics under.
+        result_queue (Queue): A multiprocessing-safe queue where timing results are pushed.
+        dim_items (int): The dimension of the square tensor used for the all_reduce operation.
+        loops (int): The number of all_reduce iterations to run for benchmarking.
+        master_addr(str): new internal addr of the process group,
+        master_port(str): port used for the process group,
+        world_size(int): number of processes expected in process_group,
+        rank(int): RANK of the current process in the process_group,
+        rail (Optional[int], optional): The CUDA device ID to use for testing. If None, defaults to the current device.
+        mode (str, optional): Mode label to tag the MLflow metrics. Defaults to "single".
+
+    Returns:
+        None
+    """
+    dist.init_process_group(  # type: ignore[possibly-unbound-attribute]
+        backend="nccl",
+        init_method=f"tcp://{master_addr}:{master_port}",
+        world_size=world_size,
+        rank=rank,
+        timeout=timedelta(minutes=PG_TIMEOUT_MIN),
+    )
+
+    if mlflow.active_run():
+        end_all_mlflow_active_runs()
+
+    mlflow.set_tracking_uri(uri=os.environ["MLFLOW_URI"])
+    mlflow.set_experiment(MLFLOW_EXPERIMENT_NAME.format(region=os.environ["AWS_DEFAULT_REGION"]))
+    mlflow.start_run(run_id=mlflow_run_id)  # type: ignore[possibly-unbound-attribute]
+
+    logging.info(f"inside run_gpu_connection_test, {dim_items}, {loops}, {device_id}")
+    timer = Timer()
+    device_id = device_id if device_id is not None else torch.cuda.current_device()
+
+    with timer("cuda"):
+        device = torch.device("cuda", device_id)
+        torch.cuda.set_device(device_id)
+        buffer = torch.ones((dim_items, dim_items), device=device, dtype=torch.float64)
+
+    # warmup
+    dist.all_reduce(buffer, op=dist.ReduceOp.AVG, async_op=False)  # type: ignore[possibly-unbound-attribute]
+
+    results = []
+    for i in range(loops):
+        with timer(f"all_reduce_{i}"):
+            with timer("send"):
+                waiter = dist.all_reduce(buffer, op=dist.ReduceOp.AVG, async_op=True)  # type: ignore[possibly-unbound-attribute]
+            with timer("sync"):
+                waiter.wait()
+                dist.barrier()  # type: ignore[possibly-unbound-attribute]
+            with timer("stat"):
+                buffer_sum = buffer.sum().item()  # noqa: F841
+
+        results.append(timer[f"all_reduce_{i}"])
+
+        mlflow.log_metric(  # type: ignore[possibly-unbound-attribute]
+            f"{PREFLIGHT_MLFLOW_METRIC_PREFIX}/{mode}_node_test/latency_gpu_{os.environ['LOCAL_RANK']}",
+            timer[f"all_reduce_{i}"],
+        )
+
+    result_queue.put(results)
+    destroy_process_group_if_initialized()

fkat/utils/cuda/preflight/health_check/gpu_stress_test.py

@@ -0,0 +1,122 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import math
+import mlflow
+import time
+from typing import Any
+import torch
+from torch.multiprocessing import Queue
+import logging
+from fkat.utils.cuda.preflight.health_check.constants import (
+    MLFLOW_EXPERIMENT_NAME,
+    PREFLIGHT_MLFLOW_METRIC_PREFIX,
+)
+from fkat.utils.cuda.preflight.health_check.logger import end_all_mlflow_active_runs
+
+
+def run_gpu_stress_test(mlflow_run_id: str, result_queue: Queue, gpu_mem: int, max_runtime: int) -> None:
+    """
+    Performs a multi-GPU stress test by executing repeated matrix multiplications and inter-GPU memory transfers.
+
+    This function:
+    - Allocates large tensors on each GPU (assuming 8 GPUs),
+    - Performs repeated `matmul` operations to stress GPU compute,
+    - Copies results across GPUs to test memory transfer integrity,
+    - Verifies data correctness after each transfer,
+    - Logs metrics to MLflow regarding correctness and loop iterations,
+    - Returns a dictionary summarizing the health of each GPU via the result queue.
+
+    Args:
+        mlflow_run_id (str): The MLflow run ID under which metrics are logged.
+        result_queue (Queue): A multiprocessing-safe queue to place GPU health results.
+        gpu_mem (int): Approximate GPU memory (in GB) to target when allocating stress test tensors.
+        max_runtime (int): Maximum runtime (in seconds) to perform the stress test.
+
+    Returns:
+        None
+    """
+    if not torch.cuda.is_available():
+        result_queue.put({})
+        return
+    mlflow.set_tracking_uri(uri=os.environ["MLFLOW_URI"])
+    mlflow.set_experiment(MLFLOW_EXPERIMENT_NAME.format(region=os.environ["AWS_DEFAULT_REGION"]))
+    if mlflow.active_run():
+        end_all_mlflow_active_runs()
+
+    mlflow.start_run(run_id=mlflow_run_id)  # type: ignore[possibly-unbound-attribute]
+
+    # Get the array size for a square array that fills 1/4 of memory with 2 byte values
+    arr_size = (((gpu_mem / 4) * 10**9) / 2) ** (1 / 2)
+    arr_size = int(math.ceil(arr_size))
+    num_gpus = torch.cuda.device_count()
+    logging.info(f"inside run_load(), num_gpus is: {num_gpus}")
+    if num_gpus != 8:
+        result_queue.put({})
+        return
+
+    gpu_health: dict[str, Any] = {str(idx): "Unknown" for idx in range(num_gpus)}
+    gpu_health["check_record"] = []
+
+    Ts = [torch.ones(arr_size, arr_size, dtype=torch.bfloat16, device=f"cuda:{gpu_num}") for gpu_num in range(num_gpus)]
+    results = [
+        torch.zeros(arr_size, arr_size, dtype=torch.bfloat16, device=f"cuda:{gpu_num}") for gpu_num in range(num_gpus)
+    ]
+    from_others = [
+        torch.zeros(arr_size, arr_size, dtype=torch.bfloat16, device=f"cuda:{gpu_num}") for gpu_num in range(num_gpus)
+    ]
+
+    torch.manual_seed(12345)
+
+    start_time = time.time()
+    curr_loop_num = 0
+    while time.time() - start_time < max_runtime:
+        # Matrix multiply into result
+        # TODO: record the latency of each matmul
+        [torch.matmul(T, T, out=result) for T, result in zip(Ts, results, strict=False)]
+
+        # Move into gpu curr_loop_num away
+        for i in range(num_gpus):
+            other_gpu = (curr_loop_num % (num_gpus - 1) + i + 1) % num_gpus
+            other = from_others[other_gpu]
+            original = results[i]
+            other[:] = original
+
+        # Check values are correct
+        checks = [
+            (other == result).sum() == result.numel() for other, result in zip(from_others, results, strict=False)
+        ]
+
+        for idx, check in enumerate(checks):
+            if not check.item():
+                gpu_health[str(idx)] = "Unhealthy"
+
+        gpu_health["check_record"].append([check.item() for check in checks])
+
+        curr_loop_num += 1
+
+        mlflow.log_metric(  # type: ignore[possibly-unbound-attribute]
+            f"{PREFLIGHT_MLFLOW_METRIC_PREFIX}/gpu_stress_test/all_gpu_calculation_is_correct",
+            all(checks),
+        )
+        mlflow.log_metric(  # type: ignore[possibly-unbound-attribute]
+            f"{PREFLIGHT_MLFLOW_METRIC_PREFIX}/gpu_stress_test/num_loop_gpu_stress_test",
+            curr_loop_num,
+        )
+
+    gpu_health["num_loop"] = str(curr_loop_num)
+
+    logging.info(f"Finsihed run_gpu_stress_test. {curr_loop_num} loops ran.")
+
+    if curr_loop_num < num_gpus:
+        logging.info(f"Few loops seen, only {curr_loop_num}")
+        for idx in range(num_gpus):
+            gpu_health[str(idx)] = "Unknown"
+
+    gpu_health["all_gpus"] = all(gpu_health.get(str(i)) == "Healthy" for i in range(num_gpus))
+
+    # Free memory
+    del Ts, results, from_others
+    torch.cuda.empty_cache()
+
+    result_queue.put(gpu_health)