Skip to content

Commit bc31b94

Browse files
connor-mccarthystijntratsaertit
connor-mccarthy
authored and
stijntratsaertit
committed
feat(sdk): support Concat and IfPresent placeholder in local container component execution #localexecution (kubeflow#10348)
* feat(sdk): support Concat and IfPresent placeholder in local container component execution #localexecution * address review feedback * fix test
1 parent 2abc173 commit bc31b94

File tree

3 files changed

+496
-24
lines changed

3 files changed

+496
-24
lines changed

‎sdk/python/kfp/local/docker_task_handler_test.py

+273-5
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
import os
15+
from typing import Optional
1516
import unittest
1617
from unittest import mock
1718

@@ -159,17 +160,20 @@ def test_latest_tag(self):
159160

160161
def test_no_tag(self):
161162
actual = docker_task_handler.add_latest_tag_if_not_present(
162-
image='alpine:123')
163-
expected = 'alpine:123'
163+
image='alpine:3.19.0')
164+
expected = 'alpine:3.19.0'
164165
self.assertEqual(actual, expected)
165166

166167

167168
class TestE2E(DockerMockTestCase,
168169
testing_utilities.LocalRunnerEnvironmentTestCase):
169170

170-
def test_python(self):
171+
def setUp(self):
172+
super().setUp()
171173
local.init(runner=local.DockerRunner())
172174

175+
def test_python(self):
176+
173177
@dsl.component
174178
def artifact_maker(x: str, a: Output[Artifact]):
175179
with open(a.path, 'w') as f:
@@ -200,7 +204,6 @@ def artifact_maker(x: str, a: Output[Artifact]):
200204
self.assertEqual(kwargs['volumes'][root_vol_key]['mode'], 'rw')
201205

202206
def test_empty_container_component(self):
203-
local.init(runner=local.DockerRunner())
204207

205208
@dsl.container_component
206209
def comp():
@@ -222,7 +225,6 @@ def comp():
222225
self.assertEqual(kwargs['command'], [])
223226

224227
def test_container_component(self):
225-
local.init(runner=local.DockerRunner())
226228

227229
@dsl.container_component
228230
def artifact_maker(x: str,):
@@ -258,6 +260,272 @@ def artifact_maker(x: str,):
258260
self.assertEqual(kwargs['volumes'][root_vol_key]['bind'], root_vol_key)
259261
self.assertEqual(kwargs['volumes'][root_vol_key]['mode'], 'rw')
260262

