Standalone notebook

client-scripts/Makefile

@@ -3,11 +3,11 @@ install-prerequisites:  ## Install dependencies
 	python3 -m pip install --upgrade pip
 	rm -rf chakra
 	git clone https://github.com/mlcommons/chakra.git
-	cp utils/pyproject.toml chakra
+	cp service_client_utils/pyproject.toml chakra
 	cd chakra && pip install .
 	rm -rf chakra
 
 
 .PHONY: test
 test:
-	pytest tests/test-unit -s -vvvv
+	pytest tests/test-unit -s -vvvv

client-scripts/notebooks/ns3_sample.py

@@ -17,8 +17,10 @@
 
 # %%
 import sys
-sys.path.append("../utils")
-from astra_sim import AstraSim, Collective, NetworkBackend
+
+# sys.path.append("../utils")
+# from astra_sim import AstraSim, Collective, NetworkBackend
+from service_client_utils.astra_sim import AstraSim, Collective, NetworkBackend
 
 # %% [markdown]
 # ##### Call the AstraSim client helper with the server endpoint and tag to connect to the ASTRA-sim gRPC server, initialize the SDK, and create a tagged folder for configs, results, and logs.
@@ -30,7 +32,9 @@
 # ##### Generate workload execution traces for each rank and set the required data size for AstraSim configuration
 
 # %%
-astra.configuration.common_config.workload = astra.generate_collective(collective=Collective.ALLREDUCE, coll_size= 8 *1024*1024, npu_range=[0,8])
+astra.configuration.common_config.workload = astra.generate_collective(
+    collective=Collective.ALLREDUCE, coll_size=8 * 1024 * 1024, npu_range=[0, 8]
+)
 print(astra.configuration.common_config.workload)
 
 
@@ -41,18 +45,28 @@
 astra.configuration.common_config.system.scheduling_policy = astra.configuration.common_config.system.LIFO
 astra.configuration.common_config.system.endpoint_delay = 10
 astra.configuration.common_config.system.active_chunks_per_dimension = 1
-astra.configuration.common_config.system.all_gather_implementation = [astra.configuration.common_config.system.RING]
-astra.configuration.common_config.system.all_to_all_implementation = [astra.configuration.common_config.system.DIRECT]
-astra.configuration.common_config.system.all_reduce_implementation = [astra.configuration.common_config.system.ONERING]
-astra.configuration.common_config.system.collective_optimization = astra.configuration.common_config.system.LOCALBWAWARE
+astra.configuration.common_config.system.all_gather_implementation = [
+    astra.configuration.common_config.system.RING
+]
+astra.configuration.common_config.system.all_to_all_implementation = [
+    astra.configuration.common_config.system.DIRECT
+]
+astra.configuration.common_config.system.all_reduce_implementation = [
+    astra.configuration.common_config.system.ONERING
+]
+astra.configuration.common_config.system.collective_optimization = (
+    astra.configuration.common_config.system.LOCALBWAWARE
+)
 astra.configuration.common_config.system.local_mem_bw = 1600
 print(astra.configuration.common_config.system)
 
 # %% [markdown]
 # ##### Configure ASTRA-sim remote memory configuration
 
 # %%
-astra.configuration.common_config.remote_memory.memory_type = astra.configuration.common_config.remote_memory.NO_MEMORY_EXPANSION
+astra.configuration.common_config.remote_memory.memory_type = (
+    astra.configuration.common_config.remote_memory.NO_MEMORY_EXPANSION
+)
 print(astra.configuration.common_config.remote_memory)
 
 # %% [markdown]
@@ -62,8 +76,8 @@
 # astra.configuration.network_backend.choice = astra.configuration.network_backend.NS3
 astra.configuration.network_backend.ns3.network.packet_payload_size = int(8192)
 astra.configuration.network_backend.ns3.logical_topology.logical_dimensions = [8]
-astra.configuration.network_backend.ns3.trace.trace_ids = [0, 1, 2, 3,4 ,5 ,6, 7]
-print("network backend choice set to:",astra.configuration.network_backend.ns3.topology.choice)
+astra.configuration.network_backend.ns3.trace.trace_ids = [0, 1, 2, 3, 4, 5, 6, 7]
+print("network backend choice set to:", astra.configuration.network_backend.ns3.topology.choice)
 print(astra.configuration.network_backend.ns3.network.packet_payload_size)
 print(astra.configuration.network_backend.ns3.logical_topology)
 print(astra.configuration.network_backend.ns3.trace)
@@ -91,7 +105,6 @@
 print(astra.configuration.network_backend.ns3.topology.nc_topology)
 
 
-
 # %% [markdown]
 # ##### Configure ASTRA-sim cmd parameters
 
@@ -122,6 +135,7 @@
 import pandas as pd
 import os
 from common import FileFolderUtils
+
 df = pd.read_csv(os.path.join(FileFolderUtils.get_instance().OUTPUT_DIR, "fct.csv"))
 df.head()
 df = pd.read_csv(os.path.join(FileFolderUtils.get_instance().OUTPUT_DIR, "flow_stats.csv"))

client-scripts/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "service_client_utils"
+version = "0.6.0"
+dependencies = [
+    "networkx",
+    "pandas",
+    "grpcio",
+]
+
+[tool.setuptools]
+packages = ["service_client_utils",]
+
+[tool.setuptools.package-dir]
+service_client_utils = "service_client_utils"

client-scripts/utils/astra_sim.py → client-scripts/service_client_utils/astra_sim.py

@@ -42,9 +42,9 @@
     GATHER,
 )
 
-from common import Utilities, FileFolderUtils, StatUtil
+from service_client_utils.common import Utilities, FileFolderUtils, StatUtil
 
-from astra_sim_client import AstraSimClient
+from service_client_utils.astra_sim_client import AstraSimClient
 import astra_sim_sdk.astra_sim_sdk as astra_sim_sdk
 
 pd.options.mode.chained_assignment = None

client-scripts/utils/astra_sim_client.py → client-scripts/service_client_utils/astra_sim_client.py

@@ -28,7 +28,7 @@
 import threading
 
 import astra_sim_sdk.astra_sim_sdk as astra_sim
-from common import Utilities, FileFolderUtils
+from service_client_utils.common import Utilities, FileFolderUtils
 
 
 class AstraSimClient:

client-scripts/utils/common.py → client-scripts/service_client_utils/common.py

@@ -236,7 +236,7 @@ def ns3_flow_statistics():
         """
         df = pd.read_csv(
             os.path.join(FileFolderUtils().OUTPUT_DIR, "fct.txt"),
-            delim_whitespace=True,
+            sep=r"\s+",
             header=None,
         )
 
@@ -273,7 +273,7 @@ def ns3_fct_csv():
         """
         df = pd.read_csv(
             os.path.join(FileFolderUtils().OUTPUT_DIR, "fct.txt"),
-            delim_whitespace=True,
+            sep=r"\s+",
             header=None,
         )
         df.columns = [