Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
from cluster_tools.cluster_tasks import SlurmTask, LocalTask, LSFTask
class UniqueBlockLabelsBase(luigi.Task):
""" UniqueBlockLabels base class
"""
task_name = 'unique_block_labels'
src_file = os.path.abspath(__file__)
input_path = luigi.Parameter()
input_key = luigi.Parameter()
output_path = luigi.Parameter()
output_key = luigi.Parameter()
dependency = luigi.TaskParameter()
effective_scale_factor = luigi.ListParameter()
prefix = luigi.Parameter(default='')
def requires(self):
return self.dependency
def run_impl(self):
# get the global config and init configs
shebang, _, roi_begin, roi_end = self.global_config_values()
self.init(shebang)
# read shape chunks and mulit-set from input
with vu.file_reader(self.input_path) as f:
ds = f[self.input_key]
shape = ds.shape
chunks = ds.chunks
dtype = ds.dtype
task_name = 'sub_solutions'
src_file = os.path.abspath(__file__)
# input volumes and graph
problem_path = luigi.Parameter()
scale = luigi.IntParameter()
ws_path = luigi.Parameter()
ws_key = luigi.Parameter()
output_path = luigi.Parameter()
output_key = luigi.Parameter()
# we have extra roi parameters here,
# to quickly inspect a roi for debugging independent
# of the roi of the global conifg
# this roi must be inside of the global config's roi though.
roi_begin = luigi.ListParameter(default=None)
roi_end = luigi.ListParameter(default=None)
# identifiers for the sub-result and sub-graph sub-keys
sub_graph_identifier = luigi.Parameter('sub_graphs')
sub_result_identifier = luigi.Parameter('sub_results')
dependency = luigi.TaskParameter()
def requires(self):
return self.dependency
def clean_up_for_retry(self, block_list):
super().clean_up_for_retry(block_list)
# TODO remove any output of failed blocks because it might be corrupted
def run_impl(self):
# get the global config and init configs
shebang, block_shape, global_roi_begin, global_roi_end = self.global_config_values()
self.init(shebang)
class RegionCentersBase(luigi.Task):
""" RegionCenters base class
"""
task_name = 'region_centers'
src_file = os.path.abspath(__file__)
allow_retry = False
input_path = luigi.Parameter()
input_key = luigi.Parameter()
morphology_path = luigi.Parameter()
morphology_key = luigi.Parameter()
output_path = luigi.Parameter()
output_key = luigi.Parameter()
ignore_label = luigi.Parameter(default=None)
resolution = luigi.ListParameter(default=[1, 1, 1])
#
dependency = luigi.TaskParameter()
def requires(self):
return self.dependency
def run_impl(self):
# get the global config and init configs
shebang = self.global_config_values()[0]
self.init(shebang)
# load the task config
config = self.get_task_config()
number_of_labels = int(vu.file_reader(self.input_path,
'r')[self.input_key].attrs['maxId']) + 1
# TODO should be a parameter
service_catalog.update_provisioned_product_properties(
ProvisionedProductId=provisioned_product_id,
ProvisionedProductProperties={
"OWNER": f"arn:aws:iam::{self.account_id}:role/servicecatalog-puppet/PuppetRole"
},
)
self.write_output(changes_made)
class RunDeployInSpokeTask(tasks.PuppetTask):
manifest_file_path = luigi.Parameter()
puppet_account_id = luigi.Parameter()
account_id = luigi.Parameter()
home_region = luigi.Parameter()
regions = luigi.ListParameter()
should_collect_cloudformation_events = luigi.BoolParameter()
should_forward_events_to_eventbridge = luigi.BoolParameter()
should_forward_failures_to_opscenter = luigi.BoolParameter()
def params_for_results_display(self):
return {
"manifest_file_path": self.manifest_file_path,
"puppet_account_id": self.puppet_account_id,
"account_id": self.account_id,
}
def run(self):
with betterboto_client.CrossAccountClientContextManager(
"s3",
f"arn:aws:iam::{self.puppet_account_id}:role/servicecatalog-puppet/PuppetRole",
f"s3-{self.puppet_account_id}",
import os
import luigi
from ..cluster_tasks import WorkflowBase
from ..utils import volume_utils as vu
from ..morphology import MorphologyWorkflow
from . import skeletonize as skeleton_tasks
class SkeletonWorkflow(WorkflowBase):
input_path = luigi.Parameter()
input_key = luigi.Parameter()
output_path = luigi.Parameter()
output_key = luigi.Parameter()
resolution = luigi.ListParameter()
size_threshold = luigi.IntParameter(default=None)
max_id = luigi.IntParameter(default=None)
method = luigi.Parameter(default='thinning')
def require_max_id(self):
with vu.file_reader(self.input_path) as f:
ds = f[self.input_key]
attrs_max_id = ds.attrs.get('maxId', None)
if self.max_id is None and attrs_max_id is None:
raise RuntimeError("Input dataset does not have maxId attribute, so it needs to be passed externally")
elif self.max_id is not None and attrs_max_id is None:
ds.attrs['maxId'] = self.max_id
max_id = self.max_id
elif self.max_id is not None and attrs_max_id is not None:
if self.max_id != attrs_max_id:
return dep
@staticmethod
def get_config():
configs = super(MulticutWorkflow, MulticutWorkflow).get_config()
configs.update({'solve_global': solve_tasks.SolveGlobalLocal.default_task_config()})
return configs
class SubSolutionsWorkflow(MulticutWorkflowBase):
ws_path = luigi.Parameter()
ws_key = luigi.Parameter()
output_path = luigi.Parameter()
output_key = luigi.Parameter()
roi_begin = luigi.ListParameter(default=None)
roi_end = luigi.ListParameter(default=None)
def requires(self):
sub_task = getattr(sub_tasks,
self._get_task_name('SubSolutions'))
dep = self._hierarchical_tasks(self.dependency, self.n_scales + 1)
t_sub = sub_task(tmp_folder=self.tmp_folder,
max_jobs=self.max_jobs,
config_dir=self.config_dir,
problem_path=self.problem_path,
ws_path=self.ws_path,
ws_key=self.ws_key,
output_path=self.output_path,
output_key=self.output_key,
scale=self.n_scales,
dependency=dep,
roi_begin=self.roi_begin,
class CopyVolumeBase(luigi.Task):
""" copy_volume base class
"""
task_name = 'copy_volume'
src_file = os.path.abspath(__file__)
# input and output volumes
input_path = luigi.Parameter()
input_key = luigi.Parameter()
output_path = luigi.Parameter()
output_key = luigi.Parameter()
prefix = luigi.Parameter()
dtype = luigi.Parameter(default=None)
fit_to_roi = luigi.BoolParameter(default=False)
effective_scale_factor = luigi.ListParameter(default=[])
dependency = luigi.TaskParameter(default=DummyTask())
@staticmethod
def default_task_config():
# we use this to get also get the common default config
config = LocalTask.default_task_config()
config.update({'chunks': None, 'compression': 'gzip',
'reduce_channels': None, 'map_uniform_blocks_to_background': False,
'value_list': None, 'offset': None, 'insert_mode': False})
return config
def requires(self):
return self.dependency
def clean_up_for_retry(self, block_list):
super().clean_up_for_retry(block_list)
dependency=dep)
return t_solve
@staticmethod
def get_config():
configs = super(LiftedMulticutWorkflow, LiftedMulticutWorkflow).get_config()
configs.update({'solve_lifted_global': solve_tasks.SolveLiftedGlobalLocal.default_task_config()})
return configs
class SubLiftedSolutionsWorkflow(LiftedMulticutWorkflowBase):
ws_path = luigi.Parameter()
ws_key = luigi.Parameter()
output_path = luigi.Parameter()
output_key = luigi.Parameter()
roi_begin = luigi.ListParameter(default=None)
roi_end = luigi.ListParameter(default=None)
def requires(self):
sub_task = getattr(sub_tasks,
self._get_task_name('SubSolutions'))
dep = self._hierarchical_tasks(self.dependency, self.n_scales + 1)
t_sub = sub_task(tmp_folder=self.tmp_folder,
max_jobs=self.max_jobs,
config_dir=self.config_dir,
problem_path=self.problem_path,
ws_path=self.ws_path,
ws_key=self.ws_key,
output_path=self.output_path,
output_key=self.output_key,
scale=self.n_scales,
dependency=dep,
return self.make_target(self.output_file_path)
def run(self):
data = self.load_data_frame(required_columns={self.key_column_name, self.value_column_name})
data.dropna(subset={self.key_column_name, self.value_column_name}, inplace=True)
result = data.groupby(by=self.key_column_name)[self.value_column_name].apply(list).to_dict()
self.dump(result)
class ConvertToOneHot(gokart.TaskOnKart):
"""
Convert column values of `categorical_column_names` to one-hot.
"""
task_namespace = 'redshells.data_frame_utils'
data_task = gokart.TaskInstanceParameter(description='A task outputs pd.DataFrame.')
categorical_column_names = luigi.ListParameter() # type: List[str]
output_file_path = luigi.Parameter(default='data/group_by_column_as_dict.pkl') # type: str
def requires(self):
return self.data_task
def output(self):
return self.make_target(self.output_file_path)
def run(self):
categorical_column_names = list(self.categorical_column_names)
data = self.load_data_frame(required_columns=set(categorical_column_names))
result = pd.get_dummies(data[categorical_column_names])
result = result.merge(data.drop(categorical_column_names, axis=1), left_index=True, right_index=True)
self.dump(result)
)
logger.info(f"[{self.portfolio}] {self.account_id}:{self.region} :: Finished importing")
class CreateLaunchRoleConstraintsForPortfolio(PuppetTask):
account_id = luigi.Parameter()
region = luigi.Parameter()
portfolio = luigi.Parameter()
hub_portfolio_id = luigi.Parameter()
puppet_account_id = luigi.Parameter()
launch_constraints = luigi.DictParameter()
dependencies = luigi.ListParameter(default=[])
post_actions = luigi.ListParameter()
should_use_sns = luigi.Parameter(default=False, significant=False)
def requires(self):
return {
'create_spoke_local_portfolio_task': ImportIntoSpokeLocalPortfolioTask(
account_id=self.account_id,
region=self.region,
portfolio=self.portfolio,
hub_portfolio_id=self.hub_portfolio_id,
),
'deps': [ProvisionProductTask(**dependency) for dependency in self.dependencies]
}
@property
def node_id(self):