Skip to content

Commit

Permalink
added tests for ssg manager in python
Browse files Browse the repository at this point in the history
  • Loading branch information
mdorier committed Jan 30, 2024
1 parent 8eba610 commit 524c337
Show file tree
Hide file tree
Showing 5 changed files with 23 additions and 12 deletions.
16 changes: 10 additions & 6 deletions python/mochi/bedrock/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,8 @@ def config(self):

@property
def spec(self) -> list[SSGSpec]:
return [SSGSpec.from_dict(group) for group in self._server.config]
abt_spec = self._server.margo.spec.argobots
return [SSGSpec.from_dict(group, abt_spec) for group in self.config]

def resolve(self, margo, location: str) -> pymargo.core.Address:
if isinstance(margo, pymargo.core.Engine):
Expand All @@ -177,12 +178,15 @@ def __len__(self) -> int:
def __getitem__(self, key: int|str) -> SSGGroup:
return SSGGroup(self._internal.get_group(key))

def create(self, name: str, config: str|dict|SSGSpec, pool: Pool,
bootstrap: str, group_file: str = "") -> SSGGroup:
if isinstance(config, dict):
config = json.dumps(config)
def create(self, name: str, pool: str|int|Pool = "__primary__",
config: str|dict|SSGSpec = "{}",
bootstrap: str = "init", group_file: str = "") -> SSGGroup:
if not isinstance(pool, Pool):
pool = self._server.margo.pools[pool]
if isinstance(config, str):
config = json.loads(config)
elif isinstance(config, SSGSpec):
config = config.to_json()
config = config.to_dict()
pool = pool._internal
return SSGGroup(self._internal.create_group(name, config, pool, bootstrap, group_file))

Expand Down
2 changes: 1 addition & 1 deletion python/mochi/bedrock/spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -1224,7 +1224,7 @@ def from_dict(data: dict, abt_spec: ArgobotsSpec) -> 'SSGSpec':
"""
args = data.copy()
args['pool'] = abt_spec.pools[data['pool']]
args['swim'] = SwimSpec.from_dict(**args['swim'])
args['swim'] = SwimSpec.from_dict(args['swim'])
ssg = SSGSpec(**args)
return ssg

Expand Down
2 changes: 1 addition & 1 deletion python/mochi/bedrock/test_abtio_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import mochi.bedrock.spec as spec


class TestMargoInstance(unittest.TestCase):
class TestAbtIOManager(unittest.TestCase):

def setUp(self):
config = {
Expand Down
2 changes: 1 addition & 1 deletion python/mochi/bedrock/test_margo_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import mochi.bedrock.spec as spec


class TestMargoInstance(unittest.TestCase):
class TestMargoManager(unittest.TestCase):

def setUp(self):
config = {
Expand Down
13 changes: 10 additions & 3 deletions src/SSGManager.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -214,22 +214,29 @@ std::shared_ptr<NamedDependency> SSGManager::getGroup(const std::string& group_n
#ifdef ENABLE_SSG
auto it = std::find_if(self->m_ssg_groups.begin(), self->m_ssg_groups.end(),
[&](auto& g) { return g->getName() == group_name; });
if (it == self->m_ssg_groups.end())
if (it == self->m_ssg_groups.end()) {
throw DETAILED_EXCEPTION("Could not find SSG group with name \"{}\"", group_name);
return nullptr;
else
} else {
return *it;
}
#else
(void)group_name;
throw DETAILED_EXCEPTION("Bedrock was not compiler with SSG support");
return nullptr;
#endif
}

std::shared_ptr<NamedDependency> SSGManager::getGroup(uint32_t group_index) const {
#ifdef ENABLE_SSG
if(group_index >= self->m_ssg_groups.size()) return nullptr;
if(group_index >= self->m_ssg_groups.size()) {
throw DETAILED_EXCEPTION("Could not find SSG group at index {}", group_index);
return nullptr;
}
return self->m_ssg_groups[group_index];
#else
(void)group_index;
throw DETAILED_EXCEPTION("Bedrock was not compiler with SSG support");
return nullptr;
#endif
}
Expand Down

0 comments on commit 524c337

Please sign in to comment.