Skip to content

Commit d73e73d

Browse files
committed
Create custom AW/RC names
1 parent 096e3df commit d73e73d

File tree

2 files changed

+15
-9
lines changed

2 files changed

+15
-9
lines changed

src/codeflare_sdk/cluster/cluster.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ def __init__(self, config: ClusterConfiguration):
1616
self.app_wrapper_yaml = self.create_app_wrapper()
1717

1818
def create_app_wrapper(self):
19+
name=self.config.name
1920
min_cpu=self.config.min_cpus
2021
max_cpu=self.config.max_cpus
2122
min_memory=self.config.min_memory
@@ -27,7 +28,7 @@ def create_app_wrapper(self):
2728
instascale=self.config.instascale
2829
instance_types=self.config.machine_types
2930
env=self.config.envs
30-
return generate_appwrapper(min_cpu=min_cpu, max_cpu=max_cpu, min_memory=min_memory,
31+
return generate_appwrapper(name=name, min_cpu=min_cpu, max_cpu=max_cpu, min_memory=min_memory,
3132
max_memory=max_memory, gpu=gpu, workers=workers, template=template,
3233
image=image, instascale=instascale, instance_types=instance_types, env=env)
3334

src/codeflare_sdk/utils/generate_yaml.py

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,14 @@ def read_template(template):
1010
except yaml.YAMLError as exc:
1111
print(exc)
1212

13-
def gen_names():
14-
gen_id = str(uuid.uuid4())
15-
appwrapper_name = "appwrapper-" + gen_id
16-
cluster_name = "cluster-" + gen_id
17-
return appwrapper_name, cluster_name
13+
def gen_names(name):
14+
if not name:
15+
gen_id = str(uuid.uuid4())
16+
appwrapper_name = "appwrapper-" + gen_id
17+
cluster_name = "cluster-" + gen_id
18+
return appwrapper_name, cluster_name
19+
else:
20+
return name, name
1821

1922
def update_names(yaml, item, appwrapper_name, cluster_name):
2023
metadata = yaml.get("metadata")
@@ -122,9 +125,9 @@ def write_user_appwrapper(user_yaml, output_file_name):
122125
yaml.dump(user_yaml, outfile, default_flow_style=False)
123126
print(f"Written to: {output_file_name}")
124127

125-
def generate_appwrapper(min_cpu, max_cpu, min_memory, max_memory, gpu, workers, template, image, instascale, instance_types, env):
128+
def generate_appwrapper(name, min_cpu, max_cpu, min_memory, max_memory, gpu, workers, template, image, instascale, instance_types, env):
126129
user_yaml = read_template(template)
127-
appwrapper_name, cluster_name = gen_names()
130+
appwrapper_name, cluster_name = gen_names(name)
128131
resources = user_yaml.get("spec","resources")
129132
item = resources["resources"].get("GenericItems")[0]
130133
update_names(user_yaml, item, appwrapper_name, cluster_name)
@@ -138,6 +141,7 @@ def generate_appwrapper(min_cpu, max_cpu, min_memory, max_memory, gpu, workers,
138141

139142
def main():
140143
parser = argparse.ArgumentParser(description='Generate user AppWrapper')
144+
parser.add_argument("--name", required=False, default="", help="User selected name for AppWrapper and Ray Cluster (auto-generated if not provided)")
141145
parser.add_argument("--min-cpu", type=int, required=True, help="min number of CPU(s) in a worker required for running job")
142146
parser.add_argument("--max-cpu", type=int, required=True, help="max number of CPU(s) in a worker required for running job")
143147
parser.add_argument("--min-memory", type=int, required=True, help="min RAM required in a worker for running job, in GB")
@@ -150,6 +154,7 @@ def main():
150154
parser.add_argument("--instance-types", type=str, nargs='+', default=[], required=False, help="Head,worker instance types (space separated)")
151155

152156
args = parser.parse_args()
157+
name = args.name
153158
min_cpu = args.min_cpu
154159
max_cpu = args.max_cpu
155160
min_memory = args.min_memory
@@ -162,7 +167,7 @@ def main():
162167
instance_types = args.instance_types
163168
env = {}
164169

165-
outfile = generate_appwrapper(min_cpu, max_cpu, min_memory, max_memory, gpu, workers, template, image, instascale, instance_types, env)
170+
outfile = generate_appwrapper(name, min_cpu, max_cpu, min_memory, max_memory, gpu, workers, template, image, instascale, instance_types, env)
166171
return outfile
167172

168173
if __name__=="__main__":

0 commit comments

Comments
 (0)