This is an automated email from the ASF dual-hosted git repository. yasith pushed a commit to branch cybershuttle-staging in repository https://gitbox.apache.org/repos/asf/airavata.git
commit bfd2774ac8908ea7fad135534c80f7276deba6e7 Author: yasith <[email protected]> AuthorDate: Sat Apr 5 20:55:44 2025 -0400 update agent to accept libraries, pip, and mounts as input. --- modules/agent-framework/airavata-agent/Makefile | 9 ++- modules/agent-framework/airavata-agent/agent.go | 62 ++++++++++++++++--- modules/agent-framework/airavata-agent/agent.sh | 82 +++++++++++++++++++++++++ 3 files changed, 142 insertions(+), 11 deletions(-) diff --git a/modules/agent-framework/airavata-agent/Makefile b/modules/agent-framework/airavata-agent/Makefile index a4f8595f60..2c4fc8b218 100644 --- a/modules/agent-framework/airavata-agent/Makefile +++ b/modules/agent-framework/airavata-agent/Makefile @@ -10,6 +10,8 @@ build-container: run-container: docker run -it $(IMAGE_NAME) /opt/airavata-agent $(AGENT_SVC_URL) $(AGENT_ID) +deploy-all: deploy-anvil-scigap deploy-anvil-waterhub deploy-anvil-gcommunityus + deploy-anvil-scigap: ssh [email protected] srun -p shared --mem 4G -t 30 singularity pull \ --disable-cache \ @@ -29,4 +31,9 @@ deploy-anvil-gcommunityus: --disable-cache \ --force \ cybershuttle/container/remote-agent-base.sif \ - docker://$(IMAGE_NAME) \ No newline at end of file + docker://$(IMAGE_NAME) + +deploy-bootstrap-anvil: + scp agent.sh [email protected]:~/cybershuttle/agent.sh && \ + scp agent.sh [email protected]:~/cybershuttle/agent.sh && \ + scp agent.sh [email protected]:~/cybershuttle/agent.sh diff --git a/modules/agent-framework/airavata-agent/agent.go b/modules/agent-framework/airavata-agent/agent.go index 1519d92454..bd64d08ae9 100644 --- a/modules/agent-framework/airavata-agent/agent.go +++ b/modules/agent-framework/airavata-agent/agent.go @@ -6,6 +6,7 @@ import ( "bytes" "context" "encoding/json" + "flag" "fmt" "io" "log" @@ -24,19 +25,33 @@ import ( type Stream = grpc.BidiStreamingClient[protos.AgentMessage, protos.ServerMessage] +var envName = "base" var pidMap = make(map[string]int) func main() { - // get CLI args - serverUrl := os.Args[1] - agentId := os.Args[2] + // Define flags with default empty values. + serverUrl := flag.String("server", "", "Server flag (optional)") + agentId := flag.String("agent", "", "Agent flag (optional)") + lib := flag.String("lib", "", "Libraries flag (optional)") + pip := flag.String("pip", "", "Pip flag (optional)") - conn, err := grpc.NewClient(serverUrl, grpc.WithTransportCredentials(insecure.NewCredentials())) + // Parse the flags provided by the user. + flag.Parse() + + // Validate required flags + if *serverUrl == "" { + log.Fatalf("[agent.go] main() Error: --server flag is required.\n") + } + if *agentId == "" { + log.Fatalf("[agent.go] main() Error: --agent flag is required.\n") + } + + conn, err := grpc.NewClient(*serverUrl, grpc.WithTransportCredentials(insecure.NewCredentials())) if err != nil { - log.Fatalf("[agent.go] main() Did not connect to %s: %v\n", serverUrl, err) + log.Fatalf("[agent.go] main() Did not connect to %s: %v\n", *serverUrl, err) } - log.Printf("[agent.go] main() Connected to %s\n", serverUrl) + log.Printf("[agent.go] main() Connected to %s\n", *serverUrl) defer conn.Close() c := protos.NewAgentCommunicationServiceClient(conn) @@ -46,11 +61,40 @@ func main() { } log.Printf("[agent.go] main() Created stream...\n") - log.Printf("[agent.go] main() Trying to connect to %s with agent id %s\n", serverUrl, agentId) + var libList []string + if strings.TrimSpace(*lib) != "" { + libList = strings.Split(*lib, ",") + log.Printf("[agent.go] main() Installing --lib: %v\n", libList) + libCmd := exec.Command("micromamba", "install", "-n", envName, "--yes") + libCmd.Args = append(libCmd.Args, libList...) + if err := libCmd.Run(); err != nil { + log.Fatalf("[agent.go] main() Error Installing --lib: %v\n", err) + } + log.Printf("[agent.go] main() Installed --lib: %v\n", libList) + } else { + log.Printf("[agent.go] main() No --lib to install.\n") + } + + var pipList []string + if strings.TrimSpace(*pip) != "" { + pipList = strings.Split(*pip, ",") + log.Printf("[agent.go] main() Installing --pip: %v\n", pipList) + pipCmd := exec.Command("micromamba", "run", "-n", envName, "pip", "install") + pipCmd.Args = append(pipCmd.Args, pipList...) + if err := pipCmd.Run(); err != nil { + log.Fatalf("[agent.go] main() Error Installing --pip: %v\n", err) + } + log.Printf("[agent.go] main() Installed --pip: %v\n", pipList) + + } else { + log.Printf("[agent.go] main() No --pip to install.\n") + } + + log.Printf("[agent.go] main() Trying to connect to %s with agent id %s\n", *serverUrl, *agentId) msg := &protos.AgentMessage{ Message: &protos.AgentMessage_AgentPing{ AgentPing: &protos.AgentPing{ - AgentId: agentId, + AgentId: *agentId, }, }, } @@ -73,8 +117,6 @@ func main() { func startInterceptor(stream Stream, grpcStreamChannel chan struct{}) { - envName := "base" - for { in, err := stream.Recv() if err == io.EOF { diff --git a/modules/agent-framework/airavata-agent/agent.sh b/modules/agent-framework/airavata-agent/agent.sh new file mode 100644 index 0000000000..8a9db1f91e --- /dev/null +++ b/modules/agent-framework/airavata-agent/agent.sh @@ -0,0 +1,82 @@ +#!/bin/bash -x + +# ##################################################################### +# Airavata Agent Bootstrapper +# ##################################################################### +# +# ---------------------------------------------------------------------- +# CONTRIBUTORS +# ---------------------------------------------------------------------- +# * Dimuthu Wannipurage +# * Lahiru Jayathilake +# * Yasith Jayawardana +# ###################################################################### + +#----------------------------------------------------------------------- +# STEP 1 - PARSE COMMAND LINE ARGS +#----------------------------------------------------------------------- + +# Default values +CS_HOME=$HOME/cybershuttle +WORKDIR=$CS_HOME/workspace +AGENT="" +SERVER="" +CONTAINER="" +LIBRARIES="" +PIP="" +BIND_OPTS=() + +PARSED_OPTIONS=$(getopt -o '' --long server:,agent:,container:,bind: -n "$0" -- "$@") +if [ $? -ne 0 ]; then + echo "Usage: $0 \ + --server SERVER \ + --agent AGENT \ + --container CONTAINER \ + --libraries LIBRARIES \ + --pip PIP \ + --mounts MOUNTS \ + [--bind BIND] ..." + exit 1 +fi +eval set -- "$PARSED_OPTIONS" + +while true; do + case "$1" in + --server) SERVER="$2"; shift 2 ;; + --agent) AGENT="$2"; shift 2 ;; + --container) CONTAINER="$2"; shift 2 ;; + --libraries) LIBRARIES="$2"; shift 2 ;; + --pip) PIP="$2"; shift 2 ;; + --mounts) + IFS=',' read -ra MOUNTS <<< "$2" + for MOUNT in "${MOUNTS[@]}"; do + BIND_OPTS+=("--bind $MOUNT:ro") + done + shift 2 ;; + --bind) BIND_OPTS+=("--bind $2:ro"); shift 2 ;; + --) shift; break ;; + *) echo "Unexpected option: $1"; exit 1 ;; + esac +done + +# Final values +echo "CS_HOME=$CS_HOME" +echo "WORKDIR=$WORKDIR" +echo "AGENT=$AGENT" +echo "SERVER=$SERVER" +echo "CONTAINER=$CONTAINER" +echo "LIBRARIES=$LIBRARIES" +echo "PIP=$PIP" +echo "BIND_OPTS=$BIND_OPTS" + +# ---------------------------------------------------------------------- +# STEP 2 - RUN AGENT +# ---------------------------------------------------------------------- +KERNEL_SOCK=$(mktemp) singularity exec \ + --bind $WORKDIR:/scratch:ro ${BIND_OPTS[@]} \ + $CS_HOME/container/$CONTAINER /opt/airavata-agent \ + --server $SERVER:19900 \ + --agent $AGENT \ + --lib $LIBRARIES \ + --pip $PIP \ + $@
