Skip to content

Commit dba437c

Browse files
khillionkhillion
committed
refactor image pulling in classes
1 parent f7a2fc6 commit dba437c

File tree

1 file changed

+127
-49
lines changed

1 file changed

+127
-49
lines changed

‎cwl_utils/docker_extract.py‎

Lines changed: 127 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,75 +1,153 @@
11
#!/usr/bin/env python3
2-
import sys
2+
from abc import ABC, abstractmethod
3+
import argparse
4+
import logging
35
import os
4-
import cwl_utils.parser_v1_0 as cwl
56
import subprocess
6-
import argparse
7+
import sys
78

9+
import cwl_utils.parser_v1_0 as cwl
810

9-
def parse_args():
10-
parser = argparse.ArgumentParser(
11-
description='Tool to save docker images from a cwl workflow \n '
12-
'and generate udocker loading commands.')
13-
parser.add_argument('dir', help='Directory in which to save images')
14-
parser.add_argument('input', help='Input CWL workflow')
15-
parser.add_argument('-s', '--singularity', help='Use singularity to pull the image', action='store_true')
16-
return parser.parse_args()
11+
logging.basicConfig(level=logging.INFO)
12+
_LOGGER = logging.getLogger(__name__)
1713

1814

19-
def main():
20-
args = parse_args()
21-
os.makedirs(args.dir, exist_ok=True)
15+
class ImagePuller(ABC):
2216

23-
top = cwl.load_document(args.input)
17+
def __init__(self, req, save_directory):
18+
self.req = req
19+
self.save_directory = save_directory
2420

25-
for req in set(traverse(top)):
26-
if args.singularity:
27-
image_name = get_image_name_singularity(req)
28-
save_docker_image_singularity(req, image_name, args.dir)
29-
else:
30-
image_name = get_image_name(req)
31-
save_docker_image(req, image_name, args.dir)
32-
print(load_docker_image(image_name))
21+
@abstractmethod
22+
def get_image_name(self):
23+
pass
3324

25+
@abstractmethod
26+
def save_docker_image(self):
27+
pass
3428

35-
def get_image_name(req):
36-
return ''.join(req.split('/')) + '.tar'
29+
@staticmethod
30+
def _run_command_pull(cmd_pull):
31+
try:
32+
subprocess.run(cmd_pull, check=True, stdout=subprocess.PIPE,
33+
stderr=subprocess.STDOUT)
34+
except subprocess.CalledProcessError as err:
35+
if err.output:
36+
raise subprocess.SubprocessError(err.output)
37+
raise err
3738

3839

39-
def get_image_name_singularity(req):
40+
class DockerImagePuller(ImagePuller):
41+
"""
42+
Pull docker image with Docker
43+
"""
44+
45+
def get_image_name(self):
46+
return ''.join(self.req.split('/')) + '.tar'
47+
48+
def generate_udocker_loading_command(self):
49+
return f'udocker load -i {self.get_image_name()}'
50+
51+
def save_docker_image(self):
52+
_LOGGER.info(f"Pulling {self.req} with Docker...")
53+
cmd_pull = ['docker', 'pull', self.req]
54+
ImagePuller._run_command_pull(cmd_pull)
55+
cmd_save = ['docker', 'save', '-o', os.path.join(self.save_directory,
56+
self.get_image_name()),
57+
self.req]
58+
subprocess.run(cmd_save, check=True)
59+
_LOGGER.info(f"Image successfully pulled: {self.save_directory}/{self.get_image_name()}")
60+
print(self.generate_udocker_loading_command())
61+
62+
63+
class SingularityImagePuller(ImagePuller):
64+
"""
65+
Pull docker image with Singularity
66+
"""
4067
CHARS_TO_REPLACE = ['/', ':']
4168
NEW_CHAR = '-'
42-
image_name = req
43-
for char in CHARS_TO_REPLACE:
44-
image_name = image_name.replace(char, NEW_CHAR)
45-
return f'{image_name}.img'
4669

70+
def get_image_name(self):
71+
image_name = self.req
72+
for char in self.CHARS_TO_REPLACE:
73+
image_name = image_name.replace(char, self.NEW_CHAR)
74+
return f'{image_name}.img'
4775

