feat: add streaming output for dp train #307
Conversation
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughIntroduces a new run_command_streaming utility for real-time stdout/stderr streaming and log file writing, and updates the training step in dpgen2/op/run_dp_train.py to use it with train.log. Previous in-memory logging for the train step is removed; post-train “freeze” continues using existing logging. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Op as run_dp_train.py
participant RC as run_command_streaming
participant Sh as Subprocess
participant Log as train.log (optional)
Op->>RC: invoke(cmd, shell=False, log_file="train.log")
activate RC
RC->>Sh: Popen(cmd, pipes, line-buffered)
par Read stdout
RC->>Sh: spawn stdout reader thread
loop lines
Sh-->>RC: stdout line
RC-->>Op: stream to terminal
RC-->>Log: append line
end
and Read stderr
RC->>Sh: spawn stderr reader thread
loop lines
Sh-->>RC: stderr line
RC-->>Op: stream to terminal
RC-->>Log: append line
end
end
Sh-->>RC: exit code
RC-->>Op: (code, stdout_str, stderr_str)
deactivate RC
note over Op: No in-memory fplog write for train step
sequenceDiagram
autonumber
participant Op as run_dp_train.py
participant RCs as run_command_streaming (train)
participant RC as run_command (freeze)
participant Log as fplog
Op->>RCs: Train (streamed to train.log)
note right of RCs: Output handled by streaming<br/>No fplog writes for train
Op->>RC: Freeze (non-streaming)
RC-->>Op: stdout/stderr captured
Op-->>Log: write freeze stdout/stderr to fplog
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
dpgen2/op/run_dp_train.py (1)
270-274: Double-opening train.log causes truncation/overwrite and log corruptionYou open
train.login write mode (Line 270) and also pass the same path torun_command_streamingwhich opens it again (in write mode per current impl). Keeping two writers to the same file handle can:
- truncate the file multiple times,
- make later writes (e.g., “freeze” section) overwrite the beginning of the file because the earlier handle’s file pointer starts at offset 0.
Net effect: training logs can be partially or fully clobbered.
Minimal, safe fix:
- Do not open
fplogbefore training.- Use
run_command_streaming(..., log_file="train.log")for both train and freeze so there’s a single writer pattern.- Remove the manual
fplog.write(...)for freeze.- Drop
clean_before_quit()or make it a no-op.Apply this diff:
@@ - with set_directory(work_dir): - # open log - fplog = open("train.log", "w") - - def clean_before_quit(): - fplog.close() + with set_directory(work_dir): + # train/freeze logs will be written by run_command_streaming to train.log + def clean_before_quit(): + pass @@ - # Use streaming output for real-time monitoring - ret, out, err = run_command_streaming(command, log_file="train.log") + # Use streaming output for real-time monitoring + ret, out, err = run_command_streaming(command, log_file="train.log") if ret != 0: - clean_before_quit() logging.error( "".join( ( "dp train failed\n", "out msg: ", out, "\n", "err msg: ", err, "\n", ) ) ) raise FatalError("dp train failed") - # Note: output is already written to log file by run_command_streaming - # No need to write again here to avoid duplication - # fplog.write("#=================== train std out ===================\n") - # fplog.write(out) - # fplog.write("#=================== train std err ===================\n") - # fplog.write(err) @@ - else: - ret, out, err = run_command(["dp", "freeze", "-o", "frozen_model.pb"]) + else: + # Stream freeze output to the same log for consistency + ret, out, err = run_command_streaming( + ["dp", "freeze", "-o", "frozen_model.pb"], log_file="train.log" + ) if ret != 0: - clean_before_quit() logging.error( "".join( ( "dp freeze failed\n", "out msg: ", out, "\n", "err msg: ", err, "\n", ) ) ) raise FatalError("dp freeze failed") model_file = "frozen_model.pb" - fplog.write("#=================== freeze std out ===================\n") - fplog.write(out) - fplog.write("#=================== freeze std err ===================\n") - fplog.write(err) @@ - clean_before_quit() + clean_before_quit()If you prefer to keep
run_commandfor freeze, alternatively re-open the log right before writing as append (with open("train.log","a", encoding="utf-8") as fplog:) and remove the pre-training open.As a follow-up, you may also want to update
run_command_streamingto open logs with "a" (append) andutf-8encoding, as suggested in my other comment.Also applies to: 296-319, 346-349, 351-351
🧹 Nitpick comments (6)
dpgen2/utils/run_command.py (4)
27-28: Open the log file in append mode with explicit encoding; consider context manager
- Opening with "w" will truncate an existing log file. Using "a" avoids accidental truncation and plays better if multiple phases (train, then freeze) write to the same log.
- Specify encoding to avoid Unicode issues.
- # Open log file if specified - log_fp = open(log_file, "w") if log_file else None + # Open log file if specified + log_fp = open(log_file, "a", encoding="utf-8") if log_file else NoneIf you want to address the SIM115 hint, wrap the whole body in a context manager. Example (minimal pattern, no functional change):
- # Open log file if specified - log_fp = open(log_file, "w") if log_file else None - - try: + # Open log file if specified + log_fp_cm = open(log_file, "a", encoding="utf-8") if log_file else None + try: + log_fp = log_fp_cm # Start subprocess ... - finally: - if log_fp: - log_fp.close() + finally: + if log_fp_cm: + log_fp_cm.close()
31-40: Clean up Popen args: avoid unused stdin=PIPE; drop universal_newlines; set encoding/errors
stdin=subprocess.PIPEis unused and can alter child behavior; inherit parent stdin to allow interactive prompts if any.text=Truealready implies universal newlines;universal_newlines=Trueis redundant.- Set
encoding="utf-8", errors="replace"for robust decoding.- process = subprocess.Popen( + process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - stdin=subprocess.PIPE, + stdin=None, shell=shell, - text=True, - bufsize=1, # Line buffered - universal_newlines=True, + text=True, + bufsize=1, # Line buffered (effective for text mode) + encoding="utf-8", + errors="replace", )
78-78: Avoid unbounded in-memory capture for long-running trainingTraining can run for days; accumulating full stdout/stderr in lists may exhaust memory. Consider capturing only the last N lines/bytes (for error messages) while streaming everything to the log file.
Example approach (conceptual):
- Add parameters like
capture_tail_lines: int | None = 2000.- Keep
collections.deque(maxlen=capture_tail_lines)instead of plain lists.- Return
"\n".join(deque_contents).I can provide a concrete patch if you want this behavior now.
31-40: Optional: pseudo-TTY for better real-time behaviorSome CLIs buffer when stdout is a pipe. If you observe laggy streaming, consider a PTY-based variant (pexpect or pty.openpty) on POSIX to make the child think it’s attached to a TTY.
dpgen2/op/run_dp_train.py (2)
296-299: Limit error payload size in logs to avoid huge log linesOn failure, you concatenate full
out/errinto a single logging record. For multi-GB logs this will be slow and may overwhelm log collectors. Consider truncating to the last N lines/KB.Example:
def _tail(s: str, max_chars=100_000): return s[-max_chars:] if len(s) > max_chars else s logging.error("dp train failed\nout msg: %s\nerr msg: %s\n", _tail(out), _tail(err))
186-194: Use the normalizedconfiginstead ofip["config"]to avoid NoneType issuesYou normalize
config = ip["config"] if ip["config"] is not None else {}(Line 186), but then readimpl = ip["config"].get(...)(Line 187). Ifip["config"]is None, this will raise. Use the normalizedconfigconsistently.- impl = ip["config"].get("impl", "tensorflow") - dp_command = ip["config"].get("command", "dp").split() + impl = config.get("impl", "tensorflow") + dp_command = config.get("command", "dp").split()
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
dpgen2/op/run_dp_train.py(4 hunks)dpgen2/utils/run_command.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
dpgen2/op/run_dp_train.py (1)
dpgen2/utils/run_command.py (1)
run_command_streaming(17-82)
🪛 Ruff (0.12.2)
dpgen2/utils/run_command.py
27-27: Use a context manager for opening files
(SIM115)
| if isinstance(cmd, str): | ||
| cmd = cmd if shell else cmd.split() | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use shlex for robust splitting and handle shell=True with list commands
cmd.split() breaks on quoted args and spaces; also, passing a list while shell=True is unsupported/undefined. Prefer shlex.split and shlex.join to cover both cases safely.
+import shlex
@@
- if isinstance(cmd, str):
- cmd = cmd if shell else cmd.split()
+ if isinstance(cmd, str):
+ cmd = cmd if shell else shlex.split(cmd)
+ elif shell and isinstance(cmd, list):
+ # When shell=True, pass a single string; preserve quoting.
+ cmd = shlex.join(cmd)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if isinstance(cmd, str): | |
| cmd = cmd if shell else cmd.split() | |
| # at the top of dpgen2/utils/run_command.py | |
| import shlex | |
| # … later, around lines 23–25 … | |
| if isinstance(cmd, str): | |
| cmd = cmd if shell else shlex.split(cmd) | |
| elif shell and isinstance(cmd, list): | |
| # When shell=True, pass a single string; preserve quoting. | |
| cmd = shlex.join(cmd) |
🤖 Prompt for AI Agents
In dpgen2/utils/run_command.py around lines 23 to 25, the current logic uses
cmd.split() which mishandles quoted arguments and does not address the undefined
behavior of passing a list when shell=True; replace this with shlex utilities:
if shell is False and cmd is a string, call shlex.split(cmd) to produce a safe
list; if shell is True and cmd is a list, convert it to a string with
shlex.join(cmd) (or if cmd is already a string leave it as-is); ensure types are
normalized accordingly and raise a clear TypeError if an unexpected type is
provided.
| # Store output | ||
| stdout_buffer = [] | ||
| stderr_buffer = [] | ||
|
|
||
| def stream_output(pipe, buffer, is_stderr=False): | ||
| for line in iter(pipe.readline, ""): | ||
| buffer.append(line) | ||
| # Print to terminal | ||
| if is_stderr: | ||
| print(line, end="", file=sys.stderr) | ||
| else: | ||
| print(line, end="") | ||
| # Write to log file | ||
| if log_fp: | ||
| log_fp.write(line) | ||
| log_fp.flush() | ||
| pipe.close() |
There was a problem hiding this comment.
Guard concurrent writes to the same log file
stdout_thread and stderr_thread write to log_fp concurrently. File writes aren’t atomic, so lines can interleave. Use a lock to serialize writes.
# Store output
stdout_buffer = []
stderr_buffer = []
+ log_lock = threading.Lock() if log_fp else None
def stream_output(pipe, buffer, is_stderr=False):
for line in iter(pipe.readline, ""):
buffer.append(line)
# Print to terminal
if is_stderr:
print(line, end="", file=sys.stderr)
else:
print(line, end="")
# Write to log file
if log_fp:
- log_fp.write(line)
- log_fp.flush()
+ with log_lock:
+ log_fp.write(line)
+ log_fp.flush()
pipe.close()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Store output | |
| stdout_buffer = [] | |
| stderr_buffer = [] | |
| def stream_output(pipe, buffer, is_stderr=False): | |
| for line in iter(pipe.readline, ""): | |
| buffer.append(line) | |
| # Print to terminal | |
| if is_stderr: | |
| print(line, end="", file=sys.stderr) | |
| else: | |
| print(line, end="") | |
| # Write to log file | |
| if log_fp: | |
| log_fp.write(line) | |
| log_fp.flush() | |
| pipe.close() | |
| # Store output | |
| stdout_buffer = [] | |
| stderr_buffer = [] | |
| log_lock = threading.Lock() if log_fp else None | |
| def stream_output(pipe, buffer, is_stderr=False): | |
| for line in iter(pipe.readline, ""): | |
| buffer.append(line) | |
| # Print to terminal | |
| if is_stderr: | |
| print(line, end="", file=sys.stderr) | |
| else: | |
| print(line, end="") | |
| # Write to log file | |
| if log_fp: | |
| with log_lock: | |
| log_fp.write(line) | |
| log_fp.flush() | |
| pipe.close() |
🤖 Prompt for AI Agents
In dpgen2/utils/run_command.py around lines 42 to 58, stdout_thread and
stderr_thread both write to the same log_fp concurrently which can interleave
output; create a single threading.Lock (e.g., log_lock) visible to both threads
and wrap any log_fp.write() and log_fp.flush() calls in
log_lock.acquire()/release() or preferably with log_lock as a context manager
(with log_lock: ...), ensuring both threads use that same lock before writing to
the file so writes are serialized and then close the pipe as before.
There was a problem hiding this comment.
The option print_oe of dflow.utils.run_command does exactly the same thing by virtue of selectors. You only need to add an argument to dflow_run_command. Refer to https://github.com/deepmodeling/dflow/blob/48c24cc4f494acb5a12c8d99293f1156c31342ad/src/dflow/utils.py#L657. Besides, I think we should provide an option rather than change the default behavior (which is silent except for error).
Thanks, is there any way to realize this in the dpgen2 configuration json? |
|
It is not supported currently. You can add an argument in config json to control the output. |
For trainning step, there is no output in log files. We can only see the output after the trainning is completed. However, the trainning normally lasts for days, so the streaming output is nessesary for users to examine the status.
The PR suggests one way to achieve this, however, a better way is to realize this directly in pydflow.
Summary by CodeRabbit