This is an automated email from the ASF dual-hosted git repository. lahirujayathilake pushed a commit to branch vizfold in repository https://gitbox.apache.org/repos/asf/airavata.git
commit 1e070f3c23054ad36e2cb0807eb920c3afb18286 Author: lahiruj <[email protected]> AuthorDate: Sat Oct 4 23:54:18 2025 -0400 initial data models for experiment workloads --- .../init/07-parameter-sweep-modeling.sql | 32 +++++ .../service/db/entity/JobBatchEntity.java | 108 ++++++++++++++++ .../service/db/entity/JobUnitEntity.java | 141 +++++++++++++++++++++ .../connection/service/db/repo/JobBatchRepo.java | 27 ++++ .../connection/service/db/repo/JobUnitRepo.java | 28 ++++ .../service/handlers/AgentManagementHandler.java | 2 +- .../service/models/AgentLaunchRequest.java | 10 ++ .../connection/service/models/JobBatchSpec.java | 60 +++++++++ .../connection/service/models/JobUnitStatus.java | 23 ++++ 9 files changed, 430 insertions(+), 1 deletion(-) diff --git a/.devcontainer/database_scripts/init/07-parameter-sweep-modeling.sql b/.devcontainer/database_scripts/init/07-parameter-sweep-modeling.sql new file mode 100644 index 0000000000..8bbe2e0372 --- /dev/null +++ b/.devcontainer/database_scripts/init/07-parameter-sweep-modeling.sql @@ -0,0 +1,32 @@ +USE `app_catalog`; + +-- Parent collection for one parameter sweep / batch +CREATE TABLE IF NOT EXISTS `JOB_BATCH` +( + `ID` VARCHAR(255) NOT NULL, + `EXPERIMENT_ID` VARCHAR(255) NOT NULL, + `CREATED_AT` TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + `PAYLOAD_JSON` JSON NULL, -- original request payload + `COMMAND_TEMPLATE` TEXT NOT NULL, -- application_command (template) + PRIMARY KEY (`ID`), + KEY `IDX_BATCH_EXPERIMENT` (`EXPERIMENT_ID`) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci; + +-- Units of work (one per parameter combination) +CREATE TABLE IF NOT EXISTS `JOB_UNIT` +( + `ID` VARCHAR(255) NOT NULL, + `BATCH_ID` VARCHAR(255) NOT NULL, + `EXPERIMENT_ID` VARCHAR(255) NOT NULL, + `CREATED_AT` TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + `RESOLVED_COMMAND` TEXT NOT NULL, -- fully expanded command to run + `STATUS` ENUM('PENDING','IN_PROGRESS','COMPLETED','FAILED') NOT NULL DEFAULT 'PENDING', + `AGENT_ID` VARCHAR(255) NULL, + `STARTED_AT` TIMESTAMP(6) NULL, + `COMPLETED_AT` TIMESTAMP(6) NULL, + PRIMARY KEY (`ID`), + KEY `IDX_UNIT_BATCH_STATUS_FIFO` (`BATCH_ID`, `STATUS`, `CREATED_AT`, `ID`), + KEY `IDX_UNIT_EXP_STATUS` (`EXPERIMENT_ID`, `STATUS`), + CONSTRAINT `FK_JOB_UNIT_BATCH` FOREIGN KEY (`BATCH_ID`) + REFERENCES `JOB_BATCH` (`ID`) ON DELETE CASCADE +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci; \ No newline at end of file diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/entity/JobBatchEntity.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/entity/JobBatchEntity.java new file mode 100644 index 0000000000..00c8f728f9 --- /dev/null +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/entity/JobBatchEntity.java @@ -0,0 +1,108 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.airavata.agent.connection.service.db.entity; + +import com.fasterxml.jackson.databind.JsonNode; +import jakarta.persistence.CascadeType; +import jakarta.persistence.Column; +import jakarta.persistence.Entity; +import jakarta.persistence.FetchType; +import jakarta.persistence.Id; +import jakarta.persistence.Lob; +import jakarta.persistence.OneToMany; +import jakarta.persistence.Table; +import org.hibernate.annotations.JdbcTypeCode; +import org.hibernate.type.SqlTypes; + +import java.time.Instant; +import java.util.List; + +@Entity +@Table(name = "JOB_BATCH") +public class JobBatchEntity { + + @Id + @Column(name = "ID", nullable = false) + private String id; + + @Column(name = "EXPERIMENT_ID", nullable = false) + private String experimentId; + + @Column(name = "CREATED_AT", updatable = false, insertable = false) + private Instant createdAt; + + @JdbcTypeCode(SqlTypes.JSON) + @Column(name = "PAYLOAD_JSON", columnDefinition = "json") + private JsonNode payloadJson; + + @Lob + @Column(name = "COMMAND_TEMPLATE", nullable = false) + private String commandTemplate; + + @OneToMany(mappedBy = "batch", fetch = FetchType.LAZY, cascade = CascadeType.REMOVE) + private List<JobUnitEntity> units; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getExperimentId() { + return experimentId; + } + + public void setExperimentId(String experimentId) { + this.experimentId = experimentId; + } + + public Instant getCreatedAt() { + return createdAt; + } + + public void setCreatedAt(Instant createdAt) { + this.createdAt = createdAt; + } + + public JsonNode getPayloadJson() { + return payloadJson; + } + + public void setPayloadJson(JsonNode payloadJson) { + this.payloadJson = payloadJson; + } + + public String getCommandTemplate() { + return commandTemplate; + } + + public void setCommandTemplate(String commandTemplate) { + this.commandTemplate = commandTemplate; + } + + public List<JobUnitEntity> getUnits() { + return units; + } + + public void setUnits(List<JobUnitEntity> units) { + this.units = units; + } +} diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/entity/JobUnitEntity.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/entity/JobUnitEntity.java new file mode 100644 index 0000000000..1d379abccd --- /dev/null +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/entity/JobUnitEntity.java @@ -0,0 +1,141 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.airavata.agent.connection.service.db.entity; + +import jakarta.persistence.Column; +import jakarta.persistence.Entity; +import jakarta.persistence.EnumType; +import jakarta.persistence.Enumerated; +import jakarta.persistence.FetchType; +import jakarta.persistence.Id; +import jakarta.persistence.JoinColumn; +import jakarta.persistence.Lob; +import jakarta.persistence.ManyToOne; +import jakarta.persistence.Table; +import org.apache.airavata.agent.connection.service.models.JobUnitStatus; + +import java.time.Instant; + +@Entity +@Table(name = "JOB_UNIT") +public class JobUnitEntity { + + @Id + @Column(name = "ID", nullable = false) + private String id; + + @ManyToOne(fetch = FetchType.LAZY, optional = false) + @JoinColumn(name = "BATCH_ID", nullable = false) + private JobBatchEntity batch; + + @Column(name = "EXPERIMENT_ID", nullable = false) + private String experimentId; + + @Column(name = "CREATED_AT", updatable = false, insertable = false) + private Instant createdAt; + + @Lob + @Column(name = "RESOLVED_COMMAND", nullable = false) + private String resolvedCommand; + + @Enumerated(EnumType.STRING) + @Column(name = "STATUS", length = 16, nullable = false) + private JobUnitStatus status = JobUnitStatus.PENDING; + + @Column(name = "AGENT_ID") + private String agentId; + + @Column(name = "STARTED_AT") + private Instant startedAt; + + @Column(name = "COMPLETED_AT") + private Instant completedAt; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public JobBatchEntity getBatch() { + return batch; + } + + public void setBatch(JobBatchEntity batch) { + this.batch = batch; + } + + public String getExperimentId() { + return experimentId; + } + + public void setExperimentId(String experimentId) { + this.experimentId = experimentId; + } + + public Instant getCreatedAt() { + return createdAt; + } + + public void setCreatedAt(Instant createdAt) { + this.createdAt = createdAt; + } + + public String getResolvedCommand() { + return resolvedCommand; + } + + public void setResolvedCommand(String resolvedCommand) { + this.resolvedCommand = resolvedCommand; + } + + public JobUnitStatus getStatus() { + return status; + } + + public void setStatus(JobUnitStatus status) { + this.status = status; + } + + public String getAgentId() { + return agentId; + } + + public void setAgentId(String agentId) { + this.agentId = agentId; + } + + public Instant getStartedAt() { + return startedAt; + } + + public void setStartedAt(Instant startedAt) { + this.startedAt = startedAt; + } + + public Instant getCompletedAt() { + return completedAt; + } + + public void setCompletedAt(Instant completedAt) { + this.completedAt = completedAt; + } +} diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/repo/JobBatchRepo.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/repo/JobBatchRepo.java new file mode 100644 index 0000000000..719394f40c --- /dev/null +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/repo/JobBatchRepo.java @@ -0,0 +1,27 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.airavata.agent.connection.service.db.repo; + +import org.apache.airavata.agent.connection.service.db.entity.JobBatchEntity; +import org.springframework.data.jpa.repository.JpaRepository; +import org.springframework.stereotype.Repository; + +@Repository +public interface JobBatchRepo extends JpaRepository<JobBatchEntity, String> { +} diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/repo/JobUnitRepo.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/repo/JobUnitRepo.java new file mode 100644 index 0000000000..02f9261352 --- /dev/null +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/db/repo/JobUnitRepo.java @@ -0,0 +1,28 @@ + +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.airavata.agent.connection.service.db.repo; + +import org.apache.airavata.agent.connection.service.db.entity.JobUnitEntity; +import org.springframework.data.jpa.repository.JpaRepository; +import org.springframework.stereotype.Repository; + +@Repository +public interface JobUnitRepo extends JpaRepository<JobUnitEntity, String> { +} diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/handlers/AgentManagementHandler.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/handlers/AgentManagementHandler.java index 409511c0db..5cdb5a179f 100644 --- a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/handlers/AgentManagementHandler.java +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/handlers/AgentManagementHandler.java @@ -158,7 +158,7 @@ public class AgentManagementHandler { public AgentLaunchResponse createAndLaunchExperiment(AgentLaunchRequest req) { try { - String agentId = "agent_" + UUID.randomUUID().toString(); + String agentId = "agent_" + UUID.randomUUID(); String envName = generateEnvName(req.getLibraries(), req.getPip()); LOGGER.info("Creating an Airavata Experiment for {} with agent id {}", req.getExperimentName(), agentId); ExperimentModel experiment = generateExperiment(req, agentId, envName); diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/AgentLaunchRequest.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/AgentLaunchRequest.java index 81b6c26631..a034d601fa 100644 --- a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/AgentLaunchRequest.java +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/AgentLaunchRequest.java @@ -39,6 +39,8 @@ public class AgentLaunchRequest { private int nodeCount = 1; private int memory = 2048; + private JobBatchSpec jobBatchSpec; + public String getExperimentName() { return experimentName; } @@ -138,4 +140,12 @@ public class AgentLaunchRequest { public void setMounts(List<String> mounts) { this.mounts = mounts; } + + public JobBatchSpec getJobBatchSpec() { + return jobBatchSpec; + } + + public void setJobBatchSpec(JobBatchSpec jobBatchSpec) { + this.jobBatchSpec = jobBatchSpec; + } } diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/JobBatchSpec.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/JobBatchSpec.java new file mode 100644 index 0000000000..bb2c3fbb87 --- /dev/null +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/JobBatchSpec.java @@ -0,0 +1,60 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.airavata.agent.connection.service.models; + +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.List; +import java.util.Map; + +public class JobBatchSpec { + + @JsonProperty("application_command") + private String applicationCommand; + + @JsonProperty("parameter_grid") + private Map<String, List<String>> parameterGrid; + + @JsonProperty("input_files") + private List<String> inputFiles; + + public String getApplicationCommand() { + return applicationCommand; + } + + public void setApplicationCommand(String applicationCommand) { + this.applicationCommand = applicationCommand; + } + + public Map<String, List<String>> getParameterGrid() { + return parameterGrid; + } + + public void setParameterGrid(Map<String, List<String>> parameterGrid) { + this.parameterGrid = parameterGrid; + } + + public List<String> getInputFiles() { + return inputFiles; + } + + public void setInputFiles(List<String> inputFiles) { + this.inputFiles = inputFiles; + } +} \ No newline at end of file diff --git a/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/JobUnitStatus.java b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/JobUnitStatus.java new file mode 100644 index 0000000000..670978fe57 --- /dev/null +++ b/modules/agent-framework/agent-service/src/main/java/org/apache/airavata/agent/connection/service/models/JobUnitStatus.java @@ -0,0 +1,23 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * <p> + * http://www.apache.org/licenses/LICENSE-2.0 + * <p> + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.airavata.agent.connection.service.models; + +public enum JobUnitStatus { + PENDING, IN_PROGRESS, COMPLETED, FAILED +}