48-
def load_docker_image(image_name):
49-
return f'udocker load -i {image_name}'
76+
def save_docker_image(self):
77+
_LOGGER.info(f"Pulling {self.req} with Singularity...")
78+
cmd_pull = ['singularity', 'pull', os.path.join(self.save_directory,
79+
self.get_image_name()),
80+
f'docker:{self.req}']
81+
ImagePuller._run_command_pull(cmd_pull)
82+
_LOGGER.info(f"Image successfully pulled: {self.save_directory}/{self.get_image_name()}")
5083

5184

52-
def _run_command_pull(cmd_pull):
53-
try:
54-
subprocess.run(cmd_pull, check=True, stdout=subprocess.PIPE,
55-
stderr=subprocess.STDOUT)
56-
except subprocess.CalledProcessError as err:
57-
if err.output:
58-
raise subprocess.SubprocessError(err.output)
59-
raise err
85+
def parse_args():
86+
parser = argparse.ArgumentParser(
87+
description='Tool to save docker images from a cwl workflow.')
88+
parser.add_argument('dir', help='Directory in which to save images')
89+
parser.add_argument('input', help='Input CWL workflow')
90+
parser.add_argument('-s', '--singularity', help='Use singularity to pull the image', action='store_true')
91+
return parser.parse_args()
6092

6193

62-
def save_docker_image(req, image_name, image_dir):
63-
cmd_pull = ['docker', 'pull', req]
64-
_run_command_pull(cmd_pull)
65-
cmd_save = ['docker', 'save', '-o', os.path.join(image_dir, image_name),
66-
req]
67-
subprocess.run(cmd_save, check=True)
94+
def main():
95+
args = parse_args()
96+
os.makedirs(args.dir, exist_ok=True)
6897

98+
top = cwl.load_document(args.input)
6999

70-
def save_docker_image_singularity(req, image_name, image_dir):
71-
cmd_pull = ['singularity', 'pull', os.path.join(image_dir, image_name), f'docker:{req}']
72-
_run_command_pull(cmd_pull)
100+
for req in set(traverse(top)):
101+
if args.singularity:
102+
image_puller = SingularityImagePuller(req, args.dir)
103+
# image_name = get_image_name_singularity(req)
104+
# save_docker_image_singularity(req, image_name, args.dir)
105+
else:
106+
image_puller = DockerImagePuller(req, args.dir)
107+
# image_name = get_image_name(req)
108+
# save_docker_image(req, image_name, args.dir)
109+
# print(load_docker_image(image_name))
110+
image_puller.save_docker_image()
111+
112+
113+
# def get_image_name(req):
114+
# return ''.join(req.split('/')) + '.tar'
115+
#
116+
#
117+
# def get_image_name_singularity(req):
118+
# CHARS_TO_REPLACE = ['/', ':']
119+
# NEW_CHAR = '-'
120+
# image_name = req
121+
# for char in CHARS_TO_REPLACE:
122+
# image_name = image_name.replace(char, NEW_CHAR)
123+
# return f'{image_name}.img'
124+
#
125+
#
126+
# def load_docker_image(image_name):
127+
# return f'udocker load -i {image_name}'
128+
129+
130+
# def _run_command_pull(cmd_pull):
131+
# try:
132+
# subprocess.run(cmd_pull, check=True, stdout=subprocess.PIPE,
133+
# stderr=subprocess.STDOUT)
134+
# except subprocess.CalledProcessError as err:
135+
# if err.output:
136+
# raise subprocess.SubprocessError(err.output)
137+
# raise err
138+
139+
140+
# def save_docker_image(req, image_name, image_dir):
141+
# cmd_pull = ['docker', 'pull', req]
142+
# _run_command_pull(cmd_pull)
143+
# cmd_save = ['docker', 'save', '-o', os.path.join(image_dir, image_name),
144+
# req]
145+
# subprocess.run(cmd_save, check=True)
146+
#
147+
#
148+
# def save_docker_image_singularity(req, image_name, image_dir):
149+
# cmd_pull = ['singularity', 'pull', os.path.join(image_dir, image_name), f'docker:{req}']
150+
# _run_command_pull(cmd_pull)
73151

74152

75153
def extract_docker_requirements(process: cwl.Process):

0 commit comments

Comments
 (0)