Skip to content

Compss+Hecuba #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cTatu wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Compss+Hecuba #9

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
20fca02
Integrate COMPSS+Hecuba for SOD2D
cTatu Nov 18, 2025
f09f394
Fix step type assignment and add initialization check
cTatu Nov 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 107 additions & 42 deletions src/smartsod2d/sod_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import shutil
import random
import numpy as np
import subprocess

from tf_agents.specs import array_spec
from tf_agents.environments import py_environment
Expand All @@ -16,8 +17,46 @@
from smartsim.log import get_logger
from smartsod2d.utils import n_witness_points, n_rectangles, numpy_str, bcolors

from pycompss.api.constraint import constraint
from pycompss.api.task import task
from pycompss.api.mpi import mpi
from pycompss.api.api import compss_wait_on
from pycompss.api.on_failure import on_failure

logger = get_logger(__name__)

from hecuba import StorageStream, StorageNumpy

class mynumpyclass(StorageNumpy,StorageStream):
'''
pass
'''


@constraint(processors=[
{"processorType": "CPU", "computingUnits": "1"},
{"processorType": "GPU", "computingUnits": "1"}
])
@mpi(
working_dir="/gpfs/scratch/bsc19/bsc19959/ceec/Cristian",
binary="sod2d",
runner="mpirun",
processes=4,
flags="-x SSDB={{db_address}} -x SR_DB_TYPE=Standalone", #-x COMPSS_BINDED_GPUS
args="{{json_name}} --restart_step={{restart_step}} --f_action={{f_action}} --t_episode={{t_episode}} --t_begin_control={{t_begin_control}}"
)
@on_failure(management="IGNORE")
@task()
def compss_sod2d(db_address, json_name, restart_step, f_action, t_episode, t_begin_control):
pass


@task(is_replicated=True)
def cancel_sod2d():
subprocess.run(['pkill', 'sod2d'])



class SodEnvBase(py_environment.PyEnvironment):
"""
SOD2D environment(s) extending the tf_agents.environments.PyEnvironment class.
Expand Down Expand Up @@ -258,15 +297,20 @@ def _generate_mpmd_settings(self, restart_file):
# set SOD2D exe arguments
sod_args = {"restart_step": restart_step, "f_action": self.f_action,
"t_episode": self.t_episode, "t_begin_control": self.t_begin_control}

for i in range(self.cfd_n_envs):
exe_args = [self.sod_JSON] + [f"--{k}={v[i]}" for k,v in sod_args.items()]
run = MpirunSettings(exe=self.sod_exe, exe_args=exe_args)
compss_sod2d(self.db_address, self.sod_JSON, restart_step[i], self.f_action[i], self.t_episode[i], self.t_begin_control[i])
# exe_args = [self.sod_JSON] + [f"--{k}={v[i]}" for k,v in sod_args.items()]
# run = MpirunSettings(exe=self.sod_exe, exe_args=exe_args)
run = MpirunSettings(exe='hostname')
run.set_tasks(self.n_tasks_per_env)

if i == 0:
f_mpmd = run
else:
f_mpmd.make_mpmd(run)
return f_mpmd

return run


def _create_mpmd_ensemble(self, restart_file):
Expand All @@ -287,8 +331,11 @@ def _create_mpmd_ensemble(self, restart_file):

def _get_n_state(self):
try:
self.client.poll_tensor(self.state_size_key, 100, self.poll_time) # poll every 100ms for 1000 times (100s in total)
state_size = self.client.get_tensor(self.state_size_key)
# self.client.poll_tensor(self.state_size_key, 100, self.poll_time) # poll every 100ms for 1000 times (100s in total)
# state_size = self.client.get_tensor(self.state_size_key)
state_size_hecuba = mynumpyclass.get_by_alias(self.state_size_key.replace('.', '_'))
state_size_hecuba.poll()
state_size = state_size_hecuba.astype(np.int32)
except Exception as exc:
raise Warning(f"Could not read state size from key: {self.state_size_key}") from exc
logger.debug(f"Read state_size: {state_size}")
Expand All @@ -297,8 +344,11 @@ def _get_n_state(self):

def _get_n_action(self):
try:
self.client.poll_tensor(self.action_size_key, 100, self.poll_time)
action_size = self.client.get_tensor(self.action_size_key)
# self.client.poll_tensor(self.action_size_key, 100, self.poll_time)
# action_size = self.client.get_tensor(self.action_size_key)
action_size_hecuba = mynumpyclass.get_by_alias(self.action_size_key.replace('.', '_'))
action_size_hecuba.poll()
action_size = action_size_hecuba.astype(np.int32)
logger.debug(f"Read action_size: {action_size}")
except Exception as exc:
raise Warning(f"Could not read action size from key: {self.action_size_key}") from exc
Expand All @@ -311,11 +361,15 @@ def _get_state(self):
"""
for i in range(self.cfd_n_envs):
if self._step_type[i] > 0: # environment still running
self.client.poll_tensor(self.state_key[i], 100, self.poll_time)
# self.client.poll_tensor(self.state_key[i], 100, self.poll_time)
try:
self._state[i, :] = self.client.get_tensor(self.state_key[i])
self.client.delete_tensor(self.state_key[i])
logger.debug(f"[Env {i}] Got state: {numpy_str(self._state[i, :5])}")
logger.debug(f"[Env {i}] Trying to read state from key: {self.state_key[i].replace('.', '_')}")
state_hecuba = mynumpyclass.get_by_alias(self.state_key[i].replace('.', '_'))
state_hecuba.poll()
self._state[i, :] = state_hecuba.astype(self._state.dtype)
# self._state[i, :] = self.client.get_tensor(self.state_key[i])
# self.client.delete_tensor(self.state_key[i])
logger.debug(f"[Env {i}] Got state: {numpy_str(self._state[i, :5])} dtype: {self._state.dtype} shape {self._state.shape}")
except Exception as exc:
raise Warning(f"Could not read state from key: {self.state_key[i]}") from exc

