Coverage for src/sparkle/selector/extractor_cli.py: 92%

1#!/usr/bin/env python3 

2# -*- coding: UTF-8 -*- 

3"""Execute Feature Extractor for an instance, write features to FeatureDataFrame.""" 

4 

5import argparse 

6from pathlib import Path 

7from filelock import FileLock 

8 

9from sparkle.structures import FeatureDataFrame 

10from sparkle.selector import Extractor 

11from sparkle.instance import resolve_instance_pair 

12 

13 

14if __name__ == "__main__": 

15 # Define command line arguments 

16 parser = argparse.ArgumentParser() 

17 parser.add_argument( 

18 "--extractor", required=True, type=Path, help="path to feature extractor" 

19 ) 

20 parser.add_argument( 

21 "--instance", 

22 required=True, 

23 type=Path, 

24 nargs="+", 

25 help="path to instance file(s) to run on", 

26 ) 

27 parser.add_argument( 

28 "--feature-csv", required=True, type=Path, help="path to feature data CSV file" 

29 ) 

30 parser.add_argument( 

31 "--cutoff", 

32 required=True, 

33 type=str, 

34 help="the maximum CPU time for the extractor.", 

35 ) 

36 parser.add_argument( 

37 "--feature-group", 

38 required=False, 

39 type=str, 

40 help="the group of features to compute, if available for the " 

41 "extractor. If not available or provided, all groups will" 

42 " be computed.", 

43 ) 

44 parser.add_argument( 

45 "--log-dir", type=Path, required=True, help="path to the log directory" 

46 ) 

47 args = parser.parse_args() 

48 

49 # Process command line arguments 

50 log_dir = args.log_dir 

51 

52 # Instance agument is a list to allow for multifile instances 

53 instance_path: list[Path] = args.instance 

54 # We only receive the instance path on the CLI, but the FeatureDataFrame is keyed by 

55 # the canonical (set_name, instance_name) pair, so resolve it from the path. All files 

56 # of a multi-file instance share the same pair, so the first file resolves it. 

57 instance_set_name, instance_name = resolve_instance_pair(instance_path[0]) 

58 extractor_path = args.extractor 

59 feature_data_csv_path = args.feature_csv 

60 cutoff_extractor = args.cutoff 

61 

62 # Ensure stringifcation of path objects 

63 if isinstance(instance_path, list): 

64 instance_list = [str(filepath) for filepath in instance_path] 

65 else: 

66 instance_list = [str(instance_path)] 

67 

68 extractor = Extractor(extractor_path) 

69 if args.feature_group: 

70 print( 

71 f"Calling {extractor.name} with feature group {args.feature_group} for instance {instance_list} with cutoff {cutoff_extractor}" 

72 ) 

73 else: 

74 print( 

75 f"Calling {extractor.name} for instance {instance_list} with cutoff {cutoff_extractor}" 

76 ) 

77 

78 features = extractor.run( 

79 instance_list, 

80 feature_group=args.feature_group, 

81 cutoff_time=cutoff_extractor, 

82 log_dir=log_dir, 

83 ) 

84 

85 if features is None or len(features) == 0: 

86 raise ValueError( 

87 "No features found! This may be due to a timeout. Check extractor logs." 

88 ) 

89 

90 feature_data_per_group = {} 

91 for feature_group, feature_name, value in features: 

92 if feature_group not in feature_data_per_group: 

93 feature_data_per_group[feature_group] = [[], []] 

94 print( 

95 f"{extractor_path.name} {instance_set_name} {instance_name} {feature_group} {feature_name} | {value}" 

96 ) # For logging purposes 

97 feature_data_per_group[feature_group][0] += [feature_name] 

98 feature_data_per_group[feature_group][1] += [float(value)] 

99 

100 # Now that we have our result, we write it to the FeatureDataCSV with a FileLock 

101 lock = FileLock(f"{feature_data_csv_path}.lock") 

102 if features is not None: 

103 print("Writing features to file...") 

104 with lock.acquire(timeout=600): 

105 feature_data = FeatureDataFrame(feature_data_csv_path) 

106 for feature_group, ( 

107 feature_names, 

108 feature_values, 

109 ) in feature_data_per_group.items(): 

110 # for feature_group, feature_name, value in features: 

111 feature_data.set_value( 

112 instance_set_name, 

113 instance_name, 

114 extractor_path.name, 

115 feature_group, 

116 feature_names, 

117 feature_values, 

118 append_write_csv=True, 

119 ) 

120 lock.release() 

121 print("Writing successful!") 

122 else: 

123 print( 

124 "EXCEPTION during retrieving extractor results.\n" 

125 f"****** WARNING: Feature vector computation on instance {instance_path}" 

126 " failed! ******" 

127 )