Skip to content

Commit e72e24a

Browse files
Fix PEP 8 indentation and path inconsistencies from rebranding (#7)
* Initial plan * Fix indentation and path inconsistencies per review comments Co-authored-by: TerrenceZhangX <39916879+TerrenceZhangX@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: TerrenceZhangX <39916879+TerrenceZhangX@users.noreply.github.com>
1 parent bf67019 commit e72e24a

2 files changed

Lines changed: 184 additions & 181 deletions

File tree

scripts/run_profile.py

Lines changed: 183 additions & 180 deletions
Original file line numberDiff line numberDiff line change
@@ -24,195 +24,198 @@
2424

2525

2626
def wait_for_port(host: str, port: int, timeout: int = 600) -> bool:
27-
"""Wait until a TCP port becomes reachable."""
28-
deadline = time.time() + timeout
29-
while time.time() < deadline:
30-
try:
31-
with socket.create_connection((host, port), timeout=2):
32-
return True
33-
except Exception:
34-
time.sleep(1)
35-
return False
27+
"""Wait until a TCP port becomes reachable."""
28+
deadline = time.time() + timeout
29+
while time.time() < deadline:
30+
try:
31+
with socket.create_connection((host, port), timeout=2):
32+
return True
33+
except Exception:
34+
time.sleep(1)
35+
return False
3636

3737

3838
def clean_dir(path: str) -> None:
39-
"""Clean or create a directory."""
40-
if os.path.exists(path):
41-
for name in os.listdir(path):
42-
fp = os.path.join(path, name)
43-
if os.path.isfile(fp) or os.path.islink(fp):
44-
os.unlink(fp)
45-
elif os.path.isdir(fp):
46-
shutil.rmtree(fp)
47-
else:
48-
os.makedirs(path, exist_ok=True)
39+
"""Clean or create a directory."""
40+
if os.path.exists(path):
41+
for name in os.listdir(path):
42+
fp = os.path.join(path, name)
43+
if os.path.isfile(fp) or os.path.islink(fp):
44+
os.unlink(fp)
45+
elif os.path.isdir(fp):
46+
shutil.rmtree(fp)
47+
else:
48+
os.makedirs(path, exist_ok=True)
4949

5050

5151
def parse_args(argv: Optional[list] = None) -> argparse.Namespace:
52-
p = argparse.ArgumentParser(description="Run sglang profiling workload")
53-
54-
p.add_argument(
55-
"--profile-dir",
56-
default="/flowsim/server_profile",
57-
help="Directory where profiler traces (.trace.json.gz) will be written",
58-
)
59-
p.add_argument(
60-
"--log-dir",
61-
default="/flowsim/tests/test-artifacts",
62-
help="Directory to write server/client logs",
63-
)
64-
p.add_argument(
65-
"--server-opts",
66-
required=True,
67-
help=(
68-
"All options for sglang.launch_server (include --host, --port, --model-path, --tp, etc). "
69-
"Example: '--model-path /path --tp 1 --host 0.0.0.0 --port 30001 --disable-cuda-graph'"
70-
),
71-
)
72-
p.add_argument(
73-
"--bench-opts",
74-
required=True,
75-
help=(
76-
"All options for bench_serving.py (include --backend, --host, --port, --dataset-name, --profile, etc). "
77-
"Example: '--backend sglang --host 0.0.0.0 --port 30001 --dataset-name defined-len --num-prompts 16 --profile'"
78-
),
79-
)
80-
p.add_argument(
81-
"--bench-timeout",
82-
type=int,
83-
default=1200,
84-
help="Timeout in seconds for bench_serving.py",
85-
)
86-
87-
return p.parse_args(argv)
52+
p = argparse.ArgumentParser(description="Run sglang profiling workload")
53+
54+
p.add_argument(
55+
"--profile-dir",
56+
default="/flowsim/server_profile",
57+
help="Directory where profiler traces (.trace.json.gz) will be written",
58+
)
59+
p.add_argument(
60+
"--log-dir",
61+
default="/flowsim/tests/test-artifacts",
62+
help="Directory to write server/client logs",
63+
)
64+
p.add_argument(
65+
"--server-opts",
66+
required=True,
67+
help=(
68+
"All options for sglang.launch_server (include --host, --port, --model-path, --tp, etc). "
69+
"Example: '--model-path /path --tp 1 --host 0.0.0.0 --port 30001 --disable-cuda-graph'"
70+
),
71+
)
72+
p.add_argument(
73+
"--bench-opts",
74+
required=True,
75+
help=(
76+
"All options for bench_serving.py (include --backend, --host, --port, --dataset-name, --profile, etc). "
77+
"Example: '--backend sglang --host 0.0.0.0 --port 30001 --dataset-name defined-len --num-prompts 16 --profile'"
78+
),
79+
)
80+
p.add_argument(
81+
"--bench-timeout",
82+
type=int,
83+
default=1200,
84+
help="Timeout in seconds for bench_serving.py",
85+
)
86+
87+
return p.parse_args(argv)
8888

8989

9090
def main(argv: Optional[list] = None) -> int:
91-
args = parse_args(argv)
92-
93-
profile_dir = args.profile_dir
94-
log_dir = args.log_dir
95-
96-
clean_dir(profile_dir)
97-
os.makedirs(log_dir, exist_ok=True)
98-
99-
ts = int(time.time())
100-
server_stdout_path = os.path.join(log_dir, f"server_{ts}.stdout.log")
101-
server_stderr_path = os.path.join(log_dir, f"server_{ts}.stderr.log")
102-
server_stdout_f = open(server_stdout_path, "w", encoding="utf-8")
103-
server_stderr_f = open(server_stderr_path, "w", encoding="utf-8")
104-
105-
# Set profiling environment variables
106-
env = os.environ.copy()
107-
env["SGLANG_TORCH_PROFILER_DIR"] = profile_dir
108-
env["SGLANG_PROFILE_KERNELS"] = "1"
109-
env["SGLANG_PROFILE_DEBUG"] = "1"
110-
env["SGLANG_SET_CPU_AFFINITY"] = "1"
111-
112-
# Extract host and port from server-opts for connection check
113-
server_args = shlex.split(args.server_opts)
114-
host = "0.0.0.0"
115-
port = 30001
116-
try:
117-
if "--host" in server_args:
118-
host = server_args[server_args.index("--host") + 1]
119-
if "--port" in server_args:
120-
port = int(server_args[server_args.index("--port") + 1])
121-
except (ValueError, IndexError):
122-
pass
123-
124-
# Build server command
125-
launch_cmd = [
126-
sys.executable,
127-
"-m",
128-
"sglang.launch_server",
129-
] + server_args
130-
131-
print("[INFO] Starting sglang server:", " ".join(launch_cmd), flush=True)
132-
preexec = getattr(os, "setsid", None)
133-
server_proc = subprocess.Popen(
134-
launch_cmd,
135-
cwd="/flowsim/workload/framework/sglang/python",
136-
stdout=server_stdout_f,
137-
stderr=server_stderr_f,
138-
preexec_fn=preexec,
139-
env=env,
140-
)
141-
142-
try:
143-
if not wait_for_port(host, port, timeout=600):
144-
print("[ERROR] Server did not start within timeout", file=sys.stderr)
145-
return 1
146-
147-
script = os.path.abspath(
148-
"/flowsim/workload/framework/sglang/python/sglang/bench_serving.py"
149-
)
150-
151-
bench_args = shlex.split(args.bench_opts)
152-
client_args = [sys.executable, script] + bench_args
153-
154-
print("[INFO] Running bench_serving:", " ".join(client_args), flush=True)
155-
result = subprocess.run(
156-
client_args,
157-
capture_output=True,
158-
text=True,
159-
env=env,
160-
timeout=args.bench_timeout,
161-
)
162-
163-
ts2 = int(time.time())
164-
prefix = f"bench_serving_{ts2}"
165-
client_stdout_path = os.path.join(log_dir, prefix + ".stdout.log")
166-
client_stderr_path = os.path.join(log_dir, prefix + ".stderr.log")
167-
with open(client_stdout_path, "w", encoding="utf-8") as f_out:
168-
f_out.write(result.stdout)
169-
with open(client_stderr_path, "w", encoding="utf-8") as f_err:
170-
f_err.write(result.stderr)
171-
172-
if result.returncode != 0:
173-
print(
174-
f"[ERROR] bench_serving exited with code {result.returncode}",
175-
file=sys.stderr,
176-
)
177-
return result.returncode
178-
179-
files = os.listdir(profile_dir)
180-
json_gz_files = [f for f in files if f.endswith(".trace.json.gz")]
181-
if not json_gz_files:
182-
print(
183-
f"[ERROR] No .trace.json.gz files found in {profile_dir}",
184-
file=sys.stderr,
185-
)
186-
return 1
187-
188-
print(
189-
f"[INFO] Profiling complete, found {len(json_gz_files)} trace file(s) in {profile_dir}",
190-
flush=True,
191-
)
192-
return 0
193-
finally:
194-
try:
195-
if server_proc.poll() is None:
196-
try:
197-
os.killpg(os.getpgid(server_proc.pid), signal.SIGTERM)
198-
except Exception:
199-
server_proc.terminate()
200-
server_proc.wait(timeout=30)
201-
except Exception:
202-
pass
203-
try:
204-
server_stdout_f.flush()
205-
server_stderr_f.flush()
206-
except Exception:
207-
pass
208-
try:
209-
server_stdout_f.close()
210-
server_stderr_f.close()
211-
except Exception:
212-
pass
213-
time.sleep(2)
91+
args = parse_args(argv)
21492

93+
profile_dir = args.profile_dir
94+
log_dir = args.log_dir
21595

216-
if __name__ == "__main__":
217-
raise SystemExit(main())
96+
clean_dir(profile_dir)
97+
os.makedirs(log_dir, exist_ok=True)
98+
99+
ts = int(time.time())
100+
server_stdout_path = os.path.join(log_dir, f"server_{ts}.stdout.log")
101+
server_stderr_path = os.path.join(log_dir, f"server_{ts}.stderr.log")
102+
server_stdout_f = open(server_stdout_path, "w", encoding="utf-8")
103+
server_stderr_f = open(server_stderr_path, "w", encoding="utf-8")
104+
105+
# Set profiling environment variables
106+
env = os.environ.copy()
107+
env["SGLANG_TORCH_PROFILER_DIR"] = profile_dir
108+
env["SGLANG_PROFILE_KERNELS"] = "1"
109+
env["SGLANG_PROFILE_DEBUG"] = "1"
110+
env["SGLANG_SET_CPU_AFFINITY"] = "1"
111+
112+
# Extract host and port from server-opts for connection check
113+
server_args = shlex.split(args.server_opts)
114+
host = "0.0.0.0"
115+
port = 30001
116+
try:
117+
if "--host" in server_args:
118+
host = server_args[server_args.index("--host") + 1]
119+
if "--port" in server_args:
120+
port = int(server_args[server_args.index("--port") + 1])
121+
except (ValueError, IndexError):
122+
pass
123+
124+
# Build server command
125+
launch_cmd = [
126+
sys.executable,
127+
"-m",
128+
"sglang.launch_server",
129+
] + server_args
130+
131+
print("[INFO] Starting sglang server:", " ".join(launch_cmd), flush=True)
132+
preexec = getattr(os, "setsid", None)
133+
server_proc = subprocess.Popen(
134+
launch_cmd,
135+
cwd="/flowsim/workload/framework/sglang/python",
136+
stdout=server_stdout_f,
137+
stderr=server_stderr_f,
138+
preexec_fn=preexec,
139+
env=env,
140+
)
141+
142+
try:
143+
if not wait_for_port(host, port, timeout=600):
144+
print(
145+
"[ERROR] Server did not start within timeout", file=sys.stderr
146+
)
147+
return 1
218148

149+
script = os.path.abspath(
150+
"/flowsim/workload/framework/sglang/python/sglang/bench_serving.py"
151+
)
152+
153+
bench_args = shlex.split(args.bench_opts)
154+
client_args = [sys.executable, script] + bench_args
155+
156+
print(
157+
"[INFO] Running bench_serving:", " ".join(client_args), flush=True
158+
)
159+
result = subprocess.run(
160+
client_args,
161+
capture_output=True,
162+
text=True,
163+
env=env,
164+
timeout=args.bench_timeout,
165+
)
166+
167+
ts2 = int(time.time())
168+
prefix = f"bench_serving_{ts2}"
169+
client_stdout_path = os.path.join(log_dir, prefix + ".stdout.log")
170+
client_stderr_path = os.path.join(log_dir, prefix + ".stderr.log")
171+
with open(client_stdout_path, "w", encoding="utf-8") as f_out:
172+
f_out.write(result.stdout)
173+
with open(client_stderr_path, "w", encoding="utf-8") as f_err:
174+
f_err.write(result.stderr)
175+
176+
if result.returncode != 0:
177+
print(
178+
f"[ERROR] bench_serving exited with code {result.returncode}",
179+
file=sys.stderr,
180+
)
181+
return result.returncode
182+
183+
files = os.listdir(profile_dir)
184+
json_gz_files = [f for f in files if f.endswith(".trace.json.gz")]
185+
if not json_gz_files:
186+
print(
187+
f"[ERROR] No .trace.json.gz files found in {profile_dir}",
188+
file=sys.stderr,
189+
)
190+
return 1
191+
192+
print(
193+
f"[INFO] Profiling complete, found {len(json_gz_files)} trace file(s) in {profile_dir}",
194+
flush=True,
195+
)
196+
return 0
197+
finally:
198+
try:
199+
if server_proc.poll() is None:
200+
try:
201+
os.killpg(os.getpgid(server_proc.pid), signal.SIGTERM)
202+
except Exception:
203+
server_proc.terminate()
204+
server_proc.wait(timeout=30)
205+
except Exception:
206+
pass
207+
try:
208+
server_stdout_f.flush()
209+
server_stderr_f.flush()
210+
except Exception:
211+
pass
212+
try:
213+
server_stdout_f.close()
214+
server_stderr_f.close()
215+
except Exception:
216+
pass
217+
time.sleep(2)
218+
219+
220+
if __name__ == "__main__":
221+
raise SystemExit(main())

simulator/base_parser.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -536,7 +536,7 @@ def _calibrate_communication_kernels(self) -> None:
536536
profiled_duration = comm_profile_cache[cache_key]
537537
else:
538538
profiled_duration = nb.run_nccl_all_reduce_perf(
539-
cmd_path="/workloadsim/third_party/nccl-tests/build/all_reduce_perf",
539+
cmd_path="/flowsim/third_party/nccl-tests/build/all_reduce_perf",
540540
b=str(size),
541541
e=str(size),
542542
g=str(self.tensor_parallelism),

0 commit comments

Comments
 (0)