Expand All @@ -337,15 +391,15 @@ def _get_reward(self):
raise NotImplementedError


def _get_time(self):
for i in range(self.cfd_n_envs):
self.client.poll_tensor(self.time_key[i], 100, self.poll_time)
try:
self._time[i] = self.client.get_tensor(self.time_key[i])[0]
self.client.delete_tensor(self.time_key[i])
logger.debug(f"[Env {i}] Got time: {numpy_str(self._time[i])}")
except Exception as exc:
raise Warning(f"Could not read time from key: {self.time_key[i]}") from exc
# def _get_time(self):
# for i in range(self.cfd_n_envs):
# self.client.poll_tensor(self.time_key[i], 100, self.poll_time)
# try:
# self._time[i] = self.client.get_tensor(self.time_key[i])[0]
# self.client.delete_tensor(self.time_key[i])
# logger.debug(f"[Env {i}] Got time: {numpy_str(self._time[i])}")
# except Exception as exc:
# raise Warning(f"Could not read time from key: {self.time_key[i]}") from exc


def get_status(self):
Expand All @@ -358,25 +412,35 @@ def get_status(self):
2: Running
Returns array with step_type from every environment.
"""
if not self.envs_initialised: # initialising environments - poll and wait for them to get started
for i in range(self.cfd_n_envs):
self.client.poll_tensor(self.step_type_key[i], 100, self.poll_time)
try:
self._step_type[i] = self.client.get_tensor(self.step_type_key[i])[0]
except Exception as exc:
raise Warning(f"Could not read step type from key: {self.step_type_key[i]}.") from exc
if np.any(self._step_type < 0):
raise ValueError(f"Environments could not be initialized, or initial step_type could not be read from database. \
\n step_type = {self._step_type}")
self.envs_initialised = True
else:
for i in range(self.cfd_n_envs):
try:
self._step_type[i] = self.client.get_tensor(self.step_type_key[i])[0]
except Exception as exc:
raise Warning(f"Could not read step type from key: {self.step_type_key[i]}") from exc
for i in range(self.cfd_n_envs):
step_type_hecuba = mynumpyclass.get_by_alias(self.step_type_key[i].replace('.', '_'))
step_type_hecuba.poll()
self._step_type[i] = step_type_hecuba.astype(np.int32)[0]
if np.any(self._step_type < 0):
raise ValueError(f"Environments could not be initialized, or initial step_type could not be read from database. \
\n step_type = {self._step_type}")
self.envs_initialised = True
return self._step_type

# if not self.envs_initialised: # initialising environments - poll and wait for them to get started
# for i in range(self.cfd_n_envs):
# self.client.poll_tensor(self.step_type_key[i], 100, self.poll_time)
# try:
# self._step_type[i] = self.client.get_tensor(self.step_type_key[i])[0]
# except Exception as exc:
# raise Warning(f"Could not read step type from key: {self.step_type_key[i]}.") from exc
# if np.any(self._step_type < 0):
# raise ValueError(f"Environments could not be initialized, or initial step_type could not be read from database. \
# \n step_type = {self._step_type}")
# self.envs_initialised = True
# else:
# for i in range(self.cfd_n_envs):
# try:
# self._step_type[i] = self.client.get_tensor(self.step_type_key[i])[0]
# except Exception as exc:
# raise Warning(f"Could not read step type from key: {self.step_type_key[i]}") from exc
# return self._step_type


def _set_action(self, action):
"""
Expand Down Expand Up @@ -412,12 +476,13 @@ def _stop_exp(self):
Stop SOD2D experiment (ensemble of models) with SmartSim
"""
# delete data from database
self.client.delete_tensor(self.state_size_key)
self.client.delete_tensor(self.action_size_key)
for i in range(self.cfd_n_envs):
self.client.delete_tensor(self.step_type_key[i])
# self.client.delete_tensor(self.state_size_key)
# self.client.delete_tensor(self.action_size_key)
# for i in range(self.cfd_n_envs):
# self.client.delete_tensor(self.step_type_key[i])
# stop ensemble run
self.exp.stop(self.ensemble)
# self.exp.stop(self.ensemble)
compss_wait_on(cancel_sod2d())
# move ensemble log files
for f in glob.glob("ensemble-*"):
shutil.move(f, "sod2d_exp/")
Expand Down