openssh v0.9.0-rc1 is released with a new feature native-mux.

It provides new functions like Session::connect_mux and SessionBuilder::connect_mux which communicates with the ssh multiplex master directly, through the control socket, instead of spawning a new process to communicate with it.

The advantage of this is more robust error reporting, better performance and less memory usage.

The old implementation (process-mux) checks the exit status of ssh for indication of error, then parse the output of it and the output of the ssh multiplex master to return an error.

This method is obviously not so robust as native-mux, which directly communicates with ssh multiplex master through its multiplex protocol.

The better performance and less memory usage part is mostly because we avoid creating a new process for every command you spawned on remote, every Session::check and every Session::request_port_forwarding.

The new release also add new function Session::request_port_forwarding, which supports local/remote forwarding of tcp and unix socket stream.

There are also other changes to API:

• A new type Stdio is used for setting stdin/stdout/stderr.
• ChildStd* types are now alias for tokio_pipe::{PipeRead, PipeWrite}.
• Command::spawn and Command::status now confirms to std::process::Command and tokio::process::Command, in which stdin, stdout and stderr are inherit by default.
• Command::spawn is now an async method.
• RemoteChild::wait now takes self by value.
• Error is now marked #[non_exhaustive] and new variants is added.

I know this is a huge release and upgrading it is going to be quite difficult, but I sincerely want you to try it out, as the new implementation requires feedback.

This PR attempts to gracefully shutdown SSH sessions when the user presses ctrl-c.

The overall flow of cancellation propagation is:

1. The user presses ctrl-c.
2. The handler closure (set by ctrlc::set_handler) assigns true to the variable cancelled: Arc<Mutex<bool>>.
3. The scheduling loop checks the value of cancelled at the beginning of every iteration. When true, it breaks.
4. All channel handles that are used to communicate with the SSH session tasks are explicitly droped.
5. For each SSH session task, whenever it sends or recvs from any of its channels, the channel will return an Err and the task will break out of the task execution loop.
6. For each SSH session, the session object will be dropped, terminating the SSH session.

As a side note, along the way, this PR fixes stream in Session by explicitly locking stdout. Without this, multiple calls to print! and println! are not coalesced. This lead to output lines from different commands mixing with each other.

Closes #3.

@NobodyXu Will you be interested in reviewing? Just asking since you have reviewed a lot of my code recently (and it benefited both me and Pegasus so much). Plus, I can add you as a collaborator if you'd like :)

Currently, Pegasus doesn't handle ctrl-c very well. When the user presses ctrl-c, Pegasus terminates without cancelling the SSH session tasks. So the ssh sessions remain open, .ssh-connection* directories stay as is, and in order to kill the processes spawned on remote nodes, I need to walk into each node and kill them manually.

Potential solutions

• One method would be to propagate the signal to all child ssh processes. Proper error handling inside the tasks would probably serve as an okay solution.
• Another method I'd like to explore is to see if using the native-mux feature of the coming openssh crate will make graceful shutdown any easier.

Fixed #1

Enabling feature native-mux requires once_cell to be bumped, so I bumped the Cargo.lock using cargo update.

I also disabled default-feature process-mux of openssh, which then requires me to modify src/session.rs to use the new SSHSession::connect_mux API.

openssh v0.9.0 also

• use openssh::Stdio instead of std::process::Stdio in Command::{stdin, stdout, stderr}.
• updates API of Command::spawn to an async method.

Signed-off-by: Jiahao XU [email protected]

Avoid pulling feature io-std, net, signal and fs, which added bloat to the binary that are completely unused.

Signed-off-by: Jiahao XU [email protected]

# queue.yaml
- command:
    - long_job {{ param }}
    - another_long_job {{ param }}
  param:
    - resnet50
    - ViT

Currently, Pegasus will consume the entire entry (virtually the entire file) from queue.yaml. However, consuming minimally would be nice. For instance, leaving queue.yaml in the following state would be nice:

# queue.yaml
- command:
    - another_long_job {{ param }}
   param:
    - resnet50
    - ViT

Then in the next round of get_one_job, only then queue.yaml will be empty.

Currently, given that daemon mode is not set, the scheduling loop of both broadcast mode and queue mode just terminate after sending out the last command to a free SSH session. But the time gap between sending out the last command and the final command finishing execution is vast.

In this time window, users will want to submit more jobs. Obviously since Pegasus seems to be up and running, it is more intuitive for users to assume that it will admit more jobs.

Pressing ctrl-c when Pegasus is running will also send ctrl-c to the entire foreground process group, thus also sending SIGINT to the ssh processes.

Instead of waiting for ctrl-c with the tokio ctcl-c catcher, find another way to catch the user's intent to cancel.

Possible solutions:

• Listen to stdin for something random. Like q<Enter>.
• Create a pegasus stop command. This should make sure to identify the specific pegasus process that is running on the current working directory. This will probably require a pid file.
  • It's okay to assume that there is only one instance of pegasus per directory (if not, that means more than one pegasus processes are manipulating queue.yaml and consumed.yaml).

For a large parametrized command that generates a lot of jobs and takes days, the user will want to figure out the progress of each command generated: queued, running, or done.

Internal bookkeeping is easy. Maybe keep a file that we dump the commands in progress, and add a mode in pegasus that parses and displays that file.

Use case:

• Add a node to hosts.yaml when you get access to more nodes, and the jobs you are currently running will automatically go there.
• Remove a node from hosts.yaml when someone asks you to vacate the node.

Removing a node from hosts.yaml does not kill the job. However, it is guaranteed that the removed node will not be given a new job once it's been removed from hosts.yaml.

Before entering the schedule loop, Pegasus should create and connect the SSH session to see if it connects. If not, it should gracefully terminate all previous connections and exit. The connected SSH session object can be moved into the tokio task.

Cancelling commands ran by Pegasus is very difficult. You essentially have to ssh into each node and manually figure out the PIDs of commands and kill them.

Nested commands, so to say, make things more complicated. For instance, docker exec sh -c "python train.py" will run the following commands:

• Ran by user: sh -c docker exec sh -c "python train.py"
• Ran by user: docker exec sh -c "python train.py"
• Ran by root:sh -c "python train.py"
• Ran by root: python train.py

Only killing the fourth python train.py command will truely achieve cancellation. The bottom line is, it is difficult for Pegasus to infer how to properly terminate a command.

Potential solutions

• We might ask the user for a cancellation command in queue.yaml. For example, sudo kill $(pgrep -f 'train.py'). Then the ctrl_c handler will create a new connection to the hosts and run the designated cancellation command.
• Somehow figure out the PGID of the sh process and run sudo kill -- -PGID. Can we pgrep -f with the entire command? Shell escaping might become a problem. (pgrep -f with every single word in the command and kill the intersection of all PIDs returned?)