1 // Copyright (c) 2017-2023, Lawrence Livermore National Security, LLC and other CEED contributors. 2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details. 3 // 4 // SPDX-License-Identifier: BSD-2-Clause 5 // 6 // This file is part of CEED: http://github.com/ceed 7 // Based on the instructions from https://www.craylabs.org/docs/sr_integration.html and PHASTA implementation 8 9 #include "../../include/smartsim.h" 10 11 #include "../../navierstokes.h" 12 13 PetscErrorCode SmartRedisVerifyPutTensor(void *c_client, const char *name, const size_t name_length) { 14 bool does_exist = true; 15 16 PetscFunctionBeginUser; 17 SmartRedisCall(tensor_exists(c_client, name, name_length, &does_exist)); 18 PetscCheck(does_exist, PETSC_COMM_SELF, -1, "Tensor of name '%s' was not written to the database successfully", name); 19 PetscFunctionReturn(PETSC_SUCCESS); 20 } 21 22 PetscErrorCode SmartSimTrainingSetup(User user) { 23 SmartSimData smartsim = user->smartsim; 24 PetscMPIInt rank; 25 PetscReal checkrun[2] = {1}; 26 size_t dim_2[1] = {2}; 27 PetscInt num_ranks; 28 29 PetscFunctionBeginUser; 30 PetscCallMPI(MPI_Comm_rank(user->comm, &rank)); 31 PetscCallMPI(MPI_Comm_size(user->comm, &num_ranks)); 32 33 if (rank % smartsim->collocated_database_num_ranks == 0) { 34 // -- Send array that communicates when ML is done training 35 SmartRedisCall(put_tensor(smartsim->client, "check-run", 9, checkrun, dim_2, 1, SRTensorTypeDouble, SRMemLayoutContiguous)); 36 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, "check-run", 9)); 37 } 38 39 { // -- Get minimum per-rank global vec size 40 PetscInt GlobalVecSize; 41 PetscCall(DMGetGlobalVectorInfo(user->dm, &GlobalVecSize, NULL, NULL)); 42 PetscCallMPI(MPI_Allreduce(&GlobalVecSize, &smartsim->num_tensor_nodes, 1, MPIU_INT, MPI_MIN, user->comm)); 43 smartsim->num_nodes_to_remove = GlobalVecSize - smartsim->num_tensor_nodes; 44 } 45 46 // Determine the size of the training data arrays... somehow 47 if (rank % smartsim->collocated_database_num_ranks == 0) { 48 size_t array_dims[2] = {smartsim->num_tensor_nodes, 6}, array_info_dim = 6; 49 PetscInt array_info[6] = {0}, num_features = 6; 50 51 array_info[0] = array_dims[0]; 52 array_info[1] = array_dims[1]; 53 array_info[2] = num_features; 54 array_info[3] = num_ranks; 55 array_info[4] = smartsim->collocated_database_num_ranks; 56 array_info[5] = rank; 57 58 SmartRedisCall(put_tensor(smartsim->client, "array_info", 10, array_info, &array_info_dim, 1, SRTensorTypeInt32, SRMemLayoutContiguous)); 59 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, "array_info", 10)); 60 } 61 PetscFunctionReturn(0); 62 } 63 64 PetscErrorCode SmartSimSetup(User user) { 65 PetscMPIInt rank; 66 size_t rank_id_name_len; 67 PetscInt num_orchestrator_nodes = 1; 68 69 PetscFunctionBeginUser; 70 PetscCall(PetscNew(&user->smartsim)); 71 SmartSimData smartsim = user->smartsim; 72 73 smartsim->collocated_database_num_ranks = 1; 74 PetscOptionsBegin(user->comm, NULL, "Options for SmartSim integration", NULL); 75 PetscCall(PetscOptionsInt("-smartsim_collocated_database_num_ranks", "Number of ranks per collocated database instance", NULL, 76 smartsim->collocated_database_num_ranks, &smartsim->collocated_database_num_ranks, NULL)); 77 PetscOptionsEnd(); 78 79 PetscCall(PetscStrlen(smartsim->rank_id_name, &rank_id_name_len)); 80 // Create prefix to be put on tensor names 81 PetscCallMPI(MPI_Comm_rank(user->comm, &rank)); 82 PetscCall(PetscSNPrintf(smartsim->rank_id_name, sizeof smartsim->rank_id_name, "y.%d", rank)); 83 84 SmartRedisCall(SmartRedisCClient(num_orchestrator_nodes != 1, smartsim->rank_id_name, rank_id_name_len, &smartsim->client)); 85 86 PetscCall(SmartSimTrainingSetup(user)); 87 PetscFunctionReturn(PETSC_SUCCESS); 88 } 89