@@ -463,195 +463,3 @@ def test_onnx_use_custom_allocator_with_GPU(env):
463
463
for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
464
464
env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory_access_num" ]), 11 )
465
465
466
- def tests_onnx_info (env ):
467
- if not TEST_ONNX :
468
- env .debugPrint ("skipping {} since TEST_ONNX=0" .format (sys ._getframe ().f_code .co_name ), force = True )
469
- return
470
- con = env .getConnection ()
471
-
472
- ret = con .execute_command ('AI.INFO' )
473
- env .assertEqual (6 , len (ret ))
474
-
475
- test_data_path = os .path .join (os .path .dirname (__file__ ), 'test_data' )
476
- linear_model_filename = os .path .join (test_data_path , 'linear_iris.onnx' )
477
-
478
- with open (linear_model_filename , 'rb' ) as f :
479
- model_pb = f .read ()
480
-
481
- con .execute_command ('AI.MODELSET' , 'linear{1}' , 'ONNX' , DEVICE , 'BLOB' , model_pb )
482
-
483
- ret = con .execute_command ('AI.INFO' )
484
- env .assertEqual (8 , len (ret ))
485
- env .assertEqual (b'ONNX version' , ret [6 ])
486
-
487
-
488
- def test_parallelism ():
489
- env = Env (moduleArgs = 'INTRA_OP_PARALLELISM 1 INTER_OP_PARALLELISM 1' )
490
- if not TEST_ONNX :
491
- env .debugPrint ("skipping {} since TEST_ONNX=0" .format (sys ._getframe ().f_code .co_name ), force = True )
492
- return
493
-
494
- con = env .getConnection ()
495
- test_data_path = os .path .join (os .path .dirname (__file__ ), 'test_data' )
496
- model_filename = os .path .join (test_data_path , 'mnist.onnx' )
497
- sample_filename = os .path .join (test_data_path , 'one.raw' )
498
- with open (model_filename , 'rb' ) as f :
499
- model_pb = f .read ()
500
- with open (sample_filename , 'rb' ) as f :
501
- sample_raw = f .read ()
502
-
503
- ret = con .execute_command ('AI.MODELSET' , 'm{1}' , 'ONNX' , DEVICE , 'BLOB' , model_pb )
504
- env .assertEqual (ret , b'OK' )
505
- con .execute_command ('AI.TENSORSET' , 'a{1}' , 'FLOAT' , 1 , 1 , 28 , 28 , 'BLOB' , sample_raw )
506
-
507
- con .execute_command ('AI.MODELRUN' , 'm{1}' , 'INPUTS' , 'a{1}' , 'OUTPUTS' , 'b{1}' )
508
- ensureSlaveSynced (con , env )
509
- values = con .execute_command ('AI.TENSORGET' , 'b{1}' , 'VALUES' )
510
- argmax = max (range (len (values )), key = lambda i : values [i ])
511
- env .assertEqual (argmax , 1 )
512
-
513
- load_time_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
514
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[3 ].split ()[1 :]}
515
- env .assertEqual (load_time_config ["ai_inter_op_parallelism" ], "1" )
516
- env .assertEqual (load_time_config ["ai_intra_op_parallelism" ], "1" )
517
-
518
- env = Env (moduleArgs = 'INTRA_OP_PARALLELISM 2 INTER_OP_PARALLELISM 2' )
519
- load_time_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
520
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[3 ].split ()[1 :]}
521
- env .assertEqual (load_time_config ["ai_inter_op_parallelism" ], "2" )
522
- env .assertEqual (load_time_config ["ai_intra_op_parallelism" ], "2" )
523
-
524
-
525
- def test_onnx_use_custom_allocator (env ):
526
- if not TEST_ONNX :
527
- env .debugPrint ("skipping {} since TEST_ONNX=0" .format (sys ._getframe ().f_code .co_name ), force = True )
528
- return
529
-
530
- con = env .getConnection ()
531
- test_data_path = os .path .join (os .path .dirname (__file__ ), 'test_data' )
532
- model_filename = os .path .join (test_data_path , 'mul_1.onnx' )
533
- with open (model_filename , 'rb' ) as f :
534
- model_pb = f .read ()
535
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
536
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
537
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory" ]), 0 )
538
-
539
- # Expect using the allocator during model set for allocating the model, its input name and output name:
540
- # overall 3 allocations. The model raw size is 130B ,and the names are 2B each. In practice we allocate
541
- # more than 134B as Redis allocator will use additional memory for its internal management and for the
542
- # 64-Byte alignment. When the test runs with valgrind, redis will use malloc for the allocations
543
- # (hence will not use additional memory).
544
- ret = con .execute_command ('AI.MODELSET' , 'm{1}' , 'ONNX' , 'CPU' , 'BLOB' , model_pb )
545
- env .assertEqual (ret , b'OK' )
546
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
547
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
548
-
549
- # Expect using at least 130+63+(size of an address) + 2*(2+63+(size of an address)) bytes.
550
- model_allocation_bytes_used = int (ai_memory_config ["ai_onnxruntime_memory" ])
551
- env .assertTrue (model_allocation_bytes_used > 334 )
552
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory_access_num" ]), 3 )
553
- con .execute_command ('AI.TENSORSET' , 'a_mul{1}' , 'FLOAT' , 3 , 2 , 'VALUES' , 1.0 , 2.0 , 3.0 , 4.0 , 5.0 , 6.0 )
554
-
555
- # Running the model should access the allocator 6 times: allocating+freeing input+output names,
556
- # and allocating+freeing the output as OrtValue.
557
- con .execute_command ('AI.MODELRUN' , 'm{1}' , 'INPUTS' , 'a_mul{1}' , 'OUTPUTS' , 'b{1}' )
558
- values = con .execute_command ('AI.TENSORGET' , 'b{1}' , 'VALUES' )
559
- env .assertEqual (values , [b'1' , b'4' , b'9' , b'16' , b'25' , b'36' ])
560
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
561
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
562
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory_access_num" ]), 9 )
563
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory" ]), model_allocation_bytes_used )
564
-
565
- # Expect using the allocator free function 3 times: when releasing the model, input name and output name.
566
- con .execute_command ('AI.MODELDEL' , 'm{1}' )
567
- env .assertFalse (con .execute_command ('EXISTS' , 'm{1}' ))
568
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
569
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
570
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory" ]), 0 )
571
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory_access_num" ]), 12 )
572
-
573
- # test the use of Redis allocator in model run op.
574
- model_filename = os .path .join (test_data_path , 'mnist.onnx' )
575
- sample_filename = os .path .join (test_data_path , 'one.raw' )
576
-
577
- with open (model_filename , 'rb' ) as f :
578
- model_pb = f .read ()
579
- with open (sample_filename , 'rb' ) as f :
580
- sample_raw = f .read ()
581
-
582
- ret = con .execute_command ('AI.MODELSET' , 'm{1}' , 'ONNX' , 'CPU' , 'BLOB' , model_pb )
583
- env .assertEqual (ret , b'OK' )
584
- con .execute_command ('AI.TENSORSET' , 'a{1}' , 'FLOAT' , 1 , 1 , 28 , 28 , 'BLOB' , sample_raw )
585
-
586
- # Expect 18 allocator's access from onnx during the run (in addition to the allocations that were made while
587
- # creating the model).
588
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
589
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
590
- allocator_access_num_before = ai_memory_config ["ai_onnxruntime_memory_access_num" ]
591
- con .execute_command ('AI.MODELRUN' , 'm{1}' , 'INPUTS' , 'a{1}' , 'OUTPUTS' , 'b{1}' )
592
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
593
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
594
- allocator_access_num_after = ai_memory_config ["ai_onnxruntime_memory_access_num" ]
595
- env .assertEqual (int (allocator_access_num_after ) - int (allocator_access_num_before ), 18 )
596
-
597
- values = con .execute_command ('AI.TENSORGET' , 'b{1}' , 'VALUES' )
598
- argmax = max (range (len (values )), key = lambda i : values [i ])
599
- env .assertEqual (argmax , 1 )
600
-
601
-
602
- def test_onnx_use_custom_allocator_with_GPU (env ):
603
- if not TEST_ONNX :
604
- env .debugPrint ("skipping {} since TEST_ONNX=0" .format (sys ._getframe ().f_code .co_name ), force = True )
605
- return
606
- if DEVICE == 'CPU' :
607
- env .debugPrint ("skipping {} since this test if for GPU only" .format (sys ._getframe ().f_code .co_name ), force = True )
608
- return
609
-
610
- con = env .getConnection ()
611
- test_data_path = os .path .join (os .path .dirname (__file__ ), 'test_data' )
612
- model_filename = os .path .join (test_data_path , 'mul_1.onnx' )
613
- with open (model_filename , 'rb' ) as f :
614
- model_pb = f .read ()
615
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
616
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
617
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory" ]), 0 )
618
-
619
- # Expect using the allocator during model set for allocating the model, its input name and output name:
620
- # overall 3 allocations. The model raw size is 130B ,and the names are 2B each. In practice we allocate
621
- # more than 134B as Redis allocator will use additional memory for its internal management and for the
622
- # 64-Byte alignment. When the test runs with valgrind, redis will use malloc for the allocations.
623
- ret = con .execute_command ('AI.MODELSET' , 'm_gpu{1}' , 'ONNX' , DEVICE , 'BLOB' , model_pb )
624
- env .assertEqual (ret , b'OK' )
625
-
626
- # but for GPU, expect using the allocator only for allocating input and output names (not the model itself).
627
- ret = con .execute_command ('AI.MODELSET' , 'm_cpu{1}' , 'ONNX' , 'CPU' , 'BLOB' , model_pb )
628
- env .assertEqual (ret , b'OK' )
629
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
630
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
631
-
632
- # Expect using at least 130+63+(size of an address) + 4*(2+63+(size of an address)) bytes.
633
- model_allocation_bytes_used = int (ai_memory_config ["ai_onnxruntime_memory" ])
634
- env .assertTrue (model_allocation_bytes_used > 472 )
635
- env .assertTrue (model_allocation_bytes_used < 705 )
636
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory_access_num" ]), 5 )
637
-
638
- # Make sure that allocator is not used for running and freeing the GPU model, except for
639
- # the input and output names allocations (and deallocations).
640
- con .execute_command ('AI.TENSORSET' , 'a{1}' , 'FLOAT' , 3 , 2 , 'VALUES' , 1.0 , 2.0 , 3.0 , 4.0 , 5.0 , 6.0 )
641
- con .execute_command ('AI.MODELRUN' , 'm_gpu{1}' , 'INPUTS' , 'a{1}' , 'OUTPUTS' , 'b{1}' )
642
- values = con .execute_command ('AI.TENSORGET' , 'b{1}' , 'VALUES' )
643
- env .assertEqual (values , [b'1' , b'4' , b'9' , b'16' , b'25' , b'36' ])
644
- # Expect that memory usage didn't change, and for another 4 accesses to the allocator (input and output names
645
- # allocation and free)
646
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
647
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
648
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory" ]), model_allocation_bytes_used )
649
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory_access_num" ]), 9 )
650
-
651
- # Expect only 2 more accesses in delete - for deallocating input and output names
652
- con .execute_command ('AI.MODELDEL' , 'm_gpu{1}' )
653
- env .assertFalse (con .execute_command ('EXISTS' , 'm_gpu{1}' ))
654
- ai_memory_config = {k .split (":" )[0 ]: k .split (":" )[1 ]
655
- for k in con .execute_command ("INFO MODULES" ).decode ().split ("#" )[4 ].split ()[1 :]}
656
- env .assertEqual (int (ai_memory_config ["ai_onnxruntime_memory_access_num" ]), 11 )
657
-
0 commit comments