263+
def test_if_present_with_string_omitted(self):
264+
265+
@dsl.container_component
266+
def comp(x: Optional[str] = None):
267+
return dsl.ContainerSpec(
268+
image='alpine:3.19.0',
269+
command=[
270+
dsl.IfPresentPlaceholder(
271+
input_name='x',
272+
then=['echo', x],
273+
else_=['echo', 'No input provided!'])
274+
])
275+
276+
comp()
277+
278+
run_mock = self.mocked_docker_client.containers.run
279+
run_mock.assert_called_once()
280+
kwargs = run_mock.call_args[1]
281+
self.assertEqual(
282+
kwargs['image'],
283+
'alpine:3.19.0',
284+
)
285+
self.assertEqual(kwargs['command'], [
286+
'echo',
287+
'No input provided!',
288+
])
289+
290+
def test_if_present_with_string_provided(self):
291+
292+
@dsl.container_component
293+
def comp(x: Optional[str] = None):
294+
return dsl.ContainerSpec(
295+
image='alpine:3.19.0',
296+
command=[
297+
dsl.IfPresentPlaceholder(
298+
input_name='x',
299+
then=['echo', x],
300+
else_=['echo', 'No artifact provided!'])
301+
])
302+
303+
comp(x='foo')
304+
305+
run_mock = self.mocked_docker_client.containers.run
306+
run_mock.assert_called_once()
307+
kwargs = run_mock.call_args[1]
308+
self.assertEqual(
309+
kwargs['image'],
310+
'alpine:3.19.0',
311+
)
312+
self.assertEqual(kwargs['command'], [
313+
'echo',
314+
'foo',
315+
])
316+
317+
def test_if_present_single_element_with_string_omitted(self):
318+
319+
@dsl.container_component
320+
def comp(x: Optional[str] = None):
321+
return dsl.ContainerSpec(
322+
image='alpine:3.19.0',
323+
command=[
324+
'echo',
325+
dsl.IfPresentPlaceholder(
326+
input_name='x',
327+
then=x,
328+
else_='No artifact provided!',
329+
)
330+
])
331+
332+
comp()
333+
334+
run_mock = self.mocked_docker_client.containers.run
335+
run_mock.assert_called_once()
336+
kwargs = run_mock.call_args[1]
337+
self.assertEqual(
338+
kwargs['image'],
339+
'alpine:3.19.0',
340+
)
341+
self.assertEqual(kwargs['command'], [
342+
'echo',
343+
'No artifact provided!',
344+
])
345+
346+
def test_if_present_single_element_with_string_provided(self):
347+
348+
@dsl.container_component
349+
def comp(x: Optional[str] = None):
350+
return dsl.ContainerSpec(
351+
image='alpine:3.19.0',
352+
command=[
353+
'echo',
354+
dsl.IfPresentPlaceholder(
355+
input_name='x',
356+
then=x,
357+
else_='No artifact provided!',
358+
)
359+
])
360+
361+
comp(x='foo')
362+
363+
run_mock = self.mocked_docker_client.containers.run
364+
run_mock.assert_called_once()
365+
kwargs = run_mock.call_args[1]
366+
self.assertEqual(
367+
kwargs['image'],
368+
'alpine:3.19.0',
369+
)
370+
self.assertEqual(kwargs['command'], [
371+
'echo',
372+
'foo',
373+
])
374+
375+
def test_concat_placeholder(self):
376+
377+
@dsl.container_component
378+
def comp(x: Optional[str] = None):
379+
return dsl.ContainerSpec(
380+
image='alpine',
381+
command=[dsl.ConcatPlaceholder(['prefix-', x, '-suffix'])])
382+
383+
comp()
384+
385+
run_mock = self.mocked_docker_client.containers.run
386+
run_mock.assert_called_once()
387+
kwargs = run_mock.call_args[1]
388+
self.assertEqual(
389+
kwargs['image'],
390+
'alpine:latest',
391+
)
392+
self.assertEqual(kwargs['command'], ['prefix-null-suffix'])
393+
394+
def test_nested_concat_placeholder(self):
395+
396+
@dsl.container_component
397+
def comp(x: Optional[str] = None):
398+
return dsl.ContainerSpec(
399+
image='alpine',
400+
command=[
401+
'echo',
402+
dsl.ConcatPlaceholder(
403+
['a', dsl.ConcatPlaceholder(['b', x, 'd'])])
404+
])
405+
406+
comp(x='c')
407+
408+
run_mock = self.mocked_docker_client.containers.run
409+
run_mock.assert_called_once()
410+
kwargs = run_mock.call_args[1]
411+
self.assertEqual(
412+
kwargs['image'],
413+
'alpine:latest',
414+
)
415+
self.assertEqual(kwargs['command'], ['echo', 'abcd'])
416+
417+
def test_ifpresent_in_concat_provided(self):
418+
419+
@dsl.container_component
420+
def comp(x: Optional[str] = None):
421+
return dsl.ContainerSpec(
422+
image='alpine',
423+
command=[
424+
'echo',
425+
dsl.ConcatPlaceholder([
426+
'there ',
427+
dsl.ConcatPlaceholder([
428+
'is ',
429+
dsl.IfPresentPlaceholder(
430+
input_name='x',
431+
then='one thing',
432+
else_='another thing')
433+
])
434+
])
435+
])
436+
437+
comp(x='c')
438+
439+
run_mock = self.mocked_docker_client.containers.run
440+
run_mock.assert_called_once()
441+
kwargs = run_mock.call_args[1]
442+
self.assertEqual(
443+
kwargs['image'],
444+
'alpine:latest',
445+
)
446+
self.assertEqual(kwargs['command'], ['echo', 'there is one thing'])
447+
448+
def test_ifpresent_in_concat_omitted(self):
449+
450+
@dsl.container_component
451+
def comp(x: Optional[str] = None):
452+
return dsl.ContainerSpec(
453+
image='alpine',
454+
command=[
455+
'echo',
456+
dsl.ConcatPlaceholder([
457+
'there ',
458+
dsl.ConcatPlaceholder([
459+
'is ',
460+
dsl.IfPresentPlaceholder(
461+
input_name='x',
462+
then='one thing',
463+
else_='another thing')
464+
])
465+
])
466+
])
467+
468+
comp()
469+
470+
run_mock = self.mocked_docker_client.containers.run
471+
run_mock.assert_called_once()
472+
kwargs = run_mock.call_args[1]
473+
self.assertEqual(
474+
kwargs['image'],
475+
'alpine:latest',
476+
)
477+
self.assertEqual(kwargs['command'], ['echo', 'there is another thing'])
478+
479+
def test_concat_in_ifpresent_provided(self):
480+
481+
@dsl.container_component
482+
def comp(x: Optional[str] = None):
483+
return dsl.ContainerSpec(
484+
image='alpine',
485+
command=[
486+
'echo',
487+
dsl.IfPresentPlaceholder(
488+
input_name='x',
489+
then=dsl.ConcatPlaceholder([x]),
490+
else_=dsl.ConcatPlaceholder(['something', ' ', 'else']))
491+
])
492+
493+
comp(x='something')
494+
495+
run_mock = self.mocked_docker_client.containers.run
496+
run_mock.assert_called_once()
497+
kwargs = run_mock.call_args[1]
498+
self.assertEqual(
499+
kwargs['image'],
500+
'alpine:latest',
501+
)
502+
self.assertEqual(kwargs['command'], ['echo', 'something'])
503+
504+
def test_concat_in_ifpresent_omitted(self):
505+
506+
@dsl.container_component
507+
def comp(x: Optional[str] = None):
508+
return dsl.ContainerSpec(
509+
image='alpine',
510+
command=[
511+
'echo',
512+
dsl.IfPresentPlaceholder(
513+
input_name='x',
514+
then=dsl.ConcatPlaceholder([x]),
515+
else_=dsl.ConcatPlaceholder(['another', ' ', 'thing']))
516+
])
517+
518+
comp()
519+
520+
run_mock = self.mocked_docker_client.containers.run
521+
run_mock.assert_called_once()
522+
kwargs = run_mock.call_args[1]
523+
self.assertEqual(
524+
kwargs['image'],
525+
'alpine:latest',
526+
)
527+
self.assertEqual(kwargs['command'], ['echo', 'another thing'])
528+
261529

262530
if __name__ == '__main__':
263531
unittest.main()

0 commit comments

Comments
 (